#!/usr/bin/env python3 from __future__ import annotations import argparse import json import os from pathlib import Path from typing import Any import numpy as np import pandas as pd from huggingface_hub import get_token, hf_hub_download from physical_ai_av import PhysicalAIAVDatasetInterface from PIL import Image from scipy.spatial.transform import Rotation from alpamayo_lite.config import load_config from alpamayo_lite.prompts import clean_reasoning_text CLIP_DURATION_US = 20_000_000 def parse_events(cell: object) -> list[dict[str, Any]]: if cell is None or (np.isscalar(cell) and pd.isna(cell)): return [] parsed = json.loads(cell) if isinstance(cell, str) else cell if not isinstance(parsed, (list, tuple)): return [] return [event for event in parsed if isinstance(event, dict)] def collect_candidates( reasoning_path: Path, history_seconds: float, future_seconds: float, ) -> list[dict[str, Any]]: frame = pd.read_parquet(reasoning_path) candidates: list[dict[str, Any]] = [] for clip_id, events_cell in frame["events"].items(): for event in parse_events(events_cell): if "event_start_timestamp" not in event: continue t0_us = int(event["event_start_timestamp"]) if t0_us < int(history_seconds * 1_000_000): continue if t0_us + int(future_seconds * 1_000_000) > CLIP_DURATION_US: continue # PAI command text uses `coc`; older metadata shards use `cot`. reasoning = clean_reasoning_text(str(event.get("coc") or event.get("cot") or "")) candidates.append({"clip_id": str(clip_id), "t0_us": t0_us, "reasoning": reasoning}) if not candidates: raise RuntimeError("No valid reasoning events remain after trajectory safety filtering") return candidates def load_local_trajectory( interface: PhysicalAIAVDatasetInterface, clip_id: str, t0_us: int, history_steps: int, future_steps: int, time_step: float, ) -> tuple[np.ndarray, np.ndarray]: egomotion = interface.get_clip_feature( clip_id, interface.features.LABELS.EGOMOTION, maybe_stream=True ) history_offsets = np.arange(-(history_steps - 1), 1, dtype=np.int64) history_timestamps = t0_us + history_offsets * int(time_step * 1_000_000) future_offsets = np.arange(1, future_steps + 1, dtype=np.int64) future_timestamps = t0_us + future_offsets * int(time_step * 1_000_000) history_state = egomotion(history_timestamps) future_state = egomotion(future_timestamps) history_xyz = history_state.pose.translation future_xyz = future_state.pose.translation t0_rotation = Rotation.from_quat(history_state.pose.rotation.as_quat()[-1]) inverse = t0_rotation.inv() origin = history_xyz[-1].copy() history_local = inverse.apply(history_xyz - origin).astype(np.float32) future_local = inverse.apply(future_xyz - origin).astype(np.float32) return history_local, future_local def load_camera_frame( interface: PhysicalAIAVDatasetInterface, clip_id: str, t0_us: int, camera_name: str, ) -> np.ndarray: camera_feature = getattr(interface.features.CAMERA, camera_name.upper()) reader = interface.get_clip_feature(clip_id, camera_feature, maybe_stream=True) frames, _ = reader.decode_images_from_timestamps(np.asarray([t0_us], dtype=np.int64)) return frames[0] def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", default="configs/local_overfit.yaml") args = parser.parse_args() cfg = load_config(args.config) token = os.environ.get("HF_TOKEN") or get_token() if not token: raise SystemExit( "PAI is gated. Set HF_TOKEN for an account that accepted the dataset license." ) raw_dir = cfg.resolve(cfg.paths.pai_raw_dir) raw_dir.mkdir(parents=True, exist_ok=True) cache_dir = cfg.root / ".cache" / "huggingface" reasoning_file = Path( hf_hub_download( repo_id=cfg.data.repo_id, repo_type="dataset", filename="reasoning/ood_reasoning.parquet", token=token, cache_dir=cache_dir, local_dir=raw_dir, ) ) candidates = collect_candidates( reasoning_file, history_seconds=(cfg.data.history_steps - 1) * cfg.data.time_step, future_seconds=cfg.data.source_future_steps * cfg.data.time_step, ) rng = np.random.default_rng(cfg.seed) selected_indices = rng.choice( len(candidates), size=min(cfg.data.num_clips, len(candidates)), replace=False ) selected = [candidates[int(index)] for index in selected_indices] interface = PhysicalAIAVDatasetInterface( token=token, cache_dir=cache_dir, local_dir=raw_dir, confirm_download_threshold_gb=1000.0, ) processed_dir = cfg.resolve(cfg.paths.manifest).parent image_dir = processed_dir / "images" trajectory_dir = processed_dir / "trajectories" image_dir.mkdir(parents=True, exist_ok=True) trajectory_dir.mkdir(parents=True, exist_ok=True) rows: list[dict[str, Any]] = [] for sample_index, event in enumerate(selected): clip_id = event["clip_id"] t0_us = event["t0_us"] history, future_full = load_local_trajectory( interface=interface, clip_id=clip_id, t0_us=t0_us, history_steps=cfg.data.history_steps, future_steps=cfg.data.source_future_steps, time_step=cfg.data.time_step, ) waypoint_indices = np.linspace( 0, cfg.data.source_future_steps - 1, cfg.data.expert_waypoints, dtype=np.int64, ) future = future_full[waypoint_indices] frame = load_camera_frame( interface=interface, clip_id=clip_id, t0_us=t0_us, camera_name=cfg.data.camera, ) stem = f"sample_{sample_index:02d}_{clip_id}_{t0_us}" image_path = image_dir / f"{stem}.jpg" trajectory_path = trajectory_dir / f"{stem}.npz" Image.fromarray(frame).save(image_path, quality=95) np.savez_compressed( trajectory_path, history_xyz=history, future_xyz=future, source_future_xyz=future_full, waypoint_indices=waypoint_indices, ) rows.append( { "clip_id": clip_id, "t0_us": t0_us, "reasoning": event["reasoning"], "image_path": str(image_path.relative_to(cfg.root)), "trajectory_path": str(trajectory_path.relative_to(cfg.root)), "camera": cfg.data.camera, "history_steps": cfg.data.history_steps, "expert_waypoints": cfg.data.expert_waypoints, } ) manifest_path = cfg.resolve(cfg.paths.manifest) manifest_path.write_text( "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), encoding="utf-8", ) print(json.dumps({"manifest": str(manifest_path), "samples": rows}, indent=2)) if __name__ == "__main__": main()