diff --git a/doc/changelog.rst b/doc/changelog.rst
index 950a4262..d045ca61 100644
--- a/doc/changelog.rst
+++ b/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``.
diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py
index 05bd721d..7eb1b315 100644
--- a/mjx/mujoco/mjx/_src/sensor.py
+++ b/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
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index ffe1ad5a..def0899b 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/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
diff --git a/mjx/mujoco/mjx/test_data/sensor/sensor.xml b/mjx/mujoco/mjx/test_data/sensor/sensor.xml
index 350c7bd3..0c5ca76b 100644
--- a/mjx/mujoco/mjx/test_data/sensor/sensor.xml
+++ b/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 @@
+
@@ -114,15 +117,18 @@
+
+
+