From 35339d31e53dc9b0cb4ea2a7f01230f108ab0620 Mon Sep 17 00:00:00 2001 From: Predrag Cvetkovic Date: Wed, 29 Jul 2026 12:32:03 +0200 Subject: [PATCH] fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749) * fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames Per-episode stats previously materialized every frame (and decoded up to 16 episodes in parallel), so peak memory scaled with episode length and OOM'd on large datasets (#2889). Numeric features are now read in full from the table (exact), while only image/video frames are sub-sampled per episode using the existing sample_indices heuristic. Worker count is configurable via LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior. * Update tests/datasets/test_augment_quantile_stats.py Co-authored-by: Haoming Song <1847575517@qq.com> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Haoming Song <1847575517@qq.com> --- .../scripts/augment_dataset_quantile_stats.py | 63 ++++++++--- tests/datasets/test_augment_quantile_stats.py | 104 ++++++++++++++++++ 2 files changed, 152 insertions(+), 15 deletions(-) create mode 100644 tests/datasets/test_augment_quantile_stats.py diff --git a/src/lerobot/scripts/augment_dataset_quantile_stats.py b/src/lerobot/scripts/augment_dataset_quantile_stats.py index 4ee99a541..2faca7034 100644 --- a/src/lerobot/scripts/augment_dataset_quantile_stats.py +++ b/src/lerobot/scripts/augment_dataset_quantile_stats.py @@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \ import argparse import concurrent.futures import logging +import os from pathlib import Path import numpy as np @@ -52,6 +53,7 @@ from lerobot.datasets import ( get_feature_stats, write_stats, ) +from lerobot.datasets.compute_stats import sample_indices from lerobot.utils.utils import init_logging @@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s return False -def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict: +def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict: """Process a single episode and return its statistics. Args: dataset: The LeRobot dataset episode_idx: Index of the episode to process + use_sampling: If True, sub-sample image/video frames per episode to bound + memory. If False, use every frame (exact, higher memory). Returns: Dictionary containing episode statistics @@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict: start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"] end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"] - collected_data: dict[str, list] = {} - for idx in range(start_idx, end_idx): - item = dataset[idx] - for key, value in item.items(): - if key not in dataset.features: - continue + episode_len = end_idx - start_idx - if key not in collected_data: - collected_data[key] = [] - collected_data[key].append(value) + # Images/video are the memory hog, so sub-sample those frames per episode; + # numeric columns are cheap, so read them in full (exact). + image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")] + numeric_keys = [ + k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string") + ] + + collected_data: dict[str, list] = {} + + # Numeric features: every frame, read directly from the underlying table. + if numeric_keys: + numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx] + for key in numeric_keys: + collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]] + + # Image/video features: decode only a sampled subset of frames. + if image_keys: + sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len)) + for offset in sampled_offsets: + item = dataset[start_idx + offset] + for key in image_keys: + if key in item: + collected_data.setdefault(key, []).append(item[key]) ep_stats = {} for key, data_list in collected_data.items(): @@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict: return ep_stats -def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]: +def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]: """Compute quantile statistics for all episodes in the dataset. Args: dataset: The LeRobot dataset to compute statistics for + use_sampling: If True, sub-sample image/video frames per episode to bound + memory. If False, use every frame (exact, higher memory). Returns: Dictionary containing aggregated statistics with quantiles @@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic if has_videos: logging.info("Dataset contains video keys - using sequential processing for thread safety") for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"): - ep_stats = process_single_episode(dataset, episode_idx) + ep_stats = process_single_episode(dataset, episode_idx, use_sampling) episode_stats_list.append(ep_stats) else: logging.info("Dataset has no video keys - using parallel processing for better performance") - max_workers = min(dataset.num_episodes, 16) + max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16))) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_episode = { - executor.submit(process_single_episode, dataset, episode_idx): episode_idx + executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx for episode_idx in range(dataset.num_episodes) } @@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats( repo_id: str, root: str | Path | None = None, overwrite: bool = False, + use_sampling: bool = True, ) -> None: """Augment a dataset with quantile statistics if they are missing. @@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats( repo_id: Repository ID of the dataset root: Local root directory for the dataset overwrite: Overwrite existing quantile statistics if they already exist + use_sampling: If True, sub-sample image/video frames per episode to bound + memory. If False, use every frame (exact, higher memory). """ logging.info(f"Loading dataset: {repo_id}") dataset = LeRobotDataset( @@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats( logging.info("Dataset does not contain quantile statistics. Computing them now...") - new_stats = compute_quantile_stats_for_dataset(dataset) + new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling) logging.info("Updating dataset metadata with new quantile statistics") dataset.meta.stats = new_stats @@ -248,6 +272,14 @@ def main(): action="store_true", help="Overwrite existing quantile statistics if they already exist", ) + parser.add_argument( + "--no-sampling", + action="store_true", + help=( + "Compute stats over every frame (exact, higher memory). By default, " + "image/video frames are sub-sampled per episode to bound memory." + ), + ) args = parser.parse_args() root = Path(args.root) if args.root else None @@ -258,6 +290,7 @@ def main(): repo_id=args.repo_id, root=root, overwrite=args.overwrite, + use_sampling=not args.no_sampling, ) diff --git a/tests/datasets/test_augment_quantile_stats.py b/tests/datasets/test_augment_quantile_stats.py new file mode 100644 index 000000000..e07900e02 --- /dev/null +++ b/tests/datasets/test_augment_quantile_stats.py @@ -0,0 +1,104 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.scripts.augment_dataset_quantile_stats import ( + compute_quantile_stats_for_dataset, + has_quantile_stats, +) + + +def _numeric_keys(dataset): + return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")] + + +def _image_keys(dataset): + return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")] + + +def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory): + """Sampling only touches image/video frames; numeric features are read in + full either way, so their stats must be identical with and without sampling.""" + dataset = lerobot_dataset_factory( + root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False + ) + + exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False) + sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True) + + numeric_keys = _numeric_keys(dataset) + assert numeric_keys, "fixture should expose numeric features" + for key in numeric_keys: + if key not in exact: + continue + for stat in ("mean", "std", "q01", "q50", "q99"): + if stat in exact[key]: + np.testing.assert_allclose( + sampled[key][stat], + exact[key][stat], + rtol=1e-6, + atol=1e-6, + err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling", + ) + + +def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory): + """For images, sampling should reduce the number of samples considered while + keeping the resulting statistics close to the exact ones.""" + dataset = lerobot_dataset_factory( + root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False + ) + + exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False) + sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True) + + image_keys = _image_keys(dataset) + assert image_keys, "fixture should expose at least one image feature" + for key in image_keys: + # sampling actually looked at fewer pixels + assert sampled[key]["count"][0] < exact[key]["count"][0] + # but per-channel mean stays close + np.testing.assert_allclose( + sampled[key]["mean"], + exact[key]["mean"], + rtol=0.15, + err_msg=f"image feature '{key}' mean drifted too far under sampling", + ) + + +def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory): + """With episodes shorter than the sampling floor, sampling is a no-op and + must produce exactly the same stats as the exact path.""" + dataset = lerobot_dataset_factory( + root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False + ) + + exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False) + sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True) + + for key in _image_keys(dataset): + assert sampled[key]["count"][0] == exact[key]["count"][0] + + +def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory): + """The computed stats should contain quantile keys for the dataset.""" + dataset = lerobot_dataset_factory( + root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False + ) + stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True) + assert has_quantile_stats(stats)