From d381b27e9748e5d524ad041fd92adbe9c057b92d Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Tue, 14 Jul 2026 16:55:01 +0200 Subject: [PATCH] perf(datasets): bound memory of visual episode stats via batched streaming Accumulate image/video per-episode stats into a RunningQuantileStats in frame batches instead of materialising every sampled frame at once. Peak memory is now bounded by one batch (frame_batch_size x C x H x W) regardless of episode length, preventing OOM on long, high-resolution episodes. count keeps the per-frame convention of get_feature_stats. --- src/lerobot/datasets/dataset_tools.py | 32 +++++++++++++++++++++------ 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 9beb8a6b3..405910727 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -53,11 +53,11 @@ from lerobot.utils.utils import flatten_dict from .aggregate import aggregate_datasets from .compute_stats import ( + RunningQuantileStats, aggregate_stats, auto_downsample_height_width, compute_episode_stats, compute_relative_action_stats, - get_feature_stats, sample_indices, ) from .dataset_metadata import LeRobotDatasetMetadata @@ -1633,12 +1633,18 @@ def _compute_visual_episode_stats( dataset: LeRobotDataset, ep_idx: int, visual_keys: list[str], + frame_batch_size: int = 32, ) -> dict: """Compute per-episode statistics for image/video features by sampling frames. Mirrors the image/video branch of :func:`compute_episode_stats`: per-channel stats are computed on downsampled sampled frames, then RGB stats are rescaled to [0, 1] (depth maps keep their native units). + + Frames are decoded and accumulated into a :class:`RunningQuantileStats` in batches of + ``frame_batch_size`` rather than materialising every sampled frame at once. Peak memory + is bounded by one batch (``frame_batch_size x C x H x W``) regardless of episode length, + which keeps long, high-resolution episodes from exhausting memory. """ ep_length = dataset.meta.episodes[ep_idx]["length"] frame_offsets = sample_indices(ep_length) @@ -1646,15 +1652,27 @@ def _compute_visual_episode_stats( ep_stats = {} for key in visual_keys: is_depth = key in dataset.meta.depth_keys - if dataset.meta.features[key]["dtype"] == "video": - frames = _load_episode_video_frames(dataset, key, ep_idx, frame_offsets, is_depth) - else: - frames = _load_episode_image_frames(dataset, key, ep_idx, frame_offsets, is_depth) + is_video = dataset.meta.features[key]["dtype"] == "video" - stats = get_feature_stats(frames, axis=(0, 2, 3), keepdims=True) + running = RunningQuantileStats() + for start in range(0, len(frame_offsets), frame_batch_size): + batch_offsets = frame_offsets[start : start + frame_batch_size] + if is_video: + frames = _load_episode_video_frames(dataset, key, ep_idx, batch_offsets, is_depth) + else: + frames = _load_episode_image_frames(dataset, key, ep_idx, batch_offsets, is_depth) + # (N, C, H, W) -> (N * H * W, C) so stats are accumulated per channel. + running.update(np.moveaxis(frames, 1, -1).reshape(-1, frames.shape[1])) + + stats = running.get_statistics() normalization_factor = 1.0 if is_depth else 255.0 + num_channels = stats["mean"].shape[0] + # ``count`` follows the per-frame convention of ``get_feature_stats`` (number of + # sampled frames), not the per-pixel count tracked internally by RunningQuantileStats. ep_stats[key] = { - k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0) + k: np.array([len(frame_offsets)]) + if k == "count" + else v.reshape(num_channels, 1, 1) / normalization_factor for k, v in stats.items() }