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

@@ -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,

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",
 ]

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
+    Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass the
-    the concatenation of every worker's per-episode stats. ``extra_stats`` lets callers
+    merged ``{episode_index: stats}`` mapping of every worker's per-episode stats. ``extra_stats``
-    inject feature stats computed outside the per-episode pass (e.g. relative-action 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
 
 

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
+        # recompute_stats only reads data/ and rewrites files under meta/ (stats.json, and
-        # large immutable files and copy only meta/. This avoids duplicating the dataset
+        # the episodes parquet when update_episode_stats is set), so symlink the large
-        # and works even when the source dataset is read-only.
+        # 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}")