chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Handle mocap body in bind.

Browse Source 
PiperOrigin-RevId: 759065524
Change-Id: Ic4cc737596af883f267b6e2c470ac4591ea1eead
This commit is contained in:
Alessio Quaglino
2025-05-15 03:09:30 -07:00
committed by Copybara-Service
parent e8c566bee1
commit 7ba507323d
2 changed files with 22 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+9 -1
View File
@@ -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
+13 -1
View File
@@ -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 = """