File size: 5,886 Bytes
857c2e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """
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()
|