Repository: google-deepmind/mujoco_warp
Pull request #6: Precompute valid geom collision pairs and implement contact include/exclude
Author: login: thowell
Created: 2025-03-20T20:28:23Z | Merged: 2025-04-15T18:09:56Z

this pr precomputes valid geom collision pairs with the function `collision_driver.geom_pair`. additional filtering for contact exclude is included with this function. the resulting list of geom pairs is computed in `put_model` and is available as the field `nxn_geom_pair` in `Model`.

benchmarking:
```
mjwarp-testspeed --function=collision --mjcf=collision.xml --batch_size=8192
```

this pr:
```
Summary for 8192 parallel rollouts

 Total JIT time: 0.01 s
 Total simulation time: 0.12 s
 Total steps per second: 67,267,046
 Total realtime factor: 134,534.09 x
 Total time per step: 14.87 ns
```

main:
```
Summary for 8192 parallel rollouts

 Total JIT time: 0.01 s
 Total simulation time: 0.12 s
 Total steps per second: 65,986,259
 Total realtime factor: 131,972.52 x
 Total time per step: 15.15 ns
```


Diff:
diff --git a/mujoco_warp/__init__.py b/mujoco_warp/__init__.py
index 4ab79e8c9..4bedba93b 100644
--- a/mujoco_warp/__init__.py
+++ b/mujoco_warp/__init__.py
@@ -29,6 +29,7 @@
 from ._src.forward import implicit as implicit
 from ._src.forward import rungekutta4 as rungekutta4
 from ._src.forward import step as step
+from ._src.io import geom_pair as geom_pair
 from ._src.io import get_data_into as get_data_into
 from ._src.io import make_data as make_data
 from ._src.io import put_data as put_data
diff --git a/mujoco_warp/_src/broad_phase_test.py b/mujoco_warp/_src/broad_phase_test.py
index 544398bb9..2f21df855 100644
--- a/mujoco_warp/_src/broad_phase_test.py
+++ b/mujoco_warp/_src/broad_phase_test.py
@@ -22,6 +22,7 @@
 import mujoco_warp as mjwarp
 
 from . import collision_driver
+from . import io
 from . import test_util
 
 
@@ -215,13 +216,12 @@ def test_nxn_broadphase(self):
     np.testing.assert_allclose(d3.collision_pair.numpy()[3][1], 2)
 
     # one world and zero collisions: contype and conaffinity incompatibility
-    _, _, m4, d4 = test_util.fixture(xml=_NXN_MODEL, keyframe=1)
-    m4.geom_contype = wp.array(
-      np.array(np.repeat(0, m.geom_type.shape)), dtype=wp.int32
-    )
-    m4.geom_conaffinity = wp.array(
-      np.array(np.repeat(1, m.geom_type.shape)), dtype=wp.int32
-    )
+    mjm4, mjd4, m4, d4 = test_util.fixture(xml=_NXN_MODEL, keyframe=1)
+    mjm4.geom_contype[:3] = 0
+    m4.geom_contype = wp.array(mjm4.geom_contype, dtype=wp.int32)
+    geompair, pairid = io.geom_pair(mjm4)
+    m4.nxn_geom_pair = wp.array(geompair, dtype=wp.vec2i)
+    m4.nxn_pairid = wp.array(pairid, dtype=wp.int32)
 
     collision_driver.nxn_broadphase(m4, d4)
     np.testing.assert_allclose(d4.ncollision.numpy()[0], 0)
diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py
index a7ba6270d..ae3f2849c 100644
--- a/mujoco_warp/_src/collision_box.py
+++ b/mujoco_warp/_src/collision_box.py
@@ -19,6 +19,7 @@
 
 import warp as wp
 
+from .collision_primitive import contact_params
 from .collision_primitive import write_contact
 from .math import make_frame
 from .types import Data
@@ -207,6 +208,11 @@ def box_box_kernel(
       continue
 
     worldid = d.collision_worldid[bp_idx]
+
+    geoms, margin, gap, condim, friction, solref, solreffriction, solimp = (
+      contact_params(m, d, tid)
+    )
+
     # transformations
     a_pos, b_pos = d.geom_xpos[worldid, ga], d.geom_xpos[worldid, gb]
     a_mat, b_mat = d.geom_xmat[worldid, ga], d.geom_xmat[worldid, gb]
@@ -311,7 +317,19 @@ def box_box_kernel(
       pos_glob = b_mat @ pos[i] + b_pos
       n_glob = b_mat @ sep_axis
       write_contact(
-        d, dist[i], pos_glob, make_frame(n_glob), margin, wp.vec2i(ga, gb), worldid
+        d,
+        dist[i],
+        pos_glob,
+        make_frame(n_glob),
+        margin,
+        gap,
+        condim,
+        friction,
+        solref,
+        solreffriction,
+        solimp,
+        geoms,
+        worldid,
       )
 
 
diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py
index b6b17b07e..4e383bb36 100644
--- a/mujoco_warp/_src/collision_driver.py
+++ b/mujoco_warp/_src/collision_driver.py
@@ -18,43 +18,16 @@
 from .collision_box import box_box_narrowphase
 from .collision_convex import gjk_narrowphase
 from .collision_primitive import primitive_narrowphase
-from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
 
 
 @wp.func
-def _geom_filter(m: Model, geom1: int, geom2: int, filterparent: bool) -> bool:
-  bodyid1 = m.geom_bodyid[geom1]
-  bodyid2 = m.geom_bodyid[geom2]
-  contype1 = m.geom_contype[geom1]
-  contype2 = m.geom_contype[geom2]
-  conaffinity1 = m.geom_conaffinity[geom1]
-  conaffinity2 = m.geom_conaffinity[geom2]
-  weldid1 = m.body_weldid[bodyid1]
-  weldid2 = m.body_weldid[bodyid2]
-  weld_parentid1 = m.body_weldid[m.body_parentid[weldid1]]
-  weld_parentid2 = m.body_weldid[m.body_parentid[weldid2]]
-
-  self_collision = weldid1 == weldid2
-  parent_child_collision = (
-    filterparent
-    and (weldid1 != 0)
-    and (weldid2 != 0)
-    and ((weldid1 == weld_parentid2) or (weldid2 == weld_parentid1))
-  )
-  mask = (contype1 & conaffinity2) or (contype2 & conaffinity1)
-
-  return mask and (not self_collision) and (not parent_child_collision)
-
-
-@wp.func
-def _add_geom_pair(m: Model, d: Data, geom1: int, geom2: int, worldid: int):
+def _add_geom_pair(m: Model, d: Data, geom1: int, geom2: int, worldid: int, nxnid: int):
   pairid = wp.atomic_add(d.ncollision, 0, 1)
 
   if pairid >= d.nconmax:
@@ -69,6 +42,7 @@ def _add_geom_pair(m: Model, d: Data, geom1: int, geom2: int, worldid: int):
     pair = wp.vec2i(geom1, geom2)
 
   d.collision_pair[pairid] = pair
+  d.collision_pairid[pairid] = m.nxn_pairid[nxnid]
   d.collision_worldid[pairid] = worldid
 
 
@@ -291,66 +265,21 @@ def sap_broadphase_kernel(m: Model, d: Data, num_threads: int, filter_parent: bo
     j = j % m.ngeom
 
     # geom index
-    idx1 = d.sap_sort_index[worldid, i]
-    idx2 = d.sap_sort_index[worldid, j]
+    geom1 = d.sap_sort_index[worldid, i]
+    geom2 = d.sap_sort_index[worldid, j]
+    idx = (m.ngeom - geom1 // 2) * wp.max(0, geom1 - 1) + geom2 - m.ngeom // 2 - 1
 
-    if not _geom_filter(m, idx1, idx2, filter_parent):
+    if m.nxn_pairid[idx] < -1:
       threadId += num_threads
       continue
 
     # Check if the boxes overlap
     if overlap(worldid, i, j, d.sap_geom_sort):
-      _add_geom_pair(m, d, idx1, idx2, worldid)
+      _add_geom_pair(m, d, geom1, geom2, worldid, idx)
 
     threadId += num_threads
 
 
-@wp.kernel
-def get_contact_solver_params_kernel(
-  m: Model,
-  d: Data,
-):
-  tid = wp.tid()
-
-  n_contact_pts = d.ncon[0]
-  if tid >= n_contact_pts:
-    return
-
-  geoms = d.contact.geom[tid]
-  g1 = geoms.x
-  g2 = geoms.y
-
-  margin = wp.max(m.geom_margin[g1], m.geom_margin[g2])
-  gap = wp.max(m.geom_gap[g1], m.geom_gap[g2])
-  solmix1 = m.geom_solmix[g1]
-  solmix2 = m.geom_solmix[g2]
-  mix = solmix1 / (solmix1 + solmix2)
-  mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
-  mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
-  mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
-
-  p1 = m.geom_priority[g1]
-  p2 = m.geom_priority[g2]
-  mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0))
-
-  condim1 = m.geom_condim[g1]
-  condim2 = m.geom_condim[g2]
-  condim = wp.where(
-    p1 == p2, wp.max(condim1, condim2), wp.where(p1 > p2, condim1, condim2)
-  )
-  d.contact.dim[tid] = condim
-
-  if m.geom_solref[g1].x > 0.0 and m.geom_solref[g2].x > 0.0:
-    d.contact.solref[tid] = mix * m.geom_solref[g1] + (1.0 - mix) * m.geom_solref[g2]
-  else:
-    d.contact.solref[tid] = wp.min(m.geom_solref[g1], m.geom_solref[g2])
-  d.contact.includemargin[tid] = margin - gap
-  friction_ = wp.max(m.geom_friction[g1], m.geom_friction[g2])
-  friction5 = vec5(friction_[0], friction_[0], friction_[1], friction_[2], friction_[2])
-  d.contact.friction[tid] = friction5
-  d.contact.solimp[tid] = mix * m.geom_solimp[g1] + (1.0 - mix) * m.geom_solimp[g2]
-
-
 def sap_broadphase(m: Model, d: Data):
   """Broadphase collision detection via sweep-and-prune."""
 
@@ -461,25 +390,18 @@ def sap_broadphase(m: Model, d: Data):
 
 def nxn_broadphase(m: Model, d: Data):
   """Broadphase collision detective via brute-force search."""
-  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
 
   @wp.kernel
   def _nxn_broadphase(m: Model, d: Data):
     worldid, elementid = wp.tid()
-    geom1 = (
-      m.ngeom
-      - 2
-      - int(
-        wp.sqrt(float(-8 * elementid + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5
-      )
-    )
-    geom2 = (
-      elementid
-      + geom1
-      + 1
-      - m.ngeom * (m.ngeom - 1) // 2
-      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
-    )
+
+    # check for valid geom pair
+    if m.nxn_pairid[elementid] < -1:
+      return
+
+    geom = m.nxn_geom_pair[elementid]
+    geom1 = geom[0]
+    geom2 = geom[1]
 
     margin1 = m.geom_margin[geom1]
     margin2 = m.geom_margin[geom2]
@@ -506,24 +428,11 @@ def _nxn_broadphase(m: Model, d: Data):
       dist = wp.dot(-dif, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
       bounds_filter = dist <= bound
 
-    geom_filter = _geom_filter(m, geom1, geom2, filterparent)
-
-    if bounds_filter and geom_filter:
-      _add_geom_pair(m, d, geom1, geom2, worldid)
-
-  wp.launch(
-    _nxn_broadphase, dim=(d.nworld, m.ngeom * (m.ngeom - 1) // 2), inputs=[m, d]
-  )
-
-
-def get_contact_solver_params(m: Model, d: Data):
-  wp.launch(
-    get_contact_solver_params_kernel,
-    dim=[d.nconmax],
-    inputs=[m, d],
-  )
+    if bounds_filter:
+      _add_geom_pair(m, d, geom1, geom2, worldid, elementid)
 
-  # TODO(team): do we need condim sorting, deepest penetrating contact here?
+  if m.nxn_geom_pair.shape[0]:
+    wp.launch(_nxn_broadphase, dim=(d.nworld, m.nxn_geom_pair.shape[0]), inputs=[m, d])
 
 
 @event_scope
@@ -552,8 +461,7 @@ def collision(m: Model, d: Data):
 
   # TODO(team): we should reject far-away contacts in the narrowphase instead of constraint
   #             partitioning because we can move some pressure of the atomics
+  # TODO(team) switch between collision functions and GJK/EPA here
   gjk_narrowphase(m, d)
   primitive_narrowphase(m, d)
   box_box_narrowphase(m, d)
-
-  get_contact_solver_params(m, d)
diff --git a/mujoco_warp/_src/collision_driver_test.py b/mujoco_warp/_src/collision_driver_test.py
index f785be9b7..c53f81fe2 100644
--- a/mujoco_warp/_src/collision_driver_test.py
+++ b/mujoco_warp/_src/collision_driver_test.py
@@ -21,6 +21,7 @@
 
 import mujoco_warp as mjwarp
 
+from . import io
 from . import test_util
 
 
@@ -227,6 +228,209 @@ def test_collision(self, fixture):
           break
       np.testing.assert_equal(result, True, f"Contact {i} not found in Gjk results")
 
+  def test_contact_exclude(self):
+    """Tests contact exclude."""
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body name="body1">
+            <freejoint/>
+            <geom type="sphere" size=".1"/>
+          </body>
+          <body name="body2">
+            <freejoint/>
+            <geom type="sphere" size=".1"/>
+          </body>
+          <body name="body3">
+            <freejoint/>
+            <geom type="sphere" size=".1"/>
+          </body>
+        </worldbody>
+        <contact>
+          <exclude body1="body1" body2="body2"/>
+        </contact>
+      </mujoco>
+    """)
+    geompair, pairid = io.geom_pair(mjm)
+    self.assertEqual(geompair.shape[0], 3)
+    np.testing.assert_equal(pairid, np.array([-2, -1, -1]))
+
+  def test_contact_pair(self):
+    """Tests contact pair."""
+    # no pairs
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body>
+            <freejoint/>
+            <geom type="sphere" size=".1"/>
+          </body>
+        </worldbody>
+      </mujoco>
+    """)
+    _, pairid = io.geom_pair(mjm)
+    self.assertTrue((pairid == -1).all())
+
+    # 1 pair
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body>
+            <freejoint/>
+            <geom name="geom1" type="sphere" size=".1"/>
+          </body>
+          <body>
+            <freejoint/>
+            <geom name="geom2" type="sphere" size=".1"/>
+          </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"/>
+        </contact>
+      </mujoco>
+    """)
+    _, pairid = io.geom_pair(mjm)
+    self.assertTrue((pairid == 0).all())
+
+    # generate contact
+    m = mjwarp.put_model(mjm)
+    d = mjwarp.make_data(mjm)
+    mjwarp.forward(m, d)
+
+    self.assertEqual(d.ncon.numpy()[0], 1)
+    self.assertEqual(d.contact.includemargin.numpy()[0], -1)
+    self.assertEqual(d.contact.dim.numpy()[0], 6)
+    np.testing.assert_allclose(d.contact.friction.numpy()[0], np.array([5, 4, 3, 2, 1]))
+    np.testing.assert_allclose(d.contact.solref.numpy()[0], np.array([-0.25, -0.5]))
+    np.testing.assert_allclose(
+      d.contact.solreffriction.numpy()[0], np.array([2.0, 4.0])
+    )
+    np.testing.assert_allclose(
+      d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5])
+    )
+
+    # 1 pair: override contype and conaffinity
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body name="body1">
+            <freejoint/>
+            <geom name="geom1" type="sphere" size=".1" contype="0" conaffinity="0"/>
+          </body>
+          <body name="body2">
+            <freejoint/>
+            <geom name="geom2" type="sphere" size=".1" contype="0" conaffinity="0"/>
+          </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"/>
+        </contact>
+      </mujoco>
+    """)
+    _, pairid = io.geom_pair(mjm)
+    self.assertTrue((pairid == 0).all())
+
+    # generate contact
+    m = mjwarp.put_model(mjm)
+    d = mjwarp.make_data(mjm)
+    mjwarp.forward(m, d)
+
+    self.assertEqual(d.ncon.numpy()[0], 1)
+    self.assertEqual(d.contact.includemargin.numpy()[0], -1)
+    self.assertEqual(d.contact.dim.numpy()[0], 6)
+    np.testing.assert_allclose(d.contact.friction.numpy()[0], np.array([5, 4, 3, 2, 1]))
+    np.testing.assert_allclose(d.contact.solref.numpy()[0], np.array([-0.25, -0.5]))
+    np.testing.assert_allclose(
+      d.contact.solreffriction.numpy()[0], np.array([2.0, 4.0])
+    )
+    np.testing.assert_allclose(
+      d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5])
+    )
+
+    # 1 pair: override exclude
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body name="body1">
+            <freejoint/>
+            <geom name="geom1" type="sphere" size=".1"/>
+          </body>
+          <body name="body2">
+            <freejoint/>
+            <geom name="geom2" type="sphere" size=".1"/>
+          </body>
+        </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"/>
+        </contact>
+      </mujoco>
+    """)
+    _, pairid = io.geom_pair(mjm)
+    self.assertTrue((pairid == 0).all())
+
+    # generate contact
+    m = mjwarp.put_model(mjm)
+    d = mjwarp.make_data(mjm)
+    mjwarp.forward(m, d)
+
+    self.assertEqual(d.ncon.numpy()[0], 1)
+    self.assertEqual(d.contact.includemargin.numpy()[0], -1)
+    self.assertEqual(d.contact.dim.numpy()[0], 6)
+    np.testing.assert_allclose(d.contact.friction.numpy()[0], np.array([5, 4, 3, 2, 1]))
+    np.testing.assert_allclose(d.contact.solref.numpy()[0], np.array([-0.25, -0.5]))
+    np.testing.assert_allclose(
+      d.contact.solreffriction.numpy()[0], np.array([2.0, 4.0])
+    )
+    np.testing.assert_allclose(
+      d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5])
+    )
+
+    # 1 pair 1 exclude
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <body name="body1">
+            <freejoint/>
+            <geom name="geom1" type="sphere" size=".1"/>
+          </body>
+          <body name="body2">
+            <freejoint/>
+            <geom name="geom2" type="sphere" size=".1"/>
+          </body>
+          <body name="body3">
+            <freejoint/>
+            <geom name="geom3" type="sphere" size=".1"/>
+          </body>
+        </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"/>
+        </contact>
+      </mujoco>
+    """)
+    _, pairid = io.geom_pair(mjm)
+    np.testing.assert_equal(pairid, np.array([-2, -1, 0]))
+
+    # generate contact
+    m = mjwarp.put_model(mjm)
+    d = mjwarp.make_data(mjm)
+    mjwarp.forward(m, d)
+
+    self.assertEqual(d.ncon.numpy()[0], 2)
+    self.assertEqual(d.contact.includemargin.numpy()[1], -1)
+    self.assertEqual(d.contact.dim.numpy()[1], 6)
+    np.testing.assert_allclose(d.contact.friction.numpy()[1], np.array([5, 4, 3, 2, 1]))
+    np.testing.assert_allclose(d.contact.solref.numpy()[1], np.array([-0.25, -0.5]))
+    np.testing.assert_allclose(
+      d.contact.solreffriction.numpy()[1], np.array([2.0, 4.0])
+    )
+    np.testing.assert_allclose(
+      d.contact.solimp.numpy()[1], np.array([0.1, 0.2, 0.3, 0.4, 0.5])
+    )
+
+    # TODO(team): test sap_broadphase
+
   @parameterized.parameters(
     (True, True),
     (True, False),
@@ -248,6 +452,8 @@ def test_collision_disableflags(self, constraint, contact):
 
     self.assertEqual(d.ncon.numpy()[0], mjd.ncon)
 
+  # TODO(team): test contact parameter mixing
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py
index f6d2ca996..c62afa706 100644
--- a/mujoco_warp/_src/collision_primitive.py
+++ b/mujoco_warp/_src/collision_primitive.py
@@ -19,9 +19,11 @@
 from .math import closest_segment_to_segment_points
 from .math import make_frame
 from .math import normalize_with_norm
+from .types import MJ_MINVAL
 from .types import Data
 from .types import GeomType
 from .types import Model
+from .types import vec5
 
 wp.set_module_options({"enable_backward": False})
 
@@ -67,18 +69,30 @@ def write_contact(
   pos: wp.vec3,
   frame: wp.mat33,
   margin: float,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
   geoms: wp.vec2i,
   worldid: int,
 ):
   active = (dist - margin) < 0
   if active:
-    index = wp.atomic_add(d.ncon, 0, 1)
-    if index < d.nconmax:
-      d.contact.dist[index] = dist
-      d.contact.pos[index] = pos
-      d.contact.frame[index] = frame
-      d.contact.geom[index] = geoms
-      d.contact.worldid[index] = worldid
+    cid = wp.atomic_add(d.ncon, 0, 1)
+    if cid < d.nconmax:
+      d.contact.dist[cid] = dist
+      d.contact.pos[cid] = pos
+      d.contact.frame[cid] = frame
+      d.contact.geom[cid] = geoms
+      d.contact.worldid[cid] = worldid
+      d.contact.includemargin[cid] = margin - gap
+      d.contact.dim[cid] = condim
+      d.contact.friction[cid] = friction
+      d.contact.solref[cid] = solref
+      d.contact.solreffriction[cid] = solreffriction
+      d.contact.solimp[cid] = solimp
 
 
 @wp.func
@@ -97,11 +111,31 @@ def plane_sphere(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   dist, pos = _plane_sphere(plane.normal, plane.pos, sphere.pos, sphere.size[0])
 
-  write_contact(d, dist, pos, make_frame(plane.normal), margin, geom_indices, worldid)
+  write_contact(
+    d,
+    dist,
+    pos,
+    make_frame(plane.normal),
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+    worldid,
+  )
 
 
 @wp.func
@@ -113,7 +147,13 @@ def _sphere_sphere(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   dir = pos2 - pos1
   dist = wp.length(dir)
@@ -124,7 +164,21 @@ def _sphere_sphere(
   dist = dist - (radius1 + radius2)
   pos = pos1 + n * (radius1 + 0.5 * dist)
 
-  write_contact(d, dist, pos, make_frame(n), margin, geom_indices, worldid)
+  write_contact(
+    d,
+    dist,
+    pos,
+    make_frame(n),
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+    worldid,
+  )
 
 
 @wp.func
@@ -136,7 +190,13 @@ def _sphere_sphere_ext(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
   mat1: wp.mat33,
   mat2: wp.mat33,
 ):
@@ -153,7 +213,21 @@ def _sphere_sphere_ext(
   dist = dist - (radius1 + radius2)
   pos = pos1 + n * (radius1 + 0.5 * dist)
 
-  write_contact(d, dist, pos, make_frame(n), margin, geom_indices, worldid)
+  write_contact(
+    d,
+    dist,
+    pos,
+    make_frame(n),
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+    worldid,
+  )
 
 
 @wp.func
@@ -163,7 +237,13 @@ def sphere_sphere(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   _sphere_sphere(
     sphere1.pos,
@@ -173,7 +253,13 @@ def sphere_sphere(
     worldid,
     d,
     margin,
-    geom_indices,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
   )
 
 
@@ -184,7 +270,13 @@ def sphere_capsule(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   """Calculates one contact between a sphere and a capsule."""
   axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2])
@@ -196,7 +288,20 @@ def sphere_capsule(
 
   # Treat as sphere-sphere collision between sphere and closest point
   _sphere_sphere(
-    sphere.pos, sphere.size[0], pt, cap.size[0], worldid, d, margin, geom_indices
+    sphere.pos,
+    sphere.size[0],
+    pt,
+    cap.size[0],
+    worldid,
+    d,
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
   )
 
 
@@ -207,7 +312,13 @@ def capsule_capsule(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   axis1 = wp.vec3(cap1.rot[0, 2], cap1.rot[1, 2], cap1.rot[2, 2])
   axis2 = wp.vec3(cap2.rot[0, 2], cap2.rot[1, 2], cap2.rot[2, 2])
@@ -223,7 +334,22 @@ def capsule_capsule(
     cap2.pos + seg2,
   )
 
-  _sphere_sphere(pt1, cap1.size[0], pt2, cap2.size[0], worldid, d, margin, geom_indices)
+  _sphere_sphere(
+    pt1,
+    cap1.size[0],
+    pt2,
+    cap2.size[0],
+    worldid,
+    d,
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+  )
 
 
 @wp.func
@@ -233,7 +359,13 @@ def plane_capsule(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   """Calculates two contacts between a capsule and a plane."""
   n = plane.normal
@@ -252,10 +384,38 @@ def plane_capsule(
   segment = axis * cap.size[1]
 
   dist1, pos1 = _plane_sphere(n, plane.pos, cap.pos + segment, cap.size[0])
-  write_contact(d, dist1, pos1, frame, margin, geom_indices, worldid)
+  write_contact(
+    d,
+    dist1,
+    pos1,
+    frame,
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+    worldid,
+  )
 
   dist2, pos2 = _plane_sphere(n, plane.pos, cap.pos - segment, cap.size[0])
-  write_contact(d, dist2, pos2, frame, margin, geom_indices, worldid)
+  write_contact(
+    d,
+    dist2,
+    pos2,
+    frame,
+    margin,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
+    worldid,
+  )
 
 
 @wp.func
@@ -265,7 +425,13 @@ def plane_box(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   count = int(0)
   corner = wp.vec3()
@@ -289,7 +455,21 @@ def plane_box(
     cdist = dist + ldist
     frame = make_frame(plane.normal)
     pos = corner + box.pos + (plane.normal * cdist / -2.0)
-    write_contact(d, cdist, pos, frame, margin, geom_indices, worldid)
+    write_contact(
+      d,
+      cdist,
+      pos,
+      frame,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+      worldid,
+    )
     count += 1
     if count >= 4:
       break
@@ -302,7 +482,13 @@ def sphere_cylinder(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   axis = wp.vec3(
     cylinder.rot[0, 2],
@@ -340,7 +526,13 @@ def sphere_cylinder(
       worldid,
       d,
       margin,
-      geom_indices,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
       sphere.rot,
       cylinder.rot,
     )
@@ -366,7 +558,13 @@ def sphere_cylinder(
       pos_contact,
       make_frame(plane_normal),
       margin,
-      geom_indices,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
       worldid,
     )
 
@@ -387,7 +585,13 @@ def sphere_cylinder(
     worldid,
     d,
     margin,
-    geom_indices,
+    gap,
+    condim,
+    friction,
+    solref,
+    solreffriction,
+    solimp,
+    geoms,
     sphere.rot,
     cylinder.rot,
   )
@@ -400,7 +604,13 @@ def plane_cylinder(
   worldid: int,
   d: Data,
   margin: float,
-  geom_indices: wp.vec2i,
+  gap: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2f,
+  solreffriction: wp.vec2f,
+  solimp: vec5,
+  geoms: wp.vec2i,
 ):
   """Calculates contacts between a cylinder and a plane."""
   # Extract plane normal and cylinder axis
@@ -442,7 +652,21 @@ def plane_cylinder(
   dist1 = dist0 + prjaxis + prjvec
   if dist1 <= margin:
     pos1 = cylinder.pos + vec + axis - n * (dist1 * 0.5)
-    write_contact(d, dist1, pos1, frame, margin, geom_indices, worldid)
+    write_contact(
+      d,
+      dist1,
+      pos1,
+      frame,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+      worldid,
+    )
   else:
     # If nearest point is above margin, no contacts
     return
@@ -451,7 +675,21 @@ def plane_cylinder(
   dist2 = dist0 - prjaxis + prjvec
   if dist2 <= margin:
     pos2 = cylinder.pos + vec - axis - n * (dist2 * 0.5)
-    write_contact(d, dist2, pos2, frame, margin, geom_indices, worldid)
+    write_contact(
+      d,
+      dist2,
+      pos2,
+      frame,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+      worldid,
+    )
 
   # Try triangle contact points on side closer to plane
   prjvec1 = -prjvec * 0.5
@@ -463,11 +701,98 @@ def plane_cylinder(
 
     # Add contact point A - adjust to closest side
     pos3 = cylinder.pos + vec1 + axis - vec * 0.5 - n * (dist3 * 0.5)
-    write_contact(d, dist3, pos3, frame, margin, geom_indices, worldid)
+    write_contact(
+      d,
+      dist3,
+      pos3,
+      frame,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+      worldid,
+    )
 
     # Add contact point B - adjust to closest side
     pos4 = cylinder.pos - vec1 + axis - vec * 0.5 - n * (dist3 * 0.5)
-    write_contact(d, dist3, pos4, frame, margin, geom_indices, worldid)
+    write_contact(
+      d,
+      dist3,
+      pos4,
+      frame,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+      worldid,
+    )
+
+
+@wp.func
+def contact_params(m: Model, d: Data, cid: int):
+  geoms = d.collision_pair[cid]
+  pairid = d.collision_pairid[cid]
+
+  if pairid > -1:
+    margin = m.pair_margin[pairid]
+    gap = m.pair_gap[pairid]
+    condim = m.pair_dim[pairid]
+    friction = m.pair_friction[pairid]
+    solref = m.pair_solref[pairid]
+    solreffriction = m.pair_solreffriction[pairid]
+    solimp = m.pair_solimp[pairid]
+  else:
+    g1 = geoms[0]
+    g2 = geoms[1]
+
+    p1 = m.geom_priority[g1]
+    p2 = m.geom_priority[g2]
+
+    solmix1 = m.geom_solmix[g1]
+    solmix2 = m.geom_solmix[g2]
+
+    mix = solmix1 / (solmix1 + solmix2)
+    mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
+    mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
+    mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
+    mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0))
+
+    margin = wp.max(m.geom_margin[g1], m.geom_margin[g2])
+    gap = wp.max(m.geom_gap[g1], m.geom_gap[g2])
+
+    condim1 = m.geom_condim[g1]
+    condim2 = m.geom_condim[g2]
+    condim = wp.where(
+      p1 == p2, wp.max(condim1, condim2), wp.where(p1 > p2, condim1, condim2)
+    )
+
+    geom_friction = wp.max(m.geom_friction[g1], m.geom_friction[g2])
+    friction = vec5(
+      geom_friction[0],
+      geom_friction[0],
+      geom_friction[1],
+      geom_friction[2],
+      geom_friction[2],
+    )
+
+    if m.geom_solref[g1].x > 0.0 and m.geom_solref[g2].x > 0.0:
+      solref = mix * m.geom_solref[g1] + (1.0 - mix) * m.geom_solref[g2]
+    else:
+      solref = wp.min(m.geom_solref[g1], m.geom_solref[g2])
+
+    solreffriction = wp.vec2(0.0, 0.0)
+
+    solimp = mix * m.geom_solimp[g1] + (1.0 - mix) * m.geom_solimp[g2]
+
+  return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
 
 
 @wp.kernel
@@ -480,36 +805,141 @@ def _primitive_narrowphase(
   if tid >= d.ncollision[0]:
     return
 
-  geoms = d.collision_pair[tid]
-  worldid = d.collision_worldid[tid]
-
+  geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
+    m, d, tid
+  )
   g1 = geoms[0]
   g2 = geoms[1]
-  type1 = m.geom_type[g1]
-  type2 = m.geom_type[g2]
+
+  worldid = d.collision_worldid[tid]
 
   geom1 = _geom(g1, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
   geom2 = _geom(g2, m, d.geom_xpos[worldid], d.geom_xmat[worldid])
 
-  margin = wp.max(m.geom_margin[g1], m.geom_margin[g2])
+  type1 = m.geom_type[g1]
+  type2 = m.geom_type[g2]
 
   # TODO(team): static loop unrolling to remove unnecessary branching
   if type1 == int(GeomType.PLANE.value) and type2 == int(GeomType.SPHERE.value):
-    plane_sphere(geom1, geom2, worldid, d, margin, geoms)
+    plane_sphere(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.SPHERE.value) and type2 == int(GeomType.SPHERE.value):
-    sphere_sphere(geom1, geom2, worldid, d, margin, geoms)
+    sphere_sphere(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.PLANE.value) and type2 == int(GeomType.CAPSULE.value):
-    plane_capsule(geom1, geom2, worldid, d, margin, geoms)
+    plane_capsule(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.PLANE.value) and type2 == int(GeomType.BOX.value):
-    plane_box(geom1, geom2, worldid, d, margin, geoms)
+    plane_box(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.CAPSULE.value) and type2 == int(GeomType.CAPSULE.value):
-    capsule_capsule(geom1, geom2, worldid, d, margin, geoms)
+    capsule_capsule(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.SPHERE.value) and type2 == int(GeomType.CAPSULE.value):
-    sphere_capsule(geom1, geom2, worldid, d, margin, geoms)
+    sphere_capsule(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.SPHERE.value) and type2 == int(GeomType.CYLINDER.value):
-    sphere_cylinder(geom1, geom2, worldid, d, margin, geoms)
+    sphere_cylinder(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
   elif type1 == int(GeomType.PLANE.value) and type2 == int(GeomType.CYLINDER.value):
-    plane_cylinder(geom1, geom2, worldid, d, margin, geoms)
+    plane_cylinder(
+      geom1,
+      geom2,
+      worldid,
+      d,
+      margin,
+      gap,
+      condim,
+      friction,
+      solref,
+      solreffriction,
+      solimp,
+      geoms,
+    )
 
 
 def primitive_narrowphase(m: Model, d: Data):
diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py
index 2df1f1b90..45f271536 100644
--- a/mujoco_warp/_src/io.py
+++ b/mujoco_warp/_src/io.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import mujoco
 import numpy as np
@@ -24,6 +24,52 @@
 from . import types
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
+  filterparent = not (m.opt.disableflags & types.DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+  predefined_pairs = {(m.pair_geom1[i], m.pair_geom2[i]): i for i in range(m.npair)}
+
+  tri = np.triu_indices(m.ngeom, k=1)  # k=1 to skip self collision pairs
+
+  geompairs = []
+  pairids = []
+  for geom1, geom2 in zip(*tri):
+    bodyid1 = m.geom_bodyid[geom1]
+    bodyid2 = m.geom_bodyid[geom2]
+    contype1 = m.geom_contype[geom1]
+    contype2 = m.geom_contype[geom2]
+    conaffinity1 = m.geom_conaffinity[geom1]
+    conaffinity2 = m.geom_conaffinity[geom2]
+    weldid1 = m.body_weldid[bodyid1]
+    weldid2 = m.body_weldid[bodyid2]
+    weld_parentid1 = m.body_weldid[m.body_parentid[weldid1]]
+    weld_parentid2 = m.body_weldid[m.body_parentid[weldid2]]
+
+    self_collision = weldid1 == weldid2
+    parent_child_collision = (
+      filterparent
+      and (weldid1 != 0)
+      and (weldid2 != 0)
+      and ((weldid1 == weld_parentid2) or (weldid2 == weld_parentid1))
+    )
+    mask = (contype1 & conaffinity2) or (contype2 & conaffinity1)
+    exclude = (bodyid1 << 16) + (bodyid2) in exclude_signature
+
+    if mask and (not self_collision) and (not parent_child_collision) and (not exclude):
+      pairid = -1
+    else:
+      pairid = -2
+
+    # check for predefined geom pair
+    pairid = predefined_pairs.get((geom1, geom2), pairid)
+    pairid = predefined_pairs.get((geom2, geom1), pairid)
+
+    pairids.append(pairid)
+    geompairs.append([geom1, geom2])
+
+  return np.array(geompairs), np.array(pairids)
+
+
 def put_model(mjm: mujoco.MjModel) -> types.Model:
   # check supported features
   for field, field_types, field_str in (
@@ -92,6 +138,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.nsensor = mjm.nsensor
   m.nsensordata = mjm.nsensordata
   m.nlsp = mjm.opt.ls_iterations  # TODO(team): how to set nlsp?
+  m.npair = mjm.npair
   m.nexclude = mjm.nexclude
   m.neq = mjm.neq
   m.opt.timestep = mjm.opt.timestep
@@ -439,6 +486,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     or np.any(mjm.actuator_gaintype == types.GainType.AFFINE.value)
   )
 
+  geompair, pairid = geom_pair(mjm)
+  m.nxn_geom_pair = wp.array(geompair, dtype=wp.vec2i, ndim=1)
+  m.nxn_pairid = wp.array(pairid, dtype=wp.int32, ndim=1)
+
+  # predefined collision pairs
+  m.pair_dim = wp.array(mjm.pair_dim, dtype=wp.int32, ndim=1)
+  m.pair_geom1 = wp.array(mjm.pair_geom1, dtype=wp.int32, ndim=1)
+  m.pair_geom2 = wp.array(mjm.pair_geom2, dtype=wp.int32, ndim=1)
+  m.pair_solref = wp.array(mjm.pair_solref, dtype=wp.vec2, ndim=1)
+  m.pair_solreffriction = wp.array(mjm.pair_solreffriction, dtype=wp.vec2, ndim=1)
+  m.pair_solimp = wp.array(mjm.pair_solimp, dtype=types.vec5, ndim=1)
+  m.pair_margin = wp.array(mjm.pair_margin, dtype=wp.float32, ndim=1)
+  m.pair_gap = wp.array(mjm.pair_gap, dtype=wp.float32, ndim=1)
+  m.pair_friction = wp.array(mjm.pair_friction, dtype=types.vec5, ndim=1)
   m.condim_max = np.max(mjm.geom_condim)  # TODO(team): get max after filtering
 
   # tendon
@@ -672,6 +733,7 @@ def make_data(
 
   # collision driver
   d.collision_pair = wp.empty(nconmax, dtype=wp.vec2i, ndim=1)
+  d.collision_pairid = wp.empty(nconmax, dtype=wp.int32, ndim=1)
   d.collision_worldid = wp.empty(nconmax, dtype=wp.int32, ndim=1)
   d.ncollision = wp.zeros(1, dtype=wp.int32, ndim=1)
 
@@ -939,6 +1001,7 @@ def tile(x):
 
   # collision driver
   d.collision_pair = wp.empty(nconmax, dtype=wp.vec2i, ndim=1)
+  d.collision_pairid = wp.empty(nconmax, dtype=wp.int32, ndim=1)
   d.collision_worldid = wp.empty(nconmax, dtype=wp.int32, ndim=1)
   d.ncollision = wp.zeros(1, dtype=wp.int32, ndim=1)
 
diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py
index a24f4ae61..e50ee8b93 100644
--- a/mujoco_warp/_src/types.py
+++ b/mujoco_warp/_src/types.py
@@ -532,6 +532,7 @@ class Model:
     nsensor: number of sensors                               ()
     nsensordata: number of elements in sensor data vector    ()
     nlsp: number of step sizes for parallel linsearch        ()
+    npair: number of predefined geom pairs                   ()
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                       (nq,)
@@ -672,6 +673,17 @@ class Model:
     actuator_gear: scale length and transmitted force        (nu, 6)
     exclude_signature: body1 << 16 + body2                   (nexclude,)
     actuator_affine_bias_gain: affine bias/gain present
+    nxn_geom_pair: valid collision pair geom ids             (<= ngeom * (ngeom - 1) // 2,)
+    nxn_pairid: predefined pair id, -1 if not predefined     (<= ngeom * (ngeom - 1) // 2,)
+    pair_dim: contact dimensionality                         (npair,)
+    pair_geom1: id of geom1                                  (npair,)
+    pair_geom2: id of geom2                                  (npair,)
+    pair_solref: solver reference: contact normal            (npair, mjNREF)
+    pair_solreffriction: solver reference: contact friction  (npair, mjNREF)
+    pair_solimp: solver impedance: contact                   (npair, mjNIMP)
+    pair_margin: detect contact if dist<margin               (npair,)
+    pair_gap: include in solver if dist<margin-gap           (npair,)
+    pair_friction: tangent1, 2, spin, roll1, 2               (npair, 5)
     condim_max: maximum condim for geoms
     tendon_adr: address of first object in tendon's path     (ntendon,)
     tendon_num: number of objects in tendon's path           (ntendon,)
@@ -713,6 +725,7 @@ class Model:
   nsensor: int
   nsensordata: int
   nlsp: int  # warp only
+  npair: int
   opt: Option
   stat: Statistic
   qpos0: wp.array(dtype=wp.float32, ndim=1)
@@ -853,6 +866,19 @@ class Model:
   actuator_gear: wp.array(dtype=wp.spatial_vector, ndim=1)
   exclude_signature: wp.array(dtype=wp.int32, ndim=1)
   actuator_affine_bias_gain: bool  # warp only
+  nxn_geom_pair: wp.array(dtype=wp.vec2i, ndim=1)  # warp only
+  nxn_pairid: wp.array(dtype=wp.int32, ndim=1)  # warp only
+
+  # predefined geom pairs for collision detection
+  pair_dim: wp.array(dtype=wp.int32, ndim=1)
+  pair_geom1: wp.array(dtype=wp.int32, ndim=1)
+  pair_geom2: wp.array(dtype=wp.int32, ndim=1)
+  pair_solref: wp.array(dtype=wp.vec2, ndim=1)
+  pair_solreffriction: wp.array(dtype=wp.vec2, ndim=1)
+  pair_solimp: wp.array(dtype=vec5, ndim=1)
+  pair_margin: wp.array(dtype=wp.float32, ndim=1)
+  pair_gap: wp.array(dtype=wp.float32, ndim=1)
+  pair_friction: wp.array(dtype=vec5, ndim=1)
   condim_max: int  # warp only
   tendon_adr: wp.array(dtype=wp.int32, ndim=1)
   tendon_num: wp.array(dtype=wp.int32, ndim=1)
@@ -1098,6 +1124,7 @@ class Data:
 
   # collision driver
   collision_pair: wp.array(dtype=wp.vec2i, ndim=1)
+  collision_pairid: wp.array(dtype=wp.int32, ndim=1)
   collision_worldid: wp.array(dtype=wp.int32, ndim=1)
   ncollision: wp.array(dtype=wp.int32, ndim=1)
 


Reviews:
user: login: adenzler-nvidia
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T08:22:30Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T13:28:58Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T13:29:38Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T13:32:18Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-21T13:34:20Z
user: login: erikfrey
body: OK, this is looking good.  Besides a few nits, we just need to see if we can also plug this into SAP and then we're all set to merge.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T21:42:26Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-12T14:15:57Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-12T14:16:08Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-12T14:16:17Z
user: login: erikfrey
body: OK - this ended up in a great place.  Thanks for you patience on the back and forth on this one.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-15T18:09:50Z

Review comments:
diff_hunk: @@ -24,6 +26,52 @@
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> np.array:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+
+  pairs = []
+  for i in range(m.ngeom * (m.ngeom - 1) // 2):
+    geom1 = (
+      m.ngeom
+      - 2
+      - int(wp.sqrt(float(-8 * i + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5)
+    )
+    geom2 = (
+      i
+      + geom1
+      + 1
+      - m.ngeom * (m.ngeom - 1) // 2
+      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
+    )
+
+    bodyid1 = m.geom_bodyid[geom1]
+    bodyid2 = m.geom_bodyid[geom2]
+    contype1 = m.geom_contype[geom1]
+    contype2 = m.geom_contype[geom2]
+    conaffinity1 = m.geom_conaffinity[geom1]
+    conaffinity2 = m.geom_conaffinity[geom2]
+    weldid1 = m.body_weldid[bodyid1]
+    weldid2 = m.body_weldid[bodyid2]
+    weld_parentid1 = m.body_weldid[m.body_parentid[weldid1]]
+    weld_parentid2 = m.body_weldid[m.body_parentid[weldid2]]
+
+    self_collision = weldid1 == weldid2
+    parent_child_collision = (
+      filterparent
+      and (weldid1 != 0)
+      and (weldid2 != 0)
+      and ((weldid1 == weld_parentid2) or (weldid2 == weld_parentid1))
+    )
+    mask = (contype1 & conaffinity2) or (contype2 & conaffinity1)
+    exclude = (bodyid1 << 16) + (bodyid2) in exclude_signature
+
+    if mask and (not self_collision) and (not parent_child_collision) and (not exclude):
+      pairs.append([geom1, geom2])
+
+  return np.array(pairs)
path: mujoco_warp/_src/collision_driver.py
user: login: adenzler-nvidia
body: I would have expected this function in io.py - or to be running on the Model/Data structs we have in warp. Is this a pattern you're copying from MJX?
created_at: 2025-03-21T08:15:21Z
author_association: COLLABORATOR
diff_hunk: @@ -462,20 +510,8 @@ def nxn_broadphase(m: Model, d: Data):
   @wp.kernel
   def _nxn_broadphase(m: Model, d: Data):
     worldid, elementid = wp.tid()
-    geom1 = (
-      m.ngeom
-      - 2
-      - int(
-        wp.sqrt(float(-8 * elementid + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5
-      )
-    )
-    geom2 = (
-      elementid
-      + geom1
-      + 1
-      - m.ngeom * (m.ngeom - 1) // 2
-      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
-    )
+    geom1 = m.collision_geom_pair[elementid][0]
+    geom2 = m.collision_geom_pair[elementid][1]
path: mujoco_warp/_src/collision_driver.py
user: login: adenzler-nvidia
body: Thanks - I think it makes sense to just have the lookup here. Minor nitpick - My intuition tells me that it's probably better to load the vec2 into a variable and then split from that variable, that way the compiler can lower this to a wider load instruction. Not sure if this is going to help performance here, but it certainly won't hurt.
created_at: 2025-03-21T08:17:08Z
author_association: COLLABORATOR
diff_hunk: @@ -24,6 +26,52 @@
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> np.array:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+
+  pairs = []
+  for i in range(m.ngeom * (m.ngeom - 1) // 2):
+    geom1 = (
+      m.ngeom
+      - 2
+      - int(wp.sqrt(float(-8 * i + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5)
+    )
+    geom2 = (
+      i
+      + geom1
+      + 1
+      - m.ngeom * (m.ngeom - 1) // 2
+      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
+    )
path: mujoco_warp/_src/collision_driver.py
user: login: adenzler-nvidia
body: If this isn't going to run during sim - wouldn't it make more sense to just have the double loop to enumerate the pairs? Would make this easier much easier to understand.
created_at: 2025-03-21T08:21:13Z
author_association: COLLABORATOR
diff_hunk: @@ -488,6 +488,7 @@ class Model:
     actuator_gear: scale length and transmitted force        (nu, 6)
     exclude_signature: body1 << 16 + body2                   (nexclude,)
     actuator_affine_bias_gain: affine bias/gain present
+    collision_geom_pair: valid collision pair geom ids       (<= ngeom * (ngeom - 1) // 2,)
path: mujoco_warp/_src/types.py
user: login: adenzler-nvidia
body: I think it would be useful to somehow encode in the name or at least have a comment telling that this is only for the n2 broadphase.
created_at: 2025-03-21T08:22:10Z
author_association: COLLABORATOR
diff_hunk: @@ -462,20 +510,8 @@ def nxn_broadphase(m: Model, d: Data):
   @wp.kernel
   def _nxn_broadphase(m: Model, d: Data):
     worldid, elementid = wp.tid()
-    geom1 = (
-      m.ngeom
-      - 2
-      - int(
-        wp.sqrt(float(-8 * elementid + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5
-      )
-    )
-    geom2 = (
-      elementid
-      + geom1
-      + 1
-      - m.ngeom * (m.ngeom - 1) // 2
-      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
-    )
+    geom1 = m.collision_geom_pair[elementid][0]
+    geom2 = m.collision_geom_pair[elementid][1]
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: done
created_at: 2025-03-21T13:28:58Z
author_association: COLLABORATOR
diff_hunk: @@ -488,6 +488,7 @@ class Model:
     actuator_gear: scale length and transmitted force        (nu, 6)
     exclude_signature: body1 << 16 + body2                   (nexclude,)
     actuator_affine_bias_gain: affine bias/gain present
+    collision_geom_pair: valid collision pair geom ids       (<= ngeom * (ngeom - 1) // 2,)
path: mujoco_warp/_src/types.py
user: login: thowell
body: `collision_geom_pair` -> `nxn_geom_pair`
created_at: 2025-03-21T13:29:38Z
author_association: COLLABORATOR
diff_hunk: @@ -24,6 +26,52 @@
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> np.array:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+
+  pairs = []
+  for i in range(m.ngeom * (m.ngeom - 1) // 2):
+    geom1 = (
+      m.ngeom
+      - 2
+      - int(wp.sqrt(float(-8 * i + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5)
+    )
+    geom2 = (
+      i
+      + geom1
+      + 1
+      - m.ngeom * (m.ngeom - 1) // 2
+      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
+    )
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: modified to use the indices from `numpy.triu_indices`
created_at: 2025-03-21T13:32:18Z
author_association: COLLABORATOR
diff_hunk: @@ -24,6 +26,52 @@
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> np.array:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+
+  pairs = []
+  for i in range(m.ngeom * (m.ngeom - 1) // 2):
+    geom1 = (
+      m.ngeom
+      - 2
+      - int(wp.sqrt(float(-8 * i + 4 * m.ngeom * (m.ngeom - 1) - 7)) / 2.0 - 0.5)
+    )
+    geom2 = (
+      i
+      + geom1
+      + 1
+      - m.ngeom * (m.ngeom - 1) // 2
+      + (m.ngeom - geom1) * ((m.ngeom - geom1) - 1) // 2
+    )
+
+    bodyid1 = m.geom_bodyid[geom1]
+    bodyid2 = m.geom_bodyid[geom2]
+    contype1 = m.geom_contype[geom1]
+    contype2 = m.geom_contype[geom2]
+    conaffinity1 = m.geom_conaffinity[geom1]
+    conaffinity2 = m.geom_conaffinity[geom2]
+    weldid1 = m.body_weldid[bodyid1]
+    weldid2 = m.body_weldid[bodyid2]
+    weld_parentid1 = m.body_weldid[m.body_parentid[weldid1]]
+    weld_parentid2 = m.body_weldid[m.body_parentid[weldid2]]
+
+    self_collision = weldid1 == weldid2
+    parent_child_collision = (
+      filterparent
+      and (weldid1 != 0)
+      and (weldid2 != 0)
+      and ((weldid1 == weld_parentid2) or (weldid2 == weld_parentid1))
+    )
+    mask = (contype1 & conaffinity2) or (contype2 & conaffinity1)
+    exclude = (bodyid1 << 16) + (bodyid2) in exclude_signature
+
+    if mask and (not self_collision) and (not parent_child_collision) and (not exclude):
+      pairs.append([geom1, geom2])
+
+  return np.array(pairs)
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: mjx has a similar function `geom_pairs` in collision_driver.py https://github.com/google-deepmind/mujoco/blob/dc96ed6bdb67dfb9cc277977c99340dd71ea8d2d/mjx/mujoco/mjx/_src/collision_driver.py#L121
created_at: 2025-03-21T13:34:20Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
path: mujoco_warp/_src/collision_driver.py
user: login: erikfrey
body: Although we have logic like this in C MuJoCo's collision driver, it's being used quite differently in Warp.  I'd recommend putting this in `io.py` since it's part of model creation.
created_at: 2025-04-11T19:36:41Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+  predefined_pairs = [(g1, g2) for g1, g2 in zip(m.pair_geom1, m.pair_geom2)]
+
+  tri = np.triu_indices(m.ngeom)
+  ntri = tri[0].size
+
+  pairs = []
+  pairids = []
+  for i in range(ntri):
path: mujoco_warp/_src/collision_driver.py
user: login: erikfrey
body: nit: consider instead:

```py
for geom1, geom2 in zip(*tri):
  ...
```

side benefit is you redefined `i` below which now you can avoid
created_at: 2025-04-11T19:40:00Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+  predefined_pairs = [(g1, g2) for g1, g2 in zip(m.pair_geom1, m.pair_geom2)]
path: mujoco_warp/_src/collision_driver.py
user: login: erikfrey
body: I may have been tempted to do:

```py
predefined_pairs = {(m.pair_geom1[i], m.pair_geom2[i]): i for i in range(m.npair)}
```

then down below:

```py
pairid = predefined_pairs.get((geom1, geom2), -1)
pairid = predefined_pairs.get((geom2, geom1), pairid)
```

But feel free to ignore if you think it's overkill.
created_at: 2025-04-11T19:53:03Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: done
created_at: 2025-04-12T14:15:57Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+  predefined_pairs = [(g1, g2) for g1, g2 in zip(m.pair_geom1, m.pair_geom2)]
+
+  tri = np.triu_indices(m.ngeom)
+  ntri = tri[0].size
+
+  pairs = []
+  pairids = []
+  for i in range(ntri):
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: done
created_at: 2025-04-12T14:16:08Z
author_association: COLLABORATOR
diff_hunk: @@ -13,17 +13,71 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
+import mujoco
+import numpy as np
 import warp as wp
 
 from .collision_primitive import primitive_narrowphase
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DisableBit
 from .types import Model
-from .types import vec5
 from .warp_util import event_scope
 
 
+def geom_pair(m: mujoco.MjModel) -> Tuple[np.array, np.array]:
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT.value)
+  exclude_signature = set(m.exclude_signature)
+  predefined_pairs = [(g1, g2) for g1, g2 in zip(m.pair_geom1, m.pair_geom2)]
path: mujoco_warp/_src/collision_driver.py
user: login: thowell
body: done
created_at: 2025-04-12T14:16:17Z
author_association: COLLABORATOR

PR comments:
user: login: erikfrey
created_at: 2025-03-21T14:04:54Z
body: Hmm!  This is an interesting design quesiton, I think there's a few pros and cons.  I'll list them out here later today and then we can discuss.
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-03-28T23:01:36Z
body: OK!  Very interesting!

Pros of precomputing valid geom pairs:
* We can skip ever even considering invalid geom pairs, e.g. scenes with many visual geoms will be a bit faster because broadphase never even visits that pair
* `<pair>` include and exclude is a bit easier to insert/delete statically than insert/delete on the fly.  MuJoCo efficiently merges broadphase results and pair contacts by ensuring both are sorted, then it can do a linear merge, in order to ensure we don't have duplicate contacts.  If we do this dynamically we'll have to futz with it a bit to ensure it's fast.

Cons:
* We're reifying an array of geom pairs that is ngeom^2.  Pragmatically this is probably not a big deal, even at 10,000 geoms that's only like 500MB or so on device.  So this is not an enormous con.
* If user changes contype or conaffinity, we must now re-JIT the physics step, whereas the existing version has no such requirement.

I think I could go either way.  Curious what you all think.
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-04-02T23:26:03Z
body: @thowell @adenzler-nvidia WDYT?
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-04-03T00:08:31Z
body: i think the identified cons are minor compared to the pros for the precomputed geom pairs approach.

- based on the estimate of ~500MB for 10,000 geoms, the memory requirements for the precomputed geom pairs seem reasonable
- if the model is changed (eg, contype/conaffinity modified) i think it is reasonable to expect that the physics step should be recompiled
author_association: COLLABORATOR
user: login: adenzler-nvidia
created_at: 2025-04-03T07:58:39Z
body: I would like to do a deep-dive and see how the performance profile changes with this, but that's more for my own education..

I'm with Taylor on this one. Given that we have the (somewhat arbitrary) threshold that switches to SAP broadphase for large geom counts, I think it's ok to have this list as it will not be very large in many cases.
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-04-04T18:39:13Z
body: OK!  Then I'll list two more design requests before starting to look carefully at the code:

1) Can we also add support for contact/pair in this PR?  This will help us fully address #115 and seems related enough, but let me know if you feel otherwise @thowell 

2) Can we reuse this array in the sweep and prune broadphase?  Pair and exclude are just tricky enough that we ought to reuse the same logic in both approaches if possible.
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-04-06T21:20:46Z
body: @erikfrey yes will add support for contact pair to this pr
author_association: COLLABORATOR
user: login: btaba
created_at: 2025-04-07T17:41:07Z
body: OOC why is the broadphase happening on geom pairs rather than body pairs? In MJ we do broadphase on body pairs, midphase on geom pairs arising from body pairs, etc. I did this in the CUDA impl as well, but certainly a static geom-pair list is a good improvement already since it will avoid inter-body geom pairs.
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-04-08T22:50:12Z
body: @btaba just for simplicity - you're right that in C MuJoCo, `mj_broadphase` does SAP at the body level and then compares bounding spheres for midphase - it's not clear that following these steps will be performant for GPU.  We already have a SAP broadphase but my understanding is sorting is currently a bottleneck, in practice we've found that n^2 bounding sphere test is performant for at least `ngeom` <= 100 - haven't tested higher yet.
author_association: COLLABORATOR
user: login: btaba
created_at: 2025-04-08T23:12:55Z
body: @erikfrey Thanks for context; I wasn't necessarily mentioning SAP in particular, just mentioning broadphase on bodies, midphase on geoms, rather than broadphase on geoms, no midphase. The full version I think would involve a CPU callback [something like this](https://github.com/btaba/mujoco/blob/06f5d7c0b0cfed27007cf19a2d45736edfd2e212/mjx/mujoco/mjx/_src/cuda/engine_collision_driver.cu.cc#L513-L514), but it might still be performant esp. if different parts of the mj_step graph can run in parallel (multiple streams), and it may scale better in the number of bodies/geoms

Just some food for thought
author_association: COLLABORATOR
user: login: adenzler-nvidia
created_at: 2025-04-09T08:58:09Z
body: It's always hard to have an intuition about performance, but my intuition is that you need a lot of potential geom pairs to justify the need for a midphase. Most of the time it's better to expose as much parallelism as possible, especially looking at potential future bigger GPUs. But of course only measuring this will tell us where the tradeoff is.

There are still optimizations on the table with n2 because in theory you only need to load echo geom once and can cache it in shared memory - that likely gives another nice boost.

Running with multiple streams is cool as well but I would only go there once we have a more or less complete port ready. These things are super hard to debug, and if you have good GPU utilization it likely isn't really helping a lot with performance.
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-04-09T20:19:39Z
body: @erikfrey support for contact pairs has been added
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-04-15T12:45:14Z
body: @erikfrey `io.geom_pair` now generates `ngeom * (ngeom - 1) / 2` pair ids and geom pairs:
- `id == -2`: pair should be skipped during broadphase
- `id == -1`: pair is evaluated during broadphase and contact parameters are mixed
- `id > -1`: pair is evaluated during broadphase and contact parameters are specified via contact pair

`nxn_broadphase` and `sap_broadphase` both utilize `nxn_pairid` to filter geom pairs

author_association: COLLABORATOR