src/fpgm/physics/worker_mujoco.py

#!/usr/bin/env python
"""S9's batched MuJoCo particle simulator -- runs ONLY in the isolated `mujoco` conda env.

Sibling of ``scripts/_mujoco_settle_worker.py`` (read that file's docstring first; this
one assumes the isolation rationale there and does not repeat it): same "no fpgm import,
no mujoco.viewer, no rendering, talk over files" discipline, same reason (`fpgm`'s numpy is
pinned `<2` by SAM 3.1; a plain `pip install mujoco` pulls numpy 2.x and risks a downgrade
cascade there). The one thing this worker adds on top of that precedent is **batching**:
``_mujoco_settle_worker.py`` builds and steps exactly one MJCF per subprocess invocation,
which is fine for a one-off brick drop but far too slow for S9, where a single posterior
round evaluates thousands of parameter hypotheses (particles) per episode-object. Model
*compilation* (``mujoco.MjModel.from_xml_string``, which parses the MJCF, loads the mesh,
builds the collision tree) costs milliseconds and does not depend on any of the identified
parameters in :data:`RIGID_PARAM_NAMES`/:data:`PRISMATIC_PARAM_NAMES` -- every one of them
is a plain ``MjModel`` *field* (mass, friction, solref, joint damping, ...) that can be
edited post-compile with no recompilation. So this worker compiles the MJCF **once** per
invocation (one invocation = one :class:`~fpgm.physics.types.SimSpec`, i.e. one fixed
scene *structure*) and loops every particle over the same compiled model, editing fields
and calling :func:`mujoco.mj_resetData` between particles -- amortising the one expensive
step over however many particles are in ``--particles``. This is the "reuse the compiled
MjModel where the parameter only touches model fields; recompile only when structure
changes" design the caller's docstring asks for; because a single worker invocation never
changes structure (one spec, one MJCF), that reduces here to "compile exactly once." One
field edit is NOT enough on its own, though, and this was found empirically while building
this worker, not assumed: writing ``model.body_mass``/``model.body_inertia`` alone leaves
``model.body_invweight0`` (the constraint solver's reference inverse-mass, cached at compile
time) pointing at the density-1 baseline body -- contacts still register (``ncon>0``) but are
far too soft, and a body that should rest on a support box visibly sinks through it instead.
:func:`mujoco.mj_setConst` recomputes that cache from the current field values (still no XML
reparsing, no mesh reload -- an ``O(nv)`` call, not a recompilation) and is called once per
particle, right after :func:`_apply_theta`; see that call site for the trajectory comparison
that caught this.

**What is simulated vs. what is a hypothesis.** Unlike ``_mujoco_settle_worker.py``, where
friction/restitution were fixed ASSUMPTIONS a human chose once, every entry of
:data:`RIGID_PARAM_NAMES`/:data:`PRISMATIC_PARAM_NAMES` here is a *hypothesis* carried by
one particle -- this worker's only job is to render each hypothesis into a trajectory
precisely and repeatably. Turning "does this trajectory look like what the camera saw" into
evidence for or against the hypothesis is ``fpgm.physics.likelihood``'s job (out of this
worker's scope entirely; it never sees the observed track). Anything this worker still
fixes rather than reads from a particle -- rolling friction, the contact solver's
``timeconst`` half of ``solref`` (only the damping-ratio half is identified), and a shared
mesh scale of 1.0 (i.e. the caller is trusted to have already written every hull STL, body
of interest and kinematic passengers alike, in real-world metres) -- remains a genuine
MODELLING ASSUMPTION, not a measurement, and is called out at its point of use below, in the
same spirit as ``_mujoco_settle_worker.py``'s own friction/restitution note. Bodies *beyond*
the first one in ``spec.bodies`` split into two cases (see :func:`_apply_kinematic_friction`
vs. :func:`_apply_generic_body_defaults`): a ``"kinematic"`` body (the common, expected case
-- see below) is mocap-driven off its own measured ``pose_track`` and borrows the body of
interest's *own* particle-varying ``log_mu_slide``/``log_mu_torsion`` for its friction, since
a second, per-pair friction parameter is not identifiable from one episode's worth of
evidence (see that function's docstring); any body that is neither the first one nor
``"kinematic"`` (not currently emitted by ``fpgm.physics.scene.build_sim_spec``, but the
worker must still not crash or silently mis-simulate one) falls back to fixed generic
defaults with no particle-controlled behaviour at all.

**Interface assumption: "the body of interest" is ``spec.bodies[0]``.**
:class:`~fpgm.physics.types.SimSpec` has no field naming which body a caller wants back --
:class:`~fpgm.physics.types.SimResult` carries exactly one body's ``(N, T, 7)`` trajectory
and one ``label``. The convention adopted here (and relied on by
``fpgm.physics.scene.build_sim_spec``) is: **the first entry of ``spec.bodies`` is the body
of interest**, always ``"free"`` or ``"prismatic"``; its world pose (``data.xpos``/
``data.xquat``, not raw ``qpos`` -- see :func:`_extract_pose`) is what gets recorded every
frame, and the parameter columns in ``--particles`` (``theta``/``param_names``) are
interpreted against it and it alone. Every other tracked object S6 saw in the same episode
now rides along as ``spec.bodies[1:]``, each one ``kind="kinematic"`` -- a MuJoCo mocap body
driven off its own measured ``pose_track`` (see :class:`~fpgm.physics.types.BodySpec`'s own
docstring), fully able to contact the body of interest (grasped, pushed, rested-on) but
never itself a hypothesis: nothing about *its* mass, density or inertia is read from
``theta``, because a mocap body has no dynamics for those fields to act on in the first
place. This is what turns "the brick rests on the drawer" from a fact this worker would
otherwise need to be told into one that simply falls out of both bodies being present with
their real measured poses -- see ``fpgm.physics.scene``'s module docstring for the full
motivation. Any *further* non-kinematic body beyond the first (not currently emitted by
``build_sim_spec``, kept only so the worker degrades sanely rather than crashing) gets
:data:`_GENERIC_BODY_*` physical defaults (see :func:`_apply_generic_body_defaults`) since
there is no parameter channel for it either.

**Why mocap fingers, not a real gripper joint (restated from physics/types.py, because
it drives a modelling choice below).** :class:`~fpgm.physics.types.GripperSpec` gives two
kinematically-driven mocap boxes -- ``data.mocap_pos``/``data.mocap_quat`` are set to
``finger_poses[t]`` every frame and MuJoCo treats a mocap body as having effectively
infinite mass, unaffected by contact forces but fully able to exert them. Nothing in the
model "knows" the object is grasped: whether it stays between the fingers is decided purely
by contact + ``log_mu_gripper``, which is exactly what makes one model cover grasp, slip,
push, place and release with no phase special-casing (see ``GripperSpec``'s own docstring
for the full rationale). One consequence worth being explicit about: the recorded
``finger_poses`` are **piecewise-linearly (position) / slerp (orientation) interpolated**
across the ``substeps`` physics steps between two consecutive output frames (see
:func:`_lerp`/:func:`_slerp`), not held constant at the frame-``t`` target throughout --
holding it constant would introduce a full output-frame's worth of positional lag on fast
finger motion; interpolating keeps the driven contact geometry close to the FK-measured
motion at every substep, at the cost of assuming the true motion is well-approximated by a
straight line between two consecutive 60 Hz(-ish) FK samples, which is a mild and standard
assumption for typical DROID gripper speeds.

**Priority, not friction mixing.** MuJoCo's default contact-parameter combination rule
mixes two colliding geoms' ``friction``/``solref`` (elementwise max for friction, a
weighted rule for solref) rather than letting either one win outright. That would dilute
exactly the parameters this stage exists to identify: if the heightfield/finger geoms carried
any non-trivial default friction, ``log_mu_slide``/``log_mu_gripper`` would only ever be
observed through ``max(identified, fixed_default)``, silently floored. Every geom on the
body of interest is instead given MuJoCo's ``priority="1"`` (every other geom -- the
heightfield, safety floor, fingers -- stays at the default ``priority="0"``); per the MuJoCo
docs, the higher-priority geom's own friction/solref parameters are then used *directly* for that
contact pair, unmixed. This is the mechanism, not a fixed value, that makes each particle's
``theta`` actually control what the simulation does.

Usage (always as ``-m``, never as a bare script path -- see the note below)::

    PYTHONPATH=src /home/quang/miniconda3/envs/mujoco/bin/python \\
        -m fpgm.physics.worker_mujoco \\
        --spec sim_spec.json --particles particles.npz --out result.npz

**Why ``-m``, not a direct script path (unlike ``scripts/_mujoco_settle_worker.py``,
which IS invoked as a bare path).** This module lives at ``src/fpgm/physics/worker_mujoco.py``,
a sibling of ``physics/types.py``. Running it as a bare script path makes Python prepend
*that file's own directory* to ``sys.path``, which then shadows the standard-library
``types`` module with ``fpgm/physics/types.py`` for the rest of the process -- breaking
stdlib imports that need the real ``types`` (``enum``, ``weakref``, ...) with a confusing
circular-import error. Invoking as ``-m`` (with ``PYTHONPATH=src`` and any cwd) instead
puts the current working directory at ``sys.path[0]``, never the package's own directory,
so the collision cannot occur. ``fpgm.physics.scene.MujocoSimulator`` (the subprocess
caller, ``src/fpgm/physics/simulate.py``) always launches it this way.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import mujoco
import numpy as np

# --------------------------------------------------------------------------- #
# Parameter-space contract, DUPLICATED (not imported) from fpgm.physics.types.
#
# The `mujoco` env has no `fpgm` installed (see module docstring) -- this mirrors
# scripts/settle_after_release.py's own precedent of duplicating a handful of
# quaternion helpers across the same env boundary rather than adding a cross-env
# dependency for them. `--particles` therefore always carries its own
# `param_names` (D,) array (see :func:`_load_particles`) rather than trusting
# this list alone; a name mismatch between the two sides fails loudly (KeyError
# on a missing name) instead of silently misinterpreting a column.
# --------------------------------------------------------------------------- #
RIGID_PARAM_NAMES: tuple[str, ...] = (
    "log_density",
    "com_x",
    "com_y",
    "log_mu_slide",
    "log_mu_torsion",
    "log_solref_damping",
    "log_mu_gripper",
)
PRISMATIC_PARAM_NAMES: tuple[str, ...] = (
    "log_joint_friction",
    "log_joint_damping",
)

# -- fixed modelling assumptions (never particle-controlled) ---------------- #

#: Rolling friction is not in RIGID_PARAM_NAMES (contact identifiability for a
#: body that is dropped/pushed/grasped, never rolled, does not usefully separate
#: it from sliding friction) -- held at MuJoCo's own small default, matching
#: scripts/settle_after_release.py's `_ASSUMED_FRICTION`'s rolling term.
_DEFAULT_ROLL_FRICTION = 1.0e-4
#: Finger-vs-object torsional/rolling friction: also not identified (only the
#: sliding term, log_mu_gripper, is) -- fixed at small generic defaults since
#: gripper contact in this scene is dominated by sliding/normal force, not spin.
_DEFAULT_GRIPPER_TORSION = 5.0e-3
_DEFAULT_GRIPPER_ROLL = 1.0e-4
#: solref = (timeconst, dampratio). Only dampratio (bounce/underdamping) is
#: identified via log_solref_damping; timeconst is held at MuJoCo's own default
#: (also scripts/settle_after_release.py's `_ASSUMED_SOLREF[0]`) -- varying both
#: would make solref_damping alone not identifiable from a fixed-timeconst prior.
_SOLREF_TIMECONST = 0.02
#: Generic (non-identified) physical defaults for any body beyond spec.bodies[0]
#: -- see the module docstring's "interface assumption" note. 1000 kg/m^3 is
#: plain water density, the conventional "no better information" placeholder;
#: this path is not expected to be exercised by fpgm.physics.scene.build_sim_spec,
#: which always emits a single-body spec, but the worker must still behave
#: sanely (not crash, not silently mis-simulate) if a spec ever does carry more.
_GENERIC_BODY_DENSITY = 1000.0
_GENERIC_BODY_FRICTION = (0.5, 5.0e-3, 1.0e-4)
_GENERIC_BODY_SOLREF = (_SOLREF_TIMECONST, 1.0)
#: How far below the lowest of {the heightfield's own base, every body's initial z} the
#: numerical safety floor sits. Reaching it is a divergence signal (ok=False),
#: never a feature -- same convention as _mujoco_settle_worker.py's floor_z.
_SAFETY_FLOOR_DROP_M = 1.0
_SAFETY_FLOOR_HALF_SIZE = 3.0

#: Solid depth MuJoCo's <hfield> extends BELOW its own lowest elevation vertex (the 4th
#: `size` component -- see `_heightfield_mujoco_params`'s docstring for the empirically
#: -verified geometry). Not zero: a zero-thickness heightfield is exactly the kind of thin
#: shell DEFAULT_SUBSTEPS's docstring already found tunnels through at low substep counts;
#: this gives the field's own solid base the same order of margin as
#: the old fitted-plane support box's own fixed thickness gave it (now removed; see
#: fpgm.physics.scene's module docstring for that representation's replacement).
_HFIELD_BASE_Z_M = 0.05

#: Floor on a heightfield's own elevation dynamic range (`z_max - z_min` over the WHOLE
#: grid, including hole-sentinel cells). A perfectly flat input would make this exactly 0,
#: which MuJoCo's compiler rejects (`size` component 3 must be > 0) -- this is a numerical
#: floor only, never expected to bind on real (always slightly noisy, and always containing
#: at least one deliberately-dropped hole cell) data.
_MIN_ELEVATION_RANGE_M = 1.0e-4


class WorkerError(RuntimeError):
    """Raised for a malformed --spec/--particles input; never for a diverged rollout.

    A diverged particle is not an error -- it is `ok=False` and a recorded, real
    result (see the module docstring on why silently dropping it would be wrong).
    This is only for "the wire contract itself does not hold" cases: a missing
    parameter name, a shape mismatch, a body with no geometry.
    """


# --------------------------------------------------------------------------- #
# Small standalone quaternion/geometry helpers.
#
# Deliberately not shared with fpgm.geometry.transforms (that module lives in
# `fpgm`, unavailable here) nor re-derived ad hoc per call site -- kept as the
# minimal, tested set this worker needs, matching the precedent already set by
# scripts/settle_after_release.py's own standalone quaternion helpers (see that
# file's module-level comment on why they are duplicated, not imported).
# --------------------------------------------------------------------------- #


def _mat_to_quat(rot3x3: np.ndarray) -> np.ndarray:
    """(3, 3) proper rotation matrix -> (4,) wxyz unit quaternion, via MuJoCo's own routine."""
    quat = np.zeros(4, dtype=np.float64)
    mat9 = np.ascontiguousarray(rot3x3, dtype=np.float64).reshape(-1)
    mujoco.mju_mat2Quat(quat, mat9)
    return quat


def _lerp(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
    return a + t * (b - a)


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)


# --------------------------------------------------------------------------- #
# Heightfield: metric grid -> MuJoCo's [0, 1]-normalised <hfield> contract
# --------------------------------------------------------------------------- #


def _heightfield_mujoco_params(
    height_m: np.ndarray, *, cell_size_m: float, origin_xy: tuple[float, float],
) -> dict:
    """Convert a raw metric ``(ny, nx)`` height grid into MuJoCo's ``<hfield>`` contract.

    **MuJoCo hfield geometry, EMPIRICALLY VERIFIED against this worker's own MuJoCo
    version** (its docs describe the ``size`` attribute in prose but do not spell out the
    row/column-to-world-axis mapping or the exact elevation formula; both were confirmed
    directly with a ray-cast probe rather than assumed -- see
    ``tests/test_physics_simulate.py``'s heightfield settle tests for the same fact
    re-verified end to end through a real compile):

    * ``hfield_data`` is stored ROW-MAJOR, ``(nrow, ncol)`` flattened as ``row * ncol + col``,
      values normalised to ``[0, 1]``.
    * Column index increases with world ``+X``; row index increases with world ``+Y``.
    * The ``(nrow, ncol)`` VERTICES are evenly spaced across the geom's own
      ``size = (radius_x, radius_y, elevation_z, base_z)``: vertex ``(row, col)`` sits at
      local ``x = -radius_x + col * (2*radius_x)/(ncol-1)``, similarly for ``y``/row -- i.e.
      the field spans exactly ``[-radius_x, +radius_x]`` x ``[-radius_y, +radius_y]`` about
      the geom's own ``pos``, with (ncol-1)/(nrow-1) INTERVALS between the outermost
      vertices, not (ncol)/(nrow) -- this is a vertex grid, not a pixel grid.
    * World elevation at a vertex is ``geom_pos_z + hfield_data[row, col] * elevation_z``
      (verified: a flat field of normalised value ``v`` settles a dropped sphere at exactly
      ``geom_pos_z + v * elevation_z``, to floating-point precision, for both ``v=0`` and
      ``v=1``). ``base_z`` (measured separately) extends a SOLID base downward from there,
      giving the field real volume rather than an infinitesimally thin shell -- see
      :data:`_HFIELD_BASE_Z_M`.

    ``fpgm.physics.scene._rasterize_top_down`` already produces a grid at uniform
    ``cell_size_m`` vertex spacing in both axes, so ``radius_x``/``radius_y`` fall straight
    out of ``(nx - 1) * cell_size_m / 2`` -- no resampling is needed to satisfy MuJoCo's
    "evenly spaced" requirement.

    Args:
        height_m: ``(ny, nx)`` world-frame Z, ALREADY hole-filled (see
            ``fpgm.physics.scene._fill_unknown_cells`` -- this function has no "what is an
            unknown cell" policy of its own; every entry here is treated as a real elevation
            to normalise).
        cell_size_m: Uniform vertex spacing, metres, both axes.
        origin_xy: World ``(x, y)`` of grid vertex ``(row=0, col=0)`` -- the grid's
            minimum-XY corner.

    Returns:
        A dict with ``normalized01`` (``(ny, nx)`` float64 in ``[0, 1]``, already row-major
        in the layout ``hfield_data`` needs), ``nrow``/``ncol``, ``radius_x``/``radius_y``,
        ``elevation_z``, ``base_z``, and ``pos`` (the ``<geom pos=...>`` this field must be
        placed at for the metric round-trip above to hold).
    """
    ny, nx = height_m.shape
    z_min = float(np.min(height_m))
    z_max = float(np.max(height_m))
    elevation_z = max(z_max - z_min, _MIN_ELEVATION_RANGE_M)
    normalized01 = np.clip((height_m - z_min) / elevation_z, 0.0, 1.0).astype(np.float64)

    radius_x = (nx - 1) * cell_size_m / 2.0
    radius_y = (ny - 1) * cell_size_m / 2.0
    pos_x = float(origin_xy[0]) + radius_x
    pos_y = float(origin_xy[1]) + radius_y
    pos_z = z_min

    return {
        "normalized01": normalized01,
        "nrow": ny,
        "ncol": nx,
        "radius_x": radius_x,
        "radius_y": radius_y,
        "elevation_z": elevation_z,
        "base_z": _HFIELD_BASE_Z_M,
        "pos": (pos_x, pos_y, pos_z),
    }


def _load_heightfield_grid(grid_path: Path) -> tuple[np.ndarray, np.ndarray]:
    """``(height_m, valid)`` from the sidecar ``.npz`` :class:`~fpgm.physics.types.
    HeightfieldSpec.grid_path` points at -- see that class's docstring for why this is a
    file, not inline JSON.

    Raises:
        WorkerError: file missing, unreadable, or missing the expected arrays -- the
            heightfield is the scene's ONLY static collision geometry (see the module
            docstring); a malformed one must fail loudly, not silently simulate free fall
            through empty space again.
    """
    grid_path = Path(grid_path)
    if not grid_path.exists():
        raise WorkerError(f"heightfield grid file not found: {grid_path}")
    try:
        with np.load(grid_path) as npz:
            if "height_m" not in npz or "valid" not in npz:
                raise WorkerError(
                    f"{grid_path}: missing 'height_m'/'valid' array (got {list(npz.files)})"
                )
            height_m = np.asarray(npz["height_m"], dtype=np.float64)
            valid = np.asarray(npz["valid"], dtype=bool)
    except (OSError, ValueError) as exc:
        raise WorkerError(
            f"{grid_path}: could not be read as an .npz heightfield grid: {exc}"
        ) from exc
    if height_m.shape != valid.shape or height_m.ndim != 2:
        raise WorkerError(
            f"{grid_path}: height_m/valid shape mismatch {height_m.shape}/{valid.shape}"
        )
    if not np.all(np.isfinite(height_m)):
        raise WorkerError(f"{grid_path}: height_m has non-finite entries -- hole-filling failed")
    return height_m, valid


# --------------------------------------------------------------------------- #
# MJCF assembly
# --------------------------------------------------------------------------- #


def _mesh_asset_xml(label: str, mesh_path: str) -> str:
    return f'    <mesh name="{label}_hull" file="{mesh_path}"/>\n'


def _body_xml(body: dict, *, priority: int) -> str:
    label = body["label"]
    px, py, pz = body["init_pose_pos"]
    qw, qx, qy, qz = body["init_pose_quat"]
    kind = body["kind"]

    if kind == "kinematic":
        # A mocap body, exactly like the two gripper fingers (see the module
        # docstring's "why mocap fingers" note) but carrying the tracked object's
        # OWN convex hull instead of a box, and driven from `pose_track` rather
        # than `finger_poses` -- see `_set_body_mocap`/`_set_body_mocap_interp`.
        # `priority` is intentionally unused here: this geom's own friction is set
        # every particle (see `_apply_kinematic_friction`), and MuJoCo's contact
        # priority rule (see the module docstring) already makes the body-of-
        # interest's own priority="1" geom win any contact against this one
        # regardless of what this geom's friction says.
        return (
            f'    <body name="{label}" mocap="true" '
            f'pos="{px:.8g} {py:.8g} {pz:.8g}" quat="{qw:.8g} {qx:.8g} {qy:.8g} {qz:.8g}">\n'
            f'      <geom name="{label}_geom" type="mesh" mesh="{label}_hull" '
            f'friction="0.5 0.005 0.0001" solref="{_SOLREF_TIMECONST:.8g} 1" '
            f'rgba="0.3 0.5 0.7 1"/>\n'
            f"    </body>\n"
        )

    if kind == "free":
        joint_xml = f'      <freejoint name="{label}_free"/>\n'
    else:
        ax, ay, az = body["joint_axis_local"]
        ox, oy, oz = body["joint_origin_local"]
        lo, hi = body["joint_range"]
        joint_xml = (
            f'      <joint name="{label}_slide" type="slide" '
            f'pos="{ox:.8g} {oy:.8g} {oz:.8g}" axis="{ax:.8g} {ay:.8g} {az:.8g}" '
            f'range="{lo:.8g} {hi:.8g}"/>\n'
        )

    return (
        f'    <body name="{label}" pos="{px:.8g} {py:.8g} {pz:.8g}" '
        f'quat="{qw:.8g} {qx:.8g} {qy:.8g} {qz:.8g}">\n'
        f"{joint_xml}"
        f'      <geom name="{label}_geom" type="mesh" mesh="{label}_hull" '
        f'density="1" priority="{priority}" '
        f'friction="0.5 0.005 0.0001" solref="{_SOLREF_TIMECONST:.8g} 1" '
        f'rgba="0.7 0.3 0.3 1"/>\n'
        f"    </body>\n"
    )


def _finger_xml(name: str, half_extent: np.ndarray) -> str:
    hx, hy, hz = half_extent
    return (
        f'    <body name="{name}" mocap="true" pos="0 0 -10" quat="1 0 0 0">\n'
        f'      <geom name="{name}_geom" type="box" size="{hx:.8g} {hy:.8g} {hz:.8g}" '
        f'friction="0.5 0.005 0.0001" solref="{_SOLREF_TIMECONST:.8g} 1" '
        f'rgba="0.2 0.2 0.2 0.6" contype="1" conaffinity="1"/>\n'
        f"    </body>\n"
    )


def build_mjcf(spec: dict) -> str:
    """Assemble the MJCF for one :class:`~fpgm.physics.types.SimSpec`.

    One f-string, no XML library -- same rationale as
    scripts/_mujoco_settle_worker.py's `_build_mjcf`: every element here is a
    handful of numeric attributes, easy to eyeball in a debugger if a compile
    fails, and a full XML tree would be pure overhead for that.
    """
    gx, gy, gz = spec["gravity"]
    sim_dt = float(spec["dt"]) / int(spec["substeps"])

    bodies = spec["bodies"]
    mesh_assets = "".join(_mesh_asset_xml(b["label"], b["mesh_path"]) for b in bodies)
    body_defs = "".join(
        _body_xml(b, priority=(1 if i == 0 else 0)) for i, b in enumerate(bodies)
    )

    gripper_xml = ""
    if spec.get("gripper") is not None:
        half_extent = np.asarray(spec["gripper"]["finger_size"], dtype=np.float64)
        gripper_xml = _finger_xml("finger_L", half_extent) + _finger_xml("finger_R", half_extent)

    hfield_asset_xml = ""
    hfield_geom_xml = ""
    if spec.get("heightfield") is not None:
        # Populated by `_load_spec` (see that function) -- the ONE place this worker reads
        # the sidecar grid file and converts it to MuJoCo's [0, 1] contract; see
        # `_heightfield_mujoco_params`'s docstring for the exact geometry this XML relies on.
        p = spec["heightfield"]["_mujoco"]
        px, py, pz = p["pos"]
        hfield_asset_xml = (
            f'    <hfield name="terrain" nrow="{p["nrow"]}" ncol="{p["ncol"]}" '
            f'size="{p["radius_x"]:.8g} {p["radius_y"]:.8g} {p["elevation_z"]:.8g} '
            f'{p["base_z"]:.8g}"/>\n'
        )
        hfield_geom_xml = (
            f'    <geom name="heightfield_surface" type="hfield" hfield="terrain" '
            f'pos="{px:.8g} {py:.8g} {pz:.8g}" '
            f'friction="0.5 0.005 0.0001" rgba="0.6 0.5 0.35 1"/>\n'
        )
        safety_z = pz - _SAFETY_FLOOR_DROP_M
    else:
        min_z = min(float(b["init_pose_pos"][2]) for b in bodies)
        safety_z = min_z - _SAFETY_FLOOR_DROP_M

    safety_xml = (
        f'    <geom name="safety_floor" type="plane" pos="0 0 {safety_z:.8g}" '
        f'size="{_SAFETY_FLOOR_HALF_SIZE:.8g} {_SAFETY_FLOOR_HALF_SIZE:.8g} 0.1" '
        f'friction="0.5 0.005 0.0001" rgba="0.3 0.3 0.3 0.3"/>\n'
    )

    return f"""
<mujoco model="s9_particle_batch">
  <compiler angle="radian"/>
  <option timestep="{sim_dt:.8g}" gravity="{gx:.8g} {gy:.8g} {gz:.8g}"
          integrator="implicitfast" cone="elliptic"/>
  <asset>
{mesh_assets}{hfield_asset_xml}  </asset>
  <worldbody>
{body_defs}{gripper_xml}{hfield_geom_xml}{safety_xml}  </worldbody>
</mujoco>
"""


# --------------------------------------------------------------------------- #
# Spec / particle loading
# --------------------------------------------------------------------------- #


def _load_spec(path: Path) -> dict:
    """Read ``--spec`` and pre-derive the local-frame quantities MJCF needs.

    Everything in :class:`~fpgm.physics.types.SimSpec.to_json_dict` is WORLD-frame
    (poses, prismatic axis/origin); MJCF's own ``<joint pos= axis=>`` for a body
    are expressed in that BODY's own frame. This function does that one rotation
    per prismatic body, once, rather than inside the (particle-)hot loop.

    Also where the heightfield's sidecar ``.npz`` (see
    :class:`~fpgm.physics.types.HeightfieldSpec.grid_path`) is actually read from disk and
    converted to MuJoCo's ``[0, 1]`` contract (:func:`_heightfield_mujoco_params`), stashed
    at ``raw["heightfield"]["_mujoco"]`` for :func:`build_mjcf`/:func:`simulate_all` -- once
    per worker invocation, structural (particle-independent), same "pre-derive, do not
    repeat in the hot loop" rationale as the prismatic local-frame conversion above.
    """
    raw = json.loads(Path(path).read_text())
    required = ("uuid", "camera_serial", "dt", "n_frames", "bodies")
    missing = [k for k in required if k not in raw]
    if missing:
        raise WorkerError(f"--spec missing required keys: {missing}")
    if not raw["bodies"]:
        raise WorkerError("--spec has an empty 'bodies' list")

    for body in raw["bodies"]:
        init_pose = np.asarray(body["init_pose"], dtype=np.float64)
        if init_pose.shape != (4, 4):
            raise WorkerError(f"{body['label']}: init_pose must be (4,4), got {init_pose.shape}")
        rot = init_pose[:3, :3]
        pos = init_pose[:3, 3]
        body["init_pose_pos"] = pos.tolist()
        body["init_pose_quat"] = _mat_to_quat(rot).tolist()
        body["init_lin_vel"] = list(body.get("init_lin_vel", (0.0, 0.0, 0.0)))
        body["init_ang_vel"] = list(body.get("init_ang_vel", (0.0, 0.0, 0.0)))
        if body["kind"] == "prismatic":
            axis_world = np.asarray(body["joint_axis_world"], dtype=np.float64)
            origin_world = np.asarray(body["joint_origin_world"], dtype=np.float64)
            body["joint_axis_local"] = (rot.T @ axis_world).tolist()
            body["joint_origin_local"] = (rot.T @ (origin_world - pos)).tolist()

    if raw.get("heightfield") is not None:
        hf = raw["heightfield"]
        for key in ("grid_path", "nx", "ny", "cell_size_m", "origin_xy"):
            if key not in hf:
                raise WorkerError(f"--spec heightfield missing required key {key!r}")
        height_m, _valid = _load_heightfield_grid(Path(hf["grid_path"]))
        if height_m.shape != (int(hf["ny"]), int(hf["nx"])):
            raise WorkerError(
                f"{hf['grid_path']}: grid shape {height_m.shape} disagrees with spec's "
                f"declared (ny, nx) = ({hf['ny']}, {hf['nx']})"
            )
        hf["_mujoco"] = _heightfield_mujoco_params(
            height_m, cell_size_m=float(hf["cell_size_m"]), origin_xy=tuple(hf["origin_xy"]),
        )
    return raw


def _load_particles(path: Path) -> tuple[np.ndarray, list[str]]:
    with np.load(path) as npz:
        if "theta" not in npz or "param_names" not in npz:
            raise WorkerError(
                f"--particles must contain 'theta' and 'param_names', got {list(npz.files)}"
            )
        theta = np.asarray(npz["theta"], dtype=np.float64)
        param_names = [str(n) for n in npz["param_names"]]
    if theta.ndim != 2 or theta.shape[1] != len(param_names):
        raise WorkerError(
            f"--particles theta shape {theta.shape} disagrees with "
            f"{len(param_names)} param_names"
        )
    return theta, param_names


# --------------------------------------------------------------------------- #
# Per-particle model editing
# --------------------------------------------------------------------------- #


class _BodyHandles:
    """Cached ids + density=1 baseline for one body, resolved once after compile."""

    def __init__(self, model: mujoco.MjModel, label: str, kind: str) -> None:
        self.label = label
        self.kind = kind
        self.body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, label)
        self.geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, f"{label}_geom")
        if self.body_id < 0 or self.geom_id < 0:
            raise WorkerError(f"MJCF compiled without expected body/geom for {label!r}")
        self.base_mass = float(model.body_mass[self.body_id])
        self.base_inertia = model.body_inertia[self.body_id].copy()
        self.base_ipos = model.body_ipos[self.body_id].copy()
        self.bounding_radius = float(model.geom_rbound[self.geom_id])
        if self.kind == "prismatic":
            joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{label}_slide")
            if joint_id < 0:
                raise WorkerError(f"MJCF compiled without expected slide joint for {label!r}")
            self.dof_id = int(model.jnt_dofadr[joint_id])
        else:
            self.dof_id = None
        # A "kinematic" body has no joint at all (it is mocap="true", exactly like
        # the gripper fingers) -- its own id into data.mocap_pos/mocap_quat, resolved
        # once here rather than by mj_name2id'ing "label" every substep the way
        # _set_finger_mocap does for the (fixed, only-ever-two) fingers.
        self.mocap_id = int(model.body_mocapid[self.body_id]) if self.kind == "kinematic" else None


def _apply_theta(
    model: mujoco.MjModel,
    handles: _BodyHandles,
    theta: dict[str, float],
    finger_geom_ids: list[int],
) -> None:
    """Edit exactly the MjModel fields RIGID_PARAM_NAMES/PRISMATIC_PARAM_NAMES own.

    No recompilation -- every field touched here (`body_mass`, `body_inertia`,
    `body_ipos`, `geom_friction`, `geom_solref`, `dof_frictionloss`,
    `dof_damping`) is a plain compiled-model array, safe to overwrite between
    particles. This is the amortisation the module docstring describes.
    """
    density = float(np.exp(theta["log_density"]))
    model.body_mass[handles.body_id] = handles.base_mass * density
    model.body_inertia[handles.body_id] = handles.base_inertia * density

    # com_x/com_y are in units of the object's own bounding radius (RIGID_PARAM_NAMES'
    # contract, mirrored from physics/types.py's RIGID_PARAMS docstring) -- see the
    # module docstring's note on this being a documented approximation (no
    # parallel-axis re-derivation of the inertia tensor for the shifted CoM).
    com_offset = np.array(
        [theta["com_x"] * handles.bounding_radius, theta["com_y"] * handles.bounding_radius, 0.0]
    )
    model.body_ipos[handles.body_id] = handles.base_ipos + com_offset

    model.geom_friction[handles.geom_id] = [
        float(np.exp(theta["log_mu_slide"])),
        float(np.exp(theta["log_mu_torsion"])),
        _DEFAULT_ROLL_FRICTION,
    ]
    damping = float(np.exp(theta["log_solref_damping"]))
    model.geom_solref[handles.geom_id] = [_SOLREF_TIMECONST, damping]

    mu_gripper = float(np.exp(theta["log_mu_gripper"]))
    for gid in finger_geom_ids:
        model.geom_friction[gid] = [mu_gripper, _DEFAULT_GRIPPER_TORSION, _DEFAULT_GRIPPER_ROLL]

    if handles.kind == "prismatic":
        model.dof_frictionloss[handles.dof_id] = float(np.exp(theta["log_joint_friction"]))
        model.dof_damping[handles.dof_id] = float(np.exp(theta["log_joint_damping"]))


def _apply_kinematic_friction(
    model: mujoco.MjModel, handles: _BodyHandles, theta: dict[str, float]
) -> None:
    """Give a kinematic (mocap-driven, "other tracked object") body a real friction.

    A contact needs friction numbers from BOTH surfaces to mean anything -- see the
    module docstring's "Priority, not friction mixing" note, which is exactly why this
    matters: MuJoCo still reads ``geom_friction`` off *this* geom even though the body
    of interest's own ``priority="1"`` geom is the one whose value ultimately wins the
    contact unmixed. Leaving this at the MJCF's compiled-in placeholder friction would
    silently make every brick-vs-drawer (or brick-vs-anything-else-tracked) contact
    behave as if the drawer were made of that placeholder material instead of
    participating in the identified physics at all.

    The honest choice with ONE episode's worth of evidence is to reuse the body of
    interest's OWN identified ``log_mu_slide``/``log_mu_torsion`` here, not to invent a
    second, per-pair friction parameter. A genuinely separate brick-vs-drawer friction
    is not identifiable from this data: the only observed trajectory in this SimSpec is
    the body of interest's, so there is no second, independent channel of evidence that
    could tell "the drawer's surface is slippery" apart from "the brick's own sliding
    friction with everything it touches is low" -- both explanations predict exactly the
    same brick trajectory. Adding a ``log_mu_slide_<other_label>`` per kinematic body
    would multiply the parameter count by however many objects the episode happened to
    track, purely to carry a value this stage's own likelihood could never separate from
    the one it already has.
    """
    model.geom_friction[handles.geom_id] = [
        float(np.exp(theta["log_mu_slide"])),
        float(np.exp(theta["log_mu_torsion"])),
        _DEFAULT_ROLL_FRICTION,
    ]


def _apply_generic_body_defaults(model: mujoco.MjModel, handles: _BodyHandles) -> None:
    """Physical parameters for any body beyond spec.bodies[0] -- see module docstring."""
    model.body_mass[handles.body_id] = handles.base_mass * _GENERIC_BODY_DENSITY
    model.body_inertia[handles.body_id] = handles.base_inertia * _GENERIC_BODY_DENSITY
    model.geom_friction[handles.geom_id] = list(_GENERIC_BODY_FRICTION)
    model.geom_solref[handles.geom_id] = list(_GENERIC_BODY_SOLREF)
    if handles.kind == "prismatic":
        model.dof_frictionloss[handles.dof_id] = 0.0
        model.dof_damping[handles.dof_id] = 0.1


def _extract_pose(data: mujoco.MjData, body_id: int) -> np.ndarray:
    """[x, y, z, qw, qx, qy, qz] -- `data.xpos`/`xquat`, not raw `qpos`.

    Works uniformly whether the body's own joint is a 7-dof freejoint or a
    1-dof slide joint (`qpos` for the latter is a scalar, not a pose) -- see
    the module docstring's "body of interest" note.
    """
    return np.concatenate([data.xpos[body_id], data.xquat[body_id]])


def _set_initial_state(
    model: mujoco.MjModel, data: mujoco.MjData, body: dict, handles: _BodyHandles
) -> None:
    if handles.kind == "kinematic":
        # A mocap body has no qpos/qvel at all -- its own DOF-free position is set
        # directly via data.mocap_pos/mocap_quat (see _set_body_mocap, called from
        # simulate_all right after this function for every body, kinematic or not,
        # frame 0's pose). Nothing to do here; init_lin_vel/init_ang_vel do not even
        # apply to a body with no velocity state.
        return
    lin_vel = np.asarray(body["init_lin_vel"], dtype=np.float64)
    ang_vel_world = np.asarray(body["init_ang_vel"], dtype=np.float64)
    if handles.kind == "free":
        jnt_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{handles.label}_free")
        qadr = int(model.jnt_qposadr[jnt_id])
        vadr = int(model.jnt_dofadr[jnt_id])
        # freejoint qpos is already init_pose_pos/quat (set at compile time via body
        # pos/quat) -- qpos stores the DEVIATION from that reference for a freejoint's
        # translational part too, but MuJoCo initialises qpos[0:3] to the body's
        # compiled pos, qpos[3:7] to its compiled quat, i.e. exactly init_pose already.
        # Only qvel needs to be set here.
        # MuJoCo freejoint qvel convention (verified empirically, see
        # scripts/_mujoco_settle_worker.py's docstring): linear in WORLD frame,
        # angular in the BODY-LOCAL frame -- rotate the world-frame FK estimate here.
        rot = np.asarray(body["init_pose"], dtype=np.float64)[:3, :3]
        ang_vel_body = rot.T @ ang_vel_world
        data.qvel[vadr : vadr + 3] = lin_vel
        data.qvel[vadr + 3 : vadr + 6] = ang_vel_body
        del qadr  # unused; kept for readability of the derivation above
    else:
        # Slide joint: qpos is a scalar deviation from the compiled reference frame
        # (= init_pose by construction, see _load_spec) -- simulation starts at the
        # observed pose exactly, i.e. qpos=0. qvel is the scalar rate along the axis.
        jnt_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{handles.label}_slide")
        qadr = int(model.jnt_qposadr[jnt_id])
        vadr = int(model.jnt_dofadr[jnt_id])
        axis_local = np.asarray(body["joint_axis_local"], dtype=np.float64)
        data.qpos[qadr] = 0.0
        data.qvel[vadr] = float(np.dot(lin_vel, axis_local))


def simulate_all(spec: dict, theta: np.ndarray, param_names: list[str]) -> dict:
    """Compile once, loop every particle, return the raw arrays that go into ``--out``.

    Returns a dict with ``poses`` (N, T, 7), ``ok`` (N,), and the timing breakdown
    keys documented in the module docstring / CLI usage.
    """
    t_build0 = time.perf_counter()
    xml = build_mjcf(spec)
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)

    bodies = spec["bodies"]
    handles = [_BodyHandles(model, b["label"], b["kind"]) for b in bodies]
    finger_geom_ids = []
    if spec.get("gripper") is not None:
        for name in ("finger_L_geom", "finger_R_geom"):
            gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, name)
            if gid < 0:
                raise WorkerError(f"MJCF compiled without expected gripper geom {name!r}")
            finger_geom_ids.append(gid)

    if spec.get("heightfield") is not None:
        # Structural, particle-independent data -- filled ONCE here, exactly like the mesh
        # assets MJCF's own `file=` attribute already loaded at compile time. There is no
        # `<hfield file=...>` equivalent for a raw float array without a lossy 8-bit PNG
        # round-trip (see fpgm.physics.types.HeightfieldSpec's docstring for why that was
        # rejected), so this is the one place the normalised grid actually reaches the
        # compiled model.
        hf_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, "terrain")
        if hf_id < 0:
            raise WorkerError("MJCF compiled without the expected hfield asset 'terrain'")
        adr = int(model.hfield_adr[hf_id])
        n = int(model.hfield_nrow[hf_id]) * int(model.hfield_ncol[hf_id])
        normalized01 = spec["heightfield"]["_mujoco"]["normalized01"]
        if normalized01.size != n:
            raise WorkerError(
                f"heightfield grid has {normalized01.size} cells, compiled hfield 'terrain' "
                f"has {n} (nrow={model.hfield_nrow[hf_id]}, ncol={model.hfield_ncol[hf_id]})"
            )
        model.hfield_data[adr : adr + n] = normalized01.reshape(-1)
    build_seconds = time.perf_counter() - t_build0

    n_frames = int(spec["n_frames"])
    substeps = int(spec["substeps"])
    finger_poses = (
        np.asarray(spec["gripper"]["finger_poses"], dtype=np.float64)
        if spec.get("gripper") is not None
        else None
    )
    # Every kinematic body's own (T, 4, 4) driving track, paired with its resolved
    # mocap_id -- computed once here (particle-independent, exactly like finger_poses
    # above), not inside the per-particle loop below.
    kinematic = [
        (h, np.asarray(b["pose_track"], dtype=np.float64))
        for h, b in zip(handles, bodies, strict=True)
        if h.kind == "kinematic"
    ]
    safety_z = _safety_floor_z(spec)

    n_particles = theta.shape[0]
    poses = np.full((n_particles, n_frames, 7), np.nan, dtype=np.float64)
    ok = np.zeros(n_particles, dtype=bool)

    t_step0 = time.perf_counter()
    for i in range(n_particles):
        row = dict(zip(param_names, theta[i].tolist(), strict=True))
        _apply_theta(model, handles[0], row, finger_geom_ids)
        for extra in handles[1:]:
            if extra.kind == "kinematic":
                _apply_kinematic_friction(model, extra, row)
            else:
                _apply_generic_body_defaults(model, extra)
        # mj_setConst refreshes the compile-time-cached constants that a raw
        # `body_mass`/`body_inertia` write does NOT update on its own --
        # `body_invweight0` chief among them, the reference inverse-mass the
        # constraint solver uses to weight contact impedance. Measured directly
        # while building this worker: without this call, a body compiled at
        # `density="1"` (see _BodyHandles' base_mass/base_inertia, "compile once,
        # scale per particle") then scaled up to a realistic mass produces
        # contacts that are detected (ncon>0) but far too SOFT -- the object
        # visibly sinks through a support box it should rest on, because the
        # solver is still weighting the contact as if the body were still
        # density=1 (a ~1e4x lighter, unrealistically-easy-to-push-through
        # object). With the call, the same edit reproduces a direct
        # `density="<value>"` compile's settle trajectory exactly. This is not
        # a recompilation (no XML reparsing, no mesh reload) -- it is the one
        # extra O(nv) call needed to make a mass/inertia field edit fully take
        # effect, and is paid once per particle, not per step.
        mujoco.mj_setConst(model, data)

        mujoco.mj_resetData(model, data)
        for h, b in zip(handles, bodies, strict=True):
            _set_initial_state(model, data, b, h)
        if finger_poses is not None:
            _set_finger_mocap(model, data, finger_poses[0])
        for kh, ktrack in kinematic:
            _set_body_mocap(data, kh.mocap_id, ktrack[0])

        mujoco.mj_forward(model, data)
        diverged = not np.all(np.isfinite(data.qpos)) or not np.all(np.isfinite(data.qvel))
        if not diverged:
            poses[i, 0] = _extract_pose(data, handles[0].body_id)

        if not diverged:
            for t in range(1, n_frames):
                for k in range(1, substeps + 1):
                    frac = k / substeps
                    if finger_poses is not None:
                        _set_finger_mocap_interp(
                            model, data, finger_poses[t - 1], finger_poses[t], frac
                        )
                    for kh, ktrack in kinematic:
                        _set_body_mocap_interp(data, kh.mocap_id, ktrack[t - 1], ktrack[t], frac)
                    mujoco.mj_step(model, data)
                if not np.all(np.isfinite(data.qpos)) or not np.all(np.isfinite(data.qvel)):
                    diverged = True
                    break
                if float(data.xpos[handles[0].body_id, 2]) < safety_z:
                    diverged = True
                    break
                poses[i, t] = _extract_pose(data, handles[0].body_id)
        ok[i] = not diverged
    sim_seconds = time.perf_counter() - t_step0

    return {
        "poses": poses,
        "ok": ok,
        "sim_seconds": sim_seconds,
        "build_seconds": build_seconds,
        "n_particles": n_particles,
    }


def _safety_floor_z(spec: dict) -> float:
    """See :data:`_SAFETY_FLOOR_DROP_M` -- this is ``build_mjcf``'s own ``safety_z``
    computation, duplicated here (not shared) because :func:`simulate_all` needs the
    number, not the XML string ``build_mjcf`` builds it into.
    """
    if spec.get("heightfield") is not None:
        pos_z = spec["heightfield"]["_mujoco"]["pos"][2]
        return float(pos_z) - _SAFETY_FLOOR_DROP_M
    return min(float(b["init_pose_pos"][2]) for b in spec["bodies"]) - _SAFETY_FLOOR_DROP_M


def _set_finger_mocap(model: mujoco.MjModel, data: mujoco.MjData, frame: np.ndarray) -> None:
    left_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "finger_L")
    right_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "finger_R")
    for body_id, pose in ((left_id, frame[0]), (right_id, frame[1])):
        mocap_id = int(model.body_mocapid[body_id])
        data.mocap_pos[mocap_id] = pose[:3, 3]
        data.mocap_quat[mocap_id] = _mat_to_quat(pose[:3, :3])


def _set_finger_mocap_interp(
    model: mujoco.MjModel, data: mujoco.MjData, frame_a: np.ndarray, frame_b: np.ndarray, t: float
) -> None:
    left_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "finger_L")
    right_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "finger_R")
    pairs = ((left_id, frame_a[0], frame_b[0]), (right_id, frame_a[1], frame_b[1]))
    for body_id, pose_a, pose_b in pairs:
        mocap_id = int(model.body_mocapid[body_id])
        data.mocap_pos[mocap_id] = _lerp(pose_a[:3, 3], pose_b[:3, 3], t)
        quat_a, quat_b = _mat_to_quat(pose_a[:3, :3]), _mat_to_quat(pose_b[:3, :3])
        data.mocap_quat[mocap_id] = _slerp(quat_a, quat_b, t)


def _set_body_mocap(data: mujoco.MjData, mocap_id: int, pose: np.ndarray) -> None:
    """Like ``_set_finger_mocap``, but for one already-resolved ``mocap_id`` (see
    ``_BodyHandles.mocap_id``) rather than re-resolving a fixed finger name -- a
    ``"kinematic"`` body's name is data-dependent (whichever other label S6 tracked),
    so there is no fixed name to look up the way the two fingers have.
    """
    data.mocap_pos[mocap_id] = pose[:3, 3]
    data.mocap_quat[mocap_id] = _mat_to_quat(pose[:3, :3])


def _set_body_mocap_interp(
    data: mujoco.MjData, mocap_id: int, pose_a: np.ndarray, pose_b: np.ndarray, t: float
) -> None:
    """Substep interpolation for a kinematic body -- same ``_lerp``/``_slerp`` pair
    ``_set_finger_mocap_interp`` uses, same rationale (module docstring's "why mocap
    fingers" note): holding a tracked object's pose constant for a whole output frame
    would introduce a full frame's worth of positional lag on a fast-moving object
    (e.g. a drawer mid-slide), exactly as it would for a fast-moving finger.
    """
    data.mocap_pos[mocap_id] = _lerp(pose_a[:3, 3], pose_b[:3, 3], t)
    quat_a, quat_b = _mat_to_quat(pose_a[:3, :3]), _mat_to_quat(pose_b[:3, :3])
    data.mocap_quat[mocap_id] = _slerp(quat_a, quat_b, t)


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


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


def main() -> int:
    args = parse_args()
    t_io0 = time.perf_counter()
    spec = _load_spec(args.spec)
    theta, param_names = _load_particles(args.particles)
    io_read_seconds = time.perf_counter() - t_io0

    result = simulate_all(spec, theta, param_names)

    t_io1 = time.perf_counter()
    np.savez(
        args.out,
        poses=result["poses"],
        ok=result["ok"],
        sim_seconds=result["sim_seconds"],
        n_substeps=int(spec["substeps"]),
        label=spec["bodies"][0]["label"],
        build_seconds=result["build_seconds"],
        io_seconds=io_read_seconds + (time.perf_counter() - t_io1),
        n_particles=result["n_particles"],
        mujoco_version=mujoco.__version__,
    )
    print(
        f"worker_mujoco: {result['n_particles']} particles x {spec['n_frames']} frames, "
        f"ok={int(result['ok'].sum())}/{result['n_particles']}, "
        f"build={result['build_seconds']:.3f}s step={result['sim_seconds']:.3f}s "
        f"({1000.0 * result['sim_seconds'] / max(result['n_particles'], 1):.3f} ms/particle, "
        f"mujoco {mujoco.__version__})"
    )
    return 0


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