From 6f0bd0a5f8298590dbae5d01840016c8f688f6c2 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 14 Feb 2025 09:27:47 -0800 Subject: [PATCH] Fix setting multidimensional arrays in MJX bind. PiperOrigin-RevId: 726956380 Change-Id: I51ab91b73a3a34938ac7bcd03e9691fd634cb296 --- mjx/mujoco/mjx/_src/support.py | 5 +++-- mjx/mujoco/mjx/_src/support_test.py | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 40bec2a9..9c7c60cd 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -476,6 +476,7 @@ class BindData(object): if name == 'sensordata': raise AttributeError('sensordata is readonly') array = getattr(self.data, self.__getname(name)) + dim = 1 if len(array.shape) == 1 else array.shape[-1] try: iter(value) except TypeError: @@ -490,10 +491,10 @@ class BindData(object): num = sum((typ == jt) * jt.dof_width() for jt in JointType) elif isinstance(self.id, list): adr = self.id - num = [1 for _ in range(len(self.id))] + num = [dim for _ in range(len(self.id))] else: adr = [self.id] - num = [1] + num = [dim] i = 0 for a, n in zip(adr, num): array = array.at[a: a + n].set(value[i: i + n]) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 965f369c..582d1c51 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -207,6 +207,9 @@ class SupportTest(parameterized.TestCase): np.testing.assert_array_equal( dx.bind(mx, s.bodies[i]).xpos, d.xpos[i, :] ) + np.testing.assert_array_equal( + dx.bind(mx, s.bodies[i]).xfrc_applied, d.xfrc_applied[i, :] + ) np.testing.assert_array_equal(mx.bind(s.geoms).size, m.geom_size) np.testing.assert_array_equal(dx.bind(mx, s.geoms).xpos, d.geom_xpos) @@ -269,6 +272,11 @@ class SupportTest(parameterized.TestCase): 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.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] + ) + # test invalid name with self.assertRaises(AttributeError): print(dx.bind(mx, s.geoms).ctrl)