Add velocity-dependent sensors to MJX. Includes: joint velocity, actuator velocity, and ball joint angular velocity.

PiperOrigin-RevId: 664461284
Change-Id: I5063f662d707fbe342493f9cd7993f009fac9c76

doc/changelog.rst

@@ -20,19 +20,20 @@ MJX
 6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``,
    ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``,
    ``CLOCK``.
-7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
+7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``.
-8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
+8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
+9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
 
 Bug fixes
 ^^^^^^^^^
-9.  Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`,
+10.  Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`,
     contribution by :github:user:`michael-ahn`).
-10. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit
+11. 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
 ^^^^^^^^^^^^^^^
-11. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`).
+12. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`).
 
 
 Version 3.2.2 (Aug 8, 2024)

mjx/mujoco/mjx/_src/sensor.py

@@ -35,11 +35,6 @@ def sensor_pos(m: Model, d: Data) -> Data:
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
 
-  # no position-dependent sensors
-  stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
-  if sum(stage_pos) == 0:
-    return d
-
   # position and orientation by object type
   objtype_data = {
       ObjType.UNKNOWN: (
@@ -60,6 +55,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
       SensorType.FRAMEZAXIS: 2,
   }
 
+  stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
   sensors, adrs = [], []
 
   for sensor_type in set(m.sensor_type[stage_pos]):
@@ -206,6 +202,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.CLOCK:
       sensor = jp.repeat(d.time, sum(idx))
     else:
+      # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
     sensors.append(sensor)
@@ -227,7 +224,37 @@ def sensor_vel(m: Model, d: Data) -> Data:
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
 
-  return d
+  stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL
+  sensors, adrs = [], []
+
+  for sensor_type in set(m.sensor_type[stage_vel]):
+    idx = m.sensor_type == sensor_type
+    objid = m.sensor_objid[idx]
+    adr = m.sensor_adr[idx]
+
+    if sensor_type == SensorType.JOINTVEL:
+      sensor = d.qvel[m.jnt_dofadr[objid]]
+    elif sensor_type == SensorType.ACTUATORVEL:
+      sensor = d.actuator_velocity[objid]
+    elif sensor_type == SensorType.BALLANGVEL:
+      jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
+      sensor = d.qvel[jnt_dotadr].reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    else:
+      # TODO(taylorhowell): raise error after adding sensor check to io.py
+      continue  # unsupported sensor typ
+
+    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)
 
 
 def sensor_acc(m: Model, d: Data) -> Data:

mjx/mujoco/mjx/_src/types.py

@@ -304,6 +304,9 @@ class SensorType(enum.IntEnum):
     FRAMEZAXIS: frame z-axis
     SUBTREECOM: subtree centor of mass
     CLOCK: simulation time
+    JOINTVEL: joint velocity
+    ACTUATORVEL: actuator velocity
+    BALLANGVEL: ball joint angular velocity
   """
   MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION
@@ -317,6 +320,9 @@ class SensorType(enum.IntEnum):
   FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS
   SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM
   CLOCK = mujoco.mjtSensor.mjSENS_CLOCK
+  JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL
+  ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL
+  BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL
 
 
 class ObjType(PyTreeNode):

mjx/mujoco/mjx/test_data/sensor.xml

@@ -14,6 +14,9 @@
 -subtreecom
 -clock
 * velocity-dependent sensors:
+-jointvel
+-actuatorvel
+-ballangvel
 * acceleration/force-dependent sensors:
 -->
 <mujoco model="sensor">
@@ -77,8 +80,11 @@
     <magnetometer name="magnetometer0" site="site0"/>
     <rangefinder name="rangefinder0" site="site_rangefinder0"/>
     <jointpos name="jointpos0" joint="hinge0"/>
+    <jointvel name="jointvel0" joint="hinge0"/>
     <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"/>
     <framezaxis name="framezaxis0" objtype="site" objname="site0"/>
     <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
@@ -87,12 +93,15 @@
     <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"/>
+    <actuatorvel name="actuatorvel1" actuator="motor1"/>
     <framepos name="framepos1" objtype="site" objname="site1"/>
     <framezaxis name="framezaxis1" objtype="site" objname="site1"/>
     <framexaxis name="framexaxis1" objtype="site" objname="site1"/>
     <ballquat name="ballquat3" joint="ballquat3"/>
     <camprojection site="frontcenter" camera="fixedcamera"/>
+    <ballangvel name="ballangvel3" joint="ballquat3"/>
     <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
     <rangefinder name="rangefinder1" site="site_rangefinder1"/>
     <subtreecom name="subtreecom1" body="body1"/>