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

Fix bind.set when using a single joint instead of a list.

Browse Source 
PiperOrigin-RevId: 738546056
Change-Id: I26ae045be9e983cb0760132b9d47eac018ca33cd
This commit is contained in:
Alessio Quaglino
2025-03-19 14:43:21 -07:00
committed by Copybara-Service
parent be64747166
commit 1f1d5287fb
2 changed files with 17 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+13 -8
View File
@@ -507,14 +507,19 @@ class BindData(object):
iter(value) iter(value)
except TypeError: except TypeError:
value = [value] value = [value]
if name == 'qpos': if name in ('qpos', 'qvel', 'qacc'):
adr = self.model.jnt_qposadr[self.id] adr = num = 0
typ = self.model.jnt_type[self.id] if name == 'qpos':
num = sum((typ == jt) * jt.qpos_width() for jt in JointType) adr = self.model.jnt_qposadr[self.id]
elif name == 'qvel' or name == 'qacc': typ = self.model.jnt_type[self.id]
adr = self.model.jnt_dofadr[self.id] num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
typ = self.model.jnt_type[self.id] elif name == 'qvel' or name == 'qacc':
num = sum((typ == jt) * jt.dof_width() for jt in JointType) 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)
if not isinstance(self.id, list):
adr = [adr]
num = [num]
elif isinstance(self.id, list): elif isinstance(self.id, list):
adr = self.id * dim adr = self.id * dim
num = [dim for _ in range(len(self.id))] num = [dim for _ in range(len(self.id))]
mjx/mujoco/mjx/_src/support_test.py
+4
View File
@@ -287,6 +287,10 @@ class SupportTest(parameterized.TestCase):
dx6 = dx.bind(mx, s.joints[::2]).set('qpos', [1, 0, 0, 0, 8]) 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_equal(dx6.bind(mx, s.joints).qpos, qpos_desired)
np.testing.assert_array_almost_equal(dx.bind(mx, s.joints).qpos, d.qpos) np.testing.assert_array_almost_equal(dx.bind(mx, s.joints).qpos, d.qpos)
dx6a = dx.bind(mx, s.joints[0]).set('qpos', qpos_desired[:4])
np.testing.assert_array_equal(
dx6a.bind(mx, s.joints[0]).qpos, qpos_desired[:4]
)
dx7 = dx.bind(mx, s.joints[::2]).set('qvel', [2.0, -1.2, 0.5, 0.3]) dx7 = dx.bind(mx, s.joints[::2]).set('qvel', [2.0, -1.2, 0.5, 0.3])
np.testing.assert_array_almost_equal( np.testing.assert_array_almost_equal(
dx7.bind(mx, s.joints).qvel, [2.0, -1.2, 0.5, 0.0, 0.3], decimal=6 dx7.bind(mx, s.joints).qvel, [2.0, -1.2, 0.5, 0.0, 0.3], decimal=6