scripts/plan_pick_place_books.py

#!/usr/bin/env python
"""Plan a pick of the LEGO brick and its place onto the book stack.

Two very different halves, deliberately kept apart:

  PICK  (real data, replayed):  ``trajectory.h5`` rows ``--start-row`` through
      ``--pick-end-row`` are taken verbatim -- ``joint_positions`` and
      ``gripper_position`` straight off disk, no IK. The user explicitly wants
      the *real recorded grasp* here, not a synthesized approach: this episode
      already contains a real, successful grasp of the brick (gripper closes
      from 0.0 to 0.286 -- "fully closed" for this thin brick -- by row 44,
      confirmed against the raw trace), so there is no reason to re-derive it.

  PLACE (synthesized, IK-driven): a lift/transit/descend/release/retract
      :class:`~fpgm.motion.types.ActionSpec`, realised by
      :class:`~fpgm.motion.trajectory.TrajectoryPlanner`.

The two halves share one clock: pick timestamps are ``k/15`` for the replayed
rows, and the planned portion's own ``t=0`` (== the pick's last replayed
state) is dropped before concatenating, so the merged ``timestamps`` array is
exactly ``k/15`` throughout with zero seam.

--- Why the placement is not perfectly flat (read before changing this) -------
The book surface sits ~0.88 m from the joint-1 axis, right at the edge of the
Panda's ~0.855 m reach -- and IK there only converges with the wrist pitched
40-55 degrees off vertical (measured separately; a near-vertical descent
leaves 36-135 mm of residual there, unusable). Because the brick is rigidly
attached to the gripper at the exact orientation it was grasped at (see
``T_obj_in_tcp`` below), the *required* pitch to reach the surface at all
propagates into the brick's placed orientation -- there is no freedom to
"undo" it without also leaving the reachable set.

This script does not fight that: it searches the validated pitch=[40, 55] deg
family (see the established-fact table in the task this script implements)
for the yaw that lands the brick as close to flat as the reach constraint
allows, uses real brick-mesh vertices (not a hard-coded half-height) to work
out where the *lowest* point of the resulting (tilted) brick lands, and rests
the brick there -- never clipping through the surface. The realized tilt
(tens of degrees) is reported honestly in the sidecar JSON and in this
script's stdout rather than papered over with a flatter but unreachable
target; see the printed verification block for the actual numbers.

Usage:
    PYTHONPATH=src python scripts/plan_pick_place_books.py \\
        --episode AUTOLab+0d4edc83+2023-10-21-19h-07m-04s \\
        --out outputs/pick_place_books.npz
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from dataclasses import dataclass
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.geometry.transforms import invert_se3, matrix_to_rpy, rpy_to_matrix  # noqa: E402
from fpgm.motion.trajectory import PlannerConfig, TrajectoryPlanner, summarize_plan  # noqa: E402
from fpgm.motion.types import ActionSpec, CartesianWaypoint, GripperAction  # noqa: E402
from fpgm.robot.ik import IKConfig, IKSolver  # noqa: E402
from fpgm.robot.kinematics import ArmKinematics  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

_spec = importlib.util.spec_from_file_location(
    "run_object_pipeline", REPO_ROOT / "scripts" / "run_object_pipeline.py"
)
rop = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = rop  # REQUIRED before exec_module, or dataclasses defined in
_spec.loader.exec_module(rop)  # run_object_pipeline fail to pickle/isinstance-check.

logger = get_logger("plan_pick_place_books")

_TRAJ_FPS = 15.0
_DEFAULT_URDF = rop._DEFAULT_URDF

#: Measured fact (see module docstring): a near-vertical descent to the book
#: surface fails IK (36-135 mm residual); pitched 40-55 deg it converges to
#: ~0.05 mm everywhere on the surface. The orientation search below never
#: strays outside this validated range.
_PLACE_PITCH_RANGE_DEG = (40.0, 55.0)
_PLACE_YAW_RANGE_DEG = (-90.0, 90.0)

#: Near-singular guard for the *chosen* placement orientation specifically
#: (independent of IKSolver's own per-iteration NEAR_SINGULAR_THRESHOLD):
#: reject any candidate whose min singular value doesn't clear this, so the
#: final choice has real conditioning margin, not just "barely converged".
_MIN_SINGULAR_VALUE_MARGIN = 0.03


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--episode", required=True, help="DROID episode uuid")
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument(
        "--brick-outputs", type=Path, default=None,
        help="brick object-pipeline stage root (default: "
        "outputs/<episode>/objects_sam3d_multi/20_31)",
    )
    p.add_argument(
        "--books-outputs", type=Path, default=None,
        help="books placement-surface stage root (default: "
        "outputs/<episode>/objects_books/0_11)",
    )
    p.add_argument("--start-row", type=int, default=20, help="first replayed pick row")
    p.add_argument(
        "--grasp-row", type=int, default=44,
        help="row the recorded grasp completes at (gripper first fully closed)",
    )
    p.add_argument(
        "--pick-end-row", type=int, default=46,
        help="last replayed pick row (a couple of rows after grasp completion, "
        "for a settled hand-off into the synthesized transport)",
    )
    p.add_argument("--clearance-m", type=float, default=0.003, help="gap above the fitted surface")
    p.add_argument("--urdf", type=Path, default=_DEFAULT_URDF)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--out", type=Path, required=True)
    return p.parse_args()


# --------------------------------------------------------------------------- #
# Data loading
# --------------------------------------------------------------------------- #


def _load_trajectory(episode: str) -> tuple[np.ndarray, np.ndarray]:
    with h5py.File(rop._load_trajectory_h5(episode), "r") as f:
        joint_positions = np.asarray(f["observation/robot_state/joint_positions"], dtype=np.float64)
        gripper = np.asarray(f["observation/robot_state/gripper_position"], dtype=np.float64)
    return joint_positions, gripper


def _load_brick(brick_outputs: Path) -> tuple[np.ndarray, np.ndarray]:
    """Real brick vertices (mesh-canonical frame) and its rest-pose alignment.

    Loaded through ``run_object_pipeline``'s own ``_load_mesh``/``_load_alignment``
    (see this module's importlib-trick comment above) so the mesh/alignment
    contract is never duplicated here.
    """
    mesh = rop._load_mesh(rop.stage_dir(brick_outputs, "mesh"))
    alignment = rop._load_alignment(rop.stage_dir(brick_outputs, "align"))
    vertices = np.asarray(mesh.vertices, dtype=np.float64)
    transform = np.asarray(alignment.transform, dtype=np.float64)
    return vertices, transform


def _load_placement_target(books_outputs: Path) -> dict:
    path = books_outputs / "objects" / "2_placement" / "placement_target.json"
    return json.loads(path.read_text())


# --------------------------------------------------------------------------- #
# Geometry helpers
# --------------------------------------------------------------------------- #


def _plane_basis(normal: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Orthonormal in-plane axes ``(u_hat, v_hat)`` perpendicular to ``normal``.

    ``u_hat`` is world +X projected onto the plane (world +Y as a fallback if
    +X is nearly parallel to ``normal``, which it never is for this
    near-vertical book-stack normal, but the fallback keeps this general).
    """
    normal = normal / np.linalg.norm(normal)
    seed = np.array([1.0, 0.0, 0.0]) if abs(normal[0]) < 0.9 else np.array([0.0, 1.0, 0.0])
    u_hat = seed - normal * np.dot(seed, normal)
    u_hat = u_hat / np.linalg.norm(u_hat)
    v_hat = np.cross(normal, u_hat)
    return u_hat, v_hat


def _place_orientation_family(yaw_rad: float, pitch_rad: float) -> np.ndarray:
    """The validated reachable family: gripper flipped to point down, then
    pitched/yawed off vertical -- see the module docstring's measured table.
    """
    flip_down = rpy_to_matrix(np.array([np.pi, 0.0, 0.0]))
    post = rpy_to_matrix(np.array([0.0, pitch_rad, yaw_rad]))
    return flip_down @ post


@dataclass
class PlaceOrientationResult:
    rotation: np.ndarray  # (3, 3) base_link -> TCP
    yaw_deg: float
    pitch_deg: float
    tilt_from_flat_deg: float
    position_error_m: float
    min_singular_value: float


def search_place_orientation(
    T_obj_in_tcp: np.ndarray,
    local_up_axis_tcp: np.ndarray,
    normal: np.ndarray,
    ik: IKSolver,
    gripper: float,
    probe_xyz: np.ndarray,
) -> PlaceOrientationResult:
    """Grid-search the validated pitch range for the least-tilted reachable placement.

    For each ``(yaw, pitch)`` in the established-reachable family, this checks
    (a) IK actually converges at ``probe_xyz`` with this orientation and a
    comfortable singular-value margin, and (b) how far the resulting brick
    orientation would sit from "flat" (its rest-pose up-axis aligned with the
    book-surface normal). It returns the reachable candidate with the least
    tilt -- never a candidate outside the measured-safe pitch band.

    ``local_up_axis_tcp`` is the brick's rest-pose "up" direction, already
    expressed in the TCP-local frame (i.e. ``T_obj_in_tcp``'s rotation applied
    to the rest pose's own up axis) -- a fixed unit vector once the grasp has
    happened, since the attachment is rigid.
    """
    best: PlaceOrientationResult | None = None
    n_reachable = 0
    n_total = 0
    for yaw_deg in np.linspace(*_PLACE_YAW_RANGE_DEG, 37):
        # Cold-start (seed_q=None) at the start of every yaw row, deliberately
        # never warm-chained across rows: a genuine IK failure at one (yaw,
        # pitch) -- e.g. an extreme yaw that pushes the TCP target out of
        # reach -- must not poison every subsequent row with a bad warm seed.
        # A warm seed gets zero restarts by default (see IKConfig.n_restarts_
        # warm's docstring), so a bad one never recovers; a fresh cold seed
        # gets IKConfig.n_restarts_cold retries and is the only way later,
        # perfectly reachable (yaw, pitch) pairs get a fair chance.
        q_seed_row = None
        for pitch_deg in np.linspace(*_PLACE_PITCH_RANGE_DEG, 16):
            n_total += 1
            rot = _place_orientation_family(np.radians(yaw_deg), np.radians(pitch_deg))
            target = np.eye(4, dtype=np.float64)
            target[:3, :3] = rot
            target[:3, 3] = probe_xyz - rot @ T_obj_in_tcp[:3, 3]
            result = ik.solve(target, gripper=gripper, seed_q=q_seed_row)
            q_seed_row = result.q
            if not (result.success and result.min_singular_value > _MIN_SINGULAR_VALUE_MARGIN):
                continue
            n_reachable += 1
            up_world = rot @ local_up_axis_tcp
            up_world = up_world / np.linalg.norm(up_world)
            tilt_deg = float(np.degrees(np.arccos(np.clip(np.dot(up_world, normal), -1.0, 1.0))))
            if best is None or tilt_deg < best.tilt_from_flat_deg:
                best = PlaceOrientationResult(
                    rotation=rot, yaw_deg=float(yaw_deg), pitch_deg=float(pitch_deg),
                    tilt_from_flat_deg=tilt_deg, position_error_m=result.position_error_m,
                    min_singular_value=result.min_singular_value,
                )
    logger.info(
        "place-orientation search: %d/%d (yaw, pitch) samples reachable with margin; "
        "best tilt-from-flat=%.2f deg at yaw=%.1f pitch=%.1f (pos_err=%.2e m, min_sv=%.4f)",
        n_reachable, n_total,
        best.tilt_from_flat_deg if best else float("nan"),
        best.yaw_deg if best else float("nan"),
        best.pitch_deg if best else float("nan"),
        best.position_error_m if best else float("nan"),
        best.min_singular_value if best else float("nan"),
    )
    if best is None:
        raise RuntimeError(
            "no orientation in the validated pitch=[40, 55] deg family reaches the "
            "placement point with a safe singular-value margin -- the placement is "
            "genuinely outside the reachable set at this point, not a solver tuning issue"
        )
    return best


@dataclass
class RestingPlacement:
    T_object_final: np.ndarray  # mesh-canonical -> world, scale baked in
    T_tcp_place: np.ndarray  # base_link -> TCP
    lowest_vertex_clearance_m: float  # should equal the requested clearance, by construction
    centroid_height_above_plane_m: float
    max_vertex_height_above_plane_m: float
    footprint_uv_half_extent_m: np.ndarray  # (2,) actual brick footprint half-extent used


def solve_resting_placement(
    vertices: np.ndarray,
    T_obj_in_tcp: np.ndarray,
    R_place: np.ndarray,
    normal: np.ndarray,
    placement_pos: np.ndarray,
    clearance_m: float,
) -> RestingPlacement:
    """Where the (possibly tilted) brick actually rests, from real mesh vertices.

    Generalizes "offset along the normal by half the brick's height" to the
    tilted case this reach constraint forces (see module docstring): rather
    than a single hard-coded half-height, every mesh vertex is rotated into
    its placement orientation and the *lowest* one (along the surface normal)
    is set to sit exactly ``clearance_m`` above the fitted plane -- i.e. the
    brick is lowered until first contact, never clipped through the surface.
    The in-plane (u, v) translation is chosen so the brick's *centroid*
    lands over the fitted placement point, matching how you'd actually lower
    an object onto a target from directly above.
    """
    u_hat, v_hat = _plane_basis(normal)
    r_obj_full = R_place @ T_obj_in_tcp[:3, :3]  # mesh -> world rotation, scale baked in
    verts_rot = vertices @ r_obj_full.T  # untranslated world-oriented vertices

    proj_n = verts_rot @ normal
    proj_min = float(np.min(proj_n))
    centroid_rot = verts_rot.mean(axis=0)

    target_n = float(np.dot(placement_pos, normal)) + clearance_m
    target_u = float(np.dot(placement_pos, u_hat))
    target_v = float(np.dot(placement_pos, v_hat))

    t_obj = (
        (target_u - float(np.dot(centroid_rot, u_hat))) * u_hat
        + (target_v - float(np.dot(centroid_rot, v_hat))) * v_hat
        + (target_n - proj_min) * normal
    )

    T_object_final = np.eye(4, dtype=np.float64)
    T_object_final[:3, :3] = r_obj_full
    T_object_final[:3, 3] = t_obj

    T_tcp_place = T_object_final @ np.linalg.inv(T_obj_in_tcp)

    verts_world = verts_rot + t_obj
    height_above_plane = verts_world @ normal - float(np.dot(placement_pos, normal))
    footprint_u = verts_world @ u_hat - target_u
    footprint_v = verts_world @ v_hat - target_v

    return RestingPlacement(
        T_object_final=T_object_final,
        T_tcp_place=T_tcp_place,
        lowest_vertex_clearance_m=float(np.min(height_above_plane)),
        centroid_height_above_plane_m=float(np.mean(height_above_plane)),
        max_vertex_height_above_plane_m=float(np.max(height_above_plane)),
        footprint_uv_half_extent_m=np.array(
            [max(abs(footprint_u.min()), abs(footprint_u.max())),
             max(abs(footprint_v.min()), abs(footprint_v.max()))]
        ),
    )


def _rest_up_axis_in_tcp_frame(T_obj: np.ndarray, T_obj_in_tcp: np.ndarray) -> np.ndarray:
    """Unit "up" direction of the brick's rest pose, expressed in the TCP-local frame.

    "Up" is whichever of the rest pose's own three (orthonormal, once scale is
    divided out) local axes has the largest world-Z component -- the axis the
    reconstruction implies was vertical while the brick sat on the bench.
    Only used to *score* candidate placement orientations by how close they
    come to flat (see :func:`search_place_orientation`); the actual placement
    geometry in :func:`solve_resting_placement` uses the full vertex cloud,
    not this single-axis approximation.
    """
    linear = T_obj[:3, :3]
    scale = float(np.mean(np.linalg.norm(linear, axis=0)))
    r_rest = linear / scale
    up_idx = int(np.argmax(np.abs(r_rest[2, :])))
    sign = np.sign(r_rest[2, up_idx]) or 1.0
    e_up = np.zeros(3)
    e_up[up_idx] = sign
    r_obj_in_tcp_rot = T_obj_in_tcp[:3, :3] / scale
    v = r_obj_in_tcp_rot @ e_up
    return v / np.linalg.norm(v)


# --------------------------------------------------------------------------- #
# Pick replay
# --------------------------------------------------------------------------- #


@dataclass
class PickReplay:
    timestamps: np.ndarray
    joint_positions: np.ndarray
    gripper: np.ndarray
    grasp_frame_idx: int  # index into the arrays above where the grasp completes


def replay_pick(
    joint_positions_full: np.ndarray,
    gripper_full: np.ndarray,
    start_row: int,
    grasp_row: int,
    pick_end_row: int,
) -> PickReplay:
    """Take ``trajectory.h5`` rows verbatim -- the real recorded grasp, no IK.

    Timestamps are ``k/15`` for ``k = 0 .. (pick_end_row - start_row)``,
    matching :class:`~fpgm.motion.types.JointTrajectory`'s own "grid index,
    never accumulated dt" convention so concatenating with the planned
    portion later introduces no clock drift.
    """
    if not (start_row <= grasp_row <= pick_end_row):
        raise ValueError(
            f"expected start_row <= grasp_row <= pick_end_row, got "
            f"{start_row} <= {grasp_row} <= {pick_end_row}"
        )
    rows = np.arange(start_row, pick_end_row + 1)
    timestamps = (rows - start_row).astype(np.float64) / _TRAJ_FPS
    return PickReplay(
        timestamps=timestamps,
        joint_positions=joint_positions_full[rows].copy(),
        gripper=gripper_full[rows].copy(),
        grasp_frame_idx=int(grasp_row - start_row),
    )


# --------------------------------------------------------------------------- #
# Transport + place action spec
# --------------------------------------------------------------------------- #


def build_transport_spec(
    tcp_pick_end: np.ndarray,
    place: RestingPlacement,
    normal: np.ndarray,
    lift_m: float = 0.15,
    transit_hover_m: float = 0.05,
    hover_above_place_m: float = 0.03,
    retract_m: float = 0.05,
) -> ActionSpec:
    """Lift clear of the bench, fly to the books, descend, release, retract.

    Every waypoint near the books uses ``place.T_tcp_place``'s orientation
    (the pitched, validated-reachable family) -- see the module docstring for
    why a vertical approach is not an option here. Two things were measured
    (by probing IK independently along each candidate straight-line path, not
    guessed) and matter enough to call out:

    * **Reorienting happens as its own zero-distance waypoint, at the lift
      position** (radius ~0.66 m, comfortably reachable), rather than folded
      into the long translation to the books. Combining a large translation
      with a large simultaneous SLERP (pick orientation -> pitched place
      orientation) drove 17/73 samples of that one segment through a region
      IK could not converge in (0.06-135 mm; ``min_singular_value`` bottoming
      at 0.0, a genuine kinematic singularity, not a solver-tuning issue).
      Splitting the rotation out fixed it completely (0 issues) -- rotating
      in place a few tens of degrees near the shoulder does not stress the
      chain the way translating 0.28 m out toward full extension while also
      rotating does.
    * **The transit/retract hover heights are capped at +0.05 m along the
      normal, not further.** The book surface sits at ~0.85 m total reach
      from the joint-1 axis (see module docstring); probing indepedently
      along the normal direction past the placement point showed IK still
      converging at +0.05 m (0/11 samples failed) but degrading by +0.08 m
      (4/11 failed) and worse by +0.10 m (5/11 failed) -- rising along a
      near-vertical normal from a point already at the edge of the reachable
      sphere adds *height*, which consumes exactly the same reach budget as
      horizontal radius. A generous "lift well clear" margin here would
      silently walk the plan back into the unreachable region this whole
      script exists to route around.
    """
    rpy_pick = matrix_to_rpy(tcp_pick_end[:3, :3])
    rpy_place = matrix_to_rpy(place.T_tcp_place[:3, :3])
    xyz_place = place.T_tcp_place[:3, 3]
    lift_xyz = tcp_pick_end[:3, 3] + np.array([0.0, 0.0, lift_m])

    spec = ActionSpec()
    spec.append(CartesianWaypoint(xyz=lift_xyz, rpy=rpy_pick, speed_mps=0.15))
    spec.append(CartesianWaypoint(xyz=lift_xyz, rpy=rpy_place, speed_mps=0.30))
    spec.append(
        CartesianWaypoint(xyz=xyz_place + normal * transit_hover_m, rpy=rpy_place, speed_mps=0.15)
    )
    spec.append(
        CartesianWaypoint(
            xyz=xyz_place + normal * hover_above_place_m, rpy=rpy_place, speed_mps=0.06
        )
    )
    spec.append(CartesianWaypoint(xyz=xyz_place, rpy=rpy_place, speed_mps=0.03))
    spec.append(GripperAction(value=0.0))
    spec.append(
        CartesianWaypoint(xyz=xyz_place + normal * retract_m, rpy=rpy_place, speed_mps=0.08)
    )
    return spec


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #


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

    brick_outputs = args.brick_outputs or (
        REPO_ROOT / "outputs" / args.episode / "objects_sam3d_multi" / "20_31"
    )
    books_outputs = args.books_outputs or (
        REPO_ROOT / "outputs" / args.episode / "objects_books" / "0_11"
    )

    robot = RobotModel(str(args.urdf), load_meshes=False)
    arm = ArmKinematics.from_robot_model(robot)
    ik = IKSolver(arm, IKConfig(seed=args.seed))
    planner = TrajectoryPlanner(arm, ik, robot, PlannerConfig(on_failure="best_effort"))

    joint_positions_full, gripper_full = _load_trajectory(args.episode)
    vertices, T_obj_rest = _load_brick(brick_outputs)
    placement = _load_placement_target(books_outputs)
    placement_pos = np.asarray(placement["position_m"], dtype=np.float64)
    normal = np.asarray(placement["normal"], dtype=np.float64)
    normal = normal / np.linalg.norm(normal)
    surface_extent = np.asarray(placement["surface_extent_m"], dtype=np.float64)

    # --- 1/2: pick replay + attachment capture -----------------------------
    pick = replay_pick(
        joint_positions_full, gripper_full, args.start_row, args.grasp_row, args.pick_end_row
    )
    q_grasp = joint_positions_full[args.grasp_row]
    g_grasp = float(gripper_full[args.grasp_row])
    T_tcp_grasp = arm.fk(q_grasp, g_grasp)
    T_obj_in_tcp = invert_se3(T_tcp_grasp) @ T_obj_rest
    logger.info(
        "grasp captured at row %d: gripper=%.4f, ||T_obj_in_tcp translation||=%.4f m",
        args.grasp_row, g_grasp, float(np.linalg.norm(T_obj_in_tcp[:3, 3])),
    )

    q_pick_end = pick.joint_positions[-1]
    g_pick_end = float(pick.gripper[-1])
    tcp_pick_end = arm.fk(q_pick_end, g_pick_end)

    # --- 3: synthesized transport + place -----------------------------------
    local_up_tcp = _rest_up_axis_in_tcp_frame(T_obj_rest, T_obj_in_tcp)
    probe_xyz = placement_pos + normal * 0.03
    orientation = search_place_orientation(
        T_obj_in_tcp, local_up_tcp, normal, ik, g_pick_end, probe_xyz
    )
    place = solve_resting_placement(
        vertices, T_obj_in_tcp, orientation.rotation, normal, placement_pos, args.clearance_m
    )
    spec = build_transport_spec(tcp_pick_end, place, normal)

    transport_traj = planner.plan(spec, q_pick_end, g_pick_end)
    plan_table = summarize_plan(transport_traj)
    logger.info("transport plan:\n%s", plan_table)
    print(plan_table)

    # --- 4: concatenate onto one continuous 15 Hz clock ---------------------
    timestamps = np.concatenate(
        [pick.timestamps, pick.timestamps[-1] + transport_traj.timestamps[1:]]
    )
    joint_positions = np.concatenate(
        [pick.joint_positions, transport_traj.joint_positions[1:]], axis=0
    )
    gripper = np.concatenate([pick.gripper, transport_traj.gripper[1:]])
    n_pick = pick.timestamps.shape[0]

    gripper_step_idx = max(
        i for i, step in enumerate(spec) if isinstance(step, GripperAction)
    )
    release_samples_in_transport = np.flatnonzero(transport_traj.segment_index == gripper_step_idx)
    release_frame_idx = (n_pick - 1) + int(release_samples_in_transport[-1])
    grasp_frame_idx = pick.grasp_frame_idx

    # --- 5: per-frame object pose --------------------------------------------
    n_frames = timestamps.shape[0]
    object_poses = np.empty((n_frames, 1, 4, 4), dtype=np.float64)
    for t in range(n_frames):
        if t < grasp_frame_idx:
            object_poses[t, 0] = T_obj_rest
        elif t <= release_frame_idx:
            t_tcp = arm.fk(joint_positions[t], float(gripper[t]))
            object_poses[t, 0] = t_tcp @ T_obj_in_tcp
        else:
            object_poses[t, 0] = object_poses[release_frame_idx, 0]

    # --------------------------------------------------------------------- #
    # Verification
    # --------------------------------------------------------------------- #
    print("\n=== verification ===")
    dt_all = np.diff(timestamps)
    dt_ok = np.allclose(dt_all, 1.0 / _TRAJ_FPS, atol=1e-9)
    print(f"timestamps: {n_frames} frames, dt exactly 1/15 everywhere: {dt_ok} "
          f"(max |dt - 1/15| = {np.max(np.abs(dt_all - 1.0/_TRAJ_FPS)):.2e} s)")

    max_ik_residual = float(np.max(transport_traj.position_error_m))
    max_joint_speed_ratio = float(np.max(transport_traj.joint_speed_ratio))
    min_sv = float(np.min(transport_traj.min_singular_value))
    print(f"planned portion: max IK position residual = {max_ik_residual:.3e} m, "
          f"max joint-speed ratio = {max_joint_speed_ratio:.3f}, min singular value = {min_sv:.4f}")
    if transport_traj.issues:
        print(f"{len(transport_traj.issues)} planner issue(s):")
        for issue in transport_traj.issues:
            print(f"  [{issue.kind.value}] segment {issue.segment_index}: {issue.message}")
    else:
        print("no planner issues (every segment converged within tolerance)")
    print(f"transport_traj.success = {transport_traj.success}")

    # final_center_world is the centroid-consistent target used to place (see
    # solve_resting_placement) -- object_poses[-1, 0]'s translation column
    # matches it exactly since the object is frozen at release.
    final_center_world = place.T_object_final[:3, 3]
    u_hat, v_hat = _plane_basis(normal)
    in_plane_off = np.array(
        [
            np.dot(final_center_world - placement_pos, u_hat),
            np.dot(final_center_world - placement_pos, v_hat),
        ]
    )
    print(
        f"brick final placement: tilt-from-flat={orientation.tilt_from_flat_deg:.1f} deg "
        f"(yaw={orientation.yaw_deg:.1f}, pitch={orientation.pitch_deg:.1f}), "
        f"lowest-vertex clearance above plane={place.lowest_vertex_clearance_m*1000:.2f} mm, "
        f"mean vertex height above plane={place.centroid_height_above_plane_m*1000:.1f} mm, "
        f"max vertex height above plane={place.max_vertex_height_above_plane_m*1000:.1f} mm"
    )
    fits_surface = bool(np.all(place.footprint_uv_half_extent_m <= surface_extent / 2))
    print(
        f"brick footprint half-extent used (u, v) = "
        f"{place.footprint_uv_half_extent_m.round(4).tolist()} m; "
        f"fitted surface half-extent = {(surface_extent / 2).round(4).tolist()} m -- "
        f"fits within surface: {fits_surface}"
    )
    print(
        f"in-plane offset of brick centre from fitted placement point (u, v) = "
        f"{in_plane_off.round(4).tolist()} m"
    )

    grasped_check_rows = [
        grasp_frame_idx, (grasp_frame_idx + release_frame_idx) // 2, release_frame_idx,
    ]
    grasped_ok = True
    for t in grasped_check_rows:
        recomputed = arm.fk(joint_positions[t], float(gripper[t])) @ T_obj_in_tcp
        ok = np.allclose(recomputed, object_poses[t, 0], atol=1e-9)
        grasped_ok &= ok
        print(f"grasped-span check frame {t}: recomputed T_tcp @ T_obj_in_tcp matches pose: {ok}")
    print(f"all grasped-span checks passed: {grasped_ok}")

    static_before = np.allclose(object_poses[:grasp_frame_idx, 0], T_obj_rest, atol=1e-12)
    static_after = np.allclose(
        object_poses[release_frame_idx:, 0], object_poses[release_frame_idx, 0], atol=1e-12
    )
    print(f"object stationary before grasp (frames 0..{grasp_frame_idx - 1}): {static_before}")
    print(
        f"object stationary after release (frames {release_frame_idx}..{n_frames - 1}): "
        f"{static_after}"
    )

    # --------------------------------------------------------------------- #
    # Output
    # --------------------------------------------------------------------- #
    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=np.array(["brick"], dtype=str),
    )

    sidecar = {
        "episode": args.episode,
        "camera": args.camera,
        "brick_outputs": str(brick_outputs),
        "books_outputs": str(books_outputs),
        "phases": {
            "pick_replay_rows": [args.start_row, args.pick_end_row],
            "grasp_row": args.grasp_row,
            "grasp_frame_idx": grasp_frame_idx,
            "n_pick_frames": n_pick,
            "release_frame_idx": release_frame_idx,
            "n_total_frames": n_frames,
        },
        "attachment": {
            "T_tcp_grasp": T_tcp_grasp.tolist(),
            "T_obj_in_tcp": T_obj_in_tcp.tolist(),
        },
        "place_orientation_search": {
            "pitch_range_deg": list(_PLACE_PITCH_RANGE_DEG),
            "yaw_range_deg": list(_PLACE_YAW_RANGE_DEG),
            "chosen_yaw_deg": orientation.yaw_deg,
            "chosen_pitch_deg": orientation.pitch_deg,
            "tilt_from_flat_deg": orientation.tilt_from_flat_deg,
            "ik_position_error_m": orientation.position_error_m,
            "min_singular_value": orientation.min_singular_value,
        },
        "placement_pose_used": {
            "T_object_final": place.T_object_final.tolist(),
            "T_tcp_place": place.T_tcp_place.tolist(),
            "lowest_vertex_clearance_m": place.lowest_vertex_clearance_m,
            "centroid_height_above_plane_m": place.centroid_height_above_plane_m,
            "max_vertex_height_above_plane_m": place.max_vertex_height_above_plane_m,
            "footprint_uv_half_extent_m": place.footprint_uv_half_extent_m.tolist(),
            "surface_uv_half_extent_m": (surface_extent / 2).tolist(),
        },
        "verification": {
            "dt_exactly_1_over_15": bool(dt_ok),
            "max_ik_position_residual_m": max_ik_residual,
            "max_joint_speed_ratio": max_joint_speed_ratio,
            "min_singular_value": min_sv,
            "planner_success": bool(transport_traj.success),
            "n_planner_issues": len(transport_traj.issues),
            "grasped_span_checks_passed": bool(grasped_ok),
            "static_before_grasp": bool(static_before),
            "static_after_release": bool(static_after),
        },
        "summarize_plan": plan_table,
    }
    sidecar_path = args.out.with_suffix(".json")
    sidecar_path.write_text(json.dumps(sidecar, indent=2))

    logger.info("wrote %s and %s", args.out, sidecar_path)
    return 0


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