Fix setting multidimensional arrays in MJX bind.

PiperOrigin-RevId: 726956380
Change-Id: I51ab91b73a3a34938ac7bcd03e9691fd634cb296

Author: quagla

mjx/mujoco/mjx/_src/support.py

@@ -476,6 +476,7 @@ def set(self, name: str, value: jax.Array) -> Data:
     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 @@ def set(self, name: str, value: jax.Array) -> Data:
       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])

mjx/mujoco/mjx/_src/support_test.py

@@ -207,6 +207,9 @@ def test_bind(self):
       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 @@ def test_bind(self):
     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)