Throw an error for margin/gap with meshes.

PiperOrigin-RevId: 615937263
Change-Id: I3289cf1b1fec93385cc78a01e53866427453db59

mjx/mujoco/mjx/__init__.py

@@ -16,6 +16,7 @@
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
+from mujoco.mjx._src.collision_driver import get_params
 from mujoco.mjx._src.collision_driver import ncon
 from mujoco.mjx._src.constraint import count_constraints
 from mujoco.mjx._src.constraint import make_constraint

mjx/mujoco/mjx/_src/collision_driver.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Collide geometries."""
 
-from typing import Callable, Dict, Optional, Sequence, Tuple, Union
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import jax
 from jax import numpy as jp
@@ -187,6 +187,31 @@ def _dynamic_params(
   return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
 
 
+def get_params(
+    m: Union[Model, mujoco.MjModel], candidates: Sequence[Candidate]
+) -> Tuple[List[int], List[int], SolverParams]:
+  """Gets solver params for a list of collision candidates."""
+  # group sol params by different candidate types
+  typ_cands = {}
+  for c in candidates:
+    typ = (c.ipair > -1, c.geomp > -1)
+    typ_cands.setdefault(typ, []).append(c)
+
+  geom1, geom2, params = [], [], []
+  for (pair, priority), candidates in typ_cands.items():
+    geom1.extend([c.geom1 for c in candidates])
+    geom2.extend([c.geom2 for c in candidates])
+    if pair:
+      params.append(_pair_params(m, candidates))
+    elif priority:
+      params.append(_priority_params(m, candidates))
+    else:
+      params.append(_dynamic_params(m, candidates))
+
+  params = jax.tree_map(lambda *x: jp.concatenate(x), *params)
+  return geom1, geom2, params
+
+
 def _pair_info(
     m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int]
 ) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]:
@@ -276,24 +301,7 @@ def _collide_geoms(
   if not fn:
     return Contact.zero()
 
-  # group sol params by different candidate types
-  typ_cands = {}
-  for c in candidates:
-    typ = (c.ipair > -1, c.geomp > -1)
-    typ_cands.setdefault(typ, []).append(c)
-
-  geom1, geom2, params = [], [], []
-  for (pair, priority), candidates in typ_cands.items():
-    geom1.extend([c.geom1 for c in candidates])
-    geom2.extend([c.geom2 for c in candidates])
-    if pair:
-      params.append(_pair_params(m, candidates))
-    elif priority:
-      params.append(_priority_params(m, candidates))
-    else:
-      params.append(_dynamic_params(m, candidates))
-
-  params = jax.tree_map(lambda *x: jp.concatenate(x), *params)
+  geom1, geom2, params = get_params(m, candidates)
   g1, g2, in_axes = _pair_info(m, d, geom1, geom2)
 
   # Run a crude version of broadphase.

mjx/mujoco/mjx/_src/device.py

@@ -138,8 +138,7 @@ def _validate(m: mujoco.MjModel):
     raise NotImplementedError('Tendons are not supported.')
 
   # check collision geom types
-  candidate_set = collision_driver.collision_candidates(m)
-  for g1, g2, *_ in candidate_set:
+  for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
     g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
     if g1 == mujoco.mjtGeom.mjGEOM_PLANE and g2 in (
         mujoco.mjtGeom.mjGEOM_PLANE,
@@ -149,6 +148,12 @@ def _validate(m: mujoco.MjModel):
       continue
     if collision_driver.get_collision_fn((g1, g2)) is None:
       raise NotImplementedError(f'({g1}, {g2}) collisions not implemented.')
+    *_, params = collision_driver.get_params(m, c)
+    margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0)
+    if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap:
+      raise NotImplementedError(
+          f'Margin and gap not implemented for ({g1}, {g2})'
+      )
 
   # TODO(erikfrey): warn for high solver iterations, nefc, etc.
 

mjx/mujoco/mjx/_src/io.py

@@ -91,10 +91,16 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
     raise NotImplementedError('only condim=3 is supported')
 
   # check collision geom types
-  for g1, g2, *_ in collision_driver.collision_candidates(m):
+  for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
+    g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
     if collision_driver.get_collision_fn((g1, g2)) is None:
-      g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
       raise NotImplementedError(f'({g1}, {g2}) has no collision function')
+    *_, params = collision_driver.get_params(m, c)
+    margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0)
+    if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap:
+      raise NotImplementedError(
+          f'Margin and gap not implemented for ({g1}, {g2})'
+      )
 
   for enum_field, enum_type, mj_type in (
       (m.actuator_biastype, types.BiasType, mujoco.mjtBias),

mjx/mujoco/mjx/_src/io_test.py

@@ -209,6 +209,25 @@ class ModelIOTest(parameterized.TestCase):
           </worldbody>
         </mujoco>"""))
 
+  def test_margin_gap_mesh_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <asset>
+            <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 1"/>
+          </asset>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom type="mesh" mesh="box" margin="0.3"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
 
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""