Add rangefinder sensor to MJX.

PiperOrigin-RevId: 664417458 Change-Id: I1ad7720e48541349f155984584314b2b49ccbf36
2024-08-18 05:15:10 -07:00
parent 6acf40613d
commit 51733c2a8a
4 changed files with 32 additions and 3 deletions
@@ -17,8 +17,8 @@ General
 MJX
 ^^^
 5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`).
-6. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``,
-   ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``.
+6. Added position-dependent sensors: ``MAGNETOMETER``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``,
+   ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``.
 7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
 8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.

@@ -19,6 +19,7 @@ from jax import numpy as jp
 import mujoco
 # pylint: disable=g-importing-member
 from mujoco.mjx._src import math
+from mujoco.mjx._src import ray
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
@@ -71,6 +72,19 @@ def sensor_pos(m: Model, d: Data) -> Data:
          d.site_xmat[objid]
      ).reshape(-1)
      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.RANGEFINDER:
+      site_bodyid = m.site_bodyid[objid]
+      for sid in set(site_bodyid):
+        id_ = sid == site_bodyid
+        objid_ = objid[id_]
+        site_xpos = d.site_xpos[objid_]
+        site_mat = d.site_xmat[objid_].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_])
+      continue  # avoid adding to sensors/adrs list a second time
    elif sensor_type == SensorType.JOINTPOS:
      sensor = d.qpos[m.jnt_qposadr[objid]]
    elif sensor_type == SensorType.ACTUATORPOS:
@@ -293,6 +293,7 @@ class SensorType(enum.IntEnum):

  Members:
    MAGNETOMETER: magnetometer
+    RANGEFINDER: rangefinder
    JOINTPOS: joint position
    ACTUATORPOS: actuator position
    BALLQUAT: ball joint orientation
@@ -304,6 +305,7 @@ class SensorType(enum.IntEnum):
    CLOCK: simulation time
  """
  MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
+  RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER
  JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS
  ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS
  BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT
@@ -2,6 +2,7 @@

 * position-dependent sensors:
 -magnetometer
+-rangefinder
 -jointpos
 -actuatorpos
 -ballquat
@@ -15,11 +16,16 @@
 * acceleration/force-dependent sensors:
 -->
 <mujoco model="sensor">
+  <asset>
+    <material name="material"/>
+  </asset>
  <worldbody>
    <!-- tree 0 -->
    <body name="body0" pos="1 2 3">
      <joint name="hinge0" type="hinge" axis="1 0 0"/>
-      <geom size="1"/>
+      <geom size="0.1" material="material"/>
+      <site name="site_rangefinder0" pos="-1e-3 0 0.2"/>
+      <site name="site_rangefinder1" pos="-1e-3 0 0.175"/>
      <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"/>
@@ -39,6 +45,11 @@
      <joint name="ballquat3" type="ball" pos="0.1 0.2 0.3"/>
      <geom size="1"/>
    </body>
+
+    <!-- body for rangefinder -->
+    <body name="body_rangefinder" pos="1 2 4">
+      <geom size="0.01" material="material"/>
+    </body>
  </worldbody>

  <actuator>
@@ -51,6 +62,7 @@
  <sensor>
    <!-- position-dependent sensors -->
    <magnetometer name="magnetometer0" site="site0"/>
+    <rangefinder name="rangefinder0" site="site_rangefinder0"/>
    <jointpos name="jointpos0" joint="hinge0"/>
    <actuatorpos name="actuatorpos0" actuator="motor0"/>
    <ballquat name="ballquat2" joint="ballquat2"/>
@@ -67,6 +79,7 @@
    <framexaxis name="framexaxis1" objtype="site" objname="site1"/>
    <ballquat name="ballquat3" joint="ballquat3"/>
    <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
+    <rangefinder name="rangefinder1" site="site_rangefinder1"/>
    <subtreecom name="subtreecom1" body="body1"/>
    <clock/>
    <!-- velocity-dependent sensors -->