Handle mocap body in bind.

PiperOrigin-RevId: 759065524
Change-Id: Ic4cc737596af883f267b6e2c470ac4591ea1eead

mjx/mujoco/mjx/_src/support.py

@@ -473,6 +473,8 @@ class BindData(object):
         return self._slice(self.__getname(name), slice(adr, adr + num))
       else:
         return self._slice(self.__getname(name), adr)
+    elif name in ('mocap_pos', 'mocap_quat'):
+      return self._slice(self.__getname(name), self.model.body_mocapid[self.id])
     return self._slice(self.__getname(name), self.id)
 
   def set(self, name: str, value: jax.Array) -> Data:
@@ -485,7 +487,7 @@ class BindData(object):
       iter(value)
     except TypeError:
       value = [value]
-    if name in ('qpos', 'qvel', 'qacc'):
+    if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
       adr = num = 0
       if name == 'qpos':
         adr = self.model.jnt_qposadr[self.id]
@@ -495,6 +497,12 @@ class BindData(object):
         adr = self.model.jnt_dofadr[self.id]
         typ = self.model.jnt_type[self.id]
         num = sum((typ == jt) * jt.dof_width() for jt in JointType)
+      elif name == 'mocap_pos':
+        adr = self.model.body_mocapid[self.id] * 3
+        num = np.ones_like(self.id, dtype=int) * 3
+      elif name == 'mocap_quat':
+        adr = self.model.body_mocapid[self.id] * 4
+        num = np.ones_like(self.id, dtype=int) * 4
       if not isinstance(self.id, list):
         adr = [adr]
         num = [num]

mjx/mujoco/mjx/_src/support_test.py

@@ -173,6 +173,8 @@ class SupportTest(parameterized.TestCase):
             <joint axis="0 0 1" type="slide" name="joint3"/>
             <geom size="7 8 9" type="box" name="geom3"/>
           </body>
+          <body pos="100 110 120" name="body4" mocap="true"/>
+          <body pos="130 140 150" name="body5" mocap="true"/>
         </worldbody>
 
         <actuator>
@@ -318,6 +320,16 @@ class SupportTest(parameterized.TestCase):
     np.testing.assert_array_equal(
         dx10.bind(mx, s.bodies[1]).xfrc_applied, [2, 2, 2, 2, 2, 2]
     )
+    dx11 = dx.bind(mx, s.bodies[-1]).set(
+        'mocap_pos',
+        [1, 2, 3],
+    )
+    np.testing.assert_array_equal(
+        dx11.bind(mx, s.bodies[-2]).mocap_pos, [100, 110, 120]
+    )
+    np.testing.assert_array_equal(
+        dx11.bind(mx, s.bodies[-1]).mocap_pos, [1, 2, 3]
+    )
 
     # test attribute and type mismatches
     with self.assertRaisesRegex(
@@ -359,7 +371,7 @@ class SupportTest(parameterized.TestCase):
     self.assertEqual(
         str(e.exception),
         'mjSpec signature does not match mjx.Model signature:'
-        ' 17856615236057737915 != 12517827274439268436',
+        ' 15297169659434471387 != 2785811613804955188',
     )
 
   _CONTACTS = """