diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index aea5bced..ab330067 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/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] diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 38392d38..05c1e5bd 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -173,6 +173,8 @@ class SupportTest(parameterized.TestCase): + + @@ -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 = """