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

Add support for indexing qfrc fields in MJX bind() method.

Browse Source 
PiperOrigin-RevId: 763415415
Change-Id: Idb567a1711679565444779fb1062d7b65434ef13
This commit is contained in:
Tom Erez
2025-05-26 07:02:42 -07:00
committed by Copybara-Service
parent 22bdf22cc7
commit 4a102c37bc
2 changed files with 14 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+8 -4
View File
@@ -433,11 +433,13 @@ class BindData(object):
return name return name
else: else:
raise AttributeError('ctrl is not available for this type') raise AttributeError('ctrl is not available for this type')
if name == 'qpos' or name == 'qvel' or name == 'qacc': if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
if self.prefix == 'jnt_': if self.prefix == 'jnt_':
return name return name
else: else:
raise AttributeError('qpos, qvel, qacc are not available for this type') raise AttributeError(
'qpos, qvel, qacc, qfrc are not available for this type'
)
else: else:
return self.prefix + name return self.prefix + name
@@ -451,7 +453,9 @@ class BindData(object):
return var[..., idx, :] return var[..., idx, :]
def __getattr__(self, name: str): def __getattr__(self, name: str):
if name in ('sensordata', 'qpos', 'qvel', 'qacc'): if name in ('sensordata', 'qpos', 'qvel', 'qacc') or (
name.startswith('qfrc_')
):
adr = num = 0 adr = num = 0
if name == 'sensordata': if name == 'sensordata':
adr = self.model.sensor_adr[self.id] adr = self.model.sensor_adr[self.id]
@@ -460,7 +464,7 @@ class BindData(object):
adr = self.model.jnt_qposadr[self.id] adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id] typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.qpos_width() for jt in JointType) num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
elif name == 'qvel' or name == 'qacc': elif name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
adr = self.model.jnt_dofadr[self.id] adr = self.model.jnt_dofadr[self.id]
typ = self.model.jnt_type[self.id] typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.dof_width() for jt in JointType) num = sum((typ == jt) * jt.dof_width() for jt in JointType)
mjx/mujoco/mjx/_src/support_test.py
+6 -1
View File
@@ -249,6 +249,11 @@ class SupportTest(parameterized.TestCase):
dx.bind(mx, s.joints[i]).qacc, dx.bind(mx, s.joints[i]).qacc,
d.qacc[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6 d.qacc[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6
) )
np.testing.assert_array_almost_equal(
dx.bind(mx, s.joints[i]).qfrc_actuator,
d.qfrc_actuator[m.jnt_dofadr[i] : m.jnt_dofadr[i] + dofnum[i]],
decimal=6,
)
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl) np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
for i in range(m.nu): for i in range(m.nu):
@@ -344,7 +349,7 @@ class SupportTest(parameterized.TestCase):
): ):
print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3])) print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
with self.assertRaisesRegex( with self.assertRaisesRegex(
AttributeError, 'qpos, qvel, qacc are not available for this type' AttributeError, 'qpos, qvel, qacc, qfrc are not available for this type'
): ):
print(dx.bind(mx, s.geoms).qpos) print(dx.bind(mx, s.geoms).qpos)