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.

PiperOrigin-RevId: 702119449
Change-Id: I034363db41629f19ef98bea79ec89140052a7a3c

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),

mjx/mujoco/mjx/_src/io_test.py

@@ -513,6 +513,38 @@ class DataIOTest(parameterized.TestCase):
     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()

mjx/mujoco/mjx/_src/smooth.py

@@ -24,6 +24,7 @@ from mujoco.mjx._src import support
 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 rne_postconstraint(m: Model, d: Data) -> Data:
     )
 
   # 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):