Add velocimeter and gyro sensors to MJX.

PiperOrigin-RevId: 669337697
Change-Id: I5e7f245b96eb334ffddd01757a9f631d9ad862cd

mjx/mujoco/mjx/_src/sensor.py

@@ -266,7 +266,23 @@ def sensor_vel(m: Model, d: Data) -> Data:
     objid = m.sensor_objid[idx]
     adr = m.sensor_adr[idx]
 
-    if sensor_type == SensorType.JOINTVEL:
+    if sensor_type == SensorType.VELOCIMETER:
+      bodyid = m.site_bodyid[objid]
+      pos = d.site_xpos[objid]
+      rot = d.site_xmat[objid]
+      cvel = d.cvel[bodyid]
+      subtree_com = d.subtree_com[m.body_rootid[bodyid]]
+      sensor = jax.vmap(
+          lambda vec, dif, rot: rot.T @ (vec[3:] - jp.cross(dif, vec[:3]))
+      )(cvel, pos - subtree_com, rot).reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.GYRO:
+      bodyid = m.site_bodyid[objid]
+      rot = d.site_xmat[objid]
+      ang = d.cvel[bodyid, :3]
+      sensor = jax.vmap(lambda ang, rot: rot.T @ ang)(ang, rot).reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.JOINTVEL:
       sensor = d.qvel[m.jnt_dofadr[objid]]
     elif sensor_type == SensorType.ACTUATORVEL:
       sensor = d.actuator_velocity[objid]

mjx/mujoco/mjx/_src/types.py

@@ -307,6 +307,8 @@ class SensorType(enum.IntEnum):
     FRAMEQUAT: frame orientation, represented as quaternion
     SUBTREECOM: subtree centor of mass
     CLOCK: simulation time
+    VELOCIMETER: 3D linear velocity, in local frame
+    GYRO: 3D angular velocity, in local frame
     JOINTVEL: joint velocity
     ACTUATORVEL: actuator velocity
     BALLANGVEL: ball joint angular velocity
@@ -326,6 +328,8 @@ class SensorType(enum.IntEnum):
   FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT
   SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM
   CLOCK = mujoco.mjtSensor.mjSENS_CLOCK
+  VELOCIMETER = mujoco.mjtSensor.mjSENS_VELOCIMETER
+  GYRO = mujoco.mjtSensor.mjSENS_GYRO
   JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL
   ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL
   BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL

mjx/mujoco/mjx/test_data/sensor.xml

@@ -15,6 +15,8 @@
 -subtreecom
 -clock
 * velocity-dependent sensors:
+-velocimeter
+-gyro
 -jointvel
 -actuatorvel
 -ballangvel
@@ -82,6 +84,8 @@
 
   <sensor>
     <magnetometer name="magnetometer0" site="site0"/>
+    <velocimeter name="velocimeter0" site="site0"/>
+    <gyro name="gyro0" site="site0"/>
     <rangefinder name="rangefinder0" site="site_rangefinder0"/>
     <jointpos name="jointpos0" joint="hinge0"/>
     <jointvel name="jointvel0" joint="hinge0"/>
@@ -96,6 +100,8 @@
     <framequat name="framequat3" objtype="site" objname="site3"/>
     <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
     <framepos name="framepos0" objtype="site" objname="site0"/>
+    <velocimeter name="velocimeter1" site="site1"/>
+    <gyro name="gyro1" site="site1"/>
     <actuatorfrc name="actuatorfrc1" actuator="motor1"/>
     <subtreecom name="subtreecom0" body="body0"/>
     <camprojection site="frontorigin" camera="fixedcamera"/>