Make bodies below a mocap body follow the mocap pose in MJX - #3759

mjx.kinematics ignores the mocap pose for every body attached below a mocap body. The mocap body itself is placed correctly, but its descendants are placed as if the mocap body were still at its default pose, so their xpos, xquat, xmat, xanchor, xaxis, geom and site frames are all wrong, and so is everything computed from them afterwards (com, cdof, collisions, sensors).

The cause is the order in kinematics: the body tree scan runs from m.body_pos and m.body_quat, and mocap_pos / mocap_quat are written into the mocap bodies only after the scan, when the children have already been computed. MuJoCo C uses the mocap pose as the body pose in mj_kinematics1 before it reaches the children.

Repro

import jax
import mujoco
from mujoco import mjx

m = mujoco.MjModel.from_xml_string("""
<mujoco>
  <worldbody>
    <body mocap="true" pos="0 0 1">
      <geom size=".05"/>
      <body name="c" pos="0 0 -.2">
        <joint type="hinge" axis="0 1 0"/>
        <geom size=".05"/>
      </body>
    </body>
  </worldbody>
</mujoco>
""")
d = mujoco.MjData(m)
d.mocap_pos[0] = [1, 0, 1]
mujoco.mj_forward(m, d)
dx = jax.jit(mjx.forward)(mjx.put_model(m), mjx.put_data(m, d))
print(d.xpos[m.body('c').id])   # [1.  0.  0.8]
print(dx.xpos[m.body('c').id])  # [0.  0.  0.8] on main

Fix

Mocap bodies are static children of the world, so the mocap pose is their pose relative to the parent. This change substitutes mocap_pos and the normalized mocap_quat into the body_pos / body_quat inputs of the scan, so the children are computed from it. The write after the scan is no longer needed, since the scan now produces the mocap bodies' own frames directly.

Tests

test_mocap_children in smooth_test.py builds a model with two mocap bodies, one of them with a hinged child and a ball-jointed grandchild, sets random mocap poses (with unnormalized quaternions) and compares xanchor, xaxis, xpos, xquat, xmat, geom and site frames against mj_kinematics. On main it fails at xanchor with a max absolute difference of 0.94; with the fix it passes. test_smooth already randomizes the mocap pose, but pendula.xml has no bodies below its mocap body, which is why this was not caught.

smooth_test.py, forward_test.py and sensor_test.py pass locally (59 tests, JAX 0.11.1 on CPU, Linux, with the current nightly mujoco wheel).