scripts/build_libero_speed_dataset.py

#!/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())