diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx index 9eb449ca9..284a677ff 100644 --- a/docs/source/molmoact2.mdx +++ b/docs/source/molmoact2.mdx @@ -242,6 +242,17 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \ --repo-id=your_dataset ``` +Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram: + +```bash +python src/lerobot/scripts/augment_dataset_quantile_stats.py \ + --repo-id=your_dataset \ + --overwrite \ + --skip-images +``` + +`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes MolmoAct2's normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected. + Alternatively, train MolmoAct2 with mean/std normalization: ```bash diff --git a/docs/source/pi05.mdx b/docs/source/pi05.mdx index f8ca8bc74..bc3720302 100644 --- a/docs/source/pi05.mdx +++ b/docs/source/pi05.mdx @@ -127,6 +127,17 @@ lerobot-edit-dataset \ Or keep the dataset as-is and pass `--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}'`. +Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram: + +```bash +python src/lerobot/scripts/augment_dataset_quantile_stats.py \ + --repo-id=your_dataset \ + --overwrite \ + --skip-images +``` + +`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes π₀.₅'s normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected. + ### Training Command Example The same finetune with the VLM frozen: less memory, at some cost in success rate. Swap `--dataset.repo_id` for your own dataset. diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 02ecd81a4..e7f0abce9 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -613,8 +613,15 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d for q_key in quantile_keys: if all(q_key in s for s in stats_ft_list): quantile_values = np.stack([s[q_key] for s in stats_ft_list]) - weighted_quantiles = quantile_values * counts - aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count + # Exact global quantiles cannot be recovered from quantile summaries. + # Keep a conservative envelope of the available estimates: min + # for lower quantiles and max for upper quantiles. The resulting + # values are bounds across the inputs, not global quantile estimates. + q_percent = int(q_key[1:]) + if q_percent <= 50: + aggregated[q_key] = np.min(quantile_values, axis=0) + else: + aggregated[q_key] = np.max(quantile_values, axis=0) return aggregated diff --git a/src/lerobot/scripts/augment_dataset_quantile_stats.py b/src/lerobot/scripts/augment_dataset_quantile_stats.py index 2faca7034..f6578ae29 100644 --- a/src/lerobot/scripts/augment_dataset_quantile_stats.py +++ b/src/lerobot/scripts/augment_dataset_quantile_stats.py @@ -25,6 +25,11 @@ quantile statistics (q01, q10, q50, q90, q99) in their metadata. This script: 3. If missing, computes quantile statistics for all features 4. Updates the dataset metadata with the new quantile statistics +Statistics are accumulated into a single running histogram per feature across +all episodes rather than aggregating per-episode quantile summaries. The +resulting quantiles are histogram approximations, subject to discretization and +range-rebinning error; image/video frames are sampled by default. + Usage: ```bash @@ -34,9 +39,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 @@ -49,11 +52,10 @@ from lerobot.datasets import ( CODEBASE_VERSION, DEFAULT_QUANTILES, LeRobotDataset, - aggregate_stats, get_feature_stats, write_stats, ) -from lerobot.datasets.compute_stats import sample_indices +from lerobot.datasets.compute_stats import RunningQuantileStats, sample_indices from lerobot.utils.utils import init_logging @@ -79,20 +81,25 @@ 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, use_sampling: bool = True) -> dict: - """Process a single episode and return its statistics. +def collect_episode_arrays( + dataset: LeRobotDataset, + episode_idx: int, + use_sampling: bool = True, + skip_images: bool = False, +) -> dict[str, tuple[np.ndarray, int]]: + """Collect one episode's frames per feature, flattened to (num_samples, dim). 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). + episode_idx: Index of the episode to read + use_sampling: If True, sub-sample image/video frames to bound memory. + If False, use every frame (higher memory). + skip_images: If True, skip image/video features entirely. Returns: - Dictionary containing episode statistics + Mapping of feature name to that episode's values and the number of frames + they came from (which differs from the row count for image features). """ - logging.info(f"Computing stats for episode {episode_idx}") - start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"] end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"] @@ -102,7 +109,9 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli # 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") + k + for k in dataset.features + if dataset.features[k]["dtype"] not in ("image", "video", "string", "language") ] collected_data: dict[str, list] = {} @@ -114,7 +123,7 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli 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: + if image_keys and not skip_images: sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len)) for offset in sampled_offsets: item = dataset[start_idx + offset] @@ -122,87 +131,82 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli if key in item: collected_data.setdefault(key, []).append(item[key]) - ep_stats = {} + episode_arrays: dict[str, tuple[np.ndarray, int]] = {} for key, data_list in collected_data.items(): - if dataset.features[key]["dtype"] == "string": - continue - data = torch.stack(data_list).cpu().numpy() if dataset.features[key]["dtype"] in ["image", "video"]: if data.dtype == np.uint8: data = data.astype(np.float32) / 255.0 - - axes_to_reduce = (0, 2, 3) - keepdims = True + # (N, C, H, W) -> (N * H * W, C) so quantiles are computed per channel. + channels = data.shape[1] + values = data.transpose(0, 2, 3, 1).reshape(-1, channels) else: - axes_to_reduce = 0 - keepdims = data.ndim == 1 + values = data.reshape(-1, data.shape[-1]) if data.ndim > 1 else data.reshape(-1, 1) + episode_arrays[key] = (values, len(data_list)) - ep_stats[key] = get_feature_stats( - data, axis=axes_to_reduce, keepdims=keepdims, quantile_list=DEFAULT_QUANTILES - ) - - if dataset.features[key]["dtype"] in ["image", "video"]: - ep_stats[key] = { - k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items() - } - - return ep_stats + return episode_arrays -def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]: - """Compute quantile statistics for all episodes in the dataset. +def compute_quantile_stats_for_dataset( + dataset: LeRobotDataset, + use_sampling: bool = True, + skip_images: bool = False, +) -> dict[str, dict]: + """Compute whole-dataset statistics with one running histogram per feature. 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). + memory. If False, use every frame (higher memory). + skip_images: If True, skip image/video features and leave their stats untouched. Returns: - Dictionary containing aggregated statistics with quantiles + Dictionary containing statistics with histogram-based global quantile estimates Note: - Video decoding operations are not thread-safe, so we process episodes sequentially - when video keys are present. For datasets without videos, we use parallel processing - with ThreadPoolExecutor for better performance. + Episodes are accumulated sequentially because the running accumulators are + shared across all of them. """ logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes") - episode_stats_list = [] - has_videos = len(dataset.meta.video_keys) > 0 + running_stats: dict[str, RunningQuantileStats] = {} + frame_counts: dict[str, int] = {} + row_counts: dict[str, int] = {} + # Kept only while a feature has a single row, so it can still be finalized. + single_row_arrays: dict[str, np.ndarray] = {} - 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, 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, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16))) + for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"): + episode_arrays = collect_episode_arrays( + dataset, episode_idx, use_sampling=use_sampling, skip_images=skip_images + ) + for key, (array, num_frames) in episode_arrays.items(): + running_stats.setdefault(key, RunningQuantileStats()).update(array) + frame_counts[key] = frame_counts.get(key, 0) + num_frames + row_counts[key] = row_counts.get(key, 0) + len(array) + if row_counts[key] < 2: + single_row_arrays[key] = array + else: + single_row_arrays.pop(key, None) - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - future_to_episode = { - executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx - for episode_idx in range(dataset.num_episodes) - } - - episode_results = {} - with tqdm(total=dataset.num_episodes, desc="Processing episodes") as pbar: - for future in concurrent.futures.as_completed(future_to_episode): - episode_idx = future_to_episode[future] - ep_stats = future.result() - episode_results[episode_idx] = ep_stats - pbar.update(1) - - for episode_idx in range(dataset.num_episodes): - if episode_idx in episode_results: - episode_stats_list.append(episode_results[episode_idx]) - - if not episode_stats_list: + if not running_stats: raise ValueError("No episode data found for computing statistics") - logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes") - return aggregate_stats(episode_stats_list) + aggregated_stats: dict[str, dict] = {} + for key, accumulator in running_stats.items(): + if row_counts[key] < 2: + # Histograms need at least two samples; mirror get_feature_stats' basic-stats path. + stats = get_feature_stats(single_row_arrays[key], axis=0, keepdims=False) + else: + stats = accumulator.get_statistics() + if dataset.features[key]["dtype"] in ["image", "video"]: + # Image stats are stored as (C, 1, 1) to broadcast over height and width. + stats = {k: v if k == "count" else v[:, np.newaxis, np.newaxis] for k, v in stats.items()} + # `get_feature_stats` counts frames, not the per-channel rows the accumulator sees. + stats["count"] = np.array([frame_counts[key]]) + aggregated_stats[key] = stats + + logging.info(f"Computed global histogram statistics for {len(aggregated_stats)} features") + return aggregated_stats def augment_dataset_with_quantile_stats( @@ -210,6 +214,7 @@ def augment_dataset_with_quantile_stats( root: str | Path | None = None, overwrite: bool = False, use_sampling: bool = True, + skip_images: bool = False, ) -> None: """Augment a dataset with quantile statistics if they are missing. @@ -218,7 +223,8 @@ def augment_dataset_with_quantile_stats( 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). + memory. If False, use every frame (higher memory). + skip_images: If True, skip image/video features and keep their existing stats """ logging.info(f"Loading dataset: {repo_id}") dataset = LeRobotDataset( @@ -232,7 +238,13 @@ 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, use_sampling=use_sampling) + new_stats = compute_quantile_stats_for_dataset( + dataset, use_sampling=use_sampling, skip_images=skip_images + ) + + if skip_images and dataset.meta.stats: + for key, feature_stats in dataset.meta.stats.items(): + new_stats.setdefault(key, feature_stats) logging.info("Updating dataset metadata with new quantile statistics") dataset.meta.stats = new_stats @@ -276,10 +288,15 @@ def main(): "--no-sampling", action="store_true", help=( - "Compute stats over every frame (exact, higher memory). By default, " + "Compute stats over every frame (higher memory). By default, " "image/video frames are sub-sampled per episode to bound memory." ), ) + parser.add_argument( + "--skip-images", + action="store_true", + help="Skip image/video features and preserve their existing stats", + ) args = parser.parse_args() root = Path(args.root) if args.root else None @@ -291,6 +308,7 @@ def main(): root=root, overwrite=args.overwrite, use_sampling=not args.no_sampling, + skip_images=args.skip_images, ) diff --git a/tests/datasets/test_augment_quantile_stats.py b/tests/datasets/test_augment_quantile_stats.py index e07900e02..2314af7b6 100644 --- a/tests/datasets/test_augment_quantile_stats.py +++ b/tests/datasets/test_augment_quantile_stats.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace + import numpy as np import pytest +import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") @@ -24,7 +27,9 @@ from lerobot.scripts.augment_dataset_quantile_stats import ( def _numeric_keys(dataset): - return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")] + return [ + k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string", "language") + ] def _image_keys(dataset): @@ -102,3 +107,112 @@ def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory) ) stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True) assert has_quantile_stats(stats) + + +class FakeHFDataset: + """Minimal stand-in exposing the column slicing used by the augment script.""" + + def __init__(self, columns: dict[str, list]): + self._columns = columns + + def select_columns(self, keys): + return FakeHFDataset({key: self._columns[key] for key in keys}) + + def __getitem__(self, index): + return {key: values[index] for key, values in self._columns.items()} + + +def test_compute_quantile_stats_skips_language_features(): + class FakeDataset: + num_episodes = 1 + features = { + "action": {"dtype": "float32"}, + "observation.language": {"dtype": "language"}, + } + meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}]) + hf_dataset = FakeHFDataset( + { + "action": [[0.0], [1.0]], + "observation.language": [ + [{"role": "user", "content": "pick"}], + [{"role": "assistant", "content": "done"}], + ], + } + ) + + stats = compute_quantile_stats_for_dataset(FakeDataset()) + + assert set(stats) == {"action"} + + +def test_compute_quantile_stats_skip_images_avoids_decoding(): + class FakeDataset: + num_episodes = 1 + features = { + "action": {"dtype": "float32"}, + "observation.images.cam": {"dtype": "video"}, + } + meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}]) + hf_dataset = FakeHFDataset({"action": [[0.0], [1.0]]}) + + def __getitem__(self, index): + raise AssertionError(f"video frame {index} was decoded despite skip_images=True") + + stats = compute_quantile_stats_for_dataset(FakeDataset(), skip_images=True) + + assert set(stats) == {"action"} + + +def test_compute_quantile_stats_handles_single_frame(): + class FakeDataset: + num_episodes = 1 + features = {"action": {"dtype": "float32"}} + meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 1}]) + hf_dataset = FakeHFDataset({"action": [[5.0, 7.0]]}) + + stats = compute_quantile_stats_for_dataset(FakeDataset()) + + np.testing.assert_array_equal(stats["action"]["count"], np.array([1])) + for key in ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99"): + np.testing.assert_allclose(stats["action"][key], np.array([5.0, 7.0])) + + +def test_compute_quantile_stats_image_count_uses_frames(): + frames = [torch.zeros(3, 2, 2), torch.ones(3, 2, 2)] + + class FakeDataset: + num_episodes = 1 + features = {"observation.images.cam": {"dtype": "video"}} + meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}]) + hf_dataset = FakeHFDataset({}) + + def __getitem__(self, index): + return {"observation.images.cam": frames[index]} + + stats = compute_quantile_stats_for_dataset(FakeDataset(), use_sampling=False) + image_stats = stats["observation.images.cam"] + + np.testing.assert_array_equal(image_stats["count"], np.array([2])) + assert image_stats["mean"].shape == (3, 1, 1) + np.testing.assert_allclose(image_stats["mean"], np.full((3, 1, 1), 0.5)) + + +def test_compute_quantile_stats_accumulates_across_episodes(): + values = [[float(value)] for value in range(100)] + [[float(value)] for value in range(1000, 1010)] + + class FakeDataset: + num_episodes = 2 + features = {"action": {"dtype": "float32"}} + meta = SimpleNamespace( + episodes=[ + {"dataset_from_index": 0, "dataset_to_index": 100}, + {"dataset_from_index": 100, "dataset_to_index": 110}, + ] + ) + hf_dataset = FakeHFDataset({"action": values}) + + stats = compute_quantile_stats_for_dataset(FakeDataset()) + + np.testing.assert_array_equal(stats["action"]["count"], np.array([110])) + expected_q90 = np.percentile(np.asarray(values), 90, axis=0) + np.testing.assert_allclose(stats["action"]["q90"], expected_q90, atol=0.1) diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py index 9f399b85c..8d730fcb9 100644 --- a/tests/datasets/test_compute_stats.py +++ b/tests/datasets/test_compute_stats.py @@ -688,7 +688,7 @@ def test_compute_episode_stats_string_features_skipped(): def test_aggregate_feature_stats_with_quantiles(): - """Test aggregating feature stats that include quantiles.""" + """Test aggregating feature stats that include quantiles uses conservative bounds.""" stats_ft_list = [ { "min": np.array([1.0]), @@ -697,6 +697,9 @@ def test_aggregate_feature_stats_with_quantiles(): "std": np.array([2.0]), "count": np.array([100]), "q01": np.array([1.5]), + "q10": np.array([2.0]), + "q50": np.array([5.0]), + "q90": np.array([9.0]), "q99": np.array([9.5]), }, { @@ -706,22 +709,21 @@ def test_aggregate_feature_stats_with_quantiles(): "std": np.array([2.5]), "count": np.array([150]), "q01": np.array([2.5]), + "q10": np.array([3.0]), + "q50": np.array([6.0]), + "q90": np.array([11.0]), "q99": np.array([11.5]), }, ] result = aggregate_feature_stats(stats_ft_list) - # Should preserve quantiles - assert "q01" in result - assert "q99" in result - - # Verify quantile aggregation (weighted average) - expected_q01 = (1.5 * 100 + 2.5 * 150) / 250 # ≈ 2.1 - expected_q99 = (9.5 * 100 + 11.5 * 150) / 250 # ≈ 10.7 - - np.testing.assert_allclose(result["q01"], np.array([expected_q01]), atol=1e-6) - np.testing.assert_allclose(result["q99"], np.array([expected_q99]), atol=1e-6) + # Lower quantiles use min; upper quantiles use max, regardless of counts. + np.testing.assert_allclose(result["q01"], np.array([1.5]), atol=1e-6) + np.testing.assert_allclose(result["q10"], np.array([2.0]), atol=1e-6) + np.testing.assert_allclose(result["q50"], np.array([5.0]), atol=1e-6) + np.testing.assert_allclose(result["q90"], np.array([11.0]), atol=1e-6) + np.testing.assert_allclose(result["q99"], np.array([11.5]), atol=1e-6) def test_aggregate_stats_mixed_quantiles(): @@ -878,3 +880,60 @@ def test_fixed_quantiles_always_computed(): for q_key in expected_quantiles: assert q_key in episode_stats[key] assert episode_stats[key][q_key].shape == (features[key]["shape"][0],) + + +def test_aggregate_stats_incremental_resume(): + """Verify conservative bounds remain associative across incremental additions.""" + # Start with episode 1 stats (narrow distribution) + ep1_stats = { + "action": { + "min": np.array([-10.0, -5.0]), + "max": np.array([10.0, 5.0]), + "mean": np.array([0.0, 0.0]), + "std": np.array([3.0, 1.5]), + "count": np.array([500]), + "q01": np.array([-9.0, -4.5]), + "q99": np.array([9.0, 4.5]), + }, + } + + # Episode 2: wider distribution on dim 0 + ep2_stats = { + "action": { + "min": np.array([-30.0, -5.0]), + "max": np.array([40.0, 6.0]), + "mean": np.array([5.0, 0.5]), + "std": np.array([15.0, 2.0]), + "count": np.array([100]), + "q01": np.array([-25.0, -4.0]), + "q99": np.array([35.0, 5.5]), + }, + } + + # First aggregation: ep1 + ep2 (simulates save_episode for ep2) + cumulative = aggregate_stats([ep1_stats, ep2_stats]) + + # q01 should take min (conservative lower bound) + np.testing.assert_allclose(cumulative["action"]["q01"], np.array([-25.0, -4.5])) + # q99 should take max (conservative upper bound) + np.testing.assert_allclose(cumulative["action"]["q99"], np.array([35.0, 5.5])) + + # Episode 3: even wider on dim 1 + ep3_stats = { + "action": { + "min": np.array([-8.0, -20.0]), + "max": np.array([8.0, 25.0]), + "mean": np.array([0.0, 3.0]), + "std": np.array([2.0, 8.0]), + "count": np.array([50]), + "q01": np.array([-7.0, -18.0]), + "q99": np.array([7.0, 22.0]), + }, + } + + # Second aggregation: cumulative + ep3 (simulates save_episode for ep3) + cumulative2 = aggregate_stats([cumulative, ep3_stats]) + + # Bounds should widen monotonically + np.testing.assert_allclose(cumulative2["action"]["q01"], np.array([-25.0, -18.0])) + np.testing.assert_allclose(cumulative2["action"]["q99"], np.array([35.0, 22.0]))