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>

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] = {}
+    episode_len = end_idx - start_idx
-    for idx in range(start_idx, end_idx):
-        item = dataset[idx]
-        for key, value in item.items():
-            if key not in dataset.features:
-                continue
 
-            if key not in collected_data:
+    # Images/video are the memory hog, so sub-sample those frames per episode;
-                collected_data[key] = []
+    # numeric columns are cheap, so read them in full (exact).
-            collected_data[key].append(value)
+    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,
     )
 
 

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)