From 9f1651b7bd7b582d075949e701553fa4a3a3c8e1 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 12 Apr 2024 14:24:33 -0700
Subject: [PATCH] Raise error for gravcomp.

PiperOrigin-RevId: 624292351
Change-Id: Ia8cfce33f990616035772fb88c34fb1434c73220
---
 mjx/mujoco/mjx/_src/device.py  | 13 +++++++------
 mjx/mujoco/mjx/_src/io.py      |  3 +++
 mjx/mujoco/mjx/_src/io_test.py | 12 ++++++++++++
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py
index b920b2e6..0a8130e2 100644
--- a/mjx/mujoco/mjx/_src/device.py
+++ b/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)
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index cc272c41..3e90b26e 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/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)
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 83147503..36806c2e 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/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("""