Add framequat sensor to MJX.

PiperOrigin-RevId: 666254481
Change-Id: I50be8574e1e4677ef5dc23a2bdcbc296f301d7fc

doc/changelog.rst

@@ -18,8 +18,8 @@ MJX
 ^^^
 5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`).
 6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``,
-   ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``,
+   ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``,
-   ``CLOCK``.
+   ``SUBTREECOM``, ``CLOCK``.
 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``.
 8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``.
 9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.

mjx/mujoco/mjx/_src/sensor.py

@@ -61,6 +61,9 @@ def sensor_pos(m: Model, d: Data) -> Data:
   for sensor_type in set(m.sensor_type[stage_pos]):
     idx = m.sensor_type == sensor_type
     objid = m.sensor_objid[idx]
+    objtype = m.sensor_objtype[idx]
+    refid = m.sensor_refid[idx]
+    reftype = m.sensor_reftype[idx]
     adr = m.sensor_adr[idx]
 
     if sensor_type == SensorType.MAGNETOMETER:
@@ -113,7 +116,6 @@ def sensor_pos(m: Model, d: Data) -> Data:
 
         return sensor[:2]
 
-      refid = m.sensor_refid[idx]
       sensorsize = m.cam_sensorsize[refid]
       intrinsic = m.cam_intrinsic[refid]
       fovy = m.cam_fovy[refid]
@@ -131,15 +133,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.RANGEFINDER:
       site_bodyid = m.site_bodyid[objid]
       for sid in set(site_bodyid):
-        id_ = sid == site_bodyid
+        idxs = sid == site_bodyid
-        objid_ = objid[id_]
+        objids = objid[idxs]
-        site_xpos = d.site_xpos[objid_]
+        site_xpos = d.site_xpos[objids]
-        site_mat = d.site_xmat[objid_].reshape((-1, 9))[:, np.array([2, 5, 8])]
+        site_mat = d.site_xmat[objids].reshape((-1, 9))[:, np.array([2, 5, 8])]
         sensor, _ = jax.vmap(
             ray.ray, in_axes=(None, None, 0, 0, None, None, None)
         )(m, d, site_xpos, site_mat, (), True, sid)
         sensors.append(sensor)
-        adrs.append(adr[id_])
+        adrs.append(adr[idxs])
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.JOINTPOS:
       sensor = d.qpos[m.jnt_qposadr[objid]]
@@ -155,23 +157,19 @@ def sensor_pos(m: Model, d: Data) -> Data:
       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)
+        idxt = (objtype == ot) & (reftype == rt)
-        refid_ = refid[id_]
+        refidt = refid[idxt]
         xpos, _ = objtype_data[ot]
         xpos_ref, xmat_ref = objtype_data[rt]
-        xpos = xpos[objid[id_]]
+        xpos = xpos[objid[idxt]]
-        xpos_ref = xpos_ref[refid_]
+        xpos_ref = xpos_ref[refidt]
-        xmat_ref = xmat_ref[refid_]
+        xmat_ref = xmat_ref[refidt]
-        sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refid_)
+        sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt)
-        adr_ = adr[id_, None] + np.arange(3)[None]
+        adrt = adr[idxt, None] + np.arange(3)[None]
         sensors.append(sensor.reshape(-1))
-        adrs.append(adr_.reshape(-1))
+        adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type in frame_axis:
 
@@ -179,22 +177,58 @@ def sensor_pos(m: Model, d: Data) -> Data:
         axis = xmat[:, frame_axis[sensor_type]]
         return jp.where(refid == -1, axis, xmat_ref.T @ axis)
 
-      objtype = m.sensor_objtype[idx]
+      # evaluate for valid object and reference object type pairs
-      reftype = m.sensor_reftype[idx]
+      for ot, rt in set(zip(objtype, reftype)):
-      refid = m.sensor_refid[idx]
+        idxt = (objtype == ot) & (reftype == rt)
+        refidt = refid[idxt]
+        _, xmat = objtype_data[ot]
+        _, xmat_ref = objtype_data[rt]
+        xmat = xmat[objid[idxt]]
+        xmat_ref = xmat_ref[refidt]
+        sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt)
+        adrt = adr[idxt, None] + np.arange(3)[None]
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adrt.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type == SensorType.FRAMEQUAT:
+
+      def _quat(otype, oid):
+        if otype == ObjType.XBODY:
+          return d.xquat[oid]
+        elif otype == ObjType.BODY:
+          return jax.vmap(math.quat_mul)(d.xquat[oid], m.body_iquat[oid])
+        elif otype == ObjType.GEOM:
+          return jax.vmap(math.quat_mul)(
+              d.xquat[m.geom_bodyid[oid]], m.geom_quat[oid]
+          )
+        elif otype == ObjType.SITE:
+          return jax.vmap(math.quat_mul)(
+              d.xquat[m.site_bodyid[oid]], m.site_quat[oid]
+          )
+        elif otype == ObjType.CAMERA:
+          return jax.vmap(math.quat_mul)(
+              d.xquat[m.cam_bodyid[oid]], m.cam_quat[oid]
+          )
+        elif otype == ObjType.UNKNOWN:
+          return jp.tile(jp.array([1.0, 0.0, 0.0, 0.0]), (oid.size, 1))
+        else:
+          raise ValueError(f'Unknown object type: {otype}')
 
       # evaluate for valid object and reference object type pairs
       for ot, rt in set(zip(objtype, reftype)):
-        id_ = (objtype == ot) & (reftype == rt)
+        idxt = (objtype == ot) & (reftype == rt)
-        refid_ = refid[id_]
+        objidt = objid[idxt]
-        _, xmat = objtype_data[ot]
+        refidt = refid[idxt]
-        _, xmat_ref = objtype_data[rt]
+        quat = _quat(ot, objidt)
-        xmat = xmat[objid[id_]]
+        refquat = _quat(rt, refidt)
-        xmat_ref = xmat_ref[refid_]
+        sensor = jax.vmap(
-        sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refid_)
+            lambda q, r, rid: jp.where(
-        adr_ = adr[id_, None] + np.arange(3)[None]
+                rid == -1, q, math.quat_mul(math.quat_inv(r), q)
+            )
+        )(quat, refquat, refidt)
+        adrt = adr[idxt, None] + np.arange(4)[None]
         sensors.append(sensor.reshape(-1))
-        adrs.append(adr_.reshape(-1))
+        adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.SUBTREECOM:
       sensor = d.subtree_com[objid].reshape(-1)

mjx/mujoco/mjx/_src/types.py

@@ -304,6 +304,7 @@ class SensorType(enum.IntEnum):
     FRAMEXAXIS: frame x-axis
     FRAMEYAXIS: frame y-axis
     FRAMEZAXIS: frame z-axis
+    FRAMEQUAT: frame orientation, represented as quaternion
     SUBTREECOM: subtree centor of mass
     CLOCK: simulation time
     JOINTVEL: joint velocity
@@ -322,6 +323,7 @@ class SensorType(enum.IntEnum):
   FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS
   FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS
   FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS
+  FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT
   SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM
   CLOCK = mujoco.mjtSensor.mjSENS_CLOCK
   JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL

mjx/mujoco/mjx/test_data/sensor.xml

@@ -11,6 +11,7 @@
 -framexaxis
 -frameyaxis
 -framezaxis
+-framequat
 -subtreecom
 -clock
 * velocity-dependent sensors:
@@ -50,6 +51,7 @@
     <body name="body3" pos="-.1 -.1 -.1">
       <joint name="ballquat3" type="ball" pos="0.1 0.2 0.3"/>
       <geom size="1"/>
+      <site name="site3"/>
     </body>
 
     <!-- bodies for camera projection -->
@@ -91,6 +93,7 @@
     <framexaxis name="framexaxis0" objtype="site" objname="site0"/>
     <jointactuatorfrc name="jointactfrc1" joint="hinge1"/>
     <framezaxis name="framezaxis0" objtype="site" objname="site0"/>
+    <framequat name="framequat3" objtype="site" objname="site3"/>
     <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
     <framepos name="framepos0" objtype="site" objname="site0"/>
     <actuatorfrc name="actuatorfrc1" actuator="motor1"/>
@@ -109,6 +112,7 @@
     <camprojection site="frontcenter" camera="fixedcamera"/>
     <ballangvel name="ballangvel3" joint="ballquat3"/>
     <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
+    <framequat name="framequat0" objtype="site" objname="site0"/>
     <rangefinder name="rangefinder1" site="site_rangefinder1"/>
     <subtreecom name="subtreecom1" body="body1"/>
     <clock/>