margin and gap redesign

Author: thowell

mujoco_warp/_src/collision_convex.py

@@ -35,6 +35,7 @@
 from mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco_warp._src.types import MJ_MAXVAL
+from mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco_warp._src.types import Data
 from mujoco_warp._src.types import DisableBit
 from mujoco_warp._src.types import GeomType
@@ -773,7 +774,10 @@ def eval_ccd_write_contact(
     if is_collision_sensor:
       cutoff = 1.0e32
     else:
-      cutoff = 0.0
+      if wp.static(NEW_GAP_SEMANTICS):
+        cutoff = gap
+      else:
+        cutoff = 0.0
     dist, ncollision, w1, w2, multiccd_idx = ccd(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
@@ -793,8 +797,12 @@ def eval_ccd_write_contact(
       epa_horizon_in[ccdid],
     )
 
-    if dist >= 0.0 and pairid[1] == -1:
-      return 0
+    if wp.static(NEW_GAP_SEMANTICS):
+      if dist >= gap and pairid[1] == -1:
+        return 0
+    else:
+      if dist >= 0.0 and pairid[1] == -1:
+        return 0
 
     # CCD operates on margin-inflated shapes (support() inflates each geom by
     # 0.5 * margin).  The returned dist is therefore relative to the inflated

mujoco_warp/_src/collision_core.py

@@ -23,6 +23,7 @@
 from mujoco_warp._src.math import safe_div
 from mujoco_warp._src.types import MJ_MINMU
 from mujoco_warp._src.types import MJ_MINVAL
+from mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco_warp._src.types import ContactType
 from mujoco_warp._src.types import GeomType
 from mujoco_warp._src.types import mat63
@@ -197,14 +198,15 @@ def write_contact(
   Returns 1 if the contact is active (dist < margin), 0 otherwise.
   """
   active = dist_in < margin_in
+  detected = dist_in < margin_in + gap_in
 
   # skip contact and no collision sensor
-  if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
+  if (pairid_in[0] == -2 or not detected) and pairid_in[1] == -1:
     return 0
 
   contact_type = 0
 
-  if pairid_in[0] >= -1 and active:
+  if pairid_in[0] >= -1 and detected:
     contact_type |= ContactType.CONSTRAINT
 
   if pairid_in[1] >= 0:
@@ -217,7 +219,10 @@ def write_contact(
     contact_frame_out[cid] = frame_in
     contact_geom_out[cid] = geoms_in
     contact_worldid_out[cid] = worldid_in
-    includemargin = margin_in - gap_in
+    if wp.static(NEW_GAP_SEMANTICS):
+      includemargin = margin_in
+    else:
+      includemargin = margin_in - gap_in
     contact_includemargin_out[cid] = includemargin
     contact_dim_out[cid] = condim_in
     contact_friction_out[cid] = friction_in

mujoco_warp/_src/collision_driver.py

@@ -271,13 +271,14 @@ def _obb_filter(
   return True
 
 
-def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
+def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
   @wp.func
   def func(
     # Model:
     geom_aabb: wp.array3d[wp.vec3],
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
     geom_xmat_in: wp.array2d[wp.mat33],
@@ -299,21 +300,25 @@ def func(
     rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2]  # kernel_analyzer: ignore
     margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
     margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]  # kernel_analyzer: ignore
+    gap_id = worldid % ngeom_gap if wp.static(ngeom_gap > 1) else 0
+    gap1, gap2 = geom_gap[gap_id, geom1], geom_gap[gap_id, geom2]  # kernel_analyzer: ignore
+    effective_margin1 = margin1 + gap1
+    effective_margin2 = margin2 + gap2
     xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
     xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
 
     if rbound1 == 0.0 or rbound2 == 0.0:
       if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
-        return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
+        return _plane_filter(rbound1, rbound2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2)
     else:
       if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
-        if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
+        if not _sphere_filter(rbound1, rbound2, effective_margin1, effective_margin2, xpos1, xpos2):
           return False
       if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
-        if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
+        if not _aabb_filter(center1, center2, size1, size2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2):
           return False
       if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
-        if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
+        if not _obb_filter(center1, center2, size1, size2, effective_margin1, effective_margin2, xpos1, xpos2, xmat1, xmat2):
           return False
 
     return True
@@ -378,6 +383,7 @@ def sap_project(
     ngeom: int,
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
     nworld_in: int,
@@ -398,7 +404,7 @@ def sap_project(
       # geom is a plane
       rbound = MJ_MAXVAL
 
-    radius = rbound + geom_margin[worldid % geom_margin.shape[0], geomid]
+    radius = rbound + geom_margin[worldid % geom_margin.shape[0], geomid] + geom_gap[worldid % geom_gap.shape[0], geomid]
     center = wp.dot(direction_in, xpos)
 
     sort_index_out[worldid, geomid] = geomid
@@ -444,7 +450,7 @@ def _sap_range(
 
 
 @cache_kernel
-def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
+def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
@@ -453,6 +459,7 @@ def kernel(
     geom_aabb: wp.array3d[wp.vec3],
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
     nxn_pairid: wp.array[wp.vec2i],
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
@@ -503,8 +510,8 @@ def kernel(
         continue
 
       if (
-        wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
-          geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+        wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
+          geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
         )
         or pairid[1] >= 0
       ):
@@ -588,7 +595,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   wp.launch(
     kernel=_sap_project(m.opt.broadphase),
     dim=(d.nworld, m.ngeom),
-    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.geom_xpos, d.nworld, direction],
+    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction],
     outputs=[
       projection_lower.reshape((-1, m.ngeom)),
       projection_upper,
@@ -624,14 +631,17 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   # assumes each geom has 5 other geoms (batched over all worlds)
   nsweep = 5 * nworldgeom
   wp.launch(
-    kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
+    kernel=_sap_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=nsweep,
     inputs=[
       m.ngeom,
       m.geom_type,
       m.geom_aabb,
       m.geom_rbound,
       m.geom_margin,
+      m.geom_gap,
       m.nxn_pairid,
       d.geom_xpos,
       d.geom_xmat,
@@ -646,14 +656,15 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
 
 
 @cache_kernel
-def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
+def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
     geom_type: wp.array[int],
     geom_aabb: wp.array3d[wp.vec3],
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
     nxn_geom_pair: wp.array[wp.vec2i],
     nxn_pairid: wp.array[wp.vec2i],
     # Data in:
@@ -674,8 +685,8 @@ def kernel(
     geom2 = geom[1]
 
     if (
-      wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
-        geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+      wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
+        geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
       )
       or nxn_pairid[elementid][1] >= 0
     ):
@@ -711,13 +722,16 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
   wp.launch(
-    _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
+    _nxn_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,
       m.geom_aabb,
       m.geom_rbound,
       m.geom_margin,
+      m.geom_gap,
       m.nxn_geom_pair_filtered,
       m.nxn_pairid_filtered,
       d.geom_xpos,

mujoco_warp/_src/collision_driver_test.py

@@ -30,6 +30,7 @@
 from mujoco_warp._src.collision_driver import MJ_COLLISION_TABLE
 from mujoco_warp._src.collision_primitive import plane_convex
 from mujoco_warp._src.math import upper_trid_index
+from mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco_warp.test_data.collision_sdf.utils import register_sdf_plugins
 
 _TOLERANCE = 5e-5
@@ -702,7 +703,7 @@ def test_contact_pair(self, broadphase):
 
     # 1 pair
     _, _, m, d = test_data.fixture(
-      xml="""
+      xml=f"""
       <mujoco>
         <worldbody>
           <body>
@@ -715,7 +716,7 @@ def test_contact_pair(self, broadphase):
           </body>
         </worldbody>
         <contact>
-          <pair geom1="geom1" geom2="geom2" margin="2" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
+          <pair geom1="geom1" geom2="geom2" margin="{-1 if NEW_GAP_SEMANTICS else 2}" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
         </contact>
       </mujoco>
     """
@@ -746,7 +747,7 @@ def test_contact_pair(self, broadphase):
 
     # 1 pair: override contype and conaffinity
     _, _, m, d = test_data.fixture(
-      xml="""
+      xml=f"""
       <mujoco>
         <worldbody>
           <body name="body1">
@@ -759,7 +760,7 @@ def test_contact_pair(self, broadphase):
           </body>
         </worldbody>
         <contact>
-          <pair geom1="geom1" geom2="geom2" margin="2" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
+          <pair geom1="geom1" geom2="geom2" margin="{-1 if NEW_GAP_SEMANTICS else 2}" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
         </contact>
       </mujoco>
     """
@@ -790,7 +791,7 @@ def test_contact_pair(self, broadphase):
 
     # 1 pair: override exclude
     _, _, m, d = test_data.fixture(
-      xml="""
+      xml=f"""
       <mujoco>
         <worldbody>
           <body name="body1">
@@ -804,7 +805,7 @@ def test_contact_pair(self, broadphase):
         </worldbody>
         <contact>
           <exclude body1="body1" body2="body2"/>
-          <pair geom1="geom1" geom2="geom2" margin="2" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
+          <pair geom1="geom1" geom2="geom2" margin="{-1 if NEW_GAP_SEMANTICS else 2}" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
         </contact>
       </mujoco>
     """
@@ -835,7 +836,7 @@ def test_contact_pair(self, broadphase):
 
     # 1 pair 1 exclude
     _, _, m, d = test_data.fixture(
-      xml="""
+      xml=f"""
       <mujoco>
         <worldbody>
           <body name="body1">
@@ -853,7 +854,7 @@ def test_contact_pair(self, broadphase):
         </worldbody>
         <contact>
           <exclude body1="body1" body2="body2"/>
-          <pair geom1="geom2" geom2="geom3" margin="2" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
+          <pair geom1="geom2" geom2="geom3" margin="{-1 if NEW_GAP_SEMANTICS else 2}" gap="3" condim="6" friction="5 4 3 2 1" solref="-.25 -.5" solreffriction="2 4" solimp=".1 .2 .3 .4 .5"/>
         </contact>
       </mujoco>
     """
@@ -1126,22 +1127,23 @@ def test_sdf_volume_collision(self, fixture):
   def test_ccd_margin_dist(self):
     """Tests that CCD contact dist matches MuJoCo when margin > 0.
 
-    Two ellipsoids are placed 0.05 m apart (not touching). With margin=0.1 on
-    each geom the pair margin is 0.2, so contacts are detected within the
+    Two ellipsoids are placed 0.05 m apart (not touching). With margin=0.01
+    and gap=0.2 on each geom, the pair margin is 0.02 and pair gap is 0.4.
+    CCD inflates geometries by margin, detecting contacts within the
     speculative envelope. The reported dist must equal the true geometric
     separation (≈0.05), not the margin-biased value that the inflated
     GJK/EPA would produce.
     """
-    xml = """
+    xml = f"""
     <mujoco>
       <worldbody>
         <body pos="0 0 0">
           <freejoint/>
-          <geom type="ellipsoid" size="0.15 0.15 0.25" margin="0.1" gap="0.1"/>
+          <geom type="ellipsoid" size="0.15 0.15 0.25" margin="{0.01 if NEW_GAP_SEMANTICS else 0.1}" gap="0.2"/>
         </body>
         <body pos="0 0 0.35">
           <freejoint/>
-          <geom type="ellipsoid" size="0.1 0.1 0.05" margin="0.1" gap="0.1"/>
+          <geom type="ellipsoid" size="0.1 0.1 0.05" margin="{0.01 if NEW_GAP_SEMANTICS else 0.1}" gap="0.2"/>
         </body>
       </worldbody>
     </mujoco>
@@ -1173,7 +1175,7 @@ def test_ccd_margin_dist(self):
           break
       self.assertTrue(found, f"MJ contact {i} dist={mj_dist:.4f} not matched in MJW")
 
-    # Verify no constraint forces are generated (includemargin=0, dist > 0)
+    # dist(≈0.05) > margin(0.02): contacts are in gap zone, no constraints
     self.assertEqual(mjd.nefc, 0, "Classic MuJoCo should have no active constraints")
     self.assertEqual(d.nefc.numpy()[0], 0, "MuJoCo Warp should have no active constraints")
 

mujoco_warp/_src/constraint_test.py

@@ -26,6 +26,7 @@
 import mujoco_warp as mjw
 from mujoco_warp import ConeType
 from mujoco_warp import test_data
+from mujoco_warp._src.types import NEW_GAP_SEMANTICS
 
 # tolerance for difference between MuJoCo and MJWarp constraint calculations,
 # mostly due to float precision
@@ -270,15 +271,15 @@ def test_equality_tendon(self, jacobian):
   def test_efc_address_inactive_contacts(self):
     """Test that efc_address is -1 for inactive contacts in the gap zone."""
     # Sphere at z=0.35 with radius 0.1: dist ~ 0.15 to ground plane.
-    # margin=0.5, gap=0.4 => includemargin = 0.1.
-    # dist(0.15) < margin(0.5) => contact is detected.
-    # dist(0.15) >= includemargin(0.1) => contact is NOT active (in gap zone).
-    xml = """
+    # margin=0.1, gap=0.4 => detection at margin+gap=0.5, forces at margin=0.1.
+    # dist(0.15) < margin+gap(0.5) => contact is detected.
+    # dist(0.15) >= margin(0.1) => contact is NOT active (in gap zone).
+    xml = f"""
       <mujoco>
         <worldbody>
-          <geom type="plane" size="10 10 .001" margin="0.5" gap="0.4"/>
+          <geom type="plane" size="10 10 .001" margin="{0.1 if NEW_GAP_SEMANTICS else 0.5}" gap="0.4"/>
           <body pos="0 0 0.35">
-            <geom type="sphere" size=".1" margin="0.5" gap="0.4"/>
+            <geom type="sphere" size=".1" margin="{0.1 if NEW_GAP_SEMANTICS else 0.5}" gap="0.4"/>
             <freejoint/>
           </body>
         </worldbody>