Add framelinacc and frameangacc sensors to MJX.

PiperOrigin-RevId: 675141513
Change-Id: Ie38dba3e8be82e0554c7c9e47c6d07f3cdad29b8

doc/changelog.rst

@@ -64,7 +64,7 @@ MJX
 14. Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``,
     ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``.
 15. Added acceleration/force-dependent sensors: ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``,
-    ``JOINTACTFRC``.
+    ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``.
 16. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
 17. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
 18. Added support for :ref:`implicitfast integration<geIntegration>` for all cases except

doc/mjx.rst

@@ -220,7 +220,8 @@ The following features are **fully supported** in MJX:
      - ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``,
        ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``, ``VELOCIMETER``,
        ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``,
-       ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``, ``JOINTACTFRC``.
+       ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``,
+       ``FRAMEANGACC``.
 
 The following features are **in development** and coming soon:
 

mjx/mujoco/mjx/_src/sensor.py

@@ -419,6 +419,16 @@ def sensor_acc(m: Model, d: Data) -> Data:
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
 
+  # position and bodyid by object type
+  objtype_data = {
+      ObjType.UNKNOWN: (np.zeros((1, 3)), np.arange(1)),
+      ObjType.BODY: (d.xipos, np.arange(m.nbody)),
+      ObjType.XBODY: (d.xpos, np.arange(m.nbody)),
+      ObjType.GEOM: (d.geom_xpos, m.geom_bodyid),
+      ObjType.SITE: (d.site_xpos, m.site_bodyid),
+      ObjType.CAMERA: (d.cam_xpos, m.cam_bodyid),
+  }
+
   stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC
   sensor_types = set(m.sensor_type[stage_acc])
 
@@ -478,6 +488,41 @@ def sensor_acc(m: Model, d: Data) -> Data:
       sensor = d.actuator_force[objid]
     elif sensor_type == SensorType.JOINTACTFRC:
       sensor = d.qfrc_actuator[m.jnt_dofadr[objid]]
+    elif sensor_type in (SensorType.FRAMELINACC, SensorType.FRAMEANGACC):
+      objtype = m.sensor_objtype[idx]
+
+      for ot in set(objtype):
+        idxt = objtype == ot
+        objidt = objid[idxt]
+        pos, bodyid = objtype_data[ot]
+        pos = pos[objidt]
+        bodyid = bodyid[objidt]
+        cacc = d.cacc[bodyid]
+
+        if sensor_type == SensorType.FRAMELINACC:
+
+          @jax.vmap
+          def _framelinacc(cvel, cacc, offset):
+            ang = cvel[:3]
+            lin = cvel[3:] - jp.cross(offset, cvel[:3])
+            acc = cacc[3:] - jp.cross(offset, cacc[:3])
+            correction = jp.cross(ang, lin)
+            return acc + correction
+
+          cvel = d.cvel[bodyid]
+          offset = pos - d.subtree_com[m.body_rootid[bodyid]]
+
+          sensor = _framelinacc(cvel, cacc, offset).reshape(-1)
+        elif sensor_type == SensorType.FRAMEANGACC:
+          sensor = cacc[:, :3].reshape(-1)
+        else:
+          raise ValueError(f'Unknown sensor type: {sensor_type}')
+
+        adrt = adr[idxt, None] + np.arange(3)[None]
+
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adrt.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
     else:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type

mjx/mujoco/mjx/_src/types.py

@@ -325,6 +325,8 @@ class SensorType(enum.IntEnum):
     TORQUE: torque
     ACTUATORFRC: scalar actuator force
     JOINTACTFRC: scalar actuator force, measured at the joint
+    FRAMELINACC: 3D linear acceleration
+    FRAMEANGACC: 3D angular acceleration
   """
   MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION
@@ -353,6 +355,8 @@ class SensorType(enum.IntEnum):
   TORQUE = mujoco.mjtSensor.mjSENS_TORQUE
   ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC
   JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC
+  FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC
+  FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC
 
 
 class ObjType(PyTreeNode):

mjx/mujoco/mjx/test_data/sensor/sensor.xml

@@ -30,6 +30,8 @@
 -torque
 -actuatorfrc
 -jointactfrc
+-framelinacc
+-frameangacc
 -->
 <mujoco model="sensor">
   <asset>
@@ -144,6 +146,8 @@
     <camprojection site="frontorigin" camera="fixedcamera"/>
     <magnetometer name="magnetometer1" site="site1"/>
     <jointpos name="jointpos1" joint="hinge1"/>
+    <framelinacc name="framelinacc0" objtype="site" objname="site0"/>
+    <frameangacc name="frameangacc0" objtype="site" objname="site0"/>
     <jointvel name="jointvel1" joint="hinge1"/>
     <actuatorpos name="actuatorpos1" actuator="motor1"/>
     <subtreelinvel name="subtreelinvel1" body="body1"/>
@@ -160,9 +164,11 @@
     <frameangvel name="frameangvel0ref" objtype="site" objname="site0" reftype="site" refname="site3"/>
     <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
     <framequat name="framequat0" objtype="site" objname="site0"/>
+    <framelinacc name="framelinacc1" objtype="site" objname="site1"/>
     <rangefinder name="rangefinder1" site="site_rangefinder1"/>
     <subtreecom name="subtreecom1" body="body1"/>
     <accelerometer name="accelerometer0" site="site0"/>
+    <frameangacc name="frameangacc1" objtype="site" objname="site1"/>
     <framelinvel name="framelinvel0" objtype="site" objname="site0"/>
     <framelinvel name="framelinvel0ref" objtype="site" objname="site0" reftype="geom" refname="geom2"/>
     <framelinvel name="framelinvel0cutoff" objtype="site" objname="site0" cutoff="3e-4"/>