scripts/settle_after_release.py

#!/usr/bin/env python
"""Physically settle the brick after release, instead of leaving it hanging.

--- Why this exists ---------------------------------------------------------
The rest of this pipeline is deliberately *kinematic*: no mass, no friction,
no inertia -- a grasp is a rigid attachment and a release just... stops
updating the object's pose (see scripts/plan_pick_place_books.py:
``object_poses[t, 0] = object_poses[release_frame_idx, 0]`` for every frame
after release). That is fine everywhere the arm is in control of the object,
but it produces a visible defect at the one moment it is not: the book
surface sits at the very edge of the arm's reach and IK there only converges
with the wrist pitched 40-55 deg off vertical (see that script's docstring),
so the rigidly-attached brick is released tilted 37.26 deg from flat, and
then -- because "released" currently just means "keep the last pose forever"
-- it hangs there, its far edge floating ~5 cm above the books, forever.

Physically the brick would tip and fall flat. This script is the one place
genuine rigid-body physics is warranted, and it is confined to exactly the
post-release window: everything before ``release_frame_idx`` is left
untouched (bit-identical), and only the released trajectory is replaced by
a simulated fall-and-settle, resampled back onto the plan's own 15 Hz clock.

--- Why MuJoCo lives in its own conda env, not `fpgm` -----------------------
`pyproject.toml` pins ``numpy<2`` -- a hard constraint inherited from SAM
3.1's own dependencies, protecting the working SAM3/TAPNext/pyrender stack
already installed in the ``fpgm`` env. A plain ``pip install mujoco`` pulls
numpy 2.x and risks a downgrade cascade there. So, exactly like
``scripts/trellis_generate.py`` (TRELLIS.2, its own torch) and
``scripts/sam3d_generate.py`` (SAM 3D Objects, its own torch) before it: a
small, dependency-light worker (``scripts/_mujoco_settle_worker.py``, mujoco
+ numpy only, no fpgm import -- that env doesn't have fpgm installed) runs in
an isolated ``mujoco`` env, invoked here as a subprocess and handed JSON
files rather than shared in-process objects. This script also never imports
``mujoco`` itself and the worker never imports ``mujoco.viewer`` -- this host
has no hardware OpenGL (software OSMesa only); only MuJoCo's physics
stepping is used anywhere in this pipeline, never its renderer.

--- Why mass is not a parameter to agonise over ------------------------------
A rigid body's free-fall-and-settle trajectory under gravity, contact, and
friction is mass-INDEPENDENT: Newton's second law (F = ma) has gravity
contributing a force proportional to m, contact normal/friction forces that
(for a rigid, non-deformable contact model) also scale with the normal load
which itself scales with m, and there is no other force in this scene (no
motors, no air drag) that does not scale with m -- so m cancels out of the
equations of motion entirely. This is checked empirically at the bottom of
this script (``--check-mass-independence``, on by default): the whole
simulation is re-run at 10x the mass and the resting pose is confirmed
essentially unchanged. What DOES matter, and is NOT free, is friction and
the contact solver's softness (which stands in for restitution -- MuJoCo has
no literal coefficient-of-restitution attribute; bounciness is purely a
function of how underdamped ``solref`` is). Those are recorded plainly in
the output sidecar as ASSUMPTIONS, never as measurements.

Usage:
    PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python \\
        scripts/settle_after_release.py \\
        --npz outputs/pick_place_books/pick_place_books.npz \\
        --out outputs/pick_place_books/pick_place_books_settled.npz \\
        --out-json outputs/pick_place_books/pick_place_books_settled.json
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import numpy as np
import trimesh

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

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

# Same importlib trick used throughout this repo (see e.g.
# scripts/render_counterfactual_push.py's docstring): `sys.modules[spec.name]`
# MUST be set before `exec_module`, or `run_object_pipeline.py`'s own
# `@dataclass`-decorated classes fail isinstance checks / pickling.
_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("settle_after_release")

_TRAJ_FPS = 15.0
_DEFAULT_MUJOCO_PYTHON = Path("/home/quang/miniconda3/envs/mujoco/bin/python")
_WORKER_SCRIPT = REPO_ROOT / "scripts" / "_mujoco_settle_worker.py"

#: Friction is a genuine ASSUMPTION (never measured for this brick/book
#: contact) -- (sliding, torsional, rolling), MuJoCo's own convention. 0.5
#: sliding friction is a generic "dry, moderately rough" value (comparable to
#: wood-on-wood or cardboard-on-paper); the torsional/rolling values are
#: MuJoCo's own defaults, left alone because this is a corner-tipping brick,
#: not a rolling object, so they barely matter here.
_ASSUMED_FRICTION = (0.5, 0.005, 0.0001)
#: Restitution is likewise an ASSUMPTION: MuJoCo has no literal
#: coefficient-of-restitution attribute, bounciness is entirely a function of
#: how underdamped the contact constraint is. solref=(0.02, 1) is MuJoCo's
#: own default -- a critically damped contact, i.e. essentially NO bounce,
#: which is the physically reasonable assumption for a small, light brick
#: dropped from ~8 mm (not a superball).
_ASSUMED_SOLREF = (0.02, 1.0)
_ASSUMED_SOLIMP = (0.9, 0.95, 0.001, 0.5, 2.0)


# --------------------------------------------------------------------------- #
# Standalone quaternion <-> matrix helpers.
#
# Deliberately NOT shared with scripts/_mujoco_settle_worker.py (which uses
# mujoco.mju_mat2Quat/mju_quat2Mat directly): that worker runs in a separate
# conda env and this side of the fence has no mujoco import to reuse, so the
# handful of lines below are duplicated rather than adding a cross-env
# dependency for them. fpgm.geometry.transforms deliberately has no quaternion
# helpers (see its docstring: axis-angle only, to pin down one convention) so
# this is new, not a reinvention of an existing fpgm util.
# --------------------------------------------------------------------------- #


def _mat_to_quat_wxyz(rot: np.ndarray) -> np.ndarray:
    """(3, 3) proper rotation matrix -> (4,) unit quaternion (w, x, y, z).

    Standard branch-on-largest-diagonal-term (Shepperd's method) for numerical
    stability near all rotation angles, including near pi.
    """
    m = np.asarray(rot, dtype=np.float64)
    tr = m[0, 0] + m[1, 1] + m[2, 2]
    if tr > 0:
        s = np.sqrt(tr + 1.0) * 2.0
        w = 0.25 * s
        x = (m[2, 1] - m[1, 2]) / s
        y = (m[0, 2] - m[2, 0]) / s
        z = (m[1, 0] - m[0, 1]) / s
    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
        s = np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2]) * 2.0
        w = (m[2, 1] - m[1, 2]) / s
        x = 0.25 * s
        y = (m[0, 1] + m[1, 0]) / s
        z = (m[0, 2] + m[2, 0]) / s
    elif m[1, 1] > m[2, 2]:
        s = np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2]) * 2.0
        w = (m[0, 2] - m[2, 0]) / s
        x = (m[0, 1] + m[1, 0]) / s
        y = 0.25 * s
        z = (m[1, 2] + m[2, 1]) / s
    else:
        s = np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1]) * 2.0
        w = (m[1, 0] - m[0, 1]) / s
        x = (m[0, 2] + m[2, 0]) / s
        y = (m[1, 2] + m[2, 1]) / s
        z = 0.25 * s
    q = np.array([w, x, y, z], dtype=np.float64)
    return q / np.linalg.norm(q)


def _quat_wxyz_to_mat(quat: np.ndarray) -> np.ndarray:
    """(4,) unit quaternion (w, x, y, z) -> (3, 3) proper rotation matrix."""
    w, x, y, z = np.asarray(quat, dtype=np.float64)
    n = w * w + x * x + y * y + z * z
    s = 2.0 / n if n > 0 else 0.0
    wx, wy, wz = s * w * x, s * w * y, s * w * z
    xx, xy, xz = s * x * x, s * x * y, s * x * z
    yy, yz, zz = s * y * y, s * y * z, s * z * z
    return np.array(
        [
            [1.0 - (yy + zz), xy - wz, xz + wy],
            [xy + wz, 1.0 - (xx + zz), yz - wx],
            [xz - wy, yz + wx, 1.0 - (xx + yy)],
        ]
    )


def _slerp(q0: np.ndarray, q1: np.ndarray, t: float) -> np.ndarray:
    """Spherical linear interpolation between two wxyz unit quaternions."""
    q0 = q0 / np.linalg.norm(q0)
    q1 = q1 / np.linalg.norm(q1)
    dot = float(np.dot(q0, q1))
    if dot < 0.0:
        q1 = -q1
        dot = -dot
    dot = min(dot, 1.0)
    if dot > 0.9995:
        out = q0 + t * (q1 - q0)
        return out / np.linalg.norm(out)
    theta0 = np.arccos(dot)
    theta = theta0 * t
    q2 = q1 - q0 * dot
    q2 = q2 / np.linalg.norm(q2)
    return q0 * np.cos(theta) + q2 * np.sin(theta)


def _plane_basis(normal: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Orthonormal in-plane axes, identical convention to
    ``scripts/plan_pick_place_books.py``'s ``_plane_basis`` (world +X
    projected onto the plane, falling back to +Y), so "in-plane offset"
    numbers reported by this script are directly comparable to that script's.
    """
    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 _rest_up_axis_canonical(t_obj_rest: np.ndarray) -> np.ndarray:
    """Which mesh-canonical axis was "up" in the object's original rest pose.

    Same idea as ``plan_pick_place_books.py``'s
    ``_rest_up_axis_in_tcp_frame`` (whichever local axis has the largest
    world-Z component in the as-scanned rest pose, sign-corrected) but
    expressed directly in world/canonical terms rather than TCP-local, since
    this script has no TCP in the loop. Used purely as the fixed reference
    direction for "tilt from flat" so it is comparable, frame-for-frame, to
    the ``tilt_from_flat_deg=37.26`` figure already in the plan's sidecar
    (verified to reproduce that number to 0.001 deg before this script was
    written).
    """
    r = t_obj_rest[:3, :3]
    scale = float(np.mean(np.linalg.norm(r, axis=0)))
    r_pure = r / scale
    up_idx = int(np.argmax(np.abs(r_pure[2, :])))
    sign = np.sign(r_pure[2, up_idx]) or 1.0
    e_up = np.zeros(3, dtype=np.float64)
    e_up[up_idx] = sign
    return e_up


def tilt_from_flat_deg(r_pure: np.ndarray, e_up: np.ndarray, normal: np.ndarray) -> float:
    up_world = r_pure @ e_up
    up_world = up_world / np.linalg.norm(up_world)
    return float(np.degrees(np.arccos(np.clip(np.dot(up_world, normal), -1.0, 1.0))))


# --------------------------------------------------------------------------- #
# Velocity inheritance: finite-difference the last few attached frames.
# --------------------------------------------------------------------------- #


def estimate_release_velocity(
    object_poses: np.ndarray, release_frame_idx: int, n_frames: int, scale: float
) -> tuple[np.ndarray, np.ndarray]:
    """Linear + angular velocity of the brick at the moment of release.

    "The brick inherits the gripper's velocity at release" -- estimated by
    finite-differencing the brick's own recorded world poses over the
    ``n_frames`` samples immediately before (and including) the release
    frame, at the plan's 15 Hz clock. Linear velocity: a least-squares slope
    (more robust to per-frame FK noise than a single last-step difference).
    Angular velocity: the world-frame rotation-vector increment between each
    consecutive pair in the window, divided by dt, averaged -- returned in
    WORLD frame (the caller converts to MuJoCo's body-local qvel convention).
    """
    lo = max(0, release_frame_idx - (n_frames - 1))
    idx = np.arange(lo, release_frame_idx + 1)
    if idx.size < 2:
        raise ValueError(
            f"need >=2 frames to estimate release velocity, got {idx.size} "
            f"(release_frame_idx={release_frame_idx}, n_frames={n_frames})"
        )
    dt = 1.0 / _TRAJ_FPS
    t = (idx - idx[0]).astype(np.float64) * dt
    pos = object_poses[idx, :3, 3]
    lin_vel = np.array([np.polyfit(t, pos[:, k], 1)[0] for k in range(3)], dtype=np.float64)

    r_pure = object_poses[idx, :3, :3] / scale
    omega_samples = []
    for i in range(len(idx) - 1):
        r_rel = r_pure[i + 1] @ r_pure[i].T
        omega_samples.append(matrix_to_rotvec(r_rel) / dt)
    ang_vel_world = np.mean(np.asarray(omega_samples), axis=0) if omega_samples else np.zeros(3)
    return lin_vel, ang_vel_world


# --------------------------------------------------------------------------- #
# MuJoCo subprocess plumbing
# --------------------------------------------------------------------------- #


def _run_worker(
    mujoco_python: Path,
    hull_stl: Path,
    scale: float,
    release_pos: np.ndarray,
    release_quat: np.ndarray,
    lin_vel_world: np.ndarray,
    ang_vel_body: np.ndarray,
    mass_kg: float,
    friction: tuple[float, float, float],
    solref: tuple[float, float],
    solimp: tuple[float, float, float, float, float],
    surface_pos: np.ndarray,
    surface_quat: np.ndarray,
    surface_half_extent: np.ndarray,
    floor_z: float,
    sim_dt: float,
    max_sim_time: float,
    record_dt: float,
    settle_lin_vel_thresh: float,
    settle_ang_vel_thresh: float,
    settle_consecutive_steps: int,
    scratch_dir: Path,
    tag: str,
) -> dict:
    cfg = {
        "hull_stl_path": str(hull_stl),
        "scale": float(scale),
        "release_pos": release_pos.tolist(),
        "release_quat_wxyz": release_quat.tolist(),
        "lin_vel_world": lin_vel_world.tolist(),
        "ang_vel_body": ang_vel_body.tolist(),
        "mass_kg": float(mass_kg),
        "friction": list(friction),
        "solref": list(solref),
        "solimp": list(solimp),
        "surface_pos": surface_pos.tolist(),
        "surface_quat_wxyz": surface_quat.tolist(),
        "surface_half_extent": surface_half_extent.tolist(),
        "floor_z": float(floor_z),
        "sim_dt": float(sim_dt),
        "max_sim_time": float(max_sim_time),
        "record_dt": float(record_dt),
        "settle_lin_vel_thresh": float(settle_lin_vel_thresh),
        "settle_ang_vel_thresh": float(settle_ang_vel_thresh),
        "settle_consecutive_steps": int(settle_consecutive_steps),
    }
    in_json = scratch_dir / f"worker_in_{tag}.json"
    out_json = scratch_dir / f"worker_out_{tag}.json"
    in_json.write_text(json.dumps(cfg))

    cmd = [
        str(mujoco_python), str(_WORKER_SCRIPT),
        "--in-json", str(in_json), "--out-json", str(out_json),
    ]
    logger.info("invoking mujoco env (%s): %s", tag, " ".join(cmd))
    t0 = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True)
    logger.info(
        "mujoco subprocess (%s) finished in %.2fs, returncode=%d",
        tag, time.time() - t0, proc.returncode,
    )
    if proc.stdout:
        logger.info("[mujoco stdout %s] %s", tag, proc.stdout.strip())
    if proc.returncode != 0:
        logger.error("[mujoco stderr %s] %s", tag, proc.stderr.strip())
        raise RuntimeError(f"mujoco settle worker ({tag}) failed with code {proc.returncode}")
    return json.loads(out_json.read_text())


def _resample_trajectory(worker_out: dict, t_query: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Resample a worker's (times, pos, quat) onto arbitrary query times.

    Linear interpolation for position, slerp for orientation; times beyond
    the last recorded sample hold the final (by then, settled -- velocity is
    below threshold by construction) pose rather than extrapolating.
    """
    times = np.asarray(worker_out["times"], dtype=np.float64)
    pos = np.asarray(worker_out["pos"], dtype=np.float64)
    quat = np.asarray(worker_out["quat_wxyz"], dtype=np.float64)

    out_pos = np.empty((t_query.shape[0], 3), dtype=np.float64)
    out_quat = np.empty((t_query.shape[0], 4), dtype=np.float64)
    for i, tq in enumerate(t_query):
        if tq <= times[0]:
            out_pos[i] = pos[0]
            out_quat[i] = quat[0]
            continue
        if tq >= times[-1]:
            out_pos[i] = pos[-1]
            out_quat[i] = quat[-1]
            continue
        j = int(np.searchsorted(times, tq, side="right") - 1)
        j = min(max(j, 0), times.shape[0] - 2)
        frac = (tq - times[j]) / (times[j + 1] - times[j])
        out_pos[i] = pos[j] + frac * (pos[j + 1] - pos[j])
        out_quat[i] = _slerp(quat[j], quat[j + 1], float(frac))
    return out_pos, out_quat


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


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 (see plan_pick_place_books.py)",
    )
    p.add_argument(
        "--sidecar-json", type=Path, default=None,
        help="plan sidecar JSON (default: <npz stem>.json)",
    )
    p.add_argument(
        "--brick-outputs", type=Path, default=None, help="brick object-pipeline stage root"
    )
    p.add_argument(
        "--books-outputs", type=Path, default=None, help="books placement-surface stage root"
    )
    p.add_argument(
        "--placement-json", type=Path, default=None, help="override placement_target.json path"
    )
    p.add_argument(
        "--release-frame-idx", type=int, default=None, help="override phases.release_frame_idx"
    )
    p.add_argument(
        "--rest-frame-idx", type=int, default=0, help="frame index used as the as-scanned rest pose"
    )
    p.add_argument(
        "--n-vel-frames", type=int, default=5,
        help="window size for release-velocity finite-difference",
    )

    p.add_argument(
        "--mass-kg", type=float, default=0.03,
        help="ARBITRARY -- settle is mass-independent, see docstring",
    )
    p.add_argument("--mass-kg-check", type=float, default=None, help="default: 10x --mass-kg")
    p.add_argument(
        "--friction", type=float, nargs=3, default=list(_ASSUMED_FRICTION),
        metavar=("SLIDE", "TORSION", "ROLL"),
    )
    p.add_argument("--solref", type=float, nargs=2, default=list(_ASSUMED_SOLREF))
    p.add_argument("--solimp", type=float, nargs=5, default=list(_ASSUMED_SOLIMP))
    p.add_argument("--surface-thickness-m", type=float, default=0.02)
    p.add_argument(
        "--floor-drop-m", type=float, default=0.5, help="safety-net floor depth below the surface"
    )

    p.add_argument("--sim-dt", type=float, default=0.0005)
    p.add_argument("--record-dt", type=float, default=0.005)
    p.add_argument("--max-sim-time", type=float, default=3.0)
    p.add_argument("--settle-lin-vel-thresh", type=float, default=0.01, help="m/s")
    p.add_argument("--settle-ang-vel-thresh", type=float, default=0.05, help="rad/s")
    p.add_argument(
        "--settle-consecutive-s", type=float, default=0.05,
        help="seconds below threshold to call it settled",
    )
    p.add_argument(
        "--post-settle-buffer-s", type=float, default=0.5,
        help="extra held-still seconds appended for visibility",
    )

    p.add_argument("--check-mass-independence", action="store_true", default=True)
    p.add_argument(
        "--no-check-mass-independence", dest="check_mass_independence", action="store_false"
    )

    p.add_argument("--mujoco-python", type=Path, default=_DEFAULT_MUJOCO_PYTHON)
    p.add_argument("--scratch-dir", type=Path, default=None)

    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--out-json", type=Path, required=True)
    return p.parse_args()


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

    plan = np.load(args.npz)
    timestamps = np.asarray(plan["timestamps"], dtype=np.float64)
    joint_positions = np.asarray(plan["joint_positions"], dtype=np.float64)
    gripper = np.asarray(plan["gripper"], dtype=np.float64)
    object_poses = np.asarray(plan["object_poses"], dtype=np.float64)
    object_names = plan["object_names"]
    n_frames_in = timestamps.shape[0]

    sidecar_path = args.sidecar_json or args.npz.with_suffix(".json")
    sidecar = json.loads(sidecar_path.read_text()) if sidecar_path.exists() else {}
    release_frame_idx = args.release_frame_idx
    if release_frame_idx is None:
        release_frame_idx = int(sidecar["phases"]["release_frame_idx"])
    logger.info("release_frame_idx = %d (of %d input frames)", release_frame_idx, n_frames_in)

    brick_outputs = args.brick_outputs or Path(sidecar.get(
        "brick_outputs",
        REPO_ROOT / "outputs/AUTOLab+0d4edc83+2023-10-21-19h-07m-04s/objects_sam3d_multi/20_31",
    ))
    books_outputs = args.books_outputs or Path(sidecar.get(
        "books_outputs",
        REPO_ROOT / "outputs/AUTOLab+0d4edc83+2023-10-21-19h-07m-04s/objects_books/0_11",
    ))
    placement_json = args.placement_json or (
        books_outputs / "objects" / "2_placement" / "placement_target.json"
    )

    # --- geometry: mesh, alignment scale, convex hull ------------------------
    mesh = rop._load_mesh(rop.stage_dir(brick_outputs, "mesh"))
    alignment = rop._load_alignment(rop.stage_dir(brick_outputs, "align"))
    verts_canonical = np.asarray(mesh.vertices, dtype=np.float64)
    faces = np.asarray(mesh.faces, dtype=np.int64)
    # Instructed method: uniform scale from the alignment transform's column
    # norms (should equal the same scale baked into every object_poses[t]).
    scale_cols = np.linalg.norm(np.asarray(alignment.transform, dtype=np.float64)[:3, :3], axis=0)
    scale = float(np.mean(scale_cols))
    logger.info(
        "brick scale = %.6f m/unit (column norms %s, alignment.scale=%.6f), "
        "world extent = %s cm",
        scale, np.round(scale_cols, 6).tolist(), alignment.scale,
        np.round(mesh.extent * scale * 100.0, 2).tolist(),
    )

    tri = trimesh.Trimesh(vertices=verts_canonical, faces=faces, process=False)
    hull = tri.convex_hull
    scratch_dir = args.scratch_dir or Path(tempfile.mkdtemp(prefix="settle_after_release_"))
    scratch_dir.mkdir(parents=True, exist_ok=True)
    hull_stl = scratch_dir / "brick_hull.stl"
    hull.export(str(hull_stl))
    logger.info(
        "convex hull: %d verts / %d faces (raw mesh had %d/%d) -> %s",
        hull.vertices.shape[0], hull.faces.shape[0],
        verts_canonical.shape[0], faces.shape[0], hull_stl,
    )

    # --- release pose, up-axis reference, release velocity -------------------
    t_obj_rest = object_poses[args.rest_frame_idx, 0]
    e_up = _rest_up_axis_canonical(t_obj_rest)

    release_pose = object_poses[release_frame_idx, 0].copy()
    release_pos = release_pose[:3, 3].copy()
    r_release_pure = release_pose[:3, :3] / scale
    release_quat = _mat_to_quat_wxyz(r_release_pure)

    lin_vel_world, ang_vel_world = estimate_release_velocity(
        object_poses[:, 0], release_frame_idx, args.n_vel_frames, scale
    )
    ang_vel_body = r_release_pure.T @ ang_vel_world
    logger.info(
        "release velocity: linear = %s m/s (|v|=%.4f), angular(world) = %s rad/s (|w|=%.4f)",
        np.round(lin_vel_world, 4).tolist(), np.linalg.norm(lin_vel_world),
        np.round(ang_vel_world, 4).tolist(), np.linalg.norm(ang_vel_world),
    )

    # --- placement surface, book-top box, safety floor -----------------------
    placement = json.loads(placement_json.read_text())
    surf_pos = np.asarray(placement["position_m"], dtype=np.float64)
    surf_normal = np.asarray(placement["normal"], dtype=np.float64)
    surf_normal = surf_normal / np.linalg.norm(surf_normal)
    surf_extent = np.asarray(placement["surface_extent_m"], dtype=np.float64)
    u_hat, v_hat = _plane_basis(surf_normal)

    box_rot = np.stack([u_hat, v_hat, surf_normal], axis=1)  # columns = local axes in world frame
    box_quat = _mat_to_quat_wxyz(box_rot)
    thickness = args.surface_thickness_m
    box_center = surf_pos - surf_normal * (thickness / 2.0)
    box_half_extent = np.array([surf_extent[0] / 2.0, surf_extent[1] / 2.0, thickness / 2.0])
    floor_z = float(surf_pos[2] - args.floor_drop_m)

    tilt_at_release = tilt_from_flat_deg(r_release_pure, e_up, surf_normal)
    logger.info("tilt from flat AT RELEASE = %.2f deg (sidecar reports 37.26)", tilt_at_release)

    settle_consecutive_steps = max(1, round(args.settle_consecutive_s / args.sim_dt))
    mass_check = args.mass_kg_check if args.mass_kg_check is not None else args.mass_kg * 10.0

    common_kwargs = dict(
        mujoco_python=args.mujoco_python, hull_stl=hull_stl, scale=scale,
        release_pos=release_pos, release_quat=release_quat,
        lin_vel_world=lin_vel_world, ang_vel_body=ang_vel_body,
        friction=tuple(args.friction), solref=tuple(args.solref), solimp=tuple(args.solimp),
        surface_pos=box_center, surface_quat=box_quat, surface_half_extent=box_half_extent,
        floor_z=floor_z, sim_dt=args.sim_dt, max_sim_time=args.max_sim_time,
        record_dt=args.record_dt,
        settle_lin_vel_thresh=args.settle_lin_vel_thresh,
        settle_ang_vel_thresh=args.settle_ang_vel_thresh,
        settle_consecutive_steps=settle_consecutive_steps, scratch_dir=scratch_dir,
    )

    worker_out = _run_worker(mass_kg=args.mass_kg, tag="primary", **common_kwargs)

    mass_check_result = None
    if args.check_mass_independence:
        worker_out_check = _run_worker(mass_kg=mass_check, tag="mass_check", **common_kwargs)
        p0 = np.asarray(worker_out["final_pos"])
        p1 = np.asarray(worker_out_check["final_pos"])
        q0 = np.asarray(worker_out["final_quat_wxyz"])
        q1 = np.asarray(worker_out_check["final_quat_wxyz"])
        r0 = _quat_wxyz_to_mat(q0)
        r1 = _quat_wxyz_to_mat(q1)
        rel = r1 @ r0.T
        rel = np.clip(rel, -1.0, 1.0)
        orient_diff_deg = float(np.degrees(np.linalg.norm(matrix_to_rotvec(rel))))
        pos_diff_m = float(np.linalg.norm(p1 - p0))
        mass_check_result = {
            "mass_kg_primary": args.mass_kg,
            "mass_kg_check": mass_check,
            "final_pos_primary": p0.tolist(),
            "final_pos_check": p1.tolist(),
            "position_diff_m": pos_diff_m,
            "orientation_diff_deg": orient_diff_deg,
            "settled_primary": worker_out["settled"],
            "settled_check": worker_out_check["settled"],
        }
        logger.info(
            "mass-independence check: %.3fkg vs %.3fkg -> "
            "position diff %.2e m, orientation diff %.4f deg",
            args.mass_kg, mass_check, pos_diff_m, orient_diff_deg,
        )

    # --- resample the primary run onto the plan's 15 Hz clock ----------------
    settle_time = worker_out["settle_time_s"] if worker_out["settled"] else args.max_sim_time
    span_s = settle_time + args.post_settle_buffer_s
    n_settle_frames = int(np.ceil(span_s * _TRAJ_FPS)) + 1  # +1: index0 == release frame itself
    n_frames_out = max(n_frames_in, release_frame_idx + n_settle_frames)

    new_timestamps = np.empty(n_frames_out, dtype=np.float64)
    new_timestamps[:release_frame_idx] = timestamps[:release_frame_idx]
    dt = 1.0 / _TRAJ_FPS
    for k in range(release_frame_idx, n_frames_out):
        new_timestamps[k] = timestamps[release_frame_idx] + (k - release_frame_idx) * dt

    new_joint_positions = np.empty((n_frames_out, joint_positions.shape[1]), dtype=np.float64)
    new_gripper = np.empty(n_frames_out, dtype=np.float64)
    n_copy = min(n_frames_in, n_frames_out)
    new_joint_positions[:n_copy] = joint_positions[:n_copy]
    new_gripper[:n_copy] = gripper[:n_copy]
    if n_frames_out > n_frames_in:
        # Settle outlasts the plan: hold the robot's final joint config (it
        # has nothing to do with the brick any more -- release already
        # happened) so the extension is visually a static, already-retracted arm.
        new_joint_positions[n_frames_in:] = joint_positions[-1]
        new_gripper[n_frames_in:] = gripper[-1]

    t_query = new_timestamps[release_frame_idx:] - new_timestamps[release_frame_idx]
    sim_pos, sim_quat = _resample_trajectory(worker_out, t_query)

    new_object_poses = np.empty((n_frames_out, 1, 4, 4), dtype=np.float64)
    new_object_poses[:release_frame_idx, 0] = object_poses[:release_frame_idx, 0]
    for i, k in enumerate(range(release_frame_idx, n_frames_out)):
        r_t = _quat_wxyz_to_mat(sim_quat[i])
        mat = np.eye(4, dtype=np.float64)
        mat[:3, :3] = r_t * scale
        mat[:3, 3] = sim_pos[i]
        new_object_poses[k, 0] = mat

    # --- verification numbers --------------------------------------------
    final_r_pure = _quat_wxyz_to_mat(np.asarray(worker_out["final_quat_wxyz"]))
    final_pos = np.asarray(worker_out["final_pos"], dtype=np.float64)
    final_tilt = tilt_from_flat_deg(final_r_pure, e_up, surf_normal)

    verts_world_final = scale * (verts_canonical @ final_r_pure.T) + final_pos
    height_above_plane = verts_world_final @ surf_normal - float(np.dot(surf_pos, surf_normal))
    lowest_vertex_signed_dist = float(np.min(height_above_plane))

    # brick body origin coincides with the release-frame world origin choice
    footprint_u = verts_world_final @ u_hat - float(np.dot(surf_pos, u_hat))
    footprint_v = verts_world_final @ v_hat - float(np.dot(surf_pos, v_hat))
    centroid_u = float(np.dot(final_pos, u_hat) - np.dot(surf_pos, u_hat))
    centroid_v = float(np.dot(final_pos, v_hat) - np.dot(surf_pos, v_hat))
    footprint_half_extent = np.array(
        [
            max(abs(footprint_u.min()), abs(footprint_u.max())),
            max(abs(footprint_v.min()), abs(footprint_v.max())),
        ]
    )
    # Every vertex's (u, v) must fall within the fitted surface's half-extent box.
    within_surface = bool(
        np.all(np.abs(footprint_u) <= surf_extent[0] / 2.0 + 1e-6)
        and np.all(np.abs(footprint_v) <= surf_extent[1] / 2.0 + 1e-6)
    )

    dt_check = np.diff(new_timestamps)
    dt_ok = bool(np.allclose(dt_check, dt, atol=1e-9))
    r = release_frame_idx
    pre_release_identical = bool(
        np.array_equal(new_timestamps[:r], timestamps[:r])
        and np.array_equal(new_object_poses[:r], object_poses[:r])
        and np.array_equal(new_joint_positions[:r], joint_positions[:r])
        and np.array_equal(new_gripper[:r], gripper[:r])
    )

    logger.info("=== verification ===")
    logger.info(
        "final tilt from flat = %.2f deg (release was %.2f deg)", final_tilt, tilt_at_release
    )
    logger.info(
        "lowest vertex signed distance to surface plane = %.4f m (>=0 means no penetration)",
        lowest_vertex_signed_dist,
    )
    logger.info(
        "centroid in-plane offset from surface centre = (%.4f, %.4f) m; "
        "footprint half-extent used (u, v) = %s m vs surface half-extent %s m -> within_surface=%s",
        centroid_u, centroid_v, np.round(footprint_half_extent, 4).tolist(),
        np.round(surf_extent / 2.0, 4).tolist(), within_surface,
    )
    logger.info("dt exactly 1/15 across extended trajectory: %s", dt_ok)
    logger.info("pre-release portion bit-identical to input: %s", pre_release_identical)
    if not within_surface:
        logger.warning(
            "brick's settled footprint extends OUTSIDE the fitted book surface -- "
            "it slid/rolled off."
        )

    np.savez(
        args.out,
        timestamps=new_timestamps,
        joint_positions=new_joint_positions,
        gripper=new_gripper,
        object_poses=new_object_poses,
        object_names=object_names,
    )
    logger.info("wrote %s (%d frames, was %d)", args.out, n_frames_out, n_frames_in)

    sidecar_out = {
        "source_npz": str(args.npz),
        "source_sidecar_json": str(sidecar_path) if sidecar_path.exists() else None,
        "release_frame_idx": release_frame_idx,
        "n_frames_in": n_frames_in,
        "n_frames_out": n_frames_out,
        "release_pose": {
            "position_m": release_pos.tolist(),
            "rotation_matrix_pure": r_release_pure.tolist(),
            "quat_wxyz": release_quat.tolist(),
            "tilt_from_flat_deg": tilt_at_release,
        },
        "inherited_velocity": {
            "linear_world_mps": lin_vel_world.tolist(),
            "linear_speed_mps": float(np.linalg.norm(lin_vel_world)),
            "angular_world_radps": ang_vel_world.tolist(),
            "angular_speed_radps": float(np.linalg.norm(ang_vel_world)),
            "n_frames_used_for_finite_difference": args.n_vel_frames,
            "note": "finite-differenced from the recorded (rigidly-attached) brick poses "
            "in the few frames immediately before release -- the brick does not "
            "drop from rest, it carries the hand's motion at the instant of release.",
        },
        "physics_assumptions": {
            "mass_kg": args.mass_kg,
            "mass_is_load_bearing": False,
            "mass_note": "ARBITRARY. A rigid body's fall-and-settle trajectory under gravity "
            "+ contact + friction is mass-independent (mass cancels out of F=ma when every "
            "force in the scene, including contact/friction, itself scales with mass) -- "
            "verified empirically below by re-running at 10x mass.",
            "friction_sliding_torsional_rolling": list(args.friction),
            "friction_is_load_bearing": True,
            "friction_note": "ASSUMPTION, never measured for this brick/book contact.",
            "solref": list(args.solref),
            "solimp": list(args.solimp),
            "restitution_is_load_bearing": True,
            "restitution_note": "ASSUMPTION. MuJoCo has no literal coefficient-of-restitution "
            "attribute; bounciness is entirely a function of how underdamped `solref` is. "
            "The values used (MuJoCo's own defaults) are critically damped, i.e. essentially "
            "no bounce -- reasonable for a small, light brick, but not measured.",
        },
        "mass_independence_check": mass_check_result,
        "simulation": {
            "sim_dt_s": args.sim_dt,
            "n_sim_steps": worker_out["n_sim_steps"],
            "settled": worker_out["settled"],
            "settle_time_s": worker_out["settle_time_s"],
            "sim_wall_time_s": worker_out["sim_wall_time_s"],
            "mujoco_version": worker_out["mujoco_version"],
            "hull_stl": str(hull_stl),
            "hull_n_vertices": int(hull.vertices.shape[0]),
            "hull_n_faces": int(hull.faces.shape[0]),
        },
        "final_pose": {
            "position_m": final_pos.tolist(),
            "rotation_matrix_pure": final_r_pure.tolist(),
            "tilt_from_flat_deg": final_tilt,
            "tilt_from_flat_deg_at_release": tilt_at_release,
        },
        "verification": {
            "final_tilt_deg_lower_than_release": final_tilt < tilt_at_release,
            "lowest_vertex_signed_distance_to_surface_m": lowest_vertex_signed_dist,
            "no_penetration": lowest_vertex_signed_dist >= -1e-4,
            "centroid_in_plane_offset_uv_m": [centroid_u, centroid_v],
            "footprint_half_extent_uv_m": footprint_half_extent.tolist(),
            "surface_half_extent_uv_m": (surf_extent / 2.0).tolist(),
            "within_surface_footprint": within_surface,
            "dt_exactly_1_over_15": dt_ok,
            "pre_release_portion_bit_identical": pre_release_identical,
        },
    }
    args.out_json.write_text(json.dumps(sidecar_out, indent=2, default=str))
    logger.info("wrote %s", args.out_json)
    print(json.dumps(sidecar_out["verification"], indent=2))
    return 0


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