Throw an error for margin/gap with meshes.

PiperOrigin-RevId: 615937263
Change-Id: I3289cf1b1fec93385cc78a01e53866427453db59

doc/changelog.rst

@@ -11,6 +11,7 @@ MJX
 1. Improved performance of SAT for convex collisions.
 2. Fixed bug for sphere/capsule-convex deep penetration.
 3. Fixed bug where ``mjx.Data`` produced by ``mjx.put_data`` had different treedef than ``mjx.make_data``.
+4. Throw an error for margin/gap for convex mesh collisions, since they are not supported.
 
 Version 3.1.3 (March 5th, 2024)
 -----------------------------------

doc/mjx.rst

@@ -215,6 +215,14 @@ The following features are **in development** and coming soon:
 
    * - Category
      - Feature
+   * - :ref:`Geom <mjtGeom>`
+     - ``SDF``, ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
+   * - :ref:`Condim <coContact>`
+     - 1, 4, 6
+   * - :ref:`Constraint <mjtConstraint>`
+     - :ref:`Frictionloss <coFriction>`, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
+   * - :ref:`Integrator <mjtIntegrator>`
+     - ``IMPLICIT``, ``IMPLICITFAST``
    * - Dynamics
      - :ref:`Inverse <mj_inverse>`
    * - :ref:`Transmission <mjtTrn>`
@@ -227,16 +235,8 @@ The following features are **in development** and coming soon:
      - ``MUSCLE``
    * - :ref:`Tendon Wrapping <mjtWrap>`
      - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER``
-   * - :ref:`Geom <mjtGeom>`
-     - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
-   * - :ref:`Constraint <mjtConstraint>`
-     - :ref:`Frictionloss <coFriction>`, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
-   * - :ref:`Integrator <mjtIntegrator>`
-     - ``IMPLICIT``, ``IMPLICITFAST``
    * - :ref:`Cone <mjtCone>`
      - ``ELLIPTIC``
-   * - :ref:`Condim <coContact>`
-     - 1, 4, 6
    * - Fluid Model
      - :ref:`flEllipsoid`
    * - :ref:`Tendons <tendon>`
@@ -258,6 +258,8 @@ The following features are **unsupported**:
 
    * - Category
      - Feature
+   * - :ref:`margin<body-geom-margin>` and :ref:`gap<body-geom-gap>`
+     - Unimplemented for collisions with ``Mesh`` :ref:`Geom <mjtGeom>`.
    * - :ref:`Transmission <mjtTrn>`
      - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY``
    * - :ref:`Actuator Dynamics <mjtDyn>`
@@ -270,8 +272,6 @@ The following features are **unsupported**:
      - ``PGS``
    * - :ref:`Sensors <mjtSensor>`
      - ``PLUGIN``, ``USER``
-   * - :ref:`Geom <mjtGeom>`
-     - ``SDF``
 
 .. _MjxSharpBits:
 

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
+  geom1, geom2, params = get_params(m, candidates)
-  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)
   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, *_), c in collision_driver.collision_candidates(m).items():
-  for g1, g2, *_ in candidate_set:
     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."""