""" Post-process a LeRobot dataset recorded by the data exporter. Removes discarded episodes (flagged during collection) and stale SMPL frames (all-zero teleop.smpl_pose and frozen lead-in frames that precede them) which occur during teleop pauses or ZMQ frame drops. Can also merge multiple recording sessions into a single dataset. The script operates directly on the LeRobot v2.1 on-disk format (parquet + mp4) without any external training framework dependencies. Usage: # Clean a single dataset in-place python gear_sonic/scripts/process_dataset.py \\ --dataset-path outputs/my_dataset # Clean and write to a new directory (non-destructive) python gear_sonic/scripts/process_dataset.py \\ --dataset-path outputs/my_dataset \\ --output-path outputs/my_dataset_cleaned # Merge multiple datasets into one (validates matching script_config) python gear_sonic/scripts/process_dataset.py \\ --dataset-path outputs/session1 outputs/session2 outputs/session3 \\ --output-path outputs/merged_dataset # Merge from a text file listing dataset paths (one per line) python gear_sonic/scripts/process_dataset.py \\ --dataset-list datasets.txt \\ --output-path outputs/merged_dataset # Skip SMPL cleaning (merge only) python gear_sonic/scripts/process_dataset.py \\ --dataset-path outputs/session1 outputs/session2 \\ --output-path outputs/merged \\ --no-remove-stale-smpl # Remove discarded episodes (flagged during collection via 'x' key) python gear_sonic/scripts/process_dataset.py \\ --dataset-path outputs/my_dataset \\ --output-path outputs/my_dataset_cleaned \\ --remove-discarded """ from dataclasses import dataclass, field import json from pathlib import Path import shutil from typing import Optional import av import numpy as np import pandas as pd import tyro SMPL_POSE_COLUMN = "teleop.smpl_pose" # --------------------------------------------------------------------------- # Stale SMPL frame detection # --------------------------------------------------------------------------- def build_stale_mask(smpl_arr: np.ndarray) -> np.ndarray: """Return a boolean mask where True = frame should be removed. Marks all-zero rows AND any consecutive frozen (identical-to-next) rows that immediately precede a zero row. Frozen runs that do NOT lead into a zero row are left untouched — those occur naturally when the SMPL stream publishes at a slightly lower rate than the collection loop. """ n = len(smpl_arr) is_zero = np.all(smpl_arr == 0, axis=1) remove = is_zero.copy() diffs = np.zeros(n) diffs[1:] = np.sum(np.abs(smpl_arr[1:] - smpl_arr[:-1]), axis=1) for i in range(n): if is_zero[i]: j = i - 1 while j >= 0 and diffs[j] == 0.0 and not is_zero[j]: remove[j] = True j -= 1 return remove # --------------------------------------------------------------------------- # LeRobot on-disk helpers # --------------------------------------------------------------------------- def load_info(dataset_path: Path) -> dict: info_path = dataset_path / "meta" / "info.json" with open(info_path, encoding="utf-8") as f: return json.load(f) def load_episodes_meta(dataset_path: Path) -> list[dict]: episodes_path = dataset_path / "meta" / "episodes.jsonl" episodes = [] with open(episodes_path, encoding="utf-8") as f: for line in f: line = line.strip() if line: episodes.append(json.loads(line)) return episodes def load_tasks_meta(dataset_path: Path) -> list[dict]: tasks_path = dataset_path / "meta" / "tasks.jsonl" tasks = [] if tasks_path.exists(): with open(tasks_path, encoding="utf-8") as f: for line in f: line = line.strip() if line: tasks.append(json.loads(line)) return tasks def get_parquet_path(dataset_path: Path, info: dict, episode_index: int) -> Path: data_path_pattern = info.get("data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet") chunks_size = info.get("chunks_size", 1000) episode_chunk = episode_index // chunks_size return dataset_path / data_path_pattern.format( episode_chunk=episode_chunk, episode_index=episode_index, ) def get_video_keys(info: dict) -> list[str]: """Extract video keys from info.json, falling back to features if needed.""" keys = info.get("video_keys", []) if not keys: keys = [ k for k, v in info.get("features", {}).items() if v.get("dtype") in ("video", "image") ] return keys def get_video_paths(dataset_path: Path, info: dict, episode_index: int) -> dict[str, Path]: video_path_pattern = info.get( "video_path", "videos/{video_key}/episode_{episode_index:06d}.mp4", ) video_keys = get_video_keys(info) chunks_size = info.get("chunks_size", 1000) episode_chunk = episode_index // chunks_size paths = {} for key in video_keys: paths[key] = dataset_path / video_path_pattern.format( video_key=key, episode_index=episode_index, episode_chunk=episode_chunk, ) return paths def filter_video_frames(video_path: Path, valid_indices: np.ndarray, fps: int): """Re-encode a video keeping only frames at valid_indices.""" input_container = av.open(str(video_path)) input_stream = input_container.streams.video[0] all_frames = [] for frame in input_container.decode(input_stream): all_frames.append(frame.to_ndarray(format="rgb24")) input_container.close() if len(all_frames) == 0: return filtered = [all_frames[i] for i in valid_indices if i < len(all_frames)] if len(filtered) == 0: return tmp_path = video_path.with_suffix(".tmp.mp4") output_container = av.open(str(tmp_path), mode="w") output_stream = output_container.add_stream("h264", rate=fps) h, w = filtered[0].shape[:2] output_stream.width = w output_stream.height = h output_stream.pix_fmt = "yuv420p" for img in filtered: frame = av.VideoFrame.from_ndarray(img, format="rgb24") for packet in output_stream.encode(frame): output_container.mux(packet) for packet in output_stream.encode(): output_container.mux(packet) output_container.close() tmp_path.replace(video_path) # --------------------------------------------------------------------------- # Script config validation # --------------------------------------------------------------------------- def validate_script_configs(dataset_paths: list[Path]) -> dict | None: """Check that all datasets share the same script_config. Returns the common config if they match, or raises an error with details about which datasets differ. """ configs = {} for ds_path in dataset_paths: info = load_info(ds_path) sc = info.get("script_config") if sc is not None: configs[ds_path.name] = sc if not configs: return None canonical = json.dumps(next(iter(configs.values())), sort_keys=True) mismatched = [] for name, cfg in configs.items(): if json.dumps(cfg, sort_keys=True) != canonical: mismatched.append(name) if mismatched: print("\nERROR: script_config mismatch across datasets.") print("The following datasets have different robot configurations:\n") ref_name = next(iter(configs.keys())) print(f" Reference: {ref_name}") for name in mismatched: print(f" Differs: {name}") print( "\nDatasets recorded with different robot configurations cannot be " "merged. Verify that all sessions used the same robot setup." ) raise SystemExit(1) return next(iter(configs.values())) # --------------------------------------------------------------------------- # Core processing # --------------------------------------------------------------------------- def process_single_dataset( dataset_path: Path, remove_stale_smpl: bool, remove_discarded: bool = False, episode_index_offset: int = 0, ) -> dict: """Process one dataset: optionally clean stale SMPL frames. Returns stats dict and the list of (parquet_df, video_paths, episode_meta) tuples for merging. """ info = load_info(dataset_path) episodes_meta = load_episodes_meta(dataset_path) fps = info.get("fps", 50) discarded_indices = set(info.get("discarded_episode_indices", [])) if remove_discarded else set() stats = { "total_episodes": len(episodes_meta), "episodes_with_stale": 0, "total_frames": 0, "frames_removed": 0, "zero_frames": 0, "frozen_leadin_frames": 0, "episodes_dropped": 0, "episodes_discarded": 0, } processed_episodes = [] for ep_meta in episodes_meta: ep_idx = ep_meta["episode_index"] if ep_idx in discarded_indices: stats["episodes_discarded"] += 1 print(f" Episode {ep_idx}: discarded during collection — removing") continue parquet_path = get_parquet_path(dataset_path, info, ep_idx) video_paths = get_video_paths(dataset_path, info, ep_idx) if not parquet_path.exists(): print(f" WARNING: Missing parquet for episode {ep_idx}, skipping") continue df = pd.read_parquet(parquet_path) ep_len = len(df) stats["total_frames"] += ep_len valid_indices = None if remove_stale_smpl and SMPL_POSE_COLUMN in df.columns: smpl_arr = np.vstack( [np.asarray(x, dtype=np.float32) for x in df[SMPL_POSE_COLUMN]] ) mask = build_stale_mask(smpl_arr) n_remove = int(mask.sum()) n_zero = int(np.all(smpl_arr == 0, axis=1).sum()) n_frozen = n_remove - n_zero if n_remove > 0: stats["episodes_with_stale"] += 1 stats["frames_removed"] += n_remove stats["zero_frames"] += n_zero stats["frozen_leadin_frames"] += n_frozen pct = 100.0 * n_remove / ep_len print( f" Episode {ep_idx}: removing {n_remove}/{ep_len} frames " f"({pct:.1f}%) — {n_zero} zero + {n_frozen} frozen lead-in" ) if n_remove == ep_len: print(f" Episode {ep_idx}: ALL frames stale — dropping episode") stats["episodes_dropped"] += 1 continue valid_indices = np.where(~mask)[0] df = df.iloc[valid_indices].copy().reset_index(drop=True) if "timestamp" in df.columns: df["timestamp"] -= df["timestamp"].iloc[0] new_ep_idx = ep_idx + episode_index_offset processed_episodes.append({ "df": df, "source_video_paths": video_paths, "valid_indices": valid_indices, "episode_meta": ep_meta, "new_episode_index": new_ep_idx, "fps": fps, }) return stats, processed_episodes, info def write_output_dataset( dest_path: Path, all_episodes: list[dict], reference_info: dict, tasks_meta: list[dict], script_config: dict | None, ): """Write processed episodes to a new LeRobot dataset directory.""" dest_path.mkdir(parents=True, exist_ok=True) meta_dir = dest_path / "meta" meta_dir.mkdir(exist_ok=True) info = reference_info.copy() fps = info.get("fps", 50) chunks_size = info.get("chunks_size", 1000) if script_config is not None: info["script_config"] = script_config total_frames = 0 episodes_jsonl = [] for i, ep in enumerate(all_episodes): df = ep["df"] ep_len = len(df) df["episode_index"] = i df["index"] = range(total_frames, total_frames + ep_len) df["frame_index"] = range(ep_len) if "timestamp" in df.columns: df["timestamp"] = [j / fps for j in range(ep_len)] episode_chunk = i // chunks_size data_path_pattern = info.get( "data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", ) parquet_rel = data_path_pattern.format(episode_chunk=episode_chunk, episode_index=i) parquet_path = dest_path / parquet_rel parquet_path.parent.mkdir(parents=True, exist_ok=True) df.to_parquet(parquet_path) video_keys = get_video_keys(info) video_path_pattern = info.get( "video_path", "videos/{video_key}/episode_{episode_index:06d}.mp4", ) for vkey in video_keys: src_video = ep["source_video_paths"].get(vkey) dst_rel = video_path_pattern.format( video_key=vkey, episode_index=i, episode_chunk=episode_chunk, ) dst_video = dest_path / dst_rel dst_video.parent.mkdir(parents=True, exist_ok=True) if src_video and src_video.exists(): if ep["valid_indices"] is not None: shutil.copy2(src_video, dst_video) filter_video_frames(dst_video, ep["valid_indices"], fps) else: shutil.copy2(src_video, dst_video) ep_meta = { "episode_index": i, "tasks": ep["episode_meta"].get("tasks", []), "length": ep_len, } episodes_jsonl.append(ep_meta) total_frames += ep_len info["total_episodes"] = len(all_episodes) info["total_frames"] = total_frames info.pop("discarded_episode_indices", None) with open(meta_dir / "info.json", "w", encoding="utf-8") as f: json.dump(info, f, indent=4) with open(meta_dir / "episodes.jsonl", "w", encoding="utf-8") as f: for ep in episodes_jsonl: f.write(json.dumps(ep) + "\n") if tasks_meta: with open(meta_dir / "tasks.jsonl", "w", encoding="utf-8") as f: for task in tasks_meta: f.write(json.dumps(task) + "\n") return total_frames def copy_modality_json(source_paths: list[Path], output_path: Path): """Copy modality.json from the first source that has one.""" for src in source_paths: modality_path = src / "meta" / "modality.json" if modality_path.exists(): dst = output_path / "meta" / "modality.json" dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(modality_path, dst) return # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @dataclass class ProcessDatasetConfig: """Post-process LeRobot datasets: clean stale SMPL frames and/or merge.""" dataset_path: list[str] = field(default_factory=list) """One or more dataset directories to process.""" dataset_list: Optional[str] = None """Path to a text file listing dataset directories (one per line). Can be used instead of or in addition to --dataset-path.""" output_path: Optional[str] = None """Output directory for the processed dataset. If not specified and a single dataset is given, the dataset is modified in-place.""" remove_stale_smpl: bool = True """Remove frames where teleop.smpl_pose is all zeros (stale/dropped SMPL data) and frozen lead-in frames that precede them.""" remove_discarded: bool = True """Remove episodes that were flagged as discarded during data collection (stored in meta/info.json under discarded_episode_indices).""" def main(cfg: ProcessDatasetConfig): dataset_paths = [Path(p) for p in cfg.dataset_path] if cfg.dataset_list: list_file = Path(cfg.dataset_list) with open(list_file, encoding="utf-8") as f: for line in f: line = line.strip() if line and not line.startswith("#"): dataset_paths.append(Path(line)) if not dataset_paths: print("ERROR: No dataset paths provided. Use --dataset-path or --dataset-list.") raise SystemExit(1) for ds in dataset_paths: if not ds.exists(): print(f"ERROR: Dataset path does not exist: {ds}") raise SystemExit(1) if not (ds / "meta" / "info.json").exists(): print(f"ERROR: Not a valid LeRobot dataset (missing meta/info.json): {ds}") raise SystemExit(1) merging = len(dataset_paths) > 1 in_place = cfg.output_path is None if merging and in_place: print("ERROR: --output-path is required when merging multiple datasets.") raise SystemExit(1) output_path = Path(cfg.output_path) if cfg.output_path else dataset_paths[0] print("=" * 70) print(" LeRobot Dataset Processor") print("=" * 70) print(f" Input datasets: {len(dataset_paths)}") for ds in dataset_paths: print(f" - {ds}") print(f" Output: {output_path}{' (in-place)' if in_place else ''}") print(f" Remove stale SMPL: {cfg.remove_stale_smpl}") print(f" Remove discarded: {cfg.remove_discarded}") print("=" * 70) # Validate script configs match across all datasets if merging: print("\nValidating script_config across datasets...") script_config = validate_script_configs(dataset_paths) print(" All datasets have matching robot configurations.\n") else: info = load_info(dataset_paths[0]) script_config = info.get("script_config") # Collect tasks from all datasets (deduplicated) all_tasks_meta: list[dict] = [] seen_task_ids: set = set() for ds in dataset_paths: for task in load_tasks_meta(ds): tid = task.get("task_index", id(task)) if tid not in seen_task_ids: all_tasks_meta.append(task) seen_task_ids.add(tid) # Process each dataset all_episodes = [] total_stats = { "total_episodes": 0, "episodes_with_stale": 0, "total_frames": 0, "frames_removed": 0, "zero_frames": 0, "frozen_leadin_frames": 0, "episodes_dropped": 0, "episodes_discarded": 0, } reference_info = None for ds_path in dataset_paths: print(f"\nProcessing: {ds_path}") stats, episodes, info = process_single_dataset( ds_path, remove_stale_smpl=cfg.remove_stale_smpl, remove_discarded=cfg.remove_discarded, episode_index_offset=len(all_episodes), ) if reference_info is None: reference_info = info all_episodes.extend(episodes) for key in total_stats: total_stats[key] += stats[key] if not all_episodes: print("\nERROR: No valid episodes after processing.") raise SystemExit(1) # Write output if in_place: # In-place: rewrite parquet files and re-encode videos print(f"\nRewriting dataset in-place at {output_path}...") ds_info = load_info(output_path) fps = ds_info.get("fps", 50) # Delete files for discarded episodes if cfg.remove_discarded: discarded_indices = set(ds_info.get("discarded_episode_indices", [])) for ep_idx in discarded_indices: parquet_path = get_parquet_path(output_path, ds_info, ep_idx) if parquet_path.exists(): parquet_path.unlink() video_paths = get_video_paths(output_path, ds_info, ep_idx) for _vkey, vpath in video_paths.items(): if vpath.exists(): vpath.unlink() for ep in all_episodes: ep_idx = ep["episode_meta"]["episode_index"] parquet_path = get_parquet_path(output_path, ds_info, ep_idx) ep["df"].to_parquet(parquet_path) if ep["valid_indices"] is not None: video_paths = get_video_paths(output_path, ds_info, ep_idx) for _vkey, vpath in video_paths.items(): if vpath.exists(): filter_video_frames(vpath, ep["valid_indices"], fps) # Update episode metadata episodes_meta = [] for ep in all_episodes: meta = ep["episode_meta"].copy() meta["length"] = len(ep["df"]) episodes_meta.append(meta) with open(output_path / "meta" / "episodes.jsonl", "w", encoding="utf-8") as f: for em in episodes_meta: f.write(json.dumps(em) + "\n") ds_info["total_frames"] = sum(len(ep["df"]) for ep in all_episodes) ds_info["total_episodes"] = len(all_episodes) if cfg.remove_discarded: ds_info.pop("discarded_episode_indices", None) with open(output_path / "meta" / "info.json", "w", encoding="utf-8") as f: json.dump(ds_info, f, indent=4) else: print(f"\nWriting output dataset to {output_path}...") write_output_dataset( output_path, all_episodes, reference_info, all_tasks_meta, script_config, ) copy_modality_json(dataset_paths, output_path) # Print summary kept = total_stats["total_frames"] - total_stats["frames_removed"] kept_episodes = total_stats["total_episodes"] - total_stats["episodes_dropped"] - total_stats["episodes_discarded"] print("\n" + "=" * 70) print(" Processing complete!") print("=" * 70) print(f" Episodes: {kept_episodes} kept / {total_stats['total_episodes']} total" f" ({total_stats['episodes_dropped']} dropped, {total_stats['episodes_discarded']} discarded)") print(f" Frames: {kept} kept / {total_stats['total_frames']} total" f" ({total_stats['frames_removed']} removed)") if total_stats["frames_removed"] > 0: print(f" Zero SMPL: {total_stats['zero_frames']}") print(f" Frozen lead-in: {total_stats['frozen_leadin_frames']}") print(f" Episodes affected: {total_stats['episodes_with_stale']}") print(f" Output: {output_path}") print("=" * 70) if __name__ == "__main__": main(tyro.cli(ProcessDatasetConfig))