Add support for indexing qfrc fields in MJX bind() method.

PiperOrigin-RevId: 763415415
Change-Id: Idb567a1711679565444779fb1062d7b65434ef13

mjx/mujoco/mjx/_src/support.py

@@ -433,11 +433,13 @@ class BindData(object):
         return name
       else:
         raise AttributeError('ctrl is not available for this type')
-    if name == 'qpos' or name == 'qvel' or name == 'qacc':
+    if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
       if self.prefix == 'jnt_':
         return name
       else:
-        raise AttributeError('qpos, qvel, qacc are not available for this type')
+        raise AttributeError(
+            'qpos, qvel, qacc, qfrc are not available for this type'
+        )
     else:
       return self.prefix + name
 
@@ -451,7 +453,9 @@ class BindData(object):
     return var[..., idx, :]
 
   def __getattr__(self, name: str):
-    if name in ('sensordata', 'qpos', 'qvel', 'qacc'):
+    if name in ('sensordata', 'qpos', 'qvel', 'qacc') or (
+        name.startswith('qfrc_')
+    ):
       adr = num = 0
       if name == 'sensordata':
         adr = self.model.sensor_adr[self.id]
@@ -460,7 +464,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' or name == 'qacc':
+      elif name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
         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)

mjx/mujoco/mjx/_src/support_test.py

@@ -249,6 +249,11 @@ class SupportTest(parameterized.TestCase):
           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_almost_equal(
+          dx.bind(mx, s.joints[i]).qfrc_actuator,
+          d.qfrc_actuator[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):
@@ -344,7 +349,7 @@ class SupportTest(parameterized.TestCase):
     ):
       print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
     with self.assertRaisesRegex(
-        AttributeError, 'qpos, qvel, qacc are not available for this type'
+        AttributeError, 'qpos, qvel, qacc, qfrc are not available for this type'
     ):
       print(dx.bind(mx, s.geoms).qpos)