From 8710126d08306a9da3d6be86b4939dad5cc78780 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 11 Sep 2024 03:49:43 -0700 Subject: [PATCH] Raise NotImplementedError for unsupported sensors in MJX. Update MJX feature parity docs to include fully supported sensors. PiperOrigin-RevId: 673325233 Change-Id: Ib262bdd37937c5230765c11d1f61581e55e8c397 --- doc/mjx.rst | 5 ++ mjx/mujoco/mjx/_src/io.py | 3 +- mjx/mujoco/mjx/_src/sensor_test.py | 28 ++++++-- .../mjx/test_data/sensor/unsupported.xml | 65 ------------------- 4 files changed, 29 insertions(+), 72 deletions(-) delete mode 100644 mjx/mujoco/mjx/test_data/sensor/unsupported.xml diff --git a/doc/mjx.rst b/doc/mjx.rst index 776022cb..99fc1b8f 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -214,6 +214,11 @@ The following features are **fully supported** in MJX: - :ref:`flInertia` * - :ref:`Tendons ` - :ref:`Fixed ` + * - :ref:`Sensors ` + - ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, + ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``, ``VELOCIMETER``, + ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, + ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``ACTUATORFRC``, ``JOINTACTFRC``. The following features are **in development** and coming soon: diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 551cdc64..2d90eb33 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -116,8 +116,7 @@ def put_model( (m.actuator_gaintype, types.GainType, mujoco.mjtGain), (m.actuator_trntype, types.TrnType, mujoco.mjtTrn), (m.eq_type, types.EqType, mujoco.mjtEq), - # TODO(taylorhowell): causes Menagerie test to fail - # (m.sensor_type, types.SensorType, mujoco.mjtSensor), + (m.sensor_type, types.SensorType, mujoco.mjtSensor), (m.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = set(enum_field) - set(enum_type) diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index 42c2352a..4613b527 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/mjx/mujoco/mjx/_src/sensor_test.py @@ -93,11 +93,29 @@ class SensorTest(parameterized.TestCase): _assert_eq(random_sensor, dx.sensordata, 'sensordata') def test_unsupported_sensor(self): - """Tests MJX sensor functions do not break for unsupported sensors.""" - m = test_util.load_test_file('sensor/unsupported.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') + """Tests unsupported sensor raises error.""" + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + + + + + + + """) + with self.assertRaises(NotImplementedError): + mjx.put_model(m) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/test_data/sensor/unsupported.xml b/mjx/mujoco/mjx/test_data/sensor/unsupported.xml deleted file mode 100644 index 30ba0e70..00000000 --- a/mjx/mujoco/mjx/test_data/sensor/unsupported.xml +++ /dev/null @@ -1,65 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -