Modify MJX sensor_pos to handle unsupported sensors.

PiperOrigin-RevId: 663768186
Change-Id: I2e7600e77855a96a47f012fb64b6cc43abbf6f64

mjx/mujoco/mjx/_src/sensor.py

@@ -31,7 +31,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
   """Compute position-dependent sensors values."""
 
   # no position-dependent sensors
-  if sum(m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS) == 0:
+  stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
+  if sum(stage_pos) == 0:
     return d
 
   # position and orientation by object type
@@ -56,7 +57,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
 
   sensors, adrs = [], []
 
-  for sensor_type in set(m.sensor_type):
+  for sensor_type in set(m.sensor_type[stage_pos]):
     idx = m.sensor_type == sensor_type
     objid = m.sensor_objid[idx]
     adr = m.sensor_adr[idx]
@@ -126,10 +127,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.CLOCK:
       sensor = jp.repeat(d.time, sum(idx))
+    else:
+      continue  # unsupported sensor type
 
     sensors.append(sensor)
     adrs.append(adr)
 
+  if not adrs:
+    return d
+
   sensordata = d.sensordata.at[np.concatenate(adrs)].set(
       jp.concatenate(sensors)
   )

mjx/mujoco/mjx/_src/sensor_test.py

@@ -62,6 +62,13 @@ class SensorTest(parameterized.TestCase):
     # sensor values
     _assert_eq(d.sensordata, dx.sensordata, 'sensordata')
 
+  def test_unsupported_sensor(self):
+    """Tests MJX sensor functions do not break for unsupported sensors."""
+    m = test_util.load_test_file('unsupported_sensor.xml')
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.forward)(mx, mjx.put_data(m, mujoco.MjData(m)))
+    _assert_eq(np.zeros(m.nsensordata), dx.sensordata, 'sensordata')
+
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/test_data/unsupported_sensor.xml

@@ -0,0 +1,19 @@
+<!-- For validating model with unsupported sensor -->
+<mujoco model="unsupported_sensor">
+  <worldbody>
+    <body>
+      <joint type="hinge"/>
+      <geom name="geom0" size="1"/>
+      <site name="site"/>
+    </body>
+    <body>
+      <joint type="hinge"/>
+      <geom name="geom1" size="1"/>
+    </body>
+  </worldbody>
+  <sensor>
+    <distance name="distance" geom1="geom0" geom2="geom1"/>
+    <framelinvel name="framelinvel" objtype="site" objname="site"/>
+    <touch name="touch" site="site"/>
+  </sensor>
+</mujoco>