Add position-dependent sensors to MJX: magnetometer, ballquat, subtreecom, framepos, framexaxis, frameyaxis, framezaxis, and clock.

PiperOrigin-RevId: 663668212
Change-Id: Ifc9adc6bff4544172ab22572da04a093aebfdbd4

doc/changelog.rst

@@ -14,6 +14,8 @@ General
 MJX
 ^^^
 4. Added ``efc_pos`` to ``mjx.Data``.
+5. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``,
+   ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``.
 
 Version 3.2.2 (Aug 8, 2024)
 ---------------------------

mjx/mujoco/mjx/_src/sensor.py

@@ -15,11 +15,14 @@
 """Sensor functions."""
 
 import jax
+from jax import numpy as jp
+import mujoco
 # pylint: disable=g-importing-member
+from mujoco.mjx._src import math
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import SensorType
-from typing import Tuple
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -27,19 +30,109 @@ import numpy as np
 def sensor_pos(m: Model, d: Data) -> Data:
   """Compute position-dependent sensors values."""
 
-  sensordata = d.sensordata
+  # no position-dependent sensors
-  if np.isin(SensorType.JOINTPOS, m.sensor_type):
+  if sum(m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS) == 0:
-    # jointpos
+    return d
-    i = m.sensor_type == SensorType.JOINTPOS
+
-    objid = m.sensor_objid[i]
+  # position and orientation by object type
-    adr = m.sensor_adr[i]
+  objtype_data = {
-    sensordata = sensordata.at[adr].set(d.qpos[m.jnt_qposadr[objid]])
+      ObjType.UNKNOWN: (
-  if np.isin(SensorType.ACTUATORPOS, m.sensor_type):
+          np.expand_dims(np.eye(3), axis=0),
-    # actuatorpos
+          np.zeros((1, 3)),
-    i = m.sensor_type == SensorType.ACTUATORPOS
+      ),  # world
-    objid = m.sensor_objid[i]
+      ObjType.BODY: (d.xipos, d.ximat),
-    adr = m.sensor_adr[i]
+      ObjType.XBODY: (d.xpos, d.xmat),
-    sensordata = sensordata.at[adr].set(d.actuator_length[objid])
+      ObjType.GEOM: (d.geom_xpos, d.geom_xmat),
+      ObjType.SITE: (d.site_xpos, d.site_xmat),
+      ObjType.CAMERA: (d.cam_xpos, d.cam_xmat),
+  }
+
+  # frame axis indexing
+  frame_axis = {
+      SensorType.FRAMEXAXIS: 0,
+      SensorType.FRAMEYAXIS: 1,
+      SensorType.FRAMEZAXIS: 2,
+  }
+
+  sensors, adrs = [], []
+
+  for sensor_type in set(m.sensor_type):
+    idx = m.sensor_type == sensor_type
+    objid = m.sensor_objid[idx]
+    adr = m.sensor_adr[idx]
+
+    if sensor_type == SensorType.MAGNETOMETER:
+      sensor = jax.vmap(lambda xmat: xmat.T @ m.opt.magnetic)(
+          d.site_xmat[objid]
+      ).reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.JOINTPOS:
+      sensor = d.qpos[m.jnt_qposadr[objid]]
+    elif sensor_type == SensorType.ACTUATORPOS:
+      sensor = d.actuator_length[objid]
+    elif sensor_type == SensorType.BALLQUAT:
+      jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
+      quat = d.qpos[jnt_qposadr]
+      sensor = jax.vmap(math.normalize)(quat).reshape(-1)
+      adr = (adr[:, None] + np.arange(4)[None]).reshape(-1)
+    elif sensor_type == SensorType.FRAMEPOS:
+
+      def _framepos(xpos, xpos_ref, xmat_ref, refid):
+        return jp.where(refid == -1, xpos, xmat_ref.T @ (xpos - xpos_ref))
+
+      objtype = m.sensor_objtype[idx]
+      reftype = m.sensor_reftype[idx]
+      refid = m.sensor_refid[idx]
+
+      # evaluate for valid object and reference object type pairs
+      for ot, rt in set(zip(objtype, reftype)):
+        id_ = (objtype == ot) & (reftype == rt)
+        refid_ = refid[id_]
+        xpos, _ = objtype_data[ot]
+        xpos_ref, xmat_ref = objtype_data[rt]
+        xpos = xpos[objid[id_]]
+        xpos_ref = xpos_ref[refid_]
+        xmat_ref = xmat_ref[refid_]
+        sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refid_)
+        adr_ = adr[id_, None] + np.arange(3)[None]
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adr_.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type in frame_axis:
+
+      def _frameaxis(xmat, xmat_ref, refid):
+        axis = xmat[:, frame_axis[sensor_type]]
+        return jp.where(refid == -1, axis, xmat_ref.T @ axis)
+
+      objtype = m.sensor_objtype[idx]
+      reftype = m.sensor_reftype[idx]
+      refid = m.sensor_refid[idx]
+
+      # evaluate for valid object and reference object type pairs
+      for ot, rt in set(zip(objtype, reftype)):
+        id_ = (objtype == ot) & (reftype == rt)
+        refid_ = refid[id_]
+        _, xmat = objtype_data[ot]
+        _, xmat_ref = objtype_data[rt]
+        xmat = xmat[objid[id_]]
+        xmat_ref = xmat_ref[refid_]
+        sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refid_)
+        adr_ = adr[id_, None] + np.arange(3)[None]
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adr_.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type == SensorType.SUBTREECOM:
+      sensor = d.subtree_com[objid].reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.CLOCK:
+      sensor = jp.repeat(d.time, sum(idx))
+
+    sensors.append(sensor)
+    adrs.append(adr)
+
+  sensordata = d.sensordata.at[np.concatenate(adrs)].set(
+      jp.concatenate(sensors)
+  )
 
   return d.replace(sensordata=sensordata)
 

mjx/mujoco/mjx/_src/types.py

@@ -282,11 +282,46 @@ class SensorType(enum.IntEnum):
   """Type of sensor.
 
   Members:
+    MAGNETOMETER: magnetometer
     JOINTPOS: joint position
     ACTUATORPOS: actuator position
+    BALLQUAT: ball joint orientation
+    FRAMEPOS: frame position
+    FRAMEXAXIS: frame x-axis
+    FRAMEYAXIS: frame y-axis
+    FRAMEZAXIS: frame z-axis
+    SUBTREECOM: subtree centor of mass
+    CLOCK: simulation time
   """
+  MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS
   ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS
+  BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT
+  FRAMEPOS = mujoco.mjtSensor.mjSENS_FRAMEPOS
+  FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS
+  FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS
+  FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS
+  SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM
+  CLOCK = mujoco.mjtSensor.mjSENS_CLOCK
+
+
+class ObjType(PyTreeNode):
+  """Type of object.
+
+  Members:
+    UNKNOWN: unknown object type
+    BODY: body
+    XBODY: body, used to access regular frame instead of i-frame
+    GEOM: geom
+    SITE: site
+    CAMERA: camera
+  """
+  UNKNOWN = mujoco.mjtObj.mjOBJ_UNKNOWN
+  BODY = mujoco.mjtObj.mjOBJ_BODY
+  XBODY = mujoco.mjtObj.mjOBJ_XBODY
+  GEOM = mujoco.mjtObj.mjOBJ_GEOM
+  SITE = mujoco.mjtObj.mjOBJ_SITE
+  CAMERA = mujoco.mjtObj.mjOBJ_CAMERA
 
 
 class Option(PyTreeNode):

mjx/mujoco/mjx/test_data/sensor.xml

@@ -1,28 +1,74 @@
 <!-- For validating sensors:
 
-* position-dependent:
+* position-dependent sensors:
+-magnetometer
 -jointpos
 -actuatorpos
-* velocity-dependent:
+-ballquat
-* acceleration/force-dependent:
+-framepos
+-framexaxis
+-frameyaxis
+-framezaxis
+-subtreecom
+-clock
+* velocity-dependent sensors:
+* acceleration/force-dependent sensors:
 -->
 <mujoco model="sensor">
   <worldbody>
-    <!-- body 0 -->
+    <!-- tree 0 -->
     <body name="body0" pos="1 2 3">
       <joint name="hinge0" type="hinge" axis="1 0 0"/>
       <geom size="1"/>
+      <site name="site0" pos=".1 .2 .3"/>
+      <body name="body1" pos="0.1 0.2 0.3">
+        <joint name="hinge1" type="hinge" axis="0 1 0"/>
+        <geom size="0.25"/>
+        <site name="site1" pos=".2 .4 .6"/>
+      </body>
+    </body>
+
+    <!-- body 2 -->
+    <body name="body2" pos=".1 .1 .1">
+      <joint name="ballquat2" type="ball" pos="0.1 0.1 0.1"/>
+      <geom size="1"/>
+    </body>
+
+    <!-- body 3 -->
+    <body name="body3" pos="-.1 -.1 -.1">
+      <joint name="ballquat3" type="ball" pos="0.1 0.2 0.3"/>
+      <geom size="1"/>
     </body>
   </worldbody>
 
   <actuator>
-    <motor name="motor0" joint="hinge0" ctrlrange="-1 1" ctrllimited="true"/>
+    <motor name="motor0" joint="hinge0" ctrlrange="-1 1" gear="10"
+      ctrllimited="true"/>
+    <motor name="motor1" joint="hinge1" ctrlrange="-1 1" gear="10"
+      ctrllimited="true"/>
   </actuator>
 
   <sensor>
     <!-- position-dependent sensors -->
+    <magnetometer name="magnetometer0" site="site0"/>
     <jointpos name="jointpos0" joint="hinge0"/>
     <actuatorpos name="actuatorpos0" actuator="motor0"/>
+    <ballquat name="ballquat2" joint="ballquat2"/>
+    <framexaxis name="framexaxis0" objtype="site" objname="site0"/>
+    <framezaxis name="framezaxis0" objtype="site" objname="site0"/>
+    <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
+    <framepos name="framepos0" objtype="site" objname="site0"/>
+    <subtreecom name="subtreecom0" body="body0"/>
+    <magnetometer name="magnetometer1" site="site1"/>
+    <jointpos name="jointpos1" joint="hinge1"/>
+    <actuatorpos name="actuatorpos1" 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"/>
+    <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
+    <subtreecom name="subtreecom1" body="body1"/>
+    <clock/>
     <!-- velocity-dependent sensors -->
     <!-- acceleration/force-dependent sensors -->
   </sensor>