Add support for actuatorfrc and jointactfrc sensors in MJX.

PiperOrigin-RevId: 665804850
Change-Id: I99371fc4da055382f1564b2c67d25f7b30a0a0cf

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"/>