scripts/render_counterfactual_push.py

#!/usr/bin/env python
"""Counterfactual push: what the drawer does at push speeds that were never recorded.

Calibrated on the one push this episode actually contains, measured by
`scripts/measure_push_tracked.py`. Renders robot + brick + drawer over a frozen
initial frame, one video per case.

--- What the demo measurement actually showed (this decides the model) -------
TAPNext++ tracks over the recorded push give, across the contact window
(trajectory rows 89-108), with a 1.4-4.1 px reprojection residual on a shared
rigid-slide fit:

    drawer travel   0.190 m
    hand travel     0.168 m       -> ratio 1.13
    drawer stops    row 108
    hand stops      row 108       -> same frame

The drawer tracks the hand roughly 1:1 and stops the instant the hand does. It
never coasts. That is **sustained pushing** (the arm stays in contact and drags
a friction-dominated prismatic joint), not the impulsive strike that a
momentum-transfer model describes. The two regimes make opposite predictions
about speed, so the distinction is not academic:

  SUSTAINED (measured here)   d = kappa * hand_travel,  kappa ~ 1.13
      Displacement is set by how far the arm travels, NOT how fast. Pushing
      twice as fast covers the same distance in half the time. The mass ratio
      does not enter -- a position-controlled arm simply overpowers friction.

  STRIKE-AND-RELEASE (not in this episode)   d = beta * (alpha*v)^2
      Only here does displacement scale with speed, quadratically, via the
      momentum transfer alpha = (1+e)*m_A/(m_A+m_B) and a friction slide
      beta = 1/(2*mu*g).

`--model sustained` renders the calibrated, measured regime. `--model strike`
renders the momentum-transfer regime for comparison -- but its `mu` is an
ASSUMPTION, not a measurement, because this drawer never coasted freely and so
never revealed its friction. Every strike-mode frame is labelled accordingly;
do not read those distances as predictions this data supports.

Usage:
    PYTHONPATH=src python scripts/render_counterfactual_push.py \\
        --episode <uuid> --calibration outputs/push_calibration_tracked.json \\
        --drawer-outputs outputs/<uuid>/objects_shelf2/0_11 \\
        --brick-outputs  outputs/<uuid>/objects_sam3d_multi/20_31 \\
        --model sustained --speed-factors 0.5 1.0 2.0 --out-dir outputs/counterfactual
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from pathlib import Path

import h5py
import numpy as np
import trimesh
from PIL import Image

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.overlay import composite  # noqa: E402
from fpgm.robot.render import RobotRenderer  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402
from fpgm.viz.overlays import VideoWriter, draw_hud  # noqa: E402

import pyrender  # 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("counterfactual_push")

_TRAJ_FPS = 15.0
_G = 9.81
#: Strike-mode only, and ASSUMED (see the module docstring): a drawer on rails
#: with light binding. Never fitted -- this episode has no free-coast phase.
_ASSUMED_MU = 0.35
#: Strike-mode only, likewise assumed, used to report an implied mass ratio.
_ASSUMED_E = 0.4


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("--calibration", type=Path, required=True)
    p.add_argument("--drawer-outputs", type=Path, required=True)
    p.add_argument("--brick-outputs", type=Path, required=True)
    p.add_argument("--drawer-clip", default="0:11")
    p.add_argument("--model", choices=["sustained", "strike"], default="sustained")
    p.add_argument("--speed-factors", type=float, nargs="+", default=[0.5, 1.0, 2.0])
    p.add_argument("--contact-rows", type=int, nargs=2, default=[89, 108])
    p.add_argument("--n-frames", type=int, default=54)
    p.add_argument("--urdf", type=Path, default=None)
    p.add_argument("--out-dir", type=Path, required=True)
    return p.parse_args()


class _MultiObjectRenderer:
    """RobotRenderer plus N independently posed object meshes in one scene.

    `fpgm.objects.scene.SceneRenderer` covers exactly one object; this needs
    two (brick and drawer). Same integration pattern -- object nodes added to
    RobotRenderer's persistent scene, only ever re-posed -- minus the per-object
    mask readback, which costs an extra render pass and is unused here.
    """

    def __init__(self, link_meshes, object_meshes, width, height, **kw):
        self._robot = RobotRenderer(link_meshes, width, height, **kw)
        self._nodes = []
        for tri in object_meshes:
            node = pyrender.Node(mesh=pyrender.Mesh.from_trimesh(tri, smooth=False), matrix=np.eye(4))
            self._robot.scene.add_node(node)
            self._nodes.append(node)

    def render(self, link_poses, object_poses, camera):
        for node, pose in zip(self._nodes, object_poses, strict=True):
            self._robot.scene.set_pose(node, np.asarray(pose, dtype=np.float64))
        return self._robot.render(link_poses, camera)

    def close(self):
        self._robot.close()


def _load_object(stage_root: Path):
    mesh = rop._load_mesh(rop.stage_dir(stage_root, "mesh"))
    alignment = rop._load_alignment(rop.stage_dir(stage_root, "align"))
    # process=False: welding vertices here would desync a textured mesh's
    # per-vertex `uv` from the vertex it belongs to (see
    # scripts/run_object_pipeline.py's _save_mesh/_load_mesh docstrings).
    if mesh.uv is not None and mesh.texture is not None:
        visual = trimesh.visual.TextureVisuals(
            uv=np.asarray(mesh.uv, dtype=np.float32),
            image=Image.fromarray(np.asarray(mesh.texture, dtype=np.uint8)),
        )
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            visual=visual,
            process=False,
        )
        # No vertex_colors assignment here: it would replace `visual` with a
        # fresh ColorVisuals and silently drop the texture (trimesh visuals
        # are one-kind-at-a-time) -- the exact bug this whole fix removes.
    else:
        tri = trimesh.Trimesh(
            vertices=np.asarray(mesh.vertices, dtype=np.float64),
            faces=np.asarray(mesh.faces, dtype=np.int64),
            process=False,
        )
        if mesh.vertex_colors is not None:
            tri.visual.vertex_colors = mesh.vertex_colors
    return tri, alignment.transform


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

    cal = json.loads(args.calibration.read_text())
    disp = np.array(
        [[np.nan, np.nan] if v is None or v[0] is None else v for v in cal["displacement_xy_m"]],
        dtype=float,
    )
    hand_xyz = np.asarray(cal["hand_xyz_m"], dtype=float)
    r0, r1 = args.contact_rows

    axis2 = disp[r1] / (np.linalg.norm(disp[r1]) + 1e-12)
    axis = np.array([axis2[0], axis2[1], 0.0])
    travel = np.linalg.norm(np.nan_to_num(disp), axis=1)
    hand_along = hand_xyz @ axis

    d_drawer_demo = float(travel[r1] - travel[r0])
    d_hand_demo = float(hand_along[r1] - hand_along[r0])
    kappa = d_drawer_demo / d_hand_demo
    v_hand_demo = float(np.max(np.gradient(hand_along, 1.0 / _TRAJ_FPS)[r0:r1 + 1]))
    v_drawer_demo = float(np.max(np.gradient(travel, 1.0 / _TRAJ_FPS)[r0:r1 + 1]))
    alpha = v_drawer_demo / v_hand_demo
    beta = 1.0 / (2.0 * _ASSUMED_MU * _G)
    denom = (1.0 + _ASSUMED_E) - alpha
    mass_ratio = alpha / denom if denom > 1e-6 else float("inf")

    logger.info(
        "demo: drawer %.3f m / hand %.3f m -> kappa=%.3f ; v_hand=%.3f v_drawer=%.3f alpha=%.3f",
        d_drawer_demo, d_hand_demo, kappa, v_hand_demo, v_drawer_demo, alpha,
    )

    urdf = args.urdf or rop._DEFAULT_URDF
    robot = RobotModel(str(urdf), load_meshes=True)
    drawer_tri, drawer_T0 = _load_object(args.drawer_outputs)
    brick_tri, brick_T0 = _load_object(args.brick_outputs)

    meta = json.loads((rop.stage_dir(args.drawer_outputs, "mask") / "meta.json").read_text())
    video_w, video_h = meta["video_resolution"]
    ctx = rop._load_clip_context(args.episode, args.camera, args.drawer_clip)
    camera_K = ctx.camera_annot.rescaled(video_w, video_h).K

    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"])

    start_row = max(0, r0 - 8)
    fixed_camera = Camera(camera_K, cam2world_vector_to_world2cam(extrinsics_6d[start_row]))
    with ClipFrameSource(str(ctx.mp4_path), start_row, start_row + 1, stride=1) as frames:
        background = frames.read(0).copy()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    renderer = _MultiObjectRenderer(robot.visual_meshes(), [drawer_tri, brick_tri], video_w, video_h)
    summary = []
    try:
        for k in args.speed_factors:
            v_push = k * v_hand_demo
            if args.model == "sustained":
                # Arm replays the same recorded path (same travel), clock scaled
                # by k. Drawer follows it, scaled by the measured coupling.
                d_pred = kappa * d_hand_demo
                t_total = (r1 - start_row) / _TRAJ_FPS / k
                v_obj = k * v_drawer_demo
                note = f"MEASURED regime: d = kappa*hand_travel, kappa={kappa:.2f}"
                law = "distance set by ARM TRAVEL, not speed -> same d, {:.2f}x duration".format(1 / k)
            else:
                v_obj = alpha * v_push
                d_pred = beta * v_obj ** 2
                decel = _ASSUMED_MU * _G
                t_total = (r0 - start_row) / _TRAJ_FPS / k + (v_obj / decel if decel > 0 else 0.0)
                note = f"HYPOTHETICAL strike: mu={_ASSUMED_MU} ASSUMED (never observed coasting)"
                law = f"d = beta*(alpha*v)^2 -> {k:g}x speed gives {k*k:g}x distance"

            times = np.linspace(0.0, t_total, args.n_frames)
            out_path = args.out_dir / f"{args.model}_x{k:g}.mp4"
            with VideoWriter(out_path, fps=30.0) as writer:
                for t in times:
                    row_f = start_row + k * t * _TRAJ_FPS
                    row = int(np.clip(round(row_f), start_row, r1))
                    link_poses = robot.link_poses(J[row], float(G[row]))

                    if args.model == "sustained":
                        # Drawer rides the arm: same fraction of travel the
                        # recorded hand had reached at this row.
                        frac = np.clip((hand_along[row] - hand_along[r0]) / d_hand_demo, 0.0, 1.0)
                        s = kappa * d_hand_demo * frac
                        v_now = k * float(np.gradient(travel, 1.0 / _TRAJ_FPS)[row])
                    else:
                        t_contact = (r0 - start_row) / _TRAJ_FPS / k
                        if t <= t_contact:
                            s, v_now = 0.0, 0.0
                        else:
                            dt = t - t_contact
                            decel = _ASSUMED_MU * _G
                            v_now = max(0.0, v_obj - decel * dt)
                            s = min(d_pred, v_obj * dt - 0.5 * decel * dt * dt)

                    drawer_pose = drawer_T0.copy()
                    drawer_pose[:3, 3] = drawer_pose[:3, 3] + s * axis
                    result = renderer.render(link_poses, [drawer_pose, brick_T0], fixed_camera)
                    frame = composite(background, result, alpha=0.85)
                    frame = draw_hud(frame, [
                        f"[{args.model.upper()}]  push speed x{k:g}   v_hand = {v_push:.3f} m/s",
                        f"drawer speed = {v_obj:.3f} m/s  (now {v_now:.3f})",
                        f"drawer travel = {d_pred*100:.1f} cm  (now {s*100:.1f} cm)",
                        note,
                        law,
                    ])
                    writer.write(frame)

            summary.append({
                "model": args.model, "speed_factor": k, "v_hand_mps": v_push,
                "v_drawer_mps": v_obj, "predicted_displacement_m": d_pred,
                "duration_s": float(t_total), "video": str(out_path),
            })
            logger.info("%s x%g: v_hand=%.3f -> d=%.3f m over %.2f s",
                        args.model, k, v_push, d_pred, t_total)
    finally:
        renderer.close()

    payload = {
        "measured_demo": {
            "contact_rows": [r0, r1],
            "drawer_travel_m": d_drawer_demo,
            "hand_travel_m": d_hand_demo,
            "coupling_kappa": kappa,
            "v_hand_peak_mps": v_hand_demo,
            "v_drawer_peak_mps": v_drawer_demo,
            "alpha_speed_ratio": alpha,
            "regime": "sustained contact (drawer stops when hand stops; no free coast)",
        },
        "strike_mode_assumptions": {
            "mu_assumed": _ASSUMED_MU, "e_assumed": _ASSUMED_E,
            "implied_mass_ratio_mA_over_mB": mass_ratio,
            "warning": "not identifiable from this episode -- no free-coast phase was observed",
        },
        "cases": summary,
    }
    (args.out_dir / f"summary_{args.model}.json").write_text(json.dumps(payload, indent=2))
    print(json.dumps(payload, indent=2))
    return 0


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