From 9372f52fff0019bac0a46f381082e727a5a90a68 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 10 Jul 2026 23:35:08 +0200 Subject: [PATCH] feat(datasets): optionally rewrite per-episode stats when recomputing Add an opt-in update_episode_stats flag so recomputing stats can also rewrite the per-episode stats/* columns in the episodes parquet, keeping them consistent with meta/stats.json. compute_dataset_episode_stats now returns a {episode_index: stats} mapping so stats can be written back to the right episode and shards merge by key. Wired into recompute_stats, the lerobot-edit-dataset CLI, and the SLURM example (per-episode shards + --update-episode-stats). --- examples/dataset/slurm_recompute_stats.py | 42 ++++++++-- src/lerobot/datasets/__init__.py | 2 + src/lerobot/datasets/dataset_tools.py | 88 ++++++++++++++++++--- src/lerobot/scripts/lerobot_edit_dataset.py | 16 +++- 4 files changed, 127 insertions(+), 21 deletions(-) diff --git a/examples/dataset/slurm_recompute_stats.py b/examples/dataset/slurm_recompute_stats.py index 523b442df..eda62d265 100644 --- a/examples/dataset/slurm_recompute_stats.py +++ b/examples/dataset/slurm_recompute_stats.py @@ -27,6 +27,9 @@ Modified copy of lerobot's examples/dataset/slurm_recompute_stats.py 4. --chain-aggregate : submit ``aggregate`` with an afterok dependency on ``compute`` so it only runs once all shards exist (no manual squeue-wait, no gap/overlap race). + 5. --update-episode-stats : in ``aggregate``, also rewrite the per-episode stats in the + episodes parquet so they stay consistent with meta/stats.json + (default: only stats.json is written). Data access: no filesystem mount. Point HF_LEROBOT_HOME at a node-visible shared cache (e.g. /fsx/$USER/.cache) so the dataset downloads once and all workers read @@ -120,7 +123,16 @@ class ComputeEpisodeStatsShards(PipelineStep): class AggregateEpisodeStats(PipelineStep): """Merge all per-episode stat shards into meta/stats.json.""" - def __init__(self, repo_id, root, new_root, shard_dir, push_to_hub=False, video_backend=None): + def __init__( + self, + repo_id, + root, + new_root, + shard_dir, + push_to_hub=False, + video_backend=None, + update_episode_stats=False, + ): super().__init__() self.repo_id = repo_id self.root = root @@ -128,6 +140,7 @@ class AggregateEpisodeStats(PipelineStep): self.shard_dir = shard_dir self.push_to_hub = push_to_hub self.video_backend = video_backend + self.update_episode_stats = update_episode_stats def run(self, data=None, rank: int = 0, world_size: int = 1): # NOTE: pickled and executed on a worker; keep self-contained (see ComputeEpisodeStatsShards.run). @@ -147,10 +160,12 @@ class AggregateEpisodeStats(PipelineStep): if not shards: raise FileNotFoundError(f"No episode stat shards found in {shard_dir}") - all_episode_stats = [] + # Shards map episode_index -> stats; merging by key makes a dropped shard show up as a + # missing episode and a re-run shard overwrite rather than double-count. + all_episode_stats = {} for shard in shards: with open(shard, "rb") as f: - all_episode_stats.extend(pickle.load(f)) + all_episode_stats.update(pickle.load(f)) logging.info(f"Aggregating {len(all_episode_stats)} episode stats from {len(shards)} shards") load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {} @@ -170,23 +185,26 @@ class AggregateEpisodeStats(PipelineStep): # Frame-count check catches the case where a duplicate and a gap cancel out in the # episode count: summed per-episode frame counts must equal the dataset's total frames. + stats_values = list(all_episode_stats.values()) numeric_key = next( ( k for k, v in dataset.meta.features.items() - if v["dtype"] not in ("image", "video", "string") and all_episode_stats and k in all_episode_stats[0] + if v["dtype"] not in ("image", "video", "string") and stats_values and k in stats_values[0] ), None, ) if numeric_key is not None: - total_frames = sum(int(s[numeric_key]["count"][0]) for s in all_episode_stats) + total_frames = sum(int(s[numeric_key]["count"][0]) for s in stats_values) if total_frames != dataset.meta.total_frames: raise ValueError( f"Summed frame count from shards ({total_frames}) != dataset total_frames " f"({dataset.meta.total_frames}); episodes are double-counted or missing." ) - new_stats = aggregate_episode_stats(dataset, all_episode_stats) + new_stats = aggregate_episode_stats( + dataset, all_episode_stats, update_episode_stats=self.update_episode_stats + ) if new_stats is None: raise RuntimeError("Aggregation produced no stats") logging.info(f"Wrote stats for features: {list(new_stats.keys())} to {dataset.root}") @@ -316,10 +334,20 @@ def main(): help="After building compute, submit aggregate with an afterok dependency (single command).", ) cp.add_argument("--push-to-hub", action="store_true", help="For the chained aggregate: push after done.") + cp.add_argument( + "--update-episode-stats", + action="store_true", + help="For the chained aggregate: also rewrite per-episode stats in the episodes parquet.", + ) ap = sub.add_parser("aggregate", help="Merge shards into meta/stats.json.") _add_shared_args(ap) ap.add_argument("--push-to-hub", action="store_true", help="Push the dataset after aggregation.") + ap.add_argument( + "--update-episode-stats", + action="store_true", + help="Also rewrite per-episode stats in the episodes parquet to match stats.json.", + ) ap.add_argument( "--depends-job-id", type=str, @@ -372,6 +400,7 @@ def main(): str(args.shard_dir), args.push_to_hub, args.video_backend, + args.update_episode_stats, ) ], logs_dir=args.logs_dir, @@ -401,6 +430,7 @@ def main(): str(args.shard_dir), args.push_to_hub, args.video_backend, + args.update_episode_stats, ) ], logs_dir=args.logs_dir, diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index 9c6342898..51cef4497 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -36,6 +36,7 @@ from .dataset_tools import ( reencode_dataset, remove_feature, split_dataset, + write_episode_stats, ) from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps from .image_writer import safe_stop_image_writer @@ -103,5 +104,6 @@ __all__ = [ "resolve_delta_timestamps", "safe_stop_image_writer", "split_dataset", + "write_episode_stats", "write_stats", ] diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 92a8d6767..9beb8a6b3 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -33,6 +33,7 @@ from pathlib import Path import datasets import numpy as np import pandas as pd +import pyarrow as pa import pyarrow.parquet as pq import torch from tqdm import tqdm @@ -1665,12 +1666,12 @@ def compute_dataset_episode_stats( episode_indices: list[int] | None = None, skip_image_video: bool = True, drop_keys: list[str] | None = None, -) -> list[dict]: +) -> dict[int, dict]: """Compute per-episode statistics for a subset of episodes. This is the shardable unit of work behind :func:`recompute_stats`: distribute ``episode_indices`` across workers (e.g. ``list(range(n))[rank::world_size]``), - then combine the concatenated results with :func:`aggregate_stats`. + then combine the results with :func:`aggregate_episode_stats`. Args: dataset: The LeRobotDataset to compute stats for. @@ -1681,7 +1682,9 @@ def compute_dataset_episode_stats( in relative-action space). Returns: - A list of per-episode stat dicts, one per processed episode. + A mapping of episode index to its per-episode stat dict. Keeping the episode index + (rather than a bare list) lets callers write the stats back to the correct episode + row, and survives sharding since shards can be merged by key. """ features = dataset.meta.features meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"} @@ -1708,7 +1711,7 @@ def compute_dataset_episode_stats( for ep_idx in episode_indices: file_to_episodes.setdefault(dataset.meta.get_data_file_path(ep_idx), []).append(ep_idx) - all_episode_stats = [] + all_episode_stats = {} for src_path, eps in tqdm(sorted(file_to_episodes.items()), desc="Computing stats from data files"): df = pd.read_parquet(dataset.root / src_path) if numeric_keys else None for ep_idx in sorted(eps): @@ -1725,7 +1728,7 @@ def compute_dataset_episode_stats( ep_stats = compute_episode_stats(episode_data, features_to_compute) if visual_keys: ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys)) - all_episode_stats.append(ep_stats) + all_episode_stats[int(ep_idx)] = ep_stats return all_episode_stats @@ -1737,6 +1740,7 @@ def recompute_stats( relative_exclude_joints: list[str] | None = None, chunk_size: int = 50, num_workers: int = 0, + update_episode_stats: bool = False, ) -> LeRobotDataset: """Recompute stats.json from scratch by iterating all episodes. @@ -1746,6 +1750,11 @@ def recompute_stats( (action, state, etc.) and keep existing image/video stats unchanged. If False, image/video stats are also recomputed by sampling and decoding frames from each episode (this reads the image/video files, unlike the numeric-only path). + update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the + episodes parquet files so they stay consistent with the aggregated ``stats.json``. + Defaults to False (only ``stats.json`` is rewritten). Requires a writable + ``dataset.root``. Note that relative-action stats are aggregate-only and are not + written per-episode. relative_action: If True, compute action stats in relative space by iterating all valid action chunks and subtracting the current state. This matches the normalization distribution the model sees during @@ -1784,7 +1793,10 @@ def recompute_stats( ) new_stats = aggregate_episode_stats( - dataset, all_episode_stats, extra_stats={ACTION: relative_action_stats} if relative_action_stats else None + dataset, + all_episode_stats, + extra_stats={ACTION: relative_action_stats} if relative_action_stats else None, + update_episode_stats=update_episode_stats, ) if new_stats is None: logging.warning("No episode stats computed") @@ -1793,23 +1805,71 @@ def recompute_stats( return dataset +def write_episode_stats(dataset: LeRobotDataset, episode_stats: dict[int, dict]) -> None: + """Overwrite the per-episode ``stats/*`` columns in the episodes parquet files in place. + + Only the features present in ``episode_stats[ep_idx]`` are rewritten; stats columns for + features that were not recomputed are left untouched. Every other episode column (tasks, + length, chunk/file indices, frame ranges, …) is preserved. ``dataset.root`` must be + writable (e.g. the reference copy created for read-only sources). + """ + if not episode_stats: + return + + meta = dataset.meta + if meta.episodes is None: + meta.episodes = load_episodes(meta.root) + + # Group episodes by the parquet file that holds them so each file is rewritten once. + file_to_episodes: dict[tuple[int, int], list[int]] = {} + for ep_idx in episode_stats: + ep = meta.episodes[ep_idx] + key = (ep["meta/episodes/chunk_index"], ep["meta/episodes/file_index"]) + file_to_episodes.setdefault(key, []).append(ep_idx) + + for (chunk_idx, file_idx), eps in file_to_episodes.items(): + path = meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx) + table = pq.read_table(path) + rows = table.to_pylist() + row_by_ep = {row["episode_index"]: row for row in rows} + for ep_idx in eps: + row = row_by_ep[ep_idx] + for feature, feature_stats in episode_stats[ep_idx].items(): + for stat_name, value in feature_stats.items(): + col = f"stats/{feature}/{stat_name}" + if col in row: + row[col] = np.asarray(value).tolist() + # Reuse the source schema so the rewritten stats keep the exact on-disk types. + new_table = pa.Table.from_pylist(rows, schema=table.schema) + pq.write_table(new_table, path, compression="snappy", use_dictionary=True) + + def aggregate_episode_stats( dataset: LeRobotDataset, - all_episode_stats: list[dict], + episode_stats: dict[int, dict], extra_stats: dict | None = None, + update_episode_stats: bool = False, ) -> dict | None: """Aggregate per-episode stats, merge with existing stats, and write ``stats.json``. - Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass - the concatenation of every worker's per-episode stats. ``extra_stats`` lets callers - inject feature stats computed outside the per-episode pass (e.g. relative-action stats). + Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass the + merged ``{episode_index: stats}`` mapping of every worker's per-episode stats. ``extra_stats`` + lets callers inject feature stats computed outside the per-episode pass (e.g. relative-action + stats). + + Args: + dataset: The dataset whose ``meta/stats.json`` (and optionally episode stats) is updated. + episode_stats: Mapping of episode index to its per-episode stat dict. + extra_stats: Feature stats to inject into the aggregate (not written per-episode). + update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the + episodes parquet files via :func:`write_episode_stats`. Returns the written stats dict, or ``None`` if there was nothing to aggregate. """ - if not all_episode_stats and not extra_stats: + if not episode_stats and not extra_stats: return None - new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {} + new_stats = aggregate_stats(list(episode_stats.values())) if episode_stats else {} if extra_stats: new_stats.update(extra_stats) @@ -1820,6 +1880,10 @@ def aggregate_episode_stats( write_stats(new_stats, dataset.root) dataset.meta.stats = new_stats + + if update_episode_stats: + write_episode_stats(dataset, episode_stats) + return new_stats diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index e3bebb666..11292a559 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/src/lerobot/scripts/lerobot_edit_dataset.py @@ -186,6 +186,13 @@ Recompute stats including image/video features (samples and decodes frames from --operation.type recompute_stats \ --operation.skip_image_video false +Recompute stats and also rewrite the per-episode stats in the episodes parquet (keeps +meta/stats.json and the per-episode stats consistent): + lerobot-edit-dataset \ + --repo_id lerobot/pusht \ + --operation.type recompute_stats \ + --operation.update_episode_stats true + Recompute stats in-place (overwrites original dataset stats): lerobot-edit-dataset \ --repo_id lerobot/pusht \ @@ -333,6 +340,7 @@ class RecomputeStatsConfig(OperationConfig): relative_exclude_joints: list[str] | None = None chunk_size: int = 50 num_workers: int = 0 + update_episode_stats: bool = False overwrite: bool = False @@ -713,9 +721,10 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None: if backup_path.exists(): shutil.rmtree(backup_path) shutil.move(output_root, backup_path) - # recompute_stats only reads data/ and rewrites meta/stats.json, so symlink the - # large immutable files and copy only meta/. This avoids duplicating the dataset - # and works even when the source dataset is read-only. + # recompute_stats only reads data/ and rewrites files under meta/ (stats.json, and + # the episodes parquet when update_episode_stats is set), so symlink the large + # immutable files and copy only meta/. This avoids duplicating the dataset and works + # even when the source dataset is read-only. _reference_copy_dataset(input_root, output_root) dataset = LeRobotDataset(output_repo_id, root=output_root) @@ -733,6 +742,7 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None: relative_exclude_joints=cfg.operation.relative_exclude_joints, chunk_size=cfg.operation.chunk_size, num_workers=cfg.operation.num_workers, + update_episode_stats=cfg.operation.update_episode_stats, ) logging.info(f"Stats written to {dataset.root}")