""" Prepare an expert demonstration for each task in a LIBERO suite. The demonstration with the shortest successful frame sequence is selected as expert. Usage: python scripts/prepare_expert_demo.py \ --libero_task_suite libero_10 \ --libero_raw_data_dir /path/to/libero_10 \ --output_dir /path/to/output """ import argparse import json import shutil from pathlib import Path from typing import Dict, List, Optional, Tuple import h5py import numpy as np from PIL import Image from tqdm import tqdm def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Select and export expert LIBERO demonstrations") parser.add_argument( "--libero_task_suite", required=True, help="Name of the LIBERO task suite (used for bookkeeping in the output)", ) parser.add_argument( "--libero_raw_data_dir", required=True, type=Path, help="Directory containing the raw LIBERO HDF5 demo files", ) parser.add_argument( "--output_dir", required=True, type=Path, help="Directory where expert demo frames will be exported", ) parser.add_argument( "--frame_key", default="agentview_rgb", choices=["agentview_rgb", "eye_in_hand_rgb"], help="Observation frame key to export (eye_in_hand_rgb is useful for wrist-view references)", ) parser.add_argument( "--overwrite", action="store_true", help="Overwrite per-task export folders if they already exist", ) return parser.parse_args() def list_hdf5_files(root: Path) -> List[Path]: if not root.is_dir(): raise FileNotFoundError(f"LIBERO raw data directory not found: {root}") return sorted(p for p in root.glob("*.hdf5") if p.is_file()) def has_frames(demo_group: h5py.Group, frame_key: str) -> bool: obs_group = demo_group.get("obs") return obs_group is not None and frame_key in obs_group def demo_statistics(demo_group: h5py.Group, frame_key: str) -> Tuple[int, bool, Optional[int]]: frames_ds = demo_group[f"obs/{frame_key}"] total_frames = int(frames_ds.shape[0]) dones_ds = demo_group.get("dones") if dones_ds is None: return total_frames, True, None dones = np.asarray(dones_ds[:]) success_indices = np.where(dones == 1)[0] if success_indices.size == 0: return total_frames, False, None return total_frames, True, int(success_indices[0]) def select_expert_demo(data_group: h5py.Group, frame_key: str) -> Optional[Tuple[str, Dict[str, int]]]: best_name: Optional[str] = None best_total: Optional[int] = None best_success_index: Optional[int] = None for demo_name in sorted(data_group.keys()): demo_group = data_group[demo_name] if not has_frames(demo_group, frame_key): continue total_frames, is_successful, success_index = demo_statistics(demo_group, frame_key) if not is_successful or total_frames == 0: continue candidate_success = success_index if success_index is not None else total_frames - 1 if best_name is None: best_name = demo_name best_total = total_frames best_success_index = candidate_success continue assert best_total is not None assert best_success_index is not None if total_frames < best_total or ( total_frames == best_total and candidate_success < best_success_index ): best_name = demo_name best_total = total_frames best_success_index = candidate_success if best_name is None or best_total is None or best_success_index is None: return None return best_name, { "frame_count": best_total, "success_index": best_success_index, } def pretty_task_name(hdf5_path: Path) -> str: return hdf5_path.stem.replace("_", " ").replace("demo", "").strip() def save_frame(array: np.ndarray, file_path: Path) -> None: file_path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(array).transpose(Image.FLIP_TOP_BOTTOM).save(file_path) def export_demo_frames( demo_group: h5py.Group, output_dir: Path, output_root: Path, frame_key: str ) -> List[str]: frames_ds = demo_group[f"obs/{frame_key}"] relative_paths: List[str] = [] for frame_idx in range(frames_ds.shape[0]): array = np.asarray(frames_ds[frame_idx]) frame_path = output_dir / f"frame_{frame_idx:04d}.png" save_frame(array, frame_path) relative_paths.append(str(frame_path.relative_to(output_root))) return relative_paths def ensure_clean_directory(path: Path, overwrite: bool) -> bool: if path.exists(): if not overwrite: return False shutil.rmtree(path) path.mkdir(parents=True, exist_ok=True) return True def main() -> None: args = parse_args() raw_dir = args.libero_raw_data_dir.expanduser() output_root = args.output_dir.expanduser() suite_dir = output_root / args.libero_task_suite suite_dir.mkdir(parents=True, exist_ok=True) hdf5_files = list_hdf5_files(raw_dir) if not hdf5_files: print("No HDF5 files found. Check --libero_raw_data_dir.") return manifest: List[Dict[str, object]] = [] for hdf5_path in tqdm(hdf5_files, desc="Tasks"): with h5py.File(hdf5_path, "r") as handle: data_group = handle.get("data") if data_group is None: print(f"[skip] No data group in {hdf5_path.name}") continue selection = select_expert_demo(data_group, frame_key=args.frame_key) if selection is None: print(f"[skip] No successful demos with '{args.frame_key}' in {hdf5_path.name}") continue demo_name, stats = selection task_dir = suite_dir / hdf5_path.stem if not ensure_clean_directory(task_dir, args.overwrite): print(f"[skip] {task_dir} exists. Use --overwrite to regenerate.") continue demo_group = data_group[demo_name] frame_paths = export_demo_frames( demo_group, task_dir, output_root, frame_key=args.frame_key ) task_metadata = { "suite": args.libero_task_suite, "task_file": hdf5_path.name, "task_name": pretty_task_name(hdf5_path), "demo_name": demo_name, "frame_key": args.frame_key, "frame_count": len(frame_paths), "success_index": stats["success_index"], "frame_paths": frame_paths, } manifest.append(task_metadata) metadata_path = task_dir / "metadata.json" with metadata_path.open("w", encoding="utf-8") as fout: json.dump(task_metadata, fout, indent=2) if not manifest: print("No expert demonstrations were exported.") return manifest_path = suite_dir / "expert_manifest.json" with manifest_path.open("w", encoding="utf-8") as fout: json.dump(manifest, fout, indent=2) print(f"\nSaved {len(manifest)} expert demos to {suite_dir}") print(f"Manifest: {manifest_path}") if __name__ == "__main__": main()