#!/usr/bin/env python3 from __future__ import annotations import argparse import dataclasses import datetime import json import math from pathlib import Path import shutil import numpy as np import pandas as pd from tqdm import tqdm from datasets import Dataset from datasets import Features from datasets import Image from datasets import Sequence from datasets import Value from datasets import disable_progress_bar from datasets import enable_progress_bar from datasets import is_progress_bar_enabled from various_speed.core import SpeedTransformConfig from various_speed.core import _speed_chunk_ratio from various_speed.core import transform_episode def _episode_paths(dataset_root: Path, limit: int | None) -> list[Path]: paths = sorted((dataset_root / "data").glob("chunk-*/episode_*.parquet")) if limit is not None: paths = paths[:limit] if not paths: raise FileNotFoundError(f"No parquet episodes found under {dataset_root / 'data'}") return paths def _load_json(path: Path) -> dict: with path.open() as f: return json.load(f) def _write_json(path: Path, obj: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w") as f: json.dump(obj, f, indent=4) f.write("\n") def _read_jsonl(path: Path) -> list[dict]: if not path.exists(): return [] out = [] with path.open() as f: for raw_line in f: text = raw_line.strip() if text: out.append(json.loads(text)) return out def _write_jsonl(path: Path, rows: list[dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w") as f: for row in rows: f.write(json.dumps(row) + "\n") def _hf_value_dtype(dtype: str) -> str: if dtype in {"float32", "float64", "int8", "int16", "int32", "int64", "string", "bool"}: return dtype return "string" if dtype == "str" else dtype def _hf_feature(feature: dict): dtype = feature.get("dtype") if dtype == "image": return Image() if dtype == "video": return Value("string") shape = feature.get("shape") value = Value(_hf_value_dtype(dtype or "string")) if isinstance(shape, list) and len(shape) == 1 and shape[0] and int(shape[0]) > 1: return Sequence(value, length=int(shape[0])) return value def _features_for_columns(info: dict, columns: list[str]) -> Features: info_features = info.get("features", {}) features = {} for column in columns: if column in info_features: features[column] = _hf_feature(info_features[column]) return Features(features) if features else Features({}) def _write_episode_parquet(path: Path, df: pd.DataFrame, info: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) features = _features_for_columns(info, list(df.columns)) if features: progress_was_enabled = is_progress_bar_enabled() disable_progress_bar() try: Dataset.from_pandas(df, features=features, preserve_index=False).to_parquet(path) finally: if progress_was_enabled: enable_progress_bar() else: disable_progress_bar() else: df.to_parquet(path, index=False) def _speed_label(speed: float) -> str: text = f"{speed:g}".replace(".", "p") return f"{text}x" def _speed_phase_specs( speeds: list[float], *, chunk_aligned_observation: bool, sliding_chunk_phases: bool, ) -> list[tuple[int, float, int, int]]: if sliding_chunk_phases and not chunk_aligned_observation: raise ValueError("--sliding-chunk-phases requires --chunk-aligned-observation") specs: list[tuple[int, float, int, int]] = [] for speed_index, speed in enumerate(speeds): phase_count = _speed_chunk_ratio(speed)[0] if sliding_chunk_phases else 1 for chunk_phase in range(phase_count): specs.append((speed_index, speed, chunk_phase, phase_count)) return specs def _speed_name(speeds: list[float]) -> str: labels = [] for speed in speeds: text = f"{speed:.2f}".rstrip("0").rstrip(".") if "." not in text: text += ".0" labels.append(text.replace(".", "p")) return "_".join(labels) def _safe_name(value: str) -> str: out = [] for char in value.strip().lower(): if char.isalnum() or char in {"-", "_"}: out.append(char) else: out.append("_") return "_".join(part for part in "".join(out).split("_") if part) def _timestamp(value: str | None) -> str: if value: return _safe_name(value) return datetime.datetime.now(datetime.UTC).strftime("%Y%m%d_%H%M%S") def _resolve_output_root(dst: Path, args: argparse.Namespace, speeds: list[float], episode_count: int) -> Path: if not args.auto_name: return dst.resolve() suite = _safe_name(args.task_suite_name) tag = _safe_name(args.run_tag) episodes = f"ep{episode_count}" if args.max_episodes is not None else "full" name = f"{suite}_speed_{_speed_name(speeds)}_{episodes}_{tag}_{_timestamp(args.timestamp)}" return (dst / name).resolve() def _stack_column(df: pd.DataFrame, name: str) -> np.ndarray: return np.stack(df[name].to_numpy()).astype(np.float32) def _first_existing_column(df: pd.DataFrame, candidates: list[str]) -> str: for name in candidates: if name in df.columns: return name raise KeyError(f"Expected one of {candidates}, got columns {list(df.columns)}") def _video_keys(info: dict) -> list[str]: return [name for name, feature in info.get("features", {}).items() if feature.get("dtype") == "video"] def _image_keys(info: dict) -> list[str]: return [name for name, feature in info.get("features", {}).items() if feature.get("dtype") == "image"] def _zero_image_like(value: object, frame_index: int) -> object: if isinstance(value, dict) and isinstance(value.get("bytes"), bytes): import io from PIL import Image with Image.open(io.BytesIO(value["bytes"])) as image: zero = Image.new(image.mode, image.size) out = io.BytesIO() zero.save(out, format=image.format or "PNG") return {"bytes": out.getvalue(), "path": f"frame_{frame_index:06d}.png"} return value def _resample_embedded_images( src_df: pd.DataFrame, source_step_index: np.ndarray, observation_mask: np.ndarray, image_keys: list[str], ) -> dict[str, list[object]]: out = {} for key in image_keys: if key not in src_df: continue values = src_df[key].to_numpy() first = values[0] column = [] for frame_index, (source_step, valid) in enumerate(zip(source_step_index, observation_mask, strict=False)): if int(valid) == 0: column.append(_zero_image_like(first, frame_index)) else: column.append(values[min(int(source_step), len(values) - 1)]) out[key] = column return out def _read_video(path: Path) -> list[np.ndarray]: import imageio.v3 as iio if not path.exists(): raise FileNotFoundError(f"Missing video: {path}") return [np.asarray(frame) for frame in iio.imiter(path)] def _write_video(path: Path, frames: list[np.ndarray], fps: int) -> None: import imageio.v3 as iio path.parent.mkdir(parents=True, exist_ok=True) iio.imwrite(path, np.stack(frames), fps=fps) def _copy_or_resample_videos( src_root: Path, dst_root: Path, src_episode_index: int, dst_episode_index: int, source_frame_index: np.ndarray, observation_mask: np.ndarray, video_keys: list[str], chunks_size: int, fps: int, ) -> None: if not video_keys: return src_chunk = src_episode_index // chunks_size dst_chunk = dst_episode_index // chunks_size for key in video_keys: src = src_root / "videos" / f"chunk-{src_chunk:03d}" / key / f"episode_{src_episode_index:06d}.mp4" dst = dst_root / "videos" / f"chunk-{dst_chunk:03d}" / key / f"episode_{dst_episode_index:06d}.mp4" frames = _read_video(src) zero = np.zeros_like(frames[0]) out = [] for frame_idx, valid in zip(source_frame_index, observation_mask, strict=False): if int(valid) == 0: out.append(zero) else: out.append(frames[min(int(frame_idx), len(frames) - 1)]) _write_video(dst, out, fps=fps) def _numeric_stats(values: np.ndarray) -> dict: arr = np.asarray(values) flat = arr.reshape(-1, 1) if arr.ndim == 1 else arr.reshape(arr.shape[0], -1) return { "min": flat.min(axis=0).tolist(), "max": flat.max(axis=0).tolist(), "mean": flat.mean(axis=0).tolist(), "std": flat.std(axis=0).tolist(), "count": [int(arr.shape[0])], } def _episode_stats(df: pd.DataFrame, episode_index: int) -> dict: stats = {} for col in df.columns: first = df[col].iloc[0] if isinstance(first, np.ndarray): stats[col] = _numeric_stats(np.stack(df[col].to_numpy())) elif np.issubdtype(df[col].dtype, np.number): stats[col] = _numeric_stats(df[col].to_numpy()) return {"episode_index": int(episode_index), "stats": stats} def _aggregate_cleaning_stats(metrics_list: list[dict]) -> dict: """Aggregate near-zero cleaning counts across the source dataset. Cleaning is invariant to target speed, so we deduplicate by source_episode_index to avoid counting the same source frames once per speed. """ seen: set[int] = set() source_frames = 0 transl = 0 rot = 0 any_ = 0 both = 0 for entry in metrics_list: src = int(entry.get("source_episode_index", -1)) if src in seen: continue seen.add(src) source_frames += int(entry.get("source_frames", 0)) transl += int(entry.get("cleaned_translation_frames", 0)) rot += int(entry.get("cleaned_rotation_frames", 0)) any_ += int(entry.get("cleaned_any_frames", 0)) both += int(entry.get("cleaned_both_frames", 0)) denom = max(source_frames, 1) return { "source_episodes": len(seen), "source_frames": int(source_frames), "cleaned_translation_frames": int(transl), "cleaned_rotation_frames": int(rot), "cleaned_any_frames": int(any_), "cleaned_both_frames": int(both), "cleaned_translation_ratio": float(transl / denom), "cleaned_rotation_ratio": float(rot / denom), "cleaned_any_ratio": float(any_ / denom), "cleaned_both_ratio": float(both / denom), } def _print_cleaning_summary(summary: dict) -> None: print( "Cleaning summary: " f"{summary['cleaned_any_frames']}/{summary['source_frames']} source frames " f"({summary['cleaned_any_ratio'] * 100:.2f}%) had translation or rotation zeroed; " f"{summary['cleaned_both_frames']} ({summary['cleaned_both_ratio'] * 100:.2f}%) had both zeroed; " f"translation-only zeroed: {summary['cleaned_translation_frames']} " f"({summary['cleaned_translation_ratio'] * 100:.2f}%); " f"rotation-only zeroed: {summary['cleaned_rotation_frames']} " f"({summary['cleaned_rotation_ratio'] * 100:.2f}%)." ) def _aggregate_segment_stats(metrics_list: list[dict]) -> dict: """Aggregate segment-length distribution across the source dataset. Segments depend only on the source episode (not on target speed), so we deduplicate by ``source_episode_index`` and pool every segment from every source episode into a single distribution. """ seen: set[int] = set() seg_lens: list[float] = [] n_segments = 0 motion_counts = {0: 0, 1: 0, 2: 0, 3: 0} for entry in metrics_list: src = int(entry.get("source_episode_index", -1)) if src in seen: continue seen.add(src) # Per-episode summary statistics are pooled by treating each episode's # mean as one observation, plus its min/max for distribution range. # Per-segment lengths themselves are not stored verbatim in metrics_list # (would balloon the JSONL); instead we use the per-episode aggregates. n_segments += int(entry.get("segment_count", 0)) seg_lens.append(float(entry.get("segment_len_mean", 0.0))) for k in (0, 1, 2, 3): motion_counts[k] += int(entry.get(f"motion_class_{['still', 'translate', 'rotate', 'translate_rotate'][k]}_count", 0)) if not seg_lens: return {"source_episodes": 0, "n_segments": 0} arr = np.asarray(seg_lens, dtype=np.float64) total_motion = max(sum(motion_counts.values()), 1) return { "source_episodes": len(seen), "n_segments": int(n_segments), "segments_per_episode_mean": float(n_segments / max(len(seen), 1)), "segment_len_mean_of_means": float(arr.mean()), "segment_len_median_of_means": float(np.median(arr)), "segment_len_p10_of_means": float(np.percentile(arr, 10)), "segment_len_p90_of_means": float(np.percentile(arr, 90)), "motion_class_distribution": { "still": motion_counts[0] / total_motion, "translate": motion_counts[1] / total_motion, "rotate": motion_counts[2] / total_motion, "translate_rotate": motion_counts[3] / total_motion, }, "motion_class_counts": { "still": motion_counts[0], "translate": motion_counts[1], "rotate": motion_counts[2], "translate_rotate": motion_counts[3], }, } def _print_segment_stats(summary: dict) -> None: if summary.get("n_segments", 0) == 0: print("Segment stats: (no segments)") return mc = summary["motion_class_distribution"] print( f"Segment stats: {summary['n_segments']} segments across {summary['source_episodes']} source episodes " f"({summary['segments_per_episode_mean']:.1f} segments/episode)" ) print( f" per-episode mean segment length: " f"P10={summary['segment_len_p10_of_means']:.1f} " f"median={summary['segment_len_median_of_means']:.1f} " f"mean={summary['segment_len_mean_of_means']:.1f} " f"P90={summary['segment_len_p90_of_means']:.1f}" ) print( f" motion class distribution: " f"still={mc['still'] * 100:.1f}% translate={mc['translate'] * 100:.1f}% " f"rotate={mc['rotate'] * 100:.1f}% translate_rotate={mc['translate_rotate'] * 100:.1f}%" ) def _aggregate_replay_metrics(metrics_list: list[dict]) -> dict: """Aggregate per-(source_episode, speed) replay-fidelity numbers by target speed. Reports mean/median/max for the source-vs-replay error fields produced by compute_replay_metrics so a researcher can verify that integrated motion is preserved and that timing/gripper invariants hold after resampling. """ by_speed: dict[float, list[dict]] = {} for entry in metrics_list: speed = float(entry.get("target_speed", entry.get("speed", 0.0))) by_speed.setdefault(speed, []).append(entry) def _stat(entries: list[dict], key: str) -> dict: arr = np.asarray([float(e.get(key, 0.0)) for e in entries], dtype=np.float64) if arr.size == 0: return {"mean": 0.0, "median": 0.0, "max": 0.0} return { "mean": float(arr.mean()), "median": float(np.median(arr)), "max": float(arr.max()), } per_speed: dict[str, dict] = {} for speed, entries in sorted(by_speed.items()): per_speed[f"{speed:g}"] = { "episodes": len(entries), "actual_speed": _stat(entries, "actual_speed"), "speed_error": _stat(entries, "speed_error"), "integrated_translation_l2_error": _stat(entries, "integrated_translation_l2_error"), "integrated_rotation_l2_error": _stat(entries, "integrated_rotation_l2_error"), "translation_path_ratio": _stat(entries, "translation_path_ratio"), "rotation_path_ratio": _stat(entries, "rotation_path_ratio"), "padded_ratio": _stat(entries, "padded_ratio"), "gripper_switch_delta_sum": int(sum(int(e.get("gripper_switch_delta", 0)) for e in entries)), "gripper_switch_delta_max_abs": int( max((abs(int(e.get("gripper_switch_delta", 0))) for e in entries), default=0) ), } return {"per_speed": per_speed} def _print_replay_summary(summary: dict) -> None: print("Replay-fidelity summary (per speed):") for speed_label, stats in summary["per_speed"].items(): transl = stats["integrated_translation_l2_error"] rot = stats["integrated_rotation_l2_error"] speed_err = stats["speed_error"] actual = stats["actual_speed"] path_t = stats["translation_path_ratio"] path_r = stats["rotation_path_ratio"] padded = stats["padded_ratio"] print( f" speed={speed_label}x ep={stats['episodes']} " f"actual={actual['mean']:.3f} (max_err={speed_err['max']:.4f}) " f"transl_L2 median/max={transl['median']:.2e}/{transl['max']:.2e} " f"rot_L2 median/max={rot['median']:.2e}/{rot['max']:.2e} " f"path_ratio T/R median={path_t['median']:.4f}/{path_r['median']:.4f} " f"padded_mean={padded['mean']:.3f} " f"gripper_delta sum/maxabs={stats['gripper_switch_delta_sum']}/" f"{stats['gripper_switch_delta_max_abs']}" ) def _update_info( src_info: dict, total_episodes: int, total_frames: int, total_videos: int, chunks_size: int, fps: int, ) -> dict: info = json.loads(json.dumps(src_info)) info["total_episodes"] = int(total_episodes) info["total_frames"] = int(total_frames) info["total_videos"] = int(total_videos) info["total_chunks"] = int(math.ceil(total_episodes / chunks_size)) info["chunks_size"] = int(chunks_size) info["fps"] = int(fps) info["splits"] = {"train": f"0:{total_episodes}"} info["data_path"] = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" info["video_path"] = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" features = info.setdefault("features", {}) scalar_i64 = {"dtype": "int64", "shape": [1], "names": None} scalar_f32 = {"dtype": "float32", "shape": [1], "names": None} scalar_i8 = {"dtype": "int8", "shape": [1], "names": None} features.update( { "speed": scalar_f32, "speed_index": scalar_i64, "speed_label": {"dtype": "string", "shape": [1], "names": None}, "chunk_phase": scalar_i64, "chunk_phase_count": scalar_i64, "valid_mask": scalar_i8, "observation_mask": scalar_i8, "action_mask": scalar_i8, "is_padded": scalar_i8, "segment_id": scalar_i64, "motion_class": scalar_i64, "source_episode_index": scalar_i64, "source_frame_index": scalar_i64, "source_step_index": scalar_i64, "source_index": scalar_i64, "cleaned_translation": scalar_i8, "cleaned_rotation": scalar_i8, } ) for feature in features.values(): if feature.get("dtype") == "video": feature.setdefault("info", {})["video.fps"] = int(fps) return info def build_dataset(args: argparse.Namespace) -> None: src_root = Path(args.src).resolve() speeds = [float(x) for x in args.speeds] phase_specs = _speed_phase_specs( speeds, chunk_aligned_observation=args.chunk_aligned_observation, sliding_chunk_phases=args.sliding_chunk_phases, ) paths = _episode_paths(src_root, args.max_episodes) dst_root = _resolve_output_root(Path(args.dst), args, speeds, len(paths)) if dst_root.exists(): if not args.overwrite: raise FileExistsError(f"{dst_root} exists; pass --overwrite to replace it") shutil.rmtree(dst_root) src_info = _load_json(src_root / "meta" / "info.json") source_episode_tasks = { int(row["episode_index"]): row.get("tasks", []) for row in _read_jsonl(src_root / "meta" / "episodes.jsonl") } chunks_size = int(src_info.get("chunks_size", 1000)) fps = int(args.fps or src_info.get("fps", 20)) video_keys = _video_keys(src_info) if args.write_videos else [] image_keys = _image_keys(src_info) episode_info = _update_info( src_info, total_episodes=0, total_frames=0, total_videos=0, chunks_size=chunks_size, fps=fps, ) config = SpeedTransformConfig( transl_eps=args.segment_transl_eps, rot_eps=args.segment_rot_eps, clean_transl_eps=args.clean_transl_eps, clean_rot_eps=args.clean_rot_eps, direction_cos_threshold=args.direction_cos_threshold, min_segment_len=args.min_segment_len, keep_still_segments=not args.drop_still_segments, chunk_aligned_observation=args.chunk_aligned_observation, fps=fps, ) dst_root.mkdir(parents=True, exist_ok=True) all_episode_rows: list[dict] = [] all_stats: list[dict] = [] all_metrics: list[dict] = [] global_index = 0 dst_episode_index = 0 for path in tqdm(paths, desc="episodes"): src_df = pd.read_parquet(path) src_episode_index = int(src_df["episode_index"].iloc[0]) task_index = int(src_df["task_index"].iloc[0]) action_col = _first_existing_column(src_df, ["action", "actions"]) state_col = _first_existing_column(src_df, ["observation.state", "state"]) actions = _stack_column(src_df, action_col) states = _stack_column(src_df, state_col) source_frame_indices = src_df["frame_index"].to_numpy(dtype=np.int64) for speed_index, speed, chunk_phase, chunk_phase_count in phase_specs: phase_config = dataclasses.replace(config, chunk_phase=chunk_phase) transformed, metrics = transform_episode(actions, states, source_frame_indices, speed, phase_config) n = len(transformed["action"]) out_columns = { state_col: list(transformed["state"].astype(np.float32)), action_col: list(transformed["action"].astype(np.float32)), "timestamp": (np.arange(n, dtype=np.float32) / float(fps)).astype(np.float32), "frame_index": np.arange(n, dtype=np.int64), "episode_index": np.full(n, dst_episode_index, dtype=np.int64), "index": np.arange(global_index, global_index + n, dtype=np.int64), "task_index": np.full(n, task_index, dtype=np.int64), "speed": transformed["speed"].astype(np.float32), "speed_index": np.full(n, speed_index, dtype=np.int64), "speed_label": np.full(n, _speed_label(speed), dtype=object), "chunk_phase": np.full(n, chunk_phase, dtype=np.int64), "chunk_phase_count": np.full(n, chunk_phase_count, dtype=np.int64), "valid_mask": transformed["observation_mask"].astype(np.int8), "observation_mask": transformed["observation_mask"].astype(np.int8), "action_mask": transformed["action_mask"].astype(np.int8), "is_padded": transformed["is_padded"].astype(np.int8), "segment_id": transformed["segment_id"].astype(np.int64), "motion_class": transformed["motion_class"].astype(np.int64), "source_episode_index": np.full(n, src_episode_index, dtype=np.int64), "source_frame_index": transformed["source_frame_index"].astype(np.int64), "source_step_index": transformed["source_step_index"].astype(np.int64), "source_index": src_df["index"].to_numpy(dtype=np.int64)[ transformed["source_step_index"].astype(np.int64) ], "cleaned_translation": transformed["cleaned_translation"].astype(np.int8), "cleaned_rotation": transformed["cleaned_rotation"].astype(np.int8), } image_columns = _resample_embedded_images( src_df, transformed["source_step_index"], transformed["observation_mask"], image_keys, ) out_df = pd.DataFrame({**image_columns, **out_columns}) chunk = dst_episode_index // chunks_size out_path = dst_root / "data" / f"chunk-{chunk:03d}" / f"episode_{dst_episode_index:06d}.parquet" _write_episode_parquet(out_path, out_df, episode_info) _copy_or_resample_videos( src_root, dst_root, src_episode_index, dst_episode_index, transformed["source_frame_index"], transformed["observation_mask"], video_keys, chunks_size, fps, ) task_payload = source_episode_tasks.get(src_episode_index, []) all_episode_rows.append( { "episode_index": dst_episode_index, "tasks": task_payload, "length": int(n), "source_episode_index": int(src_episode_index), "speed": float(speed), "speed_label": _speed_label(speed), "chunk_phase": int(chunk_phase), "chunk_phase_count": int(chunk_phase_count), } ) all_stats.append(_episode_stats(out_df, dst_episode_index)) metrics.update( { "episode_index": int(dst_episode_index), "source_episode_index": int(src_episode_index), "task_index": int(task_index), "speed": float(speed), "speed_label": _speed_label(speed), "chunk_phase": int(chunk_phase), "chunk_phase_count": int(chunk_phase_count), } ) all_metrics.append(metrics) global_index += n dst_episode_index += 1 src_meta = src_root / "meta" dst_meta = dst_root / "meta" dst_meta.mkdir(parents=True, exist_ok=True) for filename in ["tasks.jsonl", "modality.json"]: src_file = src_meta / filename if src_file.exists(): shutil.copy2(src_file, dst_meta / filename) info = _update_info( src_info, total_episodes=dst_episode_index, total_frames=global_index, total_videos=dst_episode_index * len(video_keys), chunks_size=chunks_size, fps=fps, ) _write_json(dst_meta / "info.json", info) _write_jsonl(dst_meta / "episodes.jsonl", all_episode_rows) _write_jsonl(dst_meta / "episodes_stats.jsonl", all_stats) _write_jsonl(dst_meta / "speed_metrics.jsonl", all_metrics) cleaning_summary = _aggregate_cleaning_stats(all_metrics) _write_json(dst_meta / "cleaning_summary.json", cleaning_summary) replay_summary = _aggregate_replay_metrics(all_metrics) _write_json(dst_meta / "replay_summary.json", replay_summary) segment_summary = _aggregate_segment_stats(all_metrics) _write_json(dst_meta / "segment_summary.json", segment_summary) print(f"Wrote {dst_episode_index} episodes / {global_index} frames to {dst_root}") _print_cleaning_summary(cleaning_summary) _print_segment_stats(segment_summary) _print_replay_summary(replay_summary) if not args.write_videos: print("Video writing was disabled; enable --write-videos for LeRobot video loading.") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--src", required=True, help="Source LeRobot LIBERO dataset root") parser.add_argument("--dst", required=True, help="Output dataset root, or parent root with --auto-name") parser.add_argument( "--auto-name", action="store_true", help="Treat --dst as a parent directory and create a unique run subdirectory.", ) parser.add_argument("--run-tag", default="smoke", help="Run tag used by --auto-name") parser.add_argument("--task-suite-name", default="libero_spatial", help="Suite name used by --auto-name") parser.add_argument( "--timestamp", default=None, help="Optional timestamp/name suffix for --auto-name. Defaults to current YYYYMMDD_HHMMSS.", ) parser.add_argument("--speeds", nargs="+", type=float, default=[0.5, 0.75, 1.0, 1.25, 2.0]) parser.add_argument("--max-episodes", type=int, default=None) parser.add_argument("--fps", type=int, default=None) parser.add_argument("--write-videos", action="store_true") parser.add_argument("--overwrite", action="store_true") parser.add_argument("--segment-transl-eps", type=float, default=1e-4) parser.add_argument("--segment-rot-eps", type=float, default=1e-4) # Near-zero action cleaning is OFF by default (eps=0). LIBERO is clean # enough that no frame falls below 1e-4 in practice. Pass a positive value # to re-enable on noisier datasets. parser.add_argument("--clean-transl-eps", type=float, default=0.0) parser.add_argument("--clean-rot-eps", type=float, default=0.0) parser.add_argument("--direction-cos-threshold", type=float, default=-0.25) parser.add_argument("--min-segment-len", type=int, default=1) parser.add_argument("--drop-still-segments", action="store_true") parser.add_argument( "--chunk-aligned-observation", action="store_true", help=( "Divide each segment into chunks of q source frames -> p output frames " "(speed = q/p). Chunk-start outputs have mask=1 with state forced to the " "exact source state at k*q; other full-chunk outputs are padded (mask=0). " "Trailing leftover (segment_len %% q) source frames pass through 1:1 verbatim " "(each becomes one output at 1.0x, mask=1)." ), ) parser.add_argument( "--sliding-chunk-phases", action="store_true", help=( "With --chunk-aligned-observation, emit q phase-shifted episodes for " "each speed=q/p so more source frames become valid chunk starts." ), ) return parser.parse_args() if __name__ == "__main__": build_dataset(parse_args())