scripts/resample_trajectory.py

#!/usr/bin/env python
"""Resample a planner trajectory onto a different frame count / frame rate.

Written for one specific downstream constraint: Cosmos-Transfer2.5-2B consumes
control videos at **16 FPS** and states that lengths which are a multiple of
**93 frames** perform best. Our planner emits at DROID's 15 Hz, so a
pick-and-place that runs 10.27 s comes out as 155 frames -- neither the right
rate nor a multiple of 93. Rather than re-plan (which would change the motion),
this resamples the existing, already-verified trajectory onto the target grid.

The motion is *stretched*, not cropped: all ``T_src`` source samples are mapped
uniformly onto ``T_dst`` output samples, so every phase of the pick-and-place
survives and only the wall-clock duration changes. For 155 -> 186 that is
10.27 s -> 11.63 s, i.e. the action plays 1.13x slower than recorded. That is
the closest a 93-multiple gets to real time here; 93 frames would be 1.77x
*faster*, which reads as a hurried, unnatural arm.

Interpolation is deliberately linear for joints/gripper and SLERP for object
rotations:

  * Joints come from a trapezoidal-profile plan, which has curvature
    discontinuities at the accel/decel corners. A cubic spline would overshoot
    exactly there; linear cannot. The cost is chord error, which ``--validate``
    measures directly (decimate-and-restore) rather than assuming.
  * The gripper has a genuine step (release). Linear turns it into a one-frame
    ramp; a spline would ring around it.
  * Rotations must stay in SO(3). Element-wise interpolation of a 4x4 does not
    preserve orthonormality, so rotations go through SLERP and the result is
    re-checked, not trusted.

Usage:
    PYTHONPATH=src python scripts/resample_trajectory.py \\
        --npz outputs/pick_place_books/pick_place_books_settled.npz \\
        --frames 186 --fps 16 \\
        --out outputs/pick_place_books/pick_place_books_cosmos186.npz
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
from scipy.spatial.transform import Rotation, Slerp

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

from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("resample_trajectory")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--npz", type=Path, required=True, help="planner trajectory .npz")
    p.add_argument("--out", type=Path, required=True)
    p.add_argument(
        "--frames", type=int, default=186,
        help="output frame count (Cosmos-Transfer2.5 prefers multiples of 93)",
    )
    p.add_argument("--fps", type=float, default=16.0, help="output frame rate")
    p.add_argument(
        "--validate", action="store_true", default=True,
        help="measure the interpolation's own chord error before writing",
    )
    return p.parse_args()


def _interp_linear_blocks(M: np.ndarray, s_src: np.ndarray, s_dst: np.ndarray) -> np.ndarray:
    """Element-wise linear interpolation of a (T, 3, 3) stack."""
    return np.stack(
        [np.stack([np.interp(s_dst, s_src, M[:, r, c]) for c in range(3)], axis=1)
         for r in range(3)], axis=1,
    )


def _resample_poses(poses: np.ndarray, s_src: np.ndarray, s_dst: np.ndarray) -> np.ndarray:
    """(T, n_obj, 4, 4) object transforms -> (T_dst, n_obj, 4, 4).

    The 3x3 block of these matrices is **not** a bare rotation: the planner
    bakes the object mesh's own scale into it (this project's brick carries a
    uniform factor of ~0.05, i.e. the mesh is authored in units 20x the metre
    scale it is placed at). Feeding that straight to ``Rotation.from_matrix``
    silently orthonormalises the scale away and the object renders ~20x too
    big -- which is what happened the first time this ran, and is invisible in
    any check phrased as "is the result a valid rotation", because destroying
    the scale is exactly what makes it one.

    So decompose ``M = R diag(s)`` per frame, SLERP ``R``, interpolate ``s``
    linearly, and recombine. Scale is per-axis rather than assumed uniform, so
    a non-uniformly scaled asset survives too.
    """
    T_dst, n_obj = s_dst.size, poses.shape[1]
    out = np.zeros((T_dst, n_obj, 4, 4))
    out[:, :, 3, 3] = 1.0
    for j in range(n_obj):
        M = poses[:, j, :3, :3]
        scale = np.linalg.norm(M, axis=1)                     # (T, 3) per-column norms
        if not np.all(scale > 1e-12):
            raise SystemExit(f"object {j}: degenerate (zero-scale) transform, cannot resample")
        R_src = M / scale[:, None, :]

        R_dst = Slerp(s_src, Rotation.from_matrix(R_src))(s_dst).as_matrix()
        s_dst_scale = np.stack([np.interp(s_dst, s_src, scale[:, k]) for k in range(3)], axis=1)

        out[:, j, :3, :3] = R_dst * s_dst_scale[:, None, :]
        out[:, j, :3, 3] = np.stack(
            [np.interp(s_dst, s_src, poses[:, j, k, 3]) for k in range(3)], axis=1
        )
    return out


def _chord_error(values: np.ndarray, s_src: np.ndarray) -> float:
    """Bound the resample's own interpolation error, measured not assumed.

    Drop every other source sample, linearly restore the full grid from what is
    left, and report the worst discrepancy against the samples that were
    actually dropped. That is a *2x-coarser* resample than the one being
    performed, so it over-estimates the real error -- an honest upper bound.
    """
    keep = s_src[::2]
    err = 0.0
    for k in range(values.shape[1]):
        restored = np.interp(s_src, keep, values[::2, k])
        err = max(err, float(np.abs(restored - values[:, k]).max()))
    return err


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

    if args.frames % 93 != 0:
        logger.warning(
            "--frames %d is not a multiple of 93; Cosmos-Transfer2.5's model card "
            "reports multiples of 93 perform best", args.frames,
        )

    with np.load(args.npz, allow_pickle=True) as npz:
        data = {k: np.asarray(npz[k]) for k in npz.files}

    t_src = data["timestamps"].astype(np.float64)
    T_src = t_src.size
    if T_src < 2:
        raise SystemExit(f"need at least 2 source samples, got {T_src}")

    # Uniform index-space mapping: source sample i <-> s = i, output sample k
    # <-> s = k * (T_src - 1) / (T_dst - 1). Endpoints land exactly on the
    # first and last source samples, so the start and end poses are preserved
    # bit-for-bit rather than interpolated.
    s_src = np.arange(T_src, dtype=np.float64)
    s_dst = np.linspace(0.0, T_src - 1.0, args.frames)

    src_duration = float(t_src[-1] - t_src[0])
    dst_duration = (args.frames - 1) / args.fps
    stretch = dst_duration / src_duration if src_duration > 0 else float("nan")

    out = {
        "timestamps": np.arange(args.frames, dtype=np.float64) / args.fps,
    }

    joints = data["joint_positions"].astype(np.float64)
    out["joint_positions"] = np.stack(
        [np.interp(s_dst, s_src, joints[:, j]) for j in range(joints.shape[1])], axis=1
    )
    out["gripper"] = np.interp(s_dst, s_src, data["gripper"].astype(np.float64))
    out["object_poses"] = _resample_poses(data["object_poses"].astype(np.float64), s_src, s_dst)
    out["object_names"] = data["object_names"]

    # --- validation -------------------------------------------------------
    joint_err = _chord_error(joints, s_src)

    logger.info(
        "resampled %d @ %.4g Hz -> %d @ %.4g fps  (%.3f s -> %.3f s, %.3fx %s)",
        T_src, (T_src - 1) / src_duration, args.frames, args.fps,
        src_duration, dst_duration, stretch, "slower" if stretch > 1 else "faster",
    )
    logger.info(
        "joint interpolation error (2x-coarse upper bound): %.3e rad = %.4f deg",
        joint_err, np.degrees(joint_err),
    )

    # Object-transform scale is checked explicitly, and checked against the
    # *source*. The obvious formulation -- "assert the 3x3 block is a valid
    # rotation" -- is worse than useless here: it is satisfied precisely by the
    # bug of orthonormalising the asset's baked-in scale away, which renders
    # the object at the wrong size while every number looks perfect.
    src_scale = np.linalg.norm(data["object_poses"][:, :, :3, :3], axis=2)
    dst_scale = np.linalg.norm(out["object_poses"][:, :, :3, :3], axis=2)
    scale_err = float(
        np.abs(dst_scale - np.stack(
            [np.stack([np.interp(s_dst, s_src, src_scale[:, j, k]) for k in range(3)], axis=1)
             for j in range(src_scale.shape[1])], axis=1)).max()
    )
    logger.info(
        "object transform scale: source %s, resampled %s, max error %.2e",
        np.round(src_scale.reshape(-1, 3).min(axis=0), 6),
        np.round(dst_scale.reshape(-1, 3).min(axis=0), 6), scale_err,
    )
    if scale_err > 1e-12:
        raise SystemExit("object transform scale not preserved -- refusing to write")

    # Endpoints must survive exactly, or the grasp/release poses shift.
    for key, src in (("joint_positions", joints), ("gripper", data["gripper"]),
                     ("object_poses", data["object_poses"])):
        for idx in (0, -1):
            if not np.allclose(out[key][idx], src[idx], atol=1e-12):
                raise SystemExit(f"{key} endpoint {idx} not preserved")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    np.savez(args.out, **out)

    sidecar = args.out.with_suffix(".resample.json")
    sidecar.write_text(json.dumps({
        "source_npz": str(args.npz),
        "source_frames": int(T_src),
        "source_fps": float((T_src - 1) / src_duration),
        "source_duration_s": src_duration,
        "output_frames": int(args.frames),
        "output_fps": float(args.fps),
        "output_duration_s": dst_duration,
        "time_stretch_factor": stretch,
        "stretch_note": (
            "The motion is stretched uniformly onto the new frame count, not cropped or "
            "re-planned. The arm therefore moves %.1f%% %s than it does in the source "
            "trajectory; nothing about the path changes." % (
                abs(stretch - 1) * 100, "slower" if stretch > 1 else "faster")
        ),
        "interpolation": {
            "joint_positions": "linear",
            "gripper": "linear",
            "object_translation": "linear",
            "object_rotation": "SLERP on the rotation factor of M = R diag(s)",
            "object_scale": "linear on the per-axis scale factor s, then recombined",
            "why_not_spline": (
                "The joint trajectory comes from a trapezoidal velocity profile and has "
                "curvature discontinuities at the accel/decel corners; the gripper has a "
                "genuine step at release. A cubic spline overshoots at exactly those points. "
                "Linear cannot overshoot, and the measured chord error below is far under "
                "any tolerance that matters here."
            ),
        },
        "joint_interp_error_rad_upper_bound": joint_err,
        "joint_interp_error_deg_upper_bound": float(np.degrees(joint_err)),
        "object_transform_scale_max_err": scale_err,
        "object_transform_scale_source": src_scale.reshape(-1, 3).min(axis=0).tolist(),
    }, indent=2))

    print("=" * 78)
    print(f"wrote {args.out}  ({args.frames} frames @ {args.fps:g} fps = {dst_duration:.3f} s)")
    print(f"      {sidecar.name}")
    print("=" * 78)
    return 0


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