Add support for actuatorfrc and jointactfrc sensors in MJX.

PiperOrigin-RevId: 665804850
Change-Id: I99371fc4da055382f1564b2c67d25f7b30a0a0cf

doc/changelog.rst

@@ -21,21 +21,22 @@ MJX
    ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``,
    ``CLOCK``.
 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``.
-8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
+8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``.
-9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
+9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
-10. Added support for :ref:`implicitfast integration<geIntegration>` for all cases except
+10. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
+11. Added support for :ref:`implicitfast integration<geIntegration>` for all cases except
     :doc:`fluid drag <computation/fluid>`.
 
 Bug fixes
 ^^^^^^^^^
-11. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`,
+12. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`,
     contribution by :github:user:`michael-ahn`).
-12. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit
+13. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit
     integrators, wrong derivatives would be computed.
 
 Python bindings
 ^^^^^^^^^^^^^^^
-13. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`).
+14. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`).
 
 
 Version 3.2.2 (Aug 8, 2024)

mjx/mujoco/mjx/_src/sensor.py

@@ -263,4 +263,30 @@ def sensor_acc(m: Model, d: Data) -> Data:
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
 
-  return d
+  stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC
+  sensors, adrs = [], []
+
+  for sensor_type in set(m.sensor_type[stage_acc]):
+    idx = m.sensor_type == sensor_type
+    objid = m.sensor_objid[idx]
+    adr = m.sensor_adr[idx]
+
+    if sensor_type == SensorType.ACTUATORFRC:
+      sensor = d.actuator_force[objid]
+    elif sensor_type == SensorType.JOINTACTFRC:
+      sensor = d.qfrc_actuator[m.jnt_dofadr[objid]]
+    else:
+      # TODO(taylorhowell): raise error after adding sensor check to io.py
+      continue  # unsupported sensor type
+
+    sensors.append(sensor)
+    adrs.append(adr)
+
+  if not adrs:
+    return d
+
+  sensordata = d.sensordata.at[np.concatenate(adrs)].set(
+      jp.concatenate(sensors)
+  )
+
+  return d.replace(sensordata=sensordata)

mjx/mujoco/mjx/_src/types.py

@@ -309,6 +309,8 @@ class SensorType(enum.IntEnum):
     JOINTVEL: joint velocity
     ACTUATORVEL: actuator velocity
     BALLANGVEL: ball joint angular velocity
+    ACTUATORFRC: scalar actuator force
+    JOINTACTFRC: scalar actuator force, measured at the joint
   """
   MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION
@@ -325,6 +327,8 @@ class SensorType(enum.IntEnum):
   JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL
   ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL
   BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL
+  ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC
+  JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC
 
 
 class ObjType(PyTreeNode):

mjx/mujoco/mjx/test_data/sensor.xml

@@ -18,6 +18,8 @@
 -actuatorvel
 -ballangvel
 * acceleration/force-dependent sensors:
+-actuatorfrc
+-jointactfrc
 -->
 <mujoco model="sensor">
   <asset>
@@ -81,14 +83,17 @@
     <rangefinder name="rangefinder0" site="site_rangefinder0"/>
     <jointpos name="jointpos0" joint="hinge0"/>
     <jointvel name="jointvel0" joint="hinge0"/>
+    <actuatorfrc name="actuatorfrc0" actuator="motor0"/>
     <actuatorpos name="actuatorpos0" actuator="motor0"/>
     <actuatorvel name="actuatorvel0" actuator="motor0"/>
     <ballquat name="ballquat2" joint="ballquat2"/>
     <ballangvel name="ballangvel2" joint="ballquat2"/>
     <framexaxis name="framexaxis0" objtype="site" objname="site0"/>
+    <jointactuatorfrc name="jointactfrc1" joint="hinge1"/>
     <framezaxis name="framezaxis0" objtype="site" objname="site0"/>
     <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
     <framepos name="framepos0" objtype="site" objname="site0"/>
+    <actuatorfrc name="actuatorfrc1" actuator="motor1"/>
     <subtreecom name="subtreecom0" body="body0"/>
     <camprojection site="frontorigin" camera="fixedcamera"/>
     <magnetometer name="magnetometer1" site="site1"/>
@@ -97,6 +102,7 @@
     <actuatorpos name="actuatorpos1" actuator="motor1"/>
     <actuatorvel name="actuatorvel1" actuator="motor1"/>
     <framepos name="framepos1" objtype="site" objname="site1"/>
+    <jointactuatorfrc name="jointactfrc0" joint="hinge0"/>
     <framezaxis name="framezaxis1" objtype="site" objname="site1"/>
     <framexaxis name="framexaxis1" objtype="site" objname="site1"/>
     <ballquat name="ballquat3" joint="ballquat3"/>