scripts/_mujoco_settle_worker.py

#!/usr/bin/env python
"""Simulate a released rigid body falling and settling, using MuJoCo physics only.

Runs inside the isolated `mujoco` conda env (mujoco + numpy + trimesh, nothing
else) -- NEVER the `fpgm` env. `fpgm`'s numpy is pinned `<2` by `pyproject.toml`
(a hard constraint from SAM 3.1's own dependencies -- see that file's
docstring), while a plain `pip install mujoco` at the time this was written
pulls numpy 2.4.x. Installing mujoco into `fpgm` would risk a downgrade
cascade that breaks the working SAM3/TAPNext/pyrender stack there. Exactly
the same isolation pattern already used for `trellis2` and `sam3d-objects`
(see scripts/trellis_generate.py / scripts/sam3d_generate.py's docstrings):
a small, dependency-light worker script, invoked as a subprocess from the
`fpgm`-side orchestrator (scripts/settle_after_release.py) with the isolated
env's own interpreter, talking over JSON files rather than in-process
imports. This script deliberately imports nothing from `fpgm` -- the
`mujoco` env does not have it installed, and mixing envs defeats the whole
point of the isolation.

Also never imports `mujoco.viewer` or does any rendering: this host has no
hardware OpenGL (software OSMesa only), and MuJoCo's physics stepping (the
only thing this script needs) has no GL dependency at all.

Physics model (see the caller's docstring for the full rationale):
  * The brick is a single MJCF free body. Its collision geometry is the
    *convex hull* of the real reconstructed brick mesh (built by the caller
    with trimesh, passed in as an STL path) rather than a crude box --
    tipping behaviour depends on the true contact geometry (corners/edges),
    which a bounding box would get wrong.
  * The book top is a static box sized to the fitted placement-surface
    footprint (`surface_extent_m`), oriented so its top face passes through
    the fitted plane with the fitted normal.
  * A second, much larger static plane sits far below purely as a numerical
    safety net (bounds the simulation if the brick rolls or slides off the
    books) -- it is NOT part of the physical scene and reaching it is a
    signal to report honestly, not a feature.
  * Mass is whatever the caller passed in (arbitrary -- see the caller's
    docstring on why a rigid body's free-fall-and-settle trajectory does not
    depend on it). Friction and the solver's contact softness
    (`solref`/`solimp`, which stands in for restitution -- MuJoCo has no
    literal "coefficient of restitution" attribute; bounce is entirely a
    function of how underdamped the contact constraint is) are genuine
    modelling ASSUMPTIONS, recorded by the caller, not measured here.

Usage:
    /home/quang/miniconda3/envs/mujoco/bin/python scripts/_mujoco_settle_worker.py \\
        --in-json worker_input.json --out-json worker_output.json
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import mujoco
import numpy as np


def _mat_flat_to_quat(mat9: np.ndarray) -> np.ndarray:
    """(9,) row-major rotation matrix -> (4,) wxyz quaternion, via MuJoCo's own routine."""
    quat = np.zeros(4, dtype=np.float64)
    mujoco.mju_mat2Quat(quat, np.asarray(mat9, dtype=np.float64))
    return quat


def _build_mjcf(cfg: dict, hull_stl_path: str) -> str:
    """Assemble the MJCF for one settle simulation.

    Kept as a single f-string (no XML library) since every element here is a
    handful of numeric attributes -- matching the "minimal MJCF" the caller
    asked for, easy to eyeball in a debugger if something compiles wrong.
    """
    scale = float(cfg["scale"])
    px, py, pz = cfg["release_pos"]
    qw, qx, qy, qz = cfg["release_quat_wxyz"]
    sx, sy, sz = cfg["surface_pos"]
    sqw, sqx, sqy, sqz = cfg["surface_quat_wxyz"]
    shx, shy, shz = cfg["surface_half_extent"]
    fric = " ".join(f"{v:.6g}" for v in cfg["friction"])
    solref = " ".join(f"{v:.6g}" for v in cfg["solref"])
    solimp = " ".join(f"{v:.6g}" for v in cfg["solimp"])
    floor_z = float(cfg["floor_z"])
    mass = float(cfg["mass_kg"])
    dt = float(cfg["sim_dt"])

    return f"""
<mujoco model="settle_after_release">
  <compiler angle="radian"/>
  <option timestep="{dt:.8g}" gravity="0 0 -9.81" integrator="implicitfast" cone="elliptic"/>
  <asset>
    <mesh name="brick_hull" file="{hull_stl_path}" scale="{scale:.8g} {scale:.8g} {scale:.8g}"/>
  </asset>
  <worldbody>
    <light pos="0 0 2" dir="0 0 -1"/>
    <geom name="book_surface" type="box" pos="{sx:.8g} {sy:.8g} {sz:.8g}"
          quat="{sqw:.8g} {sqx:.8g} {sqy:.8g} {sqz:.8g}"
          size="{shx:.8g} {shy:.8g} {shz:.8g}"
          friction="{fric}" solref="{solref}" solimp="{solimp}" rgba="0.6 0.5 0.35 1"/>
    <geom name="safety_floor" type="plane" pos="0 0 {floor_z:.8g}" size="3 3 0.1"
          friction="{fric}" solref="{solref}" solimp="{solimp}" rgba="0.3 0.3 0.3 0.3"/>
    <body name="brick" pos="{px:.8g} {py:.8g} {pz:.8g}" quat="{qw:.8g} {qx:.8g} {qy:.8g} {qz:.8g}">
      <freejoint name="brick_free"/>
      <geom name="brick_geom" type="mesh" mesh="brick_hull" mass="{mass:.8g}"
            friction="{fric}" solref="{solref}" solimp="{solimp}" rgba="0.8 0.25 0.2 1"/>
    </body>
  </worldbody>
</mujoco>
"""


def run_settle(cfg: dict) -> dict:
    hull_stl_path = cfg["hull_stl_path"]
    xml = _build_mjcf(cfg, hull_stl_path)
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)

    # Inherited velocity: linear in world frame, angular in the BODY-LOCAL
    # frame -- MuJoCo's free-joint qvel convention (verified empirically
    # against mju_quat2Mat/mj_step before writing this: the rotational qvel
    # for a free joint integrates as R(t+dt) = R(t) @ exp([w]_x dt), i.e. w is
    # expressed in the body's own rotating frame, NOT world). The caller is
    # responsible for having already rotated a world-frame angular velocity
    # estimate into this frame using the release orientation.
    data.qvel[0:3] = cfg["lin_vel_world"]
    data.qvel[3:6] = cfg["ang_vel_body"]

    sim_dt = float(cfg["sim_dt"])
    record_dt = float(cfg["record_dt"])
    record_stride = max(1, round(record_dt / sim_dt))
    max_steps = int(round(float(cfg["max_sim_time"]) / sim_dt))
    lin_thresh = float(cfg["settle_lin_vel_thresh"])
    ang_thresh = float(cfg["settle_ang_vel_thresh"])
    consec_needed = int(cfg["settle_consecutive_steps"])

    times: list[float] = []
    pos: list[list[float]] = []
    quat: list[list[float]] = []
    lin_vel: list[list[float]] = []
    ang_vel: list[list[float]] = []

    def _record() -> None:
        times.append(float(data.time))
        pos.append(data.qpos[0:3].tolist())
        quat.append(data.qpos[3:7].tolist())
        lin_vel.append(data.qvel[0:3].tolist())
        ang_vel.append(data.qvel[3:6].tolist())

    _record()  # t=0, the release instant itself

    consec = 0
    settled = False
    settle_time: float | None = None
    step = 0
    t0 = time.time()
    while step < max_steps:
        mujoco.mj_step(model, data)
        step += 1
        lin_speed = float(np.linalg.norm(data.qvel[0:3]))
        ang_speed = float(np.linalg.norm(data.qvel[3:6]))
        if lin_speed < lin_thresh and ang_speed < ang_thresh:
            consec += 1
        else:
            consec = 0
        if step % record_stride == 0 or (consec == consec_needed and not settled):
            _record()
        if consec >= consec_needed and not settled:
            settled = True
            settle_time = float(data.time) - consec_needed * sim_dt
            break
    wall_time = time.time() - t0

    if not settled:
        _record()  # make sure the final (capped) state is in the trace

    return {
        "mujoco_version": mujoco.__version__,
        "sim_dt": sim_dt,
        "record_dt": record_dt,
        "n_sim_steps": step,
        "sim_wall_time_s": wall_time,
        "settled": settled,
        "settle_time_s": settle_time,
        "times": times,
        "pos": pos,
        "quat_wxyz": quat,
        "lin_vel": lin_vel,
        "ang_vel": ang_vel,
        "final_pos": data.qpos[0:3].tolist(),
        "final_quat_wxyz": data.qpos[3:7].tolist(),
        "final_lin_speed": float(np.linalg.norm(data.qvel[0:3])),
        "final_ang_speed": float(np.linalg.norm(data.qvel[3:6])),
    }


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--in-json", type=Path, required=True)
    p.add_argument("--out-json", type=Path, required=True)
    return p.parse_args()


def main() -> int:
    args = parse_args()
    cfg = json.loads(args.in_json.read_text())
    result = run_settle(cfg)
    args.out_json.write_text(json.dumps(result))
    print(
        f"settle worker: {result['n_sim_steps']} steps, "
        f"settled={result['settled']} at t={result['settle_time_s']}, "
        f"wall={result['sim_wall_time_s']:.2f}s (mujoco {result['mujoco_version']})"
    )
    return 0


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