Add subtreelinvel and subtreeangmom sensors to MJX.

PiperOrigin-RevId: 671347528
Change-Id: I8a32590c650ca6d0bb6df43d0347f707ffc05f12

doc/changelog.rst

@@ -53,7 +53,7 @@ MJX
   ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``,
   ``SUBTREECOM``, ``CLOCK``.
 - Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``,
-  ``FRAMELINVEL``, ``FRAMEANGVEL``.
+  ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``.
 - Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``.
 - Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
 - Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.

mjx/mujoco/mjx/_src/sensor.py

@@ -20,6 +20,7 @@ import mujoco
 # pylint: disable=g-importing-member
 from mujoco.mjx._src import math
 from mujoco.mjx._src import ray
+from mujoco.mjx._src import smooth
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
@@ -296,9 +297,13 @@ def sensor_vel(m: Model, d: Data) -> Data:
   }
 
   stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL
-  sensors, adrs = [], []
+  sensor_types = set(m.sensor_type[stage_vel])
 
-  for sensor_type in set(m.sensor_type[stage_vel]):
+  if sensor_types & {SensorType.SUBTREELINVEL, SensorType.SUBTREEANGMOM}:
+    d = smooth.subtree_vel(m, d)
+
+  sensors, adrs = [], []
+  for sensor_type in sensor_types:
     idx = m.sensor_type == sensor_type
     objid = m.sensor_objid[idx]
     adr = m.sensor_adr[idx]
@@ -329,7 +334,7 @@ def sensor_vel(m: Model, d: Data) -> Data:
       jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
       sensor = d.qvel[jnt_dotadr]
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
-    elif sensor_type in (SensorType.FRAMELINVEL, SensorType.FRAMEANGVEL):
+    elif sensor_type in {SensorType.FRAMELINVEL, SensorType.FRAMEANGVEL}:
       objtype = m.sensor_objtype[idx]
       reftype = m.sensor_reftype[idx]
       refid = m.sensor_refid[idx]
@@ -385,6 +390,12 @@ def sensor_vel(m: Model, d: Data) -> Data:
         sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type == SensorType.SUBTREELINVEL:
+      sensor = d.subtree_linvel[objid]
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.SUBTREEANGMOM:
+      sensor = d.subtree_angmom[objid]
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     else:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type

mjx/mujoco/mjx/_src/types.py

@@ -316,6 +316,8 @@ class SensorType(enum.IntEnum):
     BALLANGVEL: ball joint angular velocity
     FRAMELINVEL: 3D linear velocity
     FRAMEANGVEL: 3D angular velocity
+    SUBTREELINVEL: subtree linear velocity
+    SUBTREEANGMOM: subtree angular momentum
     ACTUATORFRC: scalar actuator force
     JOINTACTFRC: scalar actuator force, measured at the joint
   """
@@ -339,6 +341,8 @@ class SensorType(enum.IntEnum):
   BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL
   FRAMELINVEL = mujoco.mjtSensor.mjSENS_FRAMELINVEL
   FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL
+  SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL
+  SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM
   ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC
   JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC
 

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

@@ -22,6 +22,8 @@
 -ballangvel
 -framelinvel
 -frameangvel
+-subtreelinvel
+-subtreeangmom
 * acceleration/force-dependent sensors:
 -actuatorfrc
 -jointactfrc
@@ -98,6 +100,7 @@
     <actuatorfrc name="actuatorfrc0" actuator="motor0"/>
     <framelinvel name="framelinvel3" objtype="site" objname="site3"/>
     <framelinvel name="framelinvel3ref" objtype="site" objname="site3" reftype="body" refname="body2"/>
+    <subtreelinvel name="subtreelinvel0" body="body0"/>
     <actuatorpos name="actuatorpos0" actuator="motor0"/>
     <actuatorvel name="actuatorvel0" actuator="motor0"/>
     <ballquat name="ballquat2" joint="ballquat2"/>
@@ -114,15 +117,18 @@
     <frameangvel name="frameangvel3ref" objtype="site" objname="site3" reftype="body" refname="body2"/>
     <frameangvel name="frameangvel3cutoff" objtype="site" objname="site3" cutoff="1e-3"/>
     <actuatorfrc name="actuatorfrc1" actuator="motor1"/>
+    <subtreeangmom name="subtreeangmom0" body="body0"/>
     <subtreecom name="subtreecom0" body="body0"/>
     <camprojection site="frontorigin" camera="fixedcamera"/>
     <magnetometer name="magnetometer1" site="site1"/>
     <jointpos name="jointpos1" joint="hinge1"/>
     <jointvel name="jointvel1" joint="hinge1"/>
     <actuatorpos name="actuatorpos1" actuator="motor1"/>
+    <subtreelinvel name="subtreelinvel1" body="body1"/>
     <actuatorvel name="actuatorvel1" actuator="motor1"/>
     <framepos name="framepos1" objtype="site" objname="site1"/>
     <jointactuatorfrc name="jointactfrc0" joint="hinge0"/>
+    <subtreeangmom name="subtreeangmom1" body="body1"/>
     <framezaxis name="framezaxis1" objtype="site" objname="site1"/>
     <framexaxis name="framexaxis1" objtype="site" objname="site1"/>
     <ballquat name="ballquat3" joint="ballquat3"/>