Add velocimeter and gyro sensors to MJX.

PiperOrigin-RevId: 669337697
Change-Id: I5e7f245b96eb334ffddd01757a9f631d9ad862cd

doc/changelog.rst

@@ -16,10 +16,6 @@ General
      the :ref:`mjOption` struct. This is because the new convex collision detection pipeline (see below) does not use
      the MPR algorithm. The semantics of these options remain identical.
 
-- Added a new way of defining :ref:`connect<equality-connect>` equality constraints, using two sites rather than bodies.
-  The new semantic is useful when the assumption that the constraint is satisfied in the base configuration does not
-  hold. In this case the sites will "snap together" at the beginning of the simulation. Additionally, changing the site
-  positions in ``mjModel.site_pos`` at runtime can be used to modify the constraint.
 - Added the :ref:`nativeccd<option-flag-nativeccd>` flag. When this flag is enabled, general convex collision
   detection is handled natively, as opposed to using `libccd <https://github.com/danfis/libccd>`__. This feature is in
   early stages of testing.
@@ -28,7 +24,7 @@ General
 - When :ref:`attaching<meAttachment>` sub-models, :ref:`keyframes<keyframe>` will now be correctly merged into the
   parent model, but only on the first attachment.
 - Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These
-  alignments are used for computation shortcuts in :ref:`mj_kinematics`.
+  alignments are used as shortcuts in :ref:`mj_kinematics`.
 
 MJX
 ^^^
@@ -36,13 +32,12 @@ MJX
 - Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``,
   ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``,
   ``SUBTREECOM``, ``CLOCK``.
-- Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``.
+- Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``.
 - 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``.
 - Added support for :ref:`implicitfast integration<geIntegration>` for all cases except
   :doc:`fluid drag <computation/fluid>`.
-- Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
 
 Bug fixes
 ^^^^^^^^^
@@ -50,8 +45,6 @@ Bug fixes
   contribution by :github:user:`michael-ahn`).
 - 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.
-- Fixed a bug in tendon wrapping around spheres. Before this fix, tendons that wrapped around spheres with an
-  externally-placed :ref:`sidesite<spatial-geom-sidesite>` could jump inside the sphere instead of wrapping around it.
 
 Python bindings
 ^^^^^^^^^^^^^^^

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"/>