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

Add handling for joint qacc in bind

Browse Source 
PiperOrigin-RevId: 728166579
Change-Id: Ib9ecfec23a3ae1ddba198670518c48dc0af49cd1
This commit is contained in:
Silvia Cruciani
2025-02-18 05:28:05 -08:00
committed by Copybara-Service
parent 442b968e22
commit 9fb4ec95f8
2 changed files with 27 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+5 -5
View File
@@ -445,16 +445,16 @@ class BindData(object):
return name
else:
raise AttributeError('ctrl is not available for this type')
if name == 'qpos' or name == 'qvel':
if name == 'qpos' or name == 'qvel' or name == 'qacc':
if self.prefix == 'jnt_':
return name
else:
raise AttributeError('qpos and qvel are not available for this type')
raise AttributeError('qpos, qvel, qacc are not available for this type')
else:
return self.prefix + name
def __getattr__(self, name: str):
if name in ('sensordata', 'qpos', 'qvel'):
if name in ('sensordata', 'qpos', 'qvel', 'qacc'):
adr = num = 0
if name == 'sensordata':
adr = self.model.sensor_adr[self.id]
@@ -463,7 +463,7 @@ class BindData(object):
adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
elif name == 'qvel':
elif name == 'qvel' or name == 'qacc':
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)
@@ -492,7 +492,7 @@ class BindData(object):
adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
elif name == 'qvel':
elif name == 'qvel' or name == 'qacc':
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)
mjx/mujoco/mjx/_src/support_test.py
+22 -4
View File
@@ -223,7 +223,9 @@ class SupportTest(parameterized.TestCase):
np.testing.assert_array_equal(mx.bind(s.joints).axis, m.jnt_axis)
np.testing.assert_array_equal(mx.bind(s.joints).qposadr, m.jnt_qposadr)
qposnum = [4, 1, 1]
np.testing.assert_array_equal(mx.bind(s.joints).dofadr, m.jnt_dofadr)
qposnum = [4, 1, 1] # one ball joint (4) and two slide joints (1)
dofnum = [3, 1, 1] # one ball joint (3) and two slide joints (1)
for i in range(m.njnt):
np.testing.assert_array_equal(m.bind(s.joints[i]).axis, m.jnt_axis[i, :])
np.testing.assert_array_equal(mx.bind(s.joints[i]).axis, m.jnt_axis[i, :])
@@ -231,6 +233,14 @@ class SupportTest(parameterized.TestCase):
dx.bind(mx, s.joints[i]).qpos,
d.qpos[m.jnt_qposadr[i]:m.jnt_qposadr[i] + qposnum[i]], decimal=6
)
np.testing.assert_array_almost_equal(
dx.bind(mx, s.joints[i]).qvel,
d.qvel[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6
)
np.testing.assert_array_almost_equal(
dx.bind(mx, s.joints[i]).qacc,
d.qacc[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)
for i in range(m.nu):
@@ -271,10 +281,18 @@ class SupportTest(parameterized.TestCase):
dx6 = dx.bind(mx, s.joints[::2]).set('qpos', [1, 0, 0, 0, 8])
np.testing.assert_array_equal(dx6.bind(mx, s.joints).qpos, qpos_desired)
np.testing.assert_array_almost_equal(dx.bind(mx, s.joints).qpos, d.qpos)
dx7 = dx.bind(mx, s.joints[::2]).set('qvel', [2.0, -1.2, 0.5, 0.3])
np.testing.assert_array_almost_equal(
dx7.bind(mx, s.joints).qvel, [2.0, -1.2, 0.5, 0.0, 0.3], decimal=6
)
dx8 = dx.bind(mx, s.joints[::2]).set('qacc', [3.0, -2.1, 0.6, 0.4])
np.testing.assert_array_almost_equal(
dx8.bind(mx, s.joints).qacc, [3.0, -2.1, 0.6, 0.0, 0.4], decimal=6
)
dx7 = dx.bind(mx, s.bodies[1]).set('xfrc_applied', [1, 2, 3, 4, 5, 6])
dx9 = dx.bind(mx, s.bodies[1]).set('xfrc_applied', [1, 2, 3, 4, 5, 6])
np.testing.assert_array_equal(
dx7.bind(mx, s.bodies[1]).xfrc_applied, [1, 2, 3, 4, 5, 6]
dx9.bind(mx, s.bodies[1]).xfrc_applied, [1, 2, 3, 4, 5, 6]
)
for body in s.bodies[:1] + s.bodies[2:]:
np.testing.assert_array_equal(
@@ -295,7 +313,7 @@ class SupportTest(parameterized.TestCase):
):
print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
with self.assertRaises(
AttributeError, msg='qpos and qvel are not available for this type'
AttributeError, msg='qpos, qvel, qacc are not available for this type'
):
print(dx.bind(mx, s.geoms).qpos)
with self.assertRaises(KeyError, msg='invalid name: invalid_actuator_name'):