From 7035ecf9b2ed8e3dd1c9281c83c9512dda349d49 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 10 Jul 2026 16:25:40 +0200 Subject: [PATCH] fix(datasets): dequantize depth video frames when recomputing stats Depth video stats were computed on raw 12-bit codec values, leaving them in codec space instead of the recorded depth unit. Dequantize decoded frames via the feature's depth encoder config (matching DatasetReader) so recomputed stats match record-time stats. Also fix the SLURM example: the --skip-image-video flag was inverted (0 skipped visual stats), and add a --video-backend option so pyav can be used when torchcodec fails to load locally. --- examples/dataset/slurm_recompute_stats.py | 74 ++++++++++++++++++++--- src/lerobot/datasets/dataset_tools.py | 23 ++++++- 2 files changed, 87 insertions(+), 10 deletions(-) diff --git a/examples/dataset/slurm_recompute_stats.py b/examples/dataset/slurm_recompute_stats.py index b4b942338..c64b1609e 100644 --- a/examples/dataset/slurm_recompute_stats.py +++ b/examples/dataset/slurm_recompute_stats.py @@ -46,6 +46,12 @@ Usage: --repo-id someone-else/their-dataset \\ --new-root /local/writable/their-dataset_recomputed \\ --partition cpu + + # Run locally without SLURM (single process); use pyav if torchcodec won't load. + python slurm_recompute_stats.py compute \\ + --repo-id someone-else/their-dataset \\ + --new-root /local/writable/their-dataset_recomputed \\ + --skip-image-video 0 --video-backend pyav --slurm 0 """ import argparse @@ -59,7 +65,7 @@ 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): +def _load_dataset(repo_id: str, root: str | None, new_root: str | None, video_backend: str | None = None): """Load the (possibly reference-copied) dataset used for stats. When ``new_root`` differs from the source, create a read-only-safe reference copy @@ -68,21 +74,23 @@ def _load_dataset(repo_id: str, root: str | None, new_root: str | None): """ from lerobot.datasets import LeRobotDataset + kwargs = {"video_backend": video_backend} if video_backend else {} if new_root and Path(new_root).exists(): - return LeRobotDataset(repo_id, root=new_root) - return LeRobotDataset(repo_id, root=root) + return LeRobotDataset(repo_id, root=new_root, **kwargs) + return LeRobotDataset(repo_id, root=root, **kwargs) 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): + 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): import logging @@ -92,7 +100,7 @@ class ComputeEpisodeStatsShards(PipelineStep): from lerobot.utils.utils import init_logging init_logging() - dataset = _load_dataset(self.repo_id, self.root, self.new_root) + dataset = _load_dataset(self.repo_id, self.root, self.new_root, self.video_backend) my_episodes = list(range(dataset.meta.total_episodes))[rank::world_size] if not my_episodes: @@ -117,13 +125,14 @@ 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): + def __init__(self, repo_id, root, new_root, shard_dir, push_to_hub=False, video_backend=None): 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 def run(self, data=None, rank: int = 0, world_size: int = 1): import logging @@ -147,7 +156,37 @@ class AggregateEpisodeStats(PipelineStep): 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) + dataset = _load_dataset(self.repo_id, self.root, self.new_root, self.video_backend) + + # 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. + 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] + ), + None, + ) + if numeric_key is not None: + total_frames = sum(int(s[numeric_key]["count"][0]) for s in all_episode_stats) + 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) if new_stats is None: raise RuntimeError("Aggregation produced no stats") @@ -208,6 +247,13 @@ def _add_shared_args(p): 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'.") + 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.", + ) def main(): @@ -242,7 +288,12 @@ def main(): executor = _make_executor( pipeline=[ ComputeEpisodeStatsShards( - args.repo_id, args.root, args.new_root, args.skip_image_video == 0, str(args.shard_dir) + 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, @@ -260,7 +311,12 @@ def main(): executor = _make_executor( pipeline=[ AggregateEpisodeStats( - args.repo_id, args.root, args.new_root, str(args.shard_dir), args.push_to_hub + args.repo_id, + args.root, + args.new_root, + str(args.shard_dir), + args.push_to_hub, + args.video_backend, ) ], logs_dir=args.logs_dir, diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 1a4286175..92a8d6767 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -38,6 +38,7 @@ import torch from tqdm import tqdm from lerobot.configs import ( + DEFAULT_DEPTH_UNIT, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig, @@ -59,6 +60,7 @@ from .compute_stats import ( 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, @@ -1602,8 +1604,27 @@ def _load_episode_video_frames( timestamps = [from_timestamp + offset / dataset.meta.fps for offset in frame_offsets] frames = decode_video_frames( - video_path, timestamps, dataset.tolerance_s, return_uint8=not is_depth, is_depth=is_depth + 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()])