alpamayo-qwen05b-validation/scripts/materialize_pai_subset.py

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