scripts/make_standin_conditioning_npz.py

#!/usr/bin/env python
"""Synthesize a stand-in planner ``.npz`` matching render_conditioning_set.py's input contract.

The real file this substitutes for is produced by a planner running in
parallel with this script's development: ``timestamps`` (T,), ``joint_positions``
(T, 7), ``gripper`` (T,) in DROID's [0, 1]-open-at-0 convention, ``object_poses``
(T, n_obj, 4, 4) world-frame per object per frame, and ``object_names`` (n_obj,).

There is no kinematic planner here to drive that contract honestly, so this
script fakes the two halves separately from real data instead of inventing
motion from scratch:

    * the arm trajectory is the *actual* recorded DROID joint/gripper signal
      for one episode (so the robot channel exercises real, physically valid
      motion, not a synthetic wave) -- see render_robot.py for the same
      "replay a real episode" pattern used elsewhere in this repo;
    * the object is held at one *real, previously-fit* alignment transform
      (mesh -> world) for the whole clip, i.e. it does not move. This is a
      standin for "static object pose" per the task brief, not a claim that
      the object stayed exactly still during the real episode.

Neither half claims to be a physically consistent pick-and-place: the arm
never actually grasps or moves this particular object in the source episode.
This is purely a shape-and-dtype-correct fixture for exercising
render_conditioning_set.py before the real planner output exists.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import h5py
import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.types import DataError  # noqa: E402
from fpgm.utils.io import ensure_dir  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("make_standin_conditioning_npz")

_TRAJECTORY_FPS = 15.0


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--episode", required=True, help="DROID episode uuid to replay joints from")
    p.add_argument(
        "--object-name", default="brick",
        help="name written to object_names -- must match a --object-asset key given to "
        "render_conditioning_set.py",
    )
    p.add_argument(
        "--object-alignment", type=Path, default=None,
        help="path to a previously-fit alignment.npz (mesh -> world transform) whose "
        "'transform' field is held fixed for every frame. If omitted, a plausible "
        "in-front-of-the-robot pose is used instead (clearly logged, not measured).",
    )
    p.add_argument("--start-row", type=int, default=0, help="first trajectory.h5 row to use")
    p.add_argument("--n-frames", type=int, default=90, help="number of frames (rows) to take")
    p.add_argument("--out", type=Path, required=True, help="output .npz path")
    p.add_argument("--sidecar-out", type=Path, default=None,
                   help="output action-spec sidecar JSON path (default: alongside --out)")
    return p.parse_args()


def _default_object_pose() -> np.ndarray:
    """A plausible, but NOT measured, static object pose in front of the robot base.

    Used only when no real --object-alignment is supplied. Logged loudly so it
    is never mistaken for a real reconstruction.
    """
    pose = np.eye(4, dtype=np.float64)
    pose[:3, 3] = [0.55, 0.05, 0.05]
    return pose


def main() -> int:
    args = parse_args()
    setup_logging()

    traj_path = REPO_ROOT / "data" / "droid_raw" / args.episode / "trajectory.h5"
    if not traj_path.exists():
        raise DataError(f"trajectory.h5 not found for episode {args.episode!r}: {traj_path}")

    with h5py.File(traj_path, "r") as f:
        joint_positions_full = np.asarray(
            f["observation/robot_state/joint_positions"], dtype=np.float64
        )
        gripper_full = np.asarray(f["observation/robot_state/gripper_position"], dtype=np.float64)

    available = joint_positions_full.shape[0]
    end_row = min(args.start_row + args.n_frames, available)
    if end_row <= args.start_row:
        raise DataError(
            f"--start-row {args.start_row} leaves no frames in a trajectory of length {available}"
        )
    if end_row - args.start_row < args.n_frames:
        logger.warning(
            "trajectory only has %d rows; truncating stand-in from %d to %d frames",
            available, args.n_frames, end_row - args.start_row,
        )
    rows = np.arange(args.start_row, end_row)

    joint_positions = joint_positions_full[rows]
    gripper = np.clip(gripper_full[rows], 0.0, 1.0)
    timestamps = rows.astype(np.float64) / _TRAJECTORY_FPS
    timestamps = timestamps - timestamps[0]  # start at t=0, still exact multiples of 1/15s

    if args.object_alignment is not None:
        if not args.object_alignment.exists():
            raise DataError(f"--object-alignment not found: {args.object_alignment}")
        with np.load(args.object_alignment) as npz:
            object_pose = np.asarray(npz["transform"], dtype=np.float64)
        logger.info("static object pose taken from %s", args.object_alignment)
    else:
        object_pose = _default_object_pose()
        logger.warning(
            "no --object-alignment given; using an UNMEASURED placeholder pose %s "
            "for object %r -- fine for exercising the pipeline, not a real placement",
            object_pose[:3, 3].tolist(), args.object_name,
        )

    T = timestamps.shape[0]
    object_poses = np.tile(object_pose[None, None, :, :], (T, 1, 1, 1))
    object_names = np.array([args.object_name], dtype="<U32")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        args.out,
        timestamps=timestamps,
        joint_positions=joint_positions,
        gripper=gripper,
        object_poses=object_poses,
        object_names=object_names,
    )

    sidecar_out = args.sidecar_out or args.out.with_suffix(".action.json")
    ensure_dir(sidecar_out.parent)
    sidecar_out.write_text(json.dumps(
        {
            "STAND_IN": True,
            "note": (
                "This is a synthesized stand-in, not a planner action spec. "
                "joint_positions/gripper are REAL recorded DROID motion for the "
                "episode below (not a rendition of the object_names task); "
                "object_poses is a single fixed pose held for the whole clip."
            ),
            "source_episode": args.episode,
            "source_trajectory_rows": [int(rows[0]), int(rows[-1])],
            "object_name": args.object_name,
            "object_alignment_source": (
                str(args.object_alignment) if args.object_alignment else None
            ),
        },
        indent=2,
    ))

    print(f"wrote {T} frames -> {args.out.resolve()}")
    print(f"sidecar -> {sidecar_out.resolve()}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())