Release/scripts/prepare_reference_frames_pipeline.py

"""
One-command pipeline for LIBERO reference-frame preparation.

This script can:
1) optionally regenerate higher-resolution LIBERO trajectories, and
2) export expert-demo reference frames for agent view and optional wrist view.

Examples:
  # Full pipeline (regen + agentview + wrist)
  python scripts/prepare_reference_frames_pipeline.py \
    --raw-datasets-root /data/libero/raw \
    --working-root /data/libero/processed \
    --reference-root /data/libero/reference_frames \
    --include-wrist \
    --overwrite

  # Skip regeneration and use existing datasets directly
  python scripts/prepare_reference_frames_pipeline.py \
    --raw-datasets-root /data/libero/raw \
    --working-root /data/libero/processed \
    --reference-root /data/libero/reference_frames \
    --skip-regen \
    --overwrite
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path
from typing import List


DEFAULT_SUITES = ["libero_10", "libero_object", "libero_spatial", "libero_goal"]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prepare LIBERO reference frames for EVOLVE-VLA")
    parser.add_argument(
        "--raw-datasets-root",
        type=Path,
        required=True,
        help="Root directory containing raw suite folders (libero_10/libero_object/libero_spatial/libero_goal)",
    )
    parser.add_argument(
        "--working-root",
        type=Path,
        required=True,
        help="Root directory used for regenerated datasets (one subfolder per suite)",
    )
    parser.add_argument(
        "--reference-root",
        type=Path,
        required=True,
        help="Root directory where exported reference frames are written",
    )
    parser.add_argument(
        "--suites",
        type=str,
        default=",".join(DEFAULT_SUITES),
        help="Comma-separated suites to process (default: libero_10,libero_object,libero_spatial,libero_goal)",
    )
    parser.add_argument(
        "--skip-regen",
        action="store_true",
        help="Skip high-resolution regeneration and use raw datasets directly",
    )
    parser.add_argument(
        "--include-wrist",
        action="store_true",
        help="Also export wrist-view references (eye_in_hand_rgb)",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing exported frame folders",
    )
    return parser.parse_args()


def run_cmd(cmd: List[str], cwd: Path) -> None:
    print("\n[run]", " ".join(cmd))
    subprocess.run(cmd, cwd=str(cwd), check=True)


def main() -> None:
    args = parse_args()
    script_dir = Path(__file__).resolve().parent
    release_root = script_dir.parent

    suites = [s.strip() for s in args.suites.split(",") if s.strip()]
    if not suites:
        raise SystemExit("No suite provided via --suites.")

    regen_script = script_dir / "regenerate_libero_dataset.py"
    expert_script = script_dir / "prepare_expert_demo.py"
    if not regen_script.exists() or not expert_script.exists():
        raise SystemExit("Required scripts missing in Release/scripts.")

    raw_root = args.raw_datasets_root.expanduser().resolve()
    working_root = args.working_root.expanduser().resolve()
    ref_root = args.reference_root.expanduser().resolve()
    working_root.mkdir(parents=True, exist_ok=True)
    ref_root.mkdir(parents=True, exist_ok=True)

    for suite in suites:
        raw_dir = raw_root / suite
        if not raw_dir.is_dir():
            raise SystemExit(f"Raw suite directory not found: {raw_dir}")

        if args.skip_regen:
            dataset_dir = raw_dir
        else:
            dataset_dir = working_root / suite
            regen_cmd = [
                sys.executable,
                str(regen_script),
                "--libero_task_suite",
                suite,
                "--libero_raw_data_dir",
                str(raw_dir),
                "--libero_target_dir",
                str(dataset_dir),
            ]
            run_cmd(regen_cmd, cwd=release_root)

        agentview_out = ref_root / "agentview"
        agent_cmd = [
            sys.executable,
            str(expert_script),
            "--libero_task_suite",
            suite,
            "--libero_raw_data_dir",
            str(dataset_dir),
            "--output_dir",
            str(agentview_out),
        ]
        if args.overwrite:
            agent_cmd.append("--overwrite")
        run_cmd(agent_cmd, cwd=release_root)

        if args.include_wrist:
            wrist_out = ref_root / "wrist"
            wrist_cmd = [
                sys.executable,
                str(expert_script),
                "--libero_task_suite",
                suite,
                "--libero_raw_data_dir",
                str(dataset_dir),
                "--output_dir",
                str(wrist_out),
                "--frame_key",
                "eye_in_hand_rgb",
            ]
            if args.overwrite:
                wrist_cmd.append("--overwrite")
            run_cmd(wrist_cmd, cwd=release_root)

    print("\nReference-frame preparation finished.")
    print("Set EVOLVE_LIBERO_REF_* env vars to suite-specific directories, for example:")
    print(f"  EVOLVE_LIBERO_REF_LIBERO_10={ref_root}/agentview/libero_10")
    print(f"  EVOLVE_LIBERO_REF_LIBERO_OBJECT={ref_root}/agentview/libero_object")
    print(f"  EVOLVE_LIBERO_REF_LIBERO_SPATIAL={ref_root}/agentview/libero_spatial")
    print(f"  EVOLVE_LIBERO_REF_LIBERO_GOAL={ref_root}/agentview/libero_goal")
    if args.include_wrist:
        print(f"  EVOLVE_LIBERO_REF_LIBERO_OBJECT_WRIST={ref_root}/wrist/libero_object")
        print(f"  EVOLVE_LIBERO_REF_LIBERO_SPATIAL_WRIST={ref_root}/wrist/libero_spatial")
        print(f"  EVOLVE_LIBERO_REF_LIBERO_GOAL_WRIST={ref_root}/wrist/libero_goal")


if __name__ == "__main__":
    main()