mujoco-parity

# MuJoCo Parity: Comparing CPU (C) vs Accelerated (MJWarp / MJX) Physics > [!TIP] > > **Official MuJoCo documentation (`doc/*.rst`):** Read short chapters like > [`../../mjwarp/index.rst`](../../mjwarp/index.rst) (`.. _mjwDiff:`, > `.. _mjwOverflow:`, `.. _mjwPairwise:`, `.. _mjwBoxBox:`), > [`../../mjx.rst`](../../mjx.rst) (`.. _MjxFeatureParity:`), and > [`../../programming/simulation.rst`](../../programming/simulation.rst) > (`.. _siIntegrationState:`, `.. _siWarmstart:`) directly when in a MuJoCo > checkout; for large files like > [`../../computation/index.rst`](../../computation/index.rst) (`.. _Pipeline:`, > `.. _piStages:`) and > [`../../APIreference/functions.rst`](../../APIreference/functions.rst), grep > the literal `.. _Label:` anchor. Outside a checkout, the same `doc/*.rst` tree > ships offline in `pip install mujoco` at > `pathlib.Path(mujoco.__file__).parent / 'doc'`; otherwise fetch > `https://mujoco.readthedocs.io/en/latest/<path>.html` (or `/en/stable/` for > released pages; `functions.rst` renders inside > `APIreference/APIfunctions.html`, and Sphinx anchors lowercase the label, > e.g. `.. _mjwDiff:` → `#mjwdiff`). > > **Related skills:** > > - [mujoco-accelerated](../accelerated/SKILL.md) — GPU/TPU batch simulation > with MJWarp and MJX. > - [mujoco-python](../python/SKILL.md) — CPU simulation lifecycle, state > access, and stepping. > [!IMPORTANT] > > By default, MuJoCo (C) executes in `float64` with dynamic buffer allocation, > whereas MJWarp and MJX execute in `float32` with pre-allocated > contact/constraint buffers and GPU-parallel collision and solver kernels. > Comparing raw multi-step trajectories without accounting for `float32` > conditioning, contact ordering, and buffer capacities leads to false > positives and misdiagnosed root causes. For direct `float32`-to-`float32` > comparison against MJWarp/MJX (or to isolate pure precision effects by > comparing `float32` CPU against `float64` CPU), build MuJoCo C and its Python > bindings in single precision (`mjUSESINGLE`) with > `--@mujoco//:single_precision` (Bazel) or > `-DCMAKE_C_FLAGS="-DmjUSESINGLE" -DCMAKE_CXX_FLAGS="-DmjUSESINGLE"` (CMake), > and verify `mujoco.MJTNUM_DTYPE == np.float32`. ## 1. The Three-Level Parity Hierarchy Contact-rich rigid-body dynamics are chaotic: a $10^{-7}$ `float32` difference that flips a single grazing contact (`dist ≈ includemargin`) or convex collision witness point on step $k$ causes two trajectories to diverge exponentially thereafter, even when both engines are bug-free. Always diagnose parity using three levels of increasing scope: | Level | Method | What It Isolates | When to Use | | :--- | :--- | :--- | :--- | | **1. Single-State Forward** (`mj_forward` vs `mjw.forward`) | Evaluate both engines on identical states batched as worlds (`nworld = len(states)`) and compare intermediate `Data` fields in pipeline order | Pinpoints the exact pipeline stage (kinematics, tendons, mass matrix, actuation, smooth dynamics, collision, or constraint solver) and model element where a discrepancy originates | Primary unit test and root-cause diagnostic tool | | **2. Batched Single-Step Trajectory** (`mj_step` vs `mjw.step`) | Record a reference CPU trajectory of $T$ steps, then take **one** MJWarp step from all $T$ pre-step states in parallel (`nworld = T`) | Prevents error accumulation over time; a spike at step $t$ points directly to a dynamics or contact discrepancy at state $t$ rather than earlier drift | Finding state-dependent bugs (e.g., specific contact configurations) along realistic rollouts | | **3. Open-Loop Rollout with Float32 Baseline** | Step both engines from $t=0$ with identical controls alongside a **float32-perturbed CPU shadow rollout** | Distinguishes genuine parity bugs (MJWarp diverges before the `float32` CPU baseline on unambiguous steps) from benign chaotic divergence | Evaluating end-to-end rollout horizon fidelity for RL policies or controllers | --- ## 2. Pre-Flight Buffer Sizing, Overflow Detection & State Sync ### 1. Size `nconmax` and `njmax` for $2\times$ Peak CPU Counts When MJWarp exceeds `nconmax`/`naconmax` or `njmax`, excess contacts or constraint rows are dropped, which looks identical to missing collision contacts or solver failure. Size capacities to **at least twice** the maximum CPU counts (`d.ncon`, `d.nefc`) across the evaluated states so extra contacts generated by MJWarp are caught as contact mismatches rather than clipped by buffer overflow: ```python from typing import Optional, Sequence import mujoco import mujoco_warp as mjw import numpy as np import warp as wp MIN_NCONMAX = 64 MIN_NJMAX = 256 def capacities( mjds: Sequence[mujoco.MjData], nconmax: Optional[int] = None, njmax: Optional[int] = None, ) -> tuple[int, int]: if nconmax is None: nconmax = max(MIN_NCONMAX, 2 * max(d.ncon for d in mjds)) if njmax is None: njmax = max(MIN_NJMAX, 2 * max(d.nefc for d in mjds)) return nconmax, njmax def check_overflow(d: mjw.Data) -> None: nacon = int(d.nacon.numpy()[0]) if nacon > d.naconmax: raise RuntimeError( f"MJWarp contact overflow ({nacon} > naconmax {d.naconmax}); " "increase nconmax." ) nefc = int(d.nefc.numpy().max()) if nefc > d.njmax: raise RuntimeError( f"MJWarp constraint overflow ({nefc} > njmax {d.njmax}); " "increase njmax." ) buf_mask = ~(mjw.OverflowType.ITERATIONS | mjw.OverflowType.LS_ITERATIONS) overflow = int(np.bitwise_or.reduce(d.overflow.numpy()) & buf_mask) if overflow: raise RuntimeError( f"MJWarp buffer overflow: {mjw.OverflowType(overflow)!r}" ) ``` ### 2. Synchronize the Full Simulation State (`set_states`) Copying only `qpos` and `qvel` leaves actuator activations, mocap poses, applied wrenches, equality constraint toggles, delay history buffers, and solver warmstarts out of sync. In particular, unless `mjDSBL_WARMSTART` is set, MJWarp initializes its solver directly from `qacc_warmstart` (whereas CPU MuJoCo compares `qacc_warmstart` against `qacc_smooth`). Use `mujoco.mj_getState` and `mjw.set_state` with `mjSTATE_INTEGRATION` to copy the complete integration state across worlds: ```python def copy_data(mjm: mujoco.MjModel, src: mujoco.MjData) -> mujoco.MjData: dst = mujoco.MjData(mjm) mujoco.mj_copyData(dst, mjm, src) return dst def set_states( mjm: mujoco.MjModel, m: mjw.Model, d: mjw.Data, states: Sequence[mujoco.MjData], ) -> None: """Copies the full integration state of each CPU MjData into the corresponding MJWarp world.""" if len(states) != d.nworld: raise ValueError(f"Got {len(states)} states for {d.nworld} worlds.") sig = int(mujoco.mjtState.mjSTATE_INTEGRATION) size = mujoco.mj_stateSize(mjm, sig) mj_states = np.empty((len(states), size)) for i, s in enumerate(states): mujoco.mj_getState(mjm, s, mj_states[i], sig) mjw.set_state(m, d, wp.array(mj_states, dtype=float), sig) def put_batched( mjm: mujoco.MjModel, mjds: Sequence[mujoco.MjData], nconmax: Optional[int] = None, njmax: Optional[int] = None, ) -> tuple[mjw.Model, mjw.Data]: nconmax, njmax = capacities(mjds, nconmax, njmax) m = mjw.put_model(mjm) d = mjw.put_data(mjm, mjds[0], nworld=len(mjds), nconmax=nconmax, njmax=njmax) set_states(mjm, m, d, mjds) return m, d ``` --- ## 3. Stage-by-Stage `forward` Isolation & Conditioning Bounds ### 1. Pipeline Stage Order & Baseline Tolerances Each stage of `mj_forward` feeds the next. Always inspect errors in topological pipeline order—the **first** stage with normalized error $> 1$ is the root cause: | Pipeline Stage | `Data` Fields | Default `atol` | Default `rtol` | Notes | | :--- | :--- | :--- | :--- | :--- | | **1. Kinematics** | `xpos`, `xquat`, `site_xpos`, `subtree_com` | `1e-5` | `1e-4` | Forward kinematics (`mj_kinematics`, `mj_comPos`) | | **2. Tendons & Lengths** | `ten_length`, `actuator_length` | `1e-5` | `1e-4` | Spatial/fixed tendon & transmission lengths | | **3. Mass Matrix** | `M` | `1e-8` | `1e-5` | CSR lower triangle (`mjm.M_rownnz`, `mjm.M_colind`); add CRB scale bound | | **4. Velocity & Forces** | `actuator_force`, `qfrc_actuator`, `qfrc_bias`, `qfrc_passive` | `1e-3` | `1e-3` | RNE bias, passive springs/dampers, affine actuators; add stiff-actuator bound | | **5. Smooth Dynamics** | `qfrc_smooth`, `qacc_smooth` | `1e-3`, `1e-2` | `1e-3` | Unconstrained net force and $M^{-1} q_{\text{frc,smooth}}$; compare on actuation-conditioned CPU `MjData` | | **6. Constraint Dynamics** | `qfrc_constraint`, `qacc` | `1e-1` | `1e-2` | Iterative solver output (CG/Newton stop at `opt.tolerance`); compare on actuation-conditioned CPU `MjData` | Normalized error is defined as: ```text err = |mujoco - warp| / (atol + rtol * |mujoco| + extra_atol) ``` Values $\le 1.0$ are within tolerance. ### 2. Gotcha 1: Stiff Affine Actuators & Actuation Conditioning High-gain position or tendon servos (e.g. $k_p \sim 10^4\text{--}10^6$) compute force as a difference of large terms: $f = (g_0 + g_1 L + g_2 \dot{L}) u + (b_0 + b_1 L + b_2 \dot{L})$. A `float32` rounding of $L$ (`actuator_length`) or $\dot{L}$ (`actuator_velocity`) is multiplied by $g$ and $b$, producing legitimate force differences of size $O(\epsilon_{32} (|g| |u| + |b|))$ that exceed `1e-3`. When these forces act on light DOFs (such as fingers), they cause large downstream differences in `qacc_smooth`, `qfrc_constraint`, and `qacc` that masquerade as mass-matrix or constraint-solver bugs. **Two-part fix:** 1. Add a componentwise `float32` rounding bound (`actuator_rounding_atol`) to `actuator_force` and `qfrc_actuator`. 2. **Condition the CPU reference on MJWarp's actuator forces** when checking downstream fields (`qfrc_smooth`, `qacc_smooth`, `qfrc_constraint`, `qacc`) by adding `warp_qfrc_actuator - mjd.qfrc_actuator` into `mjd.qfrc_applied` and re-running `mj_forward`: ```python FLOAT32_ROUNDING_FACTOR = 8.0 ACTUATION_DOWNSTREAM_FIELDS = frozenset( {"qfrc_smooth", "qacc_smooth", "qfrc_constraint", "qacc"} ) def actuator_rounding_atol( mjm: mujoco.MjModel, mjd: mujoco.MjData ) -> dict[str, np.ndarray]: """Returns float32 rounding bounds for actuator_force and qfrc_actuator.""" if not mjm.nu or (mjm.opt.disableflags & mujoco.mjtDisableBit.mjDSBL_ACTUATION): return {} eps = float(np.finfo(np.float32).eps) length, velocity = mjd.actuator_length, mjd.actuator_velocity gain, bias = mjm.actuator_gainprm, mjm.actuator_biasprm gaintype, biastype = mjm.actuator_gaintype, mjm.actuator_biastype ctrl = mjd.ctrl.copy() if not (mjm.opt.disableflags & mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL): limited = mjm.actuator_ctrllimited.astype(bool) ctrl[limited] = np.clip( ctrl[limited], mjm.actuator_ctrlrange[limited, 0], mjm.actuator_ctrlrange[limited, 1], ) act_input = np.abs(ctrl) stateful = mjm.actuator_actnum > 0 if np.any(stateful): act_ind = mjm.actuator_actadr[stateful] + mjm.actuator_actnum[stateful] - 1 act = mjd.act[act_ind].copy() act_lim = mjm.actuator_actlimited[stateful].astype(bool) act[act_lim] = np.clip( act[act_lim], mjm.actuator_actrange[stateful][act_lim, 0], mjm.actuator_actrange[stateful][act_lim, 1], ) act_input[stateful] = np.abs(act) gain_terms = np.abs(gain[:, 0]) + np.where( gaintype == mujoco.mjtGain.mjGAIN_AFFINE, np.abs(gain[:, 1] * length) + np.abs(gain[:, 2] * velocity), 0.0, ) bias_terms = np.where(