feat(datasets): distribute stats recomputation across SLURM workers

Expose the shardable unit of work behind recompute_stats: compute_dataset_episode_stats
computes per-episode stats for an episode subset, and aggregate_episode_stats merges the
concatenated shards (count-weighted) and writes stats.json. recompute_stats now composes
these, so single-process behavior is unchanged.

Add examples/dataset/slurm_recompute_stats.py, a datatrove compute/aggregate driver that
shards episodes across workers and is read-only safe (reference-copies the source when
--new-root is given). Most useful for the expensive image/video stats path.

examples/dataset/slurm_recompute_stats.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python
+
+# 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.
+
+"""
+SLURM-distributed recomputation of a LeRobotDataset's ``meta/stats.json``.
+
+Per-episode statistics are embarrassingly parallel, so we shard episodes across
+workers, each computing stats for its subset, then a single worker aggregates all
+shards (weighted by frame counts) and writes ``meta/stats.json``. This is mostly
+useful when recomputing image/video stats (``--skip-image-video 0``), which decodes
+frames and is far more expensive than the numeric-only path.
+
+Requires: pip install 'lerobot[dataset]' datatrove
+
+Two subcommands, each a separate SLURM submission:
+
+  compute    – N workers, each writes per-episode stats for its episode shard
+  aggregate  – 1 worker, merges shards into meta/stats.json (optionally push to hub)
+
+The dataset is read-only during ``compute``. When ``--new-root`` is given, a
+lightweight reference copy is made (large files symlinked, only meta/ copied) so a
+read-only / mounted source dataset is never modified; stats land in ``--new-root``.
+
+Usage:
+    # Recompute image/video stats for a mounted, read-only dataset with 50 workers.
+    python slurm_recompute_stats.py compute \\
+        --repo-id someone-else/their-dataset \\
+        --root /path/to/mounted/repo \\
+        --new-root /local/writable/their-dataset_recomputed \\
+        --skip-image-video 0 --workers 50 --partition cpu
+
+    python slurm_recompute_stats.py aggregate \\
+        --repo-id someone-else/their-dataset \\
+        --new-root /local/writable/their-dataset_recomputed \\
+        --partition cpu
+"""
+
+import argparse
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+
+SHARD_PATTERN = "episode_stats_{rank:05d}.pkl"
+SHARD_GLOB = "episode_stats_*.pkl"
+
+
+def _load_dataset(repo_id: str, root: str | None, new_root: str | None):
+    """Load the (possibly reference-copied) dataset used for stats.
+
+    When ``new_root`` differs from the source, create a read-only-safe reference copy
+    once (only the aggregator's rank 0 or the first compute worker needs to; here every
+    rank just loads ``new_root`` if it already exists, else falls back to the source).
+    """
+    from lerobot.datasets import LeRobotDataset
+
+    if new_root and Path(new_root).exists():
+        return LeRobotDataset(repo_id, root=new_root)
+    return LeRobotDataset(repo_id, root=root)
+
+
+class ComputeEpisodeStatsShards(PipelineStep):
+    """Each worker computes per-episode stats for its ``episodes[rank::world_size]`` shard."""
+
+    def __init__(self, repo_id, root, new_root, skip_image_video, shard_dir):
+        super().__init__()
+        self.repo_id = repo_id
+        self.root = root
+        self.new_root = new_root
+        self.skip_image_video = skip_image_video
+        self.shard_dir = shard_dir
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        import logging
+        import pickle
+
+        from lerobot.datasets import compute_dataset_episode_stats
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        dataset = _load_dataset(self.repo_id, self.root, self.new_root)
+
+        my_episodes = list(range(dataset.meta.total_episodes))[rank::world_size]
+        if not my_episodes:
+            logging.info(f"Rank {rank}: no episodes assigned")
+            return
+        logging.info(f"Rank {rank}: {len(my_episodes)} / {dataset.meta.total_episodes} episodes")
+
+        episode_stats = compute_dataset_episode_stats(
+            dataset,
+            episode_indices=my_episodes,
+            skip_image_video=self.skip_image_video,
+        )
+
+        shard_dir = Path(self.shard_dir)
+        shard_dir.mkdir(parents=True, exist_ok=True)
+        out = shard_dir / SHARD_PATTERN.format(rank=rank)
+        with open(out, "wb") as f:
+            pickle.dump(episode_stats, f)
+        logging.info(f"Rank {rank}: saved {len(episode_stats)} episode stats to {out}")
+
+
+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):
+        super().__init__()
+        self.repo_id = repo_id
+        self.root = root
+        self.new_root = new_root
+        self.shard_dir = shard_dir
+        self.push_to_hub = push_to_hub
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        import logging
+        import pickle
+
+        from lerobot.datasets import aggregate_episode_stats
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        if rank != 0:
+            return
+
+        shard_dir = Path(self.shard_dir)
+        shards = sorted(shard_dir.glob(SHARD_GLOB))
+        if not shards:
+            raise FileNotFoundError(f"No episode stat shards found in {shard_dir}")
+
+        all_episode_stats = []
+        for shard in shards:
+            with open(shard, "rb") as f:
+                all_episode_stats.extend(pickle.load(f))
+        logging.info(f"Aggregating {len(all_episode_stats)} episode stats from {len(shards)} shards")
+
+        dataset = _load_dataset(self.repo_id, self.root, self.new_root)
+        new_stats = aggregate_episode_stats(dataset, all_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}")
+
+        if self.push_to_hub:
+            logging.info(f"Pushing {self.repo_id} to hub")
+            dataset.push_to_hub()
+
+
+def _make_executor(pipeline, logs_dir, job_name, slurm, workers, tasks, time, partition, cpus, mem):
+    kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)}
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": tasks,
+                "workers": workers,
+                "time": time,
+                "partition": partition,
+                "cpus_per_task": cpus,
+                "sbatch_args": {"mem-per-cpu": mem},
+            }
+        )
+        return SlurmPipelineExecutor(**kwargs)
+    kwargs.update({"tasks": tasks, "workers": 1})
+    return LocalPipelineExecutor(**kwargs)
+
+
+def _maybe_reference_copy(repo_id, root, new_root):
+    """Create the read-only-safe reference copy once, before submitting workers."""
+    if not new_root:
+        return
+    from lerobot.datasets import LeRobotDataset
+    from lerobot.scripts.lerobot_edit_dataset import _reference_copy_dataset
+
+    new_root_path = Path(new_root)
+    if new_root_path.exists():
+        return
+    src = LeRobotDataset(repo_id, root=root)
+    _reference_copy_dataset(src.root, new_root_path)
+
+
+def _add_shared_args(p):
+    p.add_argument("--repo-id", type=str, required=True, help="Dataset identifier, e.g. 'user/dataset'.")
+    p.add_argument("--root", type=str, default=None, help="Source dataset root (e.g. a mount).")
+    p.add_argument(
+        "--new-root",
+        type=str,
+        default=None,
+        help="Writable output root; a read-only-safe reference copy of --root. If omitted, stats "
+        "are written in place at --root.",
+    )
+    p.add_argument("--shard-dir", type=Path, default=Path("stats_shards"), help="Per-rank shard dir.")
+    p.add_argument("--logs-dir", type=Path, default=Path("logs"), help="datatrove logs dir.")
+    p.add_argument("--job-name", type=str, default=None, help="SLURM job name.")
+    p.add_argument("--slurm", type=int, default=1, help="1 = submit via SLURM; 0 = run locally.")
+    p.add_argument("--partition", type=str, default=None, help="SLURM partition.")
+    p.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per SLURM task.")
+    p.add_argument("--mem-per-cpu", type=str, default="4G", help="Memory per CPU, e.g. '4G'.")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="SLURM-distributed LeRobotDataset stats recomputation",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    cp = sub.add_parser("compute", help="Distribute per-episode stats across SLURM workers.")
+    _add_shared_args(cp)
+    cp.add_argument("--workers", type=int, default=50, help="Number of parallel SLURM tasks.")
+    cp.add_argument(
+        "--skip-image-video",
+        type=int,
+        default=1,
+        help="1 = numeric features only (fast); 0 = also recompute image/video stats (decodes frames).",
+    )
+
+    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.")
+
+    args = parser.parse_args()
+    slurm = args.slurm == 1
+
+    if args.command == "compute":
+        # The reference copy (if any) is created once on the submitting node so workers
+        # can all load --new-root without racing to build it.
+        _maybe_reference_copy(args.repo_id, args.root, args.new_root)
+        job_name = args.job_name or "recompute_stats_compute"
+        executor = _make_executor(
+            pipeline=[
+                ComputeEpisodeStatsShards(
+                    args.repo_id, args.root, args.new_root, args.skip_image_video == 0, str(args.shard_dir)
+                )
+            ],
+            logs_dir=args.logs_dir,
+            job_name=job_name,
+            slurm=slurm,
+            workers=args.workers,
+            tasks=args.workers,
+            time="24:00:00",
+            partition=args.partition,
+            cpus=args.cpus_per_task,
+            mem=args.mem_per_cpu,
+        )
+    else:
+        job_name = args.job_name or "recompute_stats_aggregate"
+        executor = _make_executor(
+            pipeline=[
+                AggregateEpisodeStats(
+                    args.repo_id, args.root, args.new_root, str(args.shard_dir), args.push_to_hub
+                )
+            ],
+            logs_dir=args.logs_dir,
+            job_name=job_name,
+            slurm=slurm,
+            workers=1,
+            tasks=1,
+            time="02:00:00",
+            partition=args.partition,
+            cpus=args.cpus_per_task,
+            mem=args.mem_per_cpu,
+        )
+
+    executor.run()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/datasets/__init__.py

@@ -25,6 +25,8 @@ from .compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats
 from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
 from .dataset_tools import (
     add_features,
+    aggregate_episode_stats,
+    compute_dataset_episode_stats,
     convert_image_to_video_dataset,
     delete_episodes,
     merge_datasets,
@@ -78,8 +80,10 @@ __all__ = [
     "detect_available_encoders_pyav",
     "add_features",
     "aggregate_datasets",
+    "aggregate_episode_stats",
     "aggregate_pipeline_dataset_features",
     "aggregate_stats",
+    "compute_dataset_episode_stats",
     "convert_image_to_video_dataset",
     "create_initial_features",
     "compute_sampler_state",

src/lerobot/datasets/dataset_tools.py

@@ -1639,6 +1639,76 @@ def _compute_visual_episode_stats(
     return ep_stats
 
 
+def compute_dataset_episode_stats(
+    dataset: LeRobotDataset,
+    episode_indices: list[int] | None = None,
+    skip_image_video: bool = True,
+    drop_keys: list[str] | None = None,
+) -> list[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`.
+
+    Args:
+        dataset: The LeRobotDataset to compute stats for.
+        episode_indices: Episodes to process. When ``None``, all episodes are processed.
+        skip_image_video: If True (default), only numeric features are computed. If False,
+            image/video stats are also computed by sampling and decoding frames.
+        drop_keys: Feature keys to exclude (e.g. ``action`` when it is computed separately
+            in relative-action space).
+
+    Returns:
+        A list of per-episode stat dicts, one per processed episode.
+    """
+    features = dataset.meta.features
+    meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
+    drop = set(drop_keys or [])
+    features_to_compute = {
+        k: v
+        for k, v in features.items()
+        if v["dtype"] != "string"
+        and k not in meta_keys
+        and k not in drop
+        and (not skip_image_video or v["dtype"] not in ["image", "video"])
+    }
+    numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
+    visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]]
+
+    if dataset.meta.episodes is None:
+        dataset.meta.episodes = load_episodes(dataset.meta.root)
+
+    if episode_indices is None:
+        episode_indices = list(range(dataset.meta.total_episodes))
+
+    # Group requested episodes by their data parquet file so each file is read once.
+    file_to_episodes: dict[Path, list[int]] = {}
+    for ep_idx in episode_indices:
+        file_to_episodes.setdefault(dataset.meta.get_data_file_path(ep_idx), []).append(ep_idx)
+
+    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):
+            episode_data = {}
+            if numeric_keys:
+                ep_df = df[df["episode_index"] == ep_idx]
+                for key in numeric_keys:
+                    if key in ep_df.columns:
+                        values = ep_df[key].values
+                        episode_data[key] = (
+                            np.stack(values) if hasattr(values[0], "__len__") else np.array(values)
+                        )
+
+            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)
+
+    return all_episode_stats
+
+
 def recompute_stats(
     dataset: LeRobotDataset,
     skip_image_video: bool = True,
@@ -1670,24 +1740,12 @@ def recompute_stats(
         The same dataset with updated stats.
     """
     features = dataset.meta.features
-    meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
-    numeric_features = {
-        k: v
-        for k, v in features.items()
-        if v["dtype"] not in ["image", "video", "string"] and k not in meta_keys
-    }
-
-    if skip_image_video:
-        features_to_compute = numeric_features
-    else:
-        features_to_compute = {
-            k: v for k, v in features.items() if v["dtype"] != "string" and k not in meta_keys
-        }
 
     # When relative_action is enabled, compute action stats via chunk-based sampling
     # (matching what the model sees during training) and skip action in the
     # per-episode pass below.
     relative_action_stats = None
+    drop_keys = None
     if relative_action and ACTION in features and OBS_STATE in features:
         if relative_exclude_joints is None:
             relative_exclude_joints = ["gripper"]
@@ -1698,58 +1756,50 @@ def recompute_stats(
             exclude_joints=relative_exclude_joints,
             num_workers=num_workers,
         )
-        features_to_compute.pop(ACTION, None)
+        drop_keys = [ACTION]
 
-    logging.info(f"Recomputing stats for features: {list(features_to_compute.keys())}")
+    all_episode_stats = compute_dataset_episode_stats(
+        dataset, skip_image_video=skip_image_video, drop_keys=drop_keys
+    )
 
-    data_dir = dataset.root / DATA_DIR
-    parquet_files = sorted(data_dir.glob("*/*.parquet"))
-    if not parquet_files:
-        raise ValueError(f"No parquet files found in {data_dir}")
-
-    all_episode_stats = []
-    numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
-    visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]]
-
-    for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"):
-        df = pd.read_parquet(parquet_path)
-
-        for ep_idx in sorted(df["episode_index"].unique()):
-            ep_df = df[df["episode_index"] == ep_idx]
-            episode_data = {}
-            for key in numeric_keys:
-                if key in ep_df.columns:
-                    values = ep_df[key].values
-                    if hasattr(values[0], "__len__"):
-                        episode_data[key] = np.stack(values)
-                    else:
-                        episode_data[key] = np.array(values)
-
-            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)
-
-    if features_to_compute and not all_episode_stats:
+    new_stats = aggregate_episode_stats(
+        dataset, all_episode_stats, extra_stats={ACTION: relative_action_stats} if relative_action_stats else None
+    )
+    if new_stats is None:
         logging.warning("No episode stats computed")
-        return dataset
+    else:
+        logging.info("Stats recomputed successfully")
+    return dataset
+
+
+def aggregate_episode_stats(
+    dataset: LeRobotDataset,
+    all_episode_stats: list[dict],
+    extra_stats: dict | None = None,
+) -> 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).
+
+    Returns the written stats dict, or ``None`` if there was nothing to aggregate.
+    """
+    if not all_episode_stats and not extra_stats:
+        return None
 
     new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {}
+    if extra_stats:
+        new_stats.update(extra_stats)
 
-    if relative_action_stats is not None:
-        new_stats[ACTION] = relative_action_stats
-
-    # Merge: keep existing stats for features we didn't recompute
+    # Merge: keep existing stats for features we didn't recompute.
     if dataset.meta.stats:
         for key, value in dataset.meta.stats.items():
-            if key not in new_stats:
-                new_stats[key] = value
+            new_stats.setdefault(key, value)
 
     write_stats(new_stats, dataset.root)
     dataset.meta.stats = new_stats
-
-    logging.info("Stats recomputed successfully")
-    return dataset
+    return new_stats
 
 
 def convert_image_to_video_dataset(