Set mocap_pos/quat in make_data. Fixes #2606.

PiperOrigin-RevId: 755903897
Change-Id: Ied04e9bd3febfe1bbbba24ed1f52b3c0c6e50f38

doc/changelog.rst

@@ -17,6 +17,8 @@ MJX
 ^^^
 1. Added inverse dynamics.
 2. Added tendon actuator force sensor.
+3. Fix :github:issue:`2606` such that ``make_data`` copies over ``mocap_pos`` and ``mocap_quat``
+   from ``body_pos`` and ``body_quat``.
 
 Version 3.3.1 (Apr 9, 2025)
 ----------------------------

mjx/mujoco/mjx/_src/io.py

@@ -401,6 +401,17 @@ def make_data(
       eq_active=m.eq_active0,
       **zero_fields,
   )
+
+  if m.nmocap:
+    # Set mocap_pos/quat = body_pos/quat for mocap bodies as done in C MuJoCo.
+    body_mask = m.body_mocapid >= 0
+    body_pos = m.body_pos[body_mask]
+    body_quat = m.body_quat[body_mask]
+    d = d.replace(
+        mocap_pos=body_pos[m.body_mocapid[body_mask]],
+        mocap_quat=body_quat[m.body_mocapid[body_mask]],
+    )
+
   d = jax.device_put(d, device=device)
 
   return d