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

PiperOrigin-RevId: 738546056
Change-Id: I26ae045be9e983cb0760132b9d47eac018ca33cd

Author: quagla

mjx/mujoco/mjx/_src/support.py

@@ -507,14 +507,19 @@ def set(self, name: str, value: jax.Array) -> Data:
       iter(value)
     except TypeError:
       value = [value]
-    if name == 'qpos':
-      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' 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)
+    if name in ('qpos', 'qvel', 'qacc'):
+      adr = num = 0
+      if name == 'qpos':
+        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' 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)
+      if not isinstance(self.id, list):
+        adr = [adr]
+        num = [num]
     elif isinstance(self.id, list):
       adr = self.id * dim
       num = [dim for _ in range(len(self.id))]

mjx/mujoco/mjx/_src/support_test.py

@@ -287,6 +287,10 @@ def test_bind(self):
     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)
+    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])
     np.testing.assert_array_almost_equal(
         dx7.bind(mx, s.joints).qvel, [2.0, -1.2, 0.5, 0.0, 0.3], decimal=6