scripts/measure_push_demo.py

#!/usr/bin/env python
"""Measure a real push from the demo: hand velocity in, object displacement out.

This is the calibration half of the counterfactual-push demo. It answers, from
the recorded episode alone: *when the arm pushed the drawer at speed v, how far
did the drawer actually slide?* -- which is the only honest way to fix the
transfer coefficient that `render_counterfactual_push.py` then extrapolates to
push speeds that were never recorded.

Why the displacement is measured in 2D and lifted, rather than read from depth:
PointWorld ships dense depth for exactly one frame per clip (`initial_depth`),
and its clips for this episode stop well before the drawer push happens. So the
drawer's *initial* 3D pose comes from depth (frame 0, where it is available),
and its *motion* comes from SAM 3.1 tracking the drawer's 2D mask across the
push, lifted back to metres by solving for displacement along a single
slide axis:

    u_t = project(X_0 + s_t * axis)        solve each s_t (metres) in 1D

That one-axis constraint is what makes a 2D measurement metric: a drawer is a
prismatic joint, so its motion has exactly one degree of freedom, and the axis
is recovered from the drawer's own frame-0 point cloud (PCA -- a drawer front
is a thin plate, so its thinnest direction is the slide direction).

Usage:
    PYTHONPATH=src python scripts/measure_push_demo.py \\
        --episode <uuid> --camera ext1 --mask-outputs outputs/<uuid>/objects_shelf2/0_11 \\
        --traj-start 90 --traj-end 126 --out outputs/push_calibration.json
"""

from __future__ import annotations

import argparse
import importlib.util
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.data.droid_raw import cam2world_vector_to_world2cam  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.pipeline.frames import ClipFrameSource  # 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
_spec.loader.exec_module(rop)

logger = get_logger("measure_push_demo")

#: DROID's `trajectory.h5` rows are logged at 15 Hz.
_TRAJ_FPS = 15.0
#: Video frames and trajectory rows are 1:1 for this episode (127 video frames
#: vs 128 trajectory rows). PointWorld's *annotation* clock is the half-rate one
#: (its `ClipTiming.annotation_stride == 2`, i.e. clip index 50 -> video frame
#: 100) -- that stride applies to clip keys, never to trajectory rows.
_VIDEO_STRIDE = 1


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--episode", required=True)
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument("--mask-outputs", type=Path, required=True,
                   help="stage root whose 1_mask/ holds the object's frame-0 mask (from mask_object_manual.py)")
    p.add_argument("--clip", default="0:11", help="clip the frame-0 mask/depth came from")
    p.add_argument("--traj-start", type=int, required=True, help="first trajectory row of the push window")
    p.add_argument("--traj-end", type=int, required=True, help="last trajectory row of the push window (inclusive)")
    p.add_argument("--prompt", default="open drawer", help="SAM 3.1 prompt used to re-find the object in the push window")
    p.add_argument("--device", default="cuda")
    p.add_argument("--out", type=Path, required=True)
    return p.parse_args()


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

    from fpgm.config import SegmentationConfig
    from fpgm.objects.align import object_point_cloud
    from fpgm.segmentation.sam3 import Sam3VideoSegmenter

    meta = json.loads((rop.stage_dir(args.mask_outputs, "mask") / "meta.json").read_text())
    with np.load(rop.stage_dir(args.mask_outputs, "mask") / "mask_frame0.npz") as npz:
        mask_full = npz["mask"].astype(bool)
    video_w, video_h = meta["video_resolution"]

    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)

    # --- object's frame-0 3D pose + slide axis, from depth -------------------
    depth = rop._load_initial_depth(args.episode, args.clip, meta["camera_serial"])
    adapter = rop._DepthRgbAdapter(depth, ctx.clip.initial_rgb)
    points_world, _colors = object_point_cloud(adapter, ctx.camera_annot, mask_full, (video_h, video_w))
    X0 = points_world.mean(axis=0)

    logger.info("object frame-0 centroid=%s  (n=%d points)", X0.round(3), points_world.shape[0])

    # --- robot hand kinematics over the push window -------------------------
    robot = RobotModel(str(rop._DEFAULT_URDF), load_meshes=False)
    with h5py.File(rop._load_trajectory_h5(args.episode), "r") as f:
        J = np.asarray(f["observation/robot_state/joint_positions"])
        G = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_6d = np.asarray(f[f"observation/camera_extrinsics/{meta['camera_serial']}_left"])

    rows = np.arange(args.traj_start, args.traj_end + 1)
    hand = np.array([
        (robot.link_poses(J[t], float(G[t]))["left_inner_finger"][:3, 3]
         + robot.link_poses(J[t], float(G[t]))["right_inner_finger"][:3, 3]) / 2.0
        for t in rows
    ])
    hand_vel = np.gradient(hand, 1.0 / _TRAJ_FPS, axis=0)
    hand_speed = np.linalg.norm(hand_vel, axis=1)

    # --- track the object across the push window ----------------------------
    video_start = int(rows[0]) * _VIDEO_STRIDE
    video_end = (int(rows[-1]) + 1) * _VIDEO_STRIDE

    seg_cfg = SegmentationConfig()
    segmenter = Sam3VideoSegmenter(seg_cfg, device=args.device)
    with ClipFrameSource(str(ctx.mp4_path), video_start, video_end, stride=_VIDEO_STRIDE) as frames:
        segmenter.start_session(str(frames.frame_dir))
        try:
            segmenter.add_text_prompt(0, args.prompt)
            found = segmenter.collect_masklets()
            if not found:
                raise SystemExit(f"SAM 3.1 found nothing for {args.prompt!r} in the push window")
            obj_id = segmenter.select_object(found)
            masklet = found[obj_id]
        finally:
            segmenter.close_session()

    # --- lift 2D centroid motion to metres -----------------------------------
    # Solved as a horizontal (dx, dy) displacement with dz fixed to 0, rather
    # than along an assumed slide axis: a drawer runs level, so constraining
    # dz=0 is the physical fact, whereas the in-plane *direction* is exactly
    # what the measurement should recover rather than presuppose. Two unknowns
    # from two observations (u, v) per frame -- exactly determined.
    grid = np.arange(-0.40, 0.40 + 1e-9, 0.005)
    dx, dy = np.meshgrid(grid, grid, indexing="ij")
    offsets = np.stack([dx.ravel(), dy.ravel(), np.zeros(dx.size)], axis=1)

    disp_xy: list[np.ndarray] = []
    valid_rows: list[int] = []
    residuals: list[float] = []
    for i, row in enumerate(rows):
        m = masklet.frames.get(i)
        if m is None or not m.any():
            continue
        ys, xs = np.nonzero(m)
        uv_obs = np.array([xs.mean(), ys.mean()])

        camera = Camera(
            ctx.camera_annot.rescaled(video_w, video_h).K,
            cam2world_vector_to_world2cam(extrinsics_6d[row]),
        )
        uv_cand, depth_cand = camera.project(X0[None, :] + offsets)
        ok = depth_cand > 0
        if not ok.any():
            continue
        err = np.full(offsets.shape[0], np.inf)
        err[ok] = np.linalg.norm(uv_cand[ok] - uv_obs[None, :], axis=1)
        best = int(np.argmin(err))
        disp_xy.append(offsets[best, :2].copy())
        residuals.append(float(err[best]))
        valid_rows.append(int(row))

    if len(disp_xy) < 3:
        raise SystemExit(f"only {len(disp_xy)} trackable frames -- cannot measure a push")

    disp_xy = np.asarray(disp_xy)
    disp_xy = disp_xy - disp_xy[0]  # relative to the window's first tracked frame
    vr = np.asarray(valid_rows)

    # Slide direction = principal direction of the measured horizontal motion.
    net_vec = disp_xy[-1]
    axis2 = net_vec / (np.linalg.norm(net_vec) + 1e-12)
    axis = np.array([axis2[0], axis2[1], 0.0])
    disp = disp_xy @ axis2  # signed travel along the recovered slide direction

    idx_hand_all = np.searchsorted(rows, vr)
    hand_along = (hand - hand[0]) @ axis
    hand_vel_along = hand_vel @ axis

    # --- summarise the demo push --------------------------------------------
    total_travel = float(disp[-1] - disp.min()) if abs(disp[-1]) < abs(disp.min()) else float(disp[-1])
    net_disp = float(disp[-1])
    peak_disp = float(disp[np.argmax(np.abs(disp))])
    # Contact = the hand's fastest motion along the slide axis while the drawer
    # is actually moving; that is the "push speed" the counterfactual scales.
    idx_hand = np.searchsorted(rows, vr)
    push_speed = float(np.max(np.abs(hand_vel_along[idx_hand])))
    drawer_speed = float(np.max(np.abs(np.gradient(disp, 1.0 / _TRAJ_FPS))))

    result = {
        "episode": args.episode,
        "prompt": args.prompt,
        "traj_window": [int(args.traj_start), int(args.traj_end)],
        "object_centroid_m": X0.tolist(),
        "slide_axis": axis.tolist(),
        "n_tracked_frames": int(len(disp)),
        "tracked_rows": vr.tolist(),
        "displacement_series_m": disp.tolist(),
        "displacement_xy_m": disp_xy.tolist(),
        "reprojection_residual_px": residuals,
        "hand_along_axis_m": hand_along[idx_hand].tolist(),
        "hand_speed_max_mps": float(np.max(hand_speed)),
        "push_speed_along_axis_mps": push_speed,
        "drawer_peak_speed_mps": drawer_speed,
        "drawer_net_displacement_m": net_disp,
        "drawer_peak_displacement_m": peak_disp,
        "total_travel_m": total_travel,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))
    return 0


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