Raise error in mjx.rne_postconstraint for currently unsupported connect and weld constraints and raise error in mjx.put_model for unsupported sensor and equality constraint combinations.

thowell · copybara-github · commit a793a3417aaf · 2024-12-02T16:28:23.000-08:00
PiperOrigin-RevId: 702119449
Change-Id: I034363db41629f19ef98bea79ec89140052a7a3c
diff --git a/doc/mjx.rst b/doc/mjx.rst
@@ -244,8 +244,10 @@ The following features are **fully supported** in MJX:
      - ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``TENDONPOS``, ``ACTUATORPOS``, ``BALLQUAT``,
        ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``,
        ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``TENDONVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``,
-       ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``,
-       ``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``.
+       ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACTUATORFRC``, ``JOINTACTFRC``,
+       ``FRAMELINACC``, ``FRAMEANGACC``
+     - ``ACCELEROMETER``, ``FORCE``, and ``TORQUE`` are supported if the model does not include connect or weld equality
+       constraints.
 
 The following features are **in development** and coming soon:
 
@@ -274,6 +276,9 @@ The following features are **in development** and coming soon:
      - All except ``PLUGIN``, ``USER``
    * - Lights
      - Positions and directions of lights
+   * - :ref:`Sensors <mjtSensor>`
+     - ``ACCELEROMETER``, ``FORCE``, and ``TORQUE`` for models that include connect or weld equality
+     constraints.
 
 The following features are **unsupported**:
 
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -142,6 +142,21 @@ def put_model(
           ' implemented for spatial tendons.'
       )
 
+  # check for unsupported sensor and equality constraint combinations
+  sensor_rne_postconstraint = (
+      np.any(m.sensor_type == types.SensorType.ACCELEROMETER)
+      | np.any(m.sensor_type == types.SensorType.FORCE)
+      | np.any(m.sensor_type == types.SensorType.TORQUE)
+  )
+  eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any(
+      m.eq_type == types.EqType.WELD
+  )
+  if sensor_rne_postconstraint and eq_connect_weld:
+    raise NotImplementedError(
+        'rne_postconstraint not implemented with equality constraints:'
+        ' connect, weld.'
+    )
+
   for enum_field, enum_type, mj_type in (
       (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
       (m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
@@ -513,6 +513,38 @@ def test_contact_elliptic_condim1(self):
     with self.assertRaises(NotImplementedError):
       mjx.make_data(m)
 
+  @parameterized.product(
+      sensor=['accelerometer', 'force', 'torque'], equality=['connect', 'weld']
+  )
+  def test_sensor_constraint_compatibility(self, sensor, equality):
+    """Test unsupported sensor and equality constraint combinations."""
+    equality_constraint = f'{equality} body1="body1" body2="body2"'
+    if equality == 'connect':
+      equality_constraint += ' anchor="0 0 0"'
+    m = mujoco.MjModel.from_xml_string(f"""
+        <mujoco>
+          <worldbody>
+            <body name="body1">
+              <freejoint/>
+              <geom size="0.1"/>
+              <site name="site1"/>
+            </body>
+            <body name="body2">
+              <freejoint/>
+              <geom size="0.1"/>
+            </body>
+          </worldbody>
+          <equality>
+            <{equality_constraint}/>
+          </equality>
+          <sensor>
+            <{sensor} site="site1"/>
+          </sensor>
+        </mujoco>
+      """)
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(m)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
@@ -24,6 +24,7 @@
 from mujoco.mjx._src.types import CamLightType
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
 from mujoco.mjx._src.types import TrnType
@@ -635,6 +636,10 @@ def _contact_force_to_cfrc_ext(force, pos, frame, id1, id2, com1, com2):
     )
 
   # TODO(taylorhowell): connect and weld constraints
+  if np.any(m.eq_type == EqType.CONNECT):
+    raise NotImplementedError('Connect constraints are not implemented.')
+  if np.any(m.eq_type == EqType.WELD):
+    raise NotImplementedError('Weld constraints are not implemented.')
 
   # forward pass over bodies: compute cacc, cfrc_int
   def _forward(carry, cfrc_ext, cinert, cvel, body_dofadr, body_dofnum):
