diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index b92c32e0..7dabeeb2 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/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) ) diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index d30e928e..d120e757 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/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() diff --git a/mjx/mujoco/mjx/test_data/unsupported_sensor.xml b/mjx/mujoco/mjx/test_data/unsupported_sensor.xml new file mode 100644 index 00000000..4a164bd9 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/unsupported_sensor.xml @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + +