From 9fb4ec95f8ea096e1272fa6db2c1a0ed9e6dbb97 Mon Sep 17 00:00:00 2001 From: Silvia Cruciani Date: Tue, 18 Feb 2025 05:28:05 -0800 Subject: [PATCH] Add handling for joint qacc in bind PiperOrigin-RevId: 728166579 Change-Id: Ib9ecfec23a3ae1ddba198670518c48dc0af49cd1 --- mjx/mujoco/mjx/_src/support.py | 10 +++++----- mjx/mujoco/mjx/_src/support_test.py | 26 ++++++++++++++++++++++---- 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 06fb4670..b43db087 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 34c5ba50..5fb1dcd1 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -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'):