perf(datasets): bound memory of visual episode stats via batched streaming

Accumulate image/video per-episode stats into a RunningQuantileStats in
frame batches instead of materialising every sampled frame at once. Peak
memory is now bounded by one batch (frame_batch_size x C x H x W) regardless
of episode length, preventing OOM on long, high-resolution episodes. count
keeps the per-frame convention of get_feature_stats.

examples/dataset/slurm_recompute_stats.py

@@ -0,0 +1,489 @@
+#!/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``.
+
+Modified copy of lerobot's examples/dataset/slurm_recompute_stats.py
+(feat/recompute-stats-readonly-and-visual branch) with cluster-friendly additions:
+
+  1. --qos              : pass a SLURM QoS through to every worker's sbatch.
+  2. --venv-path        : activate a venv on each worker before the python step.
+  3. --env-command      : raw shell snippet injected before the python step (e.g. to
+                          export HF_LEROBOT_HOME). Runs in addition to --venv-path.
+  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
+it. This is the download route; the source dataset is fetched from the Hub on the
+CPU workers.
+
+IMPORTANT — how to run (do NOT sbatch this file):
+  Run it as a normal python process on the LOGIN node. datatrove submits the
+  workers for you. The reference copy (--new-root) is built on the login node and
+  references the shared HF cache, so /fsx must be visible there (it is).
+
+Requires: pip install 'lerobot[dataset]' datatrove
+
+Example (single command, compute then dependent aggregate):
+
+    export HF_LEROBOT_HOME=/fsx/$USER/.cache
+
+    python slurm_recompute_stats_patched.py compute \
+        --repo-id behavior-1k/2026-challenge-demos \
+        --new-root /fsx/$USER/behavior-1k_recomputed \
+        --shard-dir /fsx/$USER/behavior-1k_recomputed/stats_shards \
+        --logs-dir /fsx/$USER/logs/recompute \
+        --skip-image-video 0 \
+        --workers 250 \
+        --partition hopper-cpu \
+        --qos normal \
+        --cpus-per-task 8 --mem-per-cpu 4G \
+        --venv-path /fsx/$USER/venvs/lerobot/bin/activate \
+        --env-command 'export HF_LEROBOT_HOME=/fsx/'"$USER"'/.cache' \
+        --chain-aggregate
+
+REHEARSE FIRST with --workers 2 --skip-image-video 1 and inspect one worker's log
+under --logs-dir to confirm QoS was accepted and a numeric stats.json is written.
+"""
+
+import argparse
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+
+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, video_backend=None):
+        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
+        self.video_backend = video_backend
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        # NOTE: this method is pickled and executed on a worker, where this script's module
+        # globals are NOT available. Keep it self-contained: import locally and don't reference
+        # module-level helpers/constants.
+        import logging
+        import pickle
+        from pathlib import Path
+
+        from lerobot.datasets import LeRobotDataset, compute_dataset_episode_stats
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {}
+        root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
+        dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
+
+        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 / f"episode_stats_{rank:05d}.pkl"
+        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,
+        video_backend=None,
+        update_episode_stats=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
+        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).
+        import logging
+        import pickle
+        from pathlib import Path
+
+        from lerobot.datasets import LeRobotDataset, 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("episode_stats_*.pkl"))
+        if not shards:
+            raise FileNotFoundError(f"No episode stat shards found in {shard_dir}")
+
+        # 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.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 {}
+        root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
+        dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
+
+        # Aggregation is order-independent, so the only way sharding changes the result is a
+        # gap (dropped shard) or an overlap (episode counted twice). Verify the shards cover
+        # every episode exactly once before writing stats.json.
+        expected_episodes = dataset.meta.total_episodes
+        if len(all_episode_stats) != expected_episodes:
+            raise ValueError(
+                f"Expected {expected_episodes} per-episode stats (one per episode) but got "
+                f"{len(all_episode_stats)} across {len(shards)} shards. A compute shard is likely "
+                "missing or was written more than once; re-run the failed shards before aggregating."
+            )
+
+        # 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 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 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, 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}")
+
+        if self.push_to_hub:
+            logging.info(f"Pushing {self.repo_id} to hub")
+            dataset.push_to_hub()
+
+
+def _mem_gb(mem: str) -> int:
+    """Parse '4G' / '4GB' / '4' into an int number of GB for datatrove's mem_per_cpu_gb."""
+    s = str(mem).strip().lower().rstrip("b").rstrip("g")
+    return int(float(s))
+
+
+def _make_executor(
+    pipeline,
+    logs_dir,
+    job_name,
+    slurm,
+    workers,
+    tasks,
+    time,
+    partition,
+    cpus,
+    mem,
+    qos=None,
+    env_command=None,
+    venv_path=None,
+    depends=None,
+):
+    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,
+                "mem_per_cpu_gb": _mem_gb(mem),  # datatrove's native field (int GB)
+                "sbatch_args": {},
+            }
+        )
+        if qos:
+            kwargs["qos"] = qos  # -> "#SBATCH --qos=<qos>" on every worker
+        if venv_path:
+            kwargs["venv_path"] = venv_path  # datatrove sources this before the python step
+        if env_command:
+            kwargs["env_command"] = env_command  # extra raw snippet before python (composes with venv_path)
+        if depends is not None:
+            kwargs["depends"] = depends  # chains --dependency=afterok:<compute jobid>
+        return SlurmPipelineExecutor(**kwargs)
+    kwargs.update({"tasks": tasks, "workers": 1})
+    return LocalPipelineExecutor(**kwargs)
+
+
+def _maybe_reference_copy(repo_id, root, new_root, download_videos):
+    """Create the read-only-safe reference copy once, before submitting workers.
+
+    Loads metadata only (to resolve the source root and revision) instead of a full
+    ``LeRobotDataset``, which would also memory-map the entire frame index just to read a
+    path. Fetches the source into the shared cache so the copy's symlinks point at real
+    files and workers don't each re-download, pulling videos only when the run needs them
+    (i.e. when image/video stats are being recomputed).
+    """
+    if not new_root:
+        return
+    from huggingface_hub import snapshot_download
+
+    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+    from lerobot.scripts.lerobot_edit_dataset import _reference_copy_dataset
+    from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
+
+    new_root_path = Path(new_root)
+    if new_root_path.exists():
+        return
+
+    meta = LeRobotDatasetMetadata(repo_id, root=Path(root) if root else None)
+    ignore_patterns = None if download_videos else "videos/"
+    if root:
+        snapshot_download(
+            repo_id,
+            repo_type="dataset",
+            revision=meta.revision,
+            local_dir=meta.root,
+            ignore_patterns=ignore_patterns,
+        )
+        src_root = Path(meta.root)
+    else:
+        src_root = Path(
+            snapshot_download(
+                repo_id,
+                repo_type="dataset",
+                revision=meta.revision,
+                cache_dir=HF_LEROBOT_HUB_CACHE,
+                ignore_patterns=ignore_patterns,
+            )
+        )
+    _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 (defaults to the Hub cache).")
+    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, e.g. 'hopper-cpu'.")
+    p.add_argument("--qos", type=str, default=None, help="SLURM QoS, e.g. 'normal'. Passed to every worker.")
+    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'.")
+    p.add_argument(
+        "--video-backend",
+        type=str,
+        default=None,
+        help="Video decoding backend (e.g. 'pyav', 'torchcodec'). Defaults to the dataset's default; "
+        "use 'pyav' if torchcodec fails to load locally.",
+    )
+    p.add_argument("--venv-path", type=str, default=None, help="venv activate script sourced on each worker.")
+    p.add_argument(
+        "--env-command",
+        type=str,
+        default=None,
+        help="Raw shell snippet injected into each worker's sbatch before the python step "
+        "(e.g. to export HF_LEROBOT_HOME). Runs in addition to --venv-path.",
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="PATCHED 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).",
+    )
+    cp.add_argument(
+        "--chain-aggregate",
+        action="store_true",
+        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,
+        default=None,
+        help="Optional SLURM job id; aggregate waits for it (afterok) before running.",
+    )
+
+    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. Videos are only fetched when
+        # image/video stats are being recomputed.
+        _maybe_reference_copy(
+            args.repo_id, args.root, args.new_root, download_videos=not bool(args.skip_image_video)
+        )
+
+        compute_exec = _make_executor(
+            pipeline=[
+                ComputeEpisodeStatsShards(
+                    args.repo_id,
+                    args.root,
+                    args.new_root,
+                    bool(args.skip_image_video),
+                    str(args.shard_dir),
+                    args.video_backend,
+                )
+            ],
+            logs_dir=args.logs_dir,
+            job_name=args.job_name or "recompute_stats_compute",
+            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,
+            qos=args.qos,
+            env_command=args.env_command,
+            venv_path=args.venv_path,
+        )
+
+        if args.chain_aggregate and slurm:
+            # Build aggregate depending on compute. datatrove launches the dependency
+            # (compute) first, then submits aggregate with --dependency=afterok:<jobid>.
+            aggregate_exec = _make_executor(
+                pipeline=[
+                    AggregateEpisodeStats(
+                        args.repo_id,
+                        args.root,
+                        args.new_root,
+                        str(args.shard_dir),
+                        args.push_to_hub,
+                        args.video_backend,
+                        args.update_episode_stats,
+                    )
+                ],
+                logs_dir=args.logs_dir,
+                job_name="recompute_stats_aggregate",
+                slurm=slurm,
+                workers=1,
+                tasks=1,
+                time="02:00:00",
+                partition=args.partition,
+                cpus=args.cpus_per_task,
+                mem=args.mem_per_cpu,
+                qos=args.qos,
+                env_command=args.env_command,
+                venv_path=args.venv_path,
+                depends=compute_exec,
+            )
+            aggregate_exec.run()
+        else:
+            compute_exec.run()
+    else:
+        aggregate_exec = _make_executor(
+            pipeline=[
+                AggregateEpisodeStats(
+                    args.repo_id,
+                    args.root,
+                    args.new_root,
+                    str(args.shard_dir),
+                    args.push_to_hub,
+                    args.video_backend,
+                    args.update_episode_stats,
+                )
+            ],
+            logs_dir=args.logs_dir,
+            job_name=args.job_name or "recompute_stats_aggregate",
+            slurm=slurm,
+            workers=1,
+            tasks=1,
+            time="02:00:00",
+            partition=args.partition,
+            cpus=args.cpus_per_task,
+            mem=args.mem_per_cpu,
+            qos=args.qos,
+            env_command=args.env_command,
+            venv_path=args.venv_path,
+        )
+        if args.depends_job_id is not None:
+            aggregate_exec.depends_job_id = args.depends_job_id
+        aggregate_exec.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,
@@ -34,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
@@ -78,8 +81,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",
@@ -99,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,11 +33,13 @@ 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
 
 from lerobot.configs import (
+    DEFAULT_DEPTH_UNIT,
     DepthEncoderConfig,
     RGBEncoderConfig,
     VideoEncoderConfig,
@@ -51,11 +53,15 @@ from lerobot.utils.utils import flatten_dict
 
 from .aggregate import aggregate_datasets
 from .compute_stats import (
+    RunningQuantileStats,
     aggregate_stats,
+    auto_downsample_height_width,
     compute_episode_stats,
     compute_relative_action_stats,
+    sample_indices,
 )
 from .dataset_metadata import LeRobotDatasetMetadata
+from .depth_utils import dequantize_depth
 from .image_writer import write_image
 from .io_utils import (
     get_parquet_file_size_in_mb,
@@ -77,6 +83,7 @@ from .utils import (
     update_chunk_file_indices,
 )
 from .video_utils import (
+    decode_video_frames,
     encode_video_frames,
     reencode_video,
 )
@@ -1559,6 +1566,191 @@ def modify_tasks(
     return dataset
 
 
+def _load_episode_image_frames(
+    dataset: LeRobotDataset,
+    key: str,
+    ep_idx: int,
+    frame_offsets: list[int],
+    is_depth: bool,
+) -> np.ndarray:
+    """Load sampled frames of an image feature for one episode as a (N, C, H, W) array."""
+    ep = dataset.meta.episodes[ep_idx]
+    from_idx = ep["dataset_from_index"]
+    column = dataset.hf_dataset.with_format(None).select_columns(key)
+
+    frames = []
+    for offset in frame_offsets:
+        img = column[from_idx + offset][key]
+        if is_depth:
+            arr = np.array(img)
+            if arr.ndim == 2:
+                arr = arr[np.newaxis, ...]
+        else:
+            arr = np.transpose(np.array(img.convert("RGB"), dtype=np.uint8), (2, 0, 1))
+        frames.append(auto_downsample_height_width(arr))
+    return np.stack(frames)
+
+
+def _load_episode_video_frames(
+    dataset: LeRobotDataset,
+    key: str,
+    ep_idx: int,
+    frame_offsets: list[int],
+    is_depth: bool,
+) -> np.ndarray:
+    """Load sampled frames of a video feature for one episode as a (N, C, H, W) array."""
+    ep = dataset.meta.episodes[ep_idx]
+    video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, key)
+    from_timestamp = ep[f"videos/{key}/from_timestamp"]
+    timestamps = [from_timestamp + offset / dataset.meta.fps for offset in frame_offsets]
+
+    frames = decode_video_frames(
+        video_path,
+        timestamps,
+        dataset.tolerance_s,
+        backend=dataset._video_backend,
+        return_uint8=not is_depth,
+        is_depth=is_depth,
+    )
+    if is_depth:
+        # ``decode_video_frames`` returns raw 12-bit codec values; dequantize back to
+        # the recorded depth unit so stats match record-time stats (which are stored in
+        # ``info.depth_unit`` and only rescaled to the output unit on read).
+        info = dataset.meta.features[key].get("info") or {}
+        depth_encoder = DepthEncoderConfig.from_video_info(info)
+        frames = dequantize_depth(
+            frames,
+            depth_min=depth_encoder.depth_min,
+            depth_max=depth_encoder.depth_max,
+            shift=depth_encoder.shift,
+            use_log=depth_encoder.use_log,
+            output_unit=info.get("depth_unit") or DEFAULT_DEPTH_UNIT,
+        )
+    return np.stack([auto_downsample_height_width(frame) for frame in frames.numpy()])
+
+
+def _compute_visual_episode_stats(
+    dataset: LeRobotDataset,
+    ep_idx: int,
+    visual_keys: list[str],
+    frame_batch_size: int = 32,
+) -> dict:
+    """Compute per-episode statistics for image/video features by sampling frames.
+
+    Mirrors the image/video branch of :func:`compute_episode_stats`: per-channel stats
+    are computed on downsampled sampled frames, then RGB stats are rescaled to [0, 1]
+    (depth maps keep their native units).
+
+    Frames are decoded and accumulated into a :class:`RunningQuantileStats` in batches of
+    ``frame_batch_size`` rather than materialising every sampled frame at once. Peak memory
+    is bounded by one batch (``frame_batch_size x C x H x W``) regardless of episode length,
+    which keeps long, high-resolution episodes from exhausting memory.
+    """
+    ep_length = dataset.meta.episodes[ep_idx]["length"]
+    frame_offsets = sample_indices(ep_length)
+
+    ep_stats = {}
+    for key in visual_keys:
+        is_depth = key in dataset.meta.depth_keys
+        is_video = dataset.meta.features[key]["dtype"] == "video"
+
+        running = RunningQuantileStats()
+        for start in range(0, len(frame_offsets), frame_batch_size):
+            batch_offsets = frame_offsets[start : start + frame_batch_size]
+            if is_video:
+                frames = _load_episode_video_frames(dataset, key, ep_idx, batch_offsets, is_depth)
+            else:
+                frames = _load_episode_image_frames(dataset, key, ep_idx, batch_offsets, is_depth)
+            # (N, C, H, W) -> (N * H * W, C) so stats are accumulated per channel.
+            running.update(np.moveaxis(frames, 1, -1).reshape(-1, frames.shape[1]))
+
+        stats = running.get_statistics()
+        normalization_factor = 1.0 if is_depth else 255.0
+        num_channels = stats["mean"].shape[0]
+        # ``count`` follows the per-frame convention of ``get_feature_stats`` (number of
+        # sampled frames), not the per-pixel count tracked internally by RunningQuantileStats.
+        ep_stats[key] = {
+            k: np.array([len(frame_offsets)])
+            if k == "count"
+            else v.reshape(num_channels, 1, 1) / normalization_factor
+            for k, v in stats.items()
+        }
+
+    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,
+) -> 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 results with :func:`aggregate_episode_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 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"}
+    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[int(ep_idx)] = ep_stats
+
+    return all_episode_stats
+
+
 def recompute_stats(
     dataset: LeRobotDataset,
     skip_image_video: bool = True,
@@ -1566,13 +1758,21 @@ 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.
 
     Args:
         dataset: The LeRobotDataset to recompute stats for.
         skip_image_video: If True (default), only recompute stats for numeric features
-            (action, state, etc.) and keep existing image/video stats unchanged.
+            (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
@@ -1588,24 +1788,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"]
@@ -1616,56 +1804,105 @@ 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
+    new_stats = aggregate_episode_stats(
-    parquet_files = sorted(data_dir.glob("*/*.parquet"))
+        dataset,
-    if not parquet_files:
+        all_episode_stats,
-        raise ValueError(f"No parquet files found in {data_dir}")
+        extra_stats={ACTION: relative_action_stats} if relative_action_stats else None,
-
+        update_episode_stats=update_episode_stats,
-    all_episode_stats = []
+    )
-    # TODO: enable image and video stats re-computation
+    if new_stats is None:
-    numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not 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)
-            all_episode_stats.append(ep_stats)
-
-    if features_to_compute and not all_episode_stats:
         logging.warning("No episode stats computed")
-        return dataset
+    else:
+        logging.info("Stats recomputed successfully")
+    return dataset
 
-    new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {}
 
-    if relative_action_stats is not None:
+def write_episode_stats(dataset: LeRobotDataset, episode_stats: dict[int, dict]) -> None:
-        new_stats[ACTION] = relative_action_stats
+    """Overwrite the per-episode ``stats/*`` columns in the episodes parquet files in place.
 
-    # Merge: keep existing stats for features we didn't recompute
+    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,
+    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
+    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 episode_stats and not extra_stats:
+        return None
+
+    new_stats = aggregate_stats(list(episode_stats.values())) if episode_stats else {}
+    if extra_stats:
+        new_stats.update(extra_stats)
+
+    # 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.setdefault(key, value)
-                new_stats[key] = value
 
     write_stats(new_stats, dataset.root)
     dataset.meta.stats = new_stats
 
-    logging.info("Stats recomputed successfully")
+    if update_episode_stats:
-    return dataset
+        write_episode_stats(dataset, episode_stats)
+
+    return new_stats
 
 
 def convert_image_to_video_dataset(

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -167,7 +167,9 @@ Show dataset information without feature details:
         --operation.type info \
         --operation.show_features false
 
-Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default):
+Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default). The source
+dataset is never modified: large files are symlinked and only meta/ is copied, so this also works
+on read-only source datasets:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
         --operation.type recompute_stats
@@ -178,6 +180,19 @@ Recompute stats and save to a specific new repo_id:
         --new_repo_id lerobot/pusht_new_stats \
         --operation.type recompute_stats
 
+Recompute stats including image/video features (samples and decodes frames from each episode):
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht \
+        --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 \
@@ -325,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
 
 
@@ -377,6 +393,30 @@ def _resolve_io_paths(
     return output_repo_id, input_path, output_path
 
 
+def _reference_copy_dataset(input_root: Path, output_root: Path) -> None:
+    """Create a lightweight copy of a dataset that never modifies the source.
+
+    The directory tree is recreated with real directories, and every file is
+    symlinked to its source counterpart so no data is duplicated and the source is
+    only ever read. Files under ``meta/`` are instead copied as real, writable files
+    so that stats/info can be rewritten without touching the original. Symlinking
+    individual files (rather than whole directories) keeps ``push_to_hub`` working,
+    since ``Path.glob`` follows file symlinks but does not descend into symlinked
+    directories. This makes the operation safe on read-only source datasets.
+    """
+    for src in input_root.rglob("*"):
+        rel = src.relative_to(input_root)
+        dst = output_root / rel
+        if src.is_dir():
+            dst.mkdir(parents=True, exist_ok=True)
+        elif rel.parts[0] == "meta":
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copyfile(src, dst)  # copyfile ignores source perms, so dst is writable
+        else:
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            dst.symlink_to(src.resolve())
+
+
 def get_output_path(
     repo_id: str,
     new_repo_id: str | None,
@@ -674,14 +714,18 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
         )
         dataset = LeRobotDataset(cfg.repo_id, root=input_root)
     else:
-        logging.info(f"Copying dataset from {input_root} to {output_root}")
+        logging.info(f"Referencing dataset from {input_root} into {output_root} (source is left untouched)")
         if output_root.exists():
             backup_path = output_root.with_name(output_root.name + "_old")
             logging.warning(f"Output directory {output_root} already exists. Moving to {backup_path}")
             if backup_path.exists():
                 shutil.rmtree(backup_path)
             shutil.move(output_root, backup_path)
-        shutil.copytree(input_root, output_root)
+        # 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)
 
     logging.info(f"Recomputing stats for {cfg.repo_id}")
@@ -698,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}")