Raise error for gravcomp.

PiperOrigin-RevId: 624292351
Change-Id: Ia8cfce33f990616035772fb88c34fb1434c73220

mjx/mujoco/mjx/_src/device.py

@@ -128,15 +128,16 @@ def _validate(m: mujoco.MjModel):
     if unsupported:
       raise NotImplementedError(f'{unsupported} not implemented.')
 
-  # check condim
-  if any(dim != 3 for dim in m.geom_condim) or any(
-      dim != 3 for dim in m.pair_dim
-  ):
-    raise NotImplementedError('Only condim=3 is supported.')
-
   if m.ntendon:
     raise NotImplementedError('Tendons are not supported.')
 
+  # check condim
+  if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
+    raise NotImplementedError('Only condim=3 is supported.')
+
+  if m.body_gravcomp.any():
+    raise NotImplementedError('gravcomp is not supported')
+
   # check collision geom types
   for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
     g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)

mjx/mujoco/mjx/_src/io.py

@@ -90,6 +90,9 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
   if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
     raise NotImplementedError('only condim=3 is supported')
 
+  if m.body_gravcomp.any():
+    raise NotImplementedError('gravcomp is not supported')
+
   # check collision geom types
   for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
     g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)

mjx/mujoco/mjx/_src/io_test.py

@@ -193,6 +193,18 @@ class ModelIOTest(parameterized.TestCase):
           </worldbody>
         </mujoco>"""))
 
+  def test_gravcomp_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body gravcomp="1">
+              <freejoint/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
   def test_cylinder_not_implemented(self):
     with self.assertRaises(NotImplementedError):
       mjx.put_model(mujoco.MjModel.from_xml_string("""