Merge pull request #74 from erikfrey/collision_driver_fixes

erikfrey · web-flow · commit a179363a54a8 · 2025-03-10T11:23:13.000-07:00
Fixes to collision driver.
diff --git a/mujoco/mjx/_src/collision_driver.py b/mujoco/mjx/_src/collision_driver.py
@@ -327,6 +327,11 @@ def get_contact_solver_params_kernel(
   p2 = m.geom_priority[g2]
   mix = where(p1 == p2, mix, where(p1 > p2, 1.0, 0.0))
 
+  condim1 = m.geom_condim[g1]
+  condim2 = m.geom_condim[g2]
+  condim = where(p1 == p2, wp.max(condim1, condim2), 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:
@@ -344,6 +349,7 @@ def group_contacts_by_type_kernel(
   d: Data,
 ):
   worldid, tid = wp.tid()
+
   if tid >= d.broadphase_result_count[worldid]:
     return
 
@@ -490,16 +496,30 @@ def _nxn_broadphase(m: Model, d: Data):
     margin2 = m.geom_margin[geom2]
     pos1 = d.geom_xpos[worldid, geom1]
     pos2 = d.geom_xpos[worldid, geom2]
+    xmat1 = d.geom_xmat[worldid, geom1]
+    xmat2 = d.geom_xmat[worldid, geom2]
     size1 = m.geom_rbound[geom1]
     size2 = m.geom_rbound[geom2]
 
     bound = size1 + size2 + wp.max(margin1, margin2)
     dif = pos2 - pos1
-    sphere_filter = wp.dot(dif, dif) <= bound * bound
+
+    if size1 != 0.0 and size2 != 0.0:
+      # neither geom is a plane
+      dist_sq = wp.dot(dif, dif)
+      bounds_filter = dist_sq <= bound * bound
+    elif size1 == 0.0:
+      # geom1 is a plane
+      dist = wp.dot(dif, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
+      bounds_filter = dist <= bound
+    else:
+      # geom2 is a plane
+      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 sphere_filter and geom_filter:
+    if bounds_filter and geom_filter:
       pairid = wp.atomic_add(d.broadphase_result_count, 0, 1)
       d.broadphase_pairs[worldid, pairid] = _geom_pair(m, geom1, geom2)
 
diff --git a/mujoco/mjx/_src/collision_driver_test.py b/mujoco/mjx/_src/collision_driver_test.py
@@ -17,25 +17,51 @@
 import mujoco
 from mujoco import mjx
 import numpy as np
+from absl.testing import absltest
 from absl.testing import parameterized
 
 
-class ConvexTest(parameterized.TestCase):
-  """Tests the convex contact functions."""
+class PrimitiveTest(parameterized.TestCase):
+  """Tests the primitive contact functions."""
 
-  _BOX_PLANE = """
+  _MJCFS = {
+    "box_plane": """
         <mujoco>
           <worldbody>
             <geom size="40 40 40" type="plane"/>
-            <body pos="0 0 0.7" euler="45 0 0">
+            <body pos="0 0 0.3" euler="45 0 0">
               <freejoint/>
               <geom size="0.5 0.5 0.5" type="box"/>
             </body>
           </worldbody>
         </mujoco>
-        """
-
-  _CAPSULE_CAPSULE = """
+      """,
+    "plane_sphere": """
+        <mujoco>
+          <worldbody>
+            <geom size="40 40 40" type="plane"/>
+            <body pos="0 0 0.2" euler="45 0 0">
+              <freejoint/>
+              <geom size="0.5" type="sphere"/>
+            </body>
+          </worldbody>
+        </mujoco>
+        """,
+    "sphere_sphere": """
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint type="free"/>
+              <geom pos="0 0 0" size="0.2" type="sphere"/>
+            </body>
+            <body >
+              <joint type="free"/>
+              <geom pos="0 0.3 0" size="0.11" type="sphere"/>
+            </body>
+          </worldbody>
+        </mujoco>
+        """,
+    "capsule_capsule": """
         <mujoco model="two_capsules">
           <worldbody>
             <body>
@@ -50,40 +76,45 @@ class ConvexTest(parameterized.TestCase):
             </body>
           </worldbody>
         </mujoco>
-        """
-
-  _SPHERE_SPHERE = """
+        """,
+    "plane_capsule": """
         <mujoco>
           <worldbody>
-            <body>
-              <joint type="free"/>
-              <geom pos="0 0 0" size="0.2" type="sphere"/>
-            </body>
-            <body >
-              <joint type="free"/>
-              <geom pos="0 0.3 0" size="0.11" type="sphere"/>
+            <geom size="40 40 40" type="plane"/>
+            <body pos="0 0 0.0" euler="30 30 0">
+              <freejoint/>
+              <geom size="0.05 0.05" type="capsule"/>
             </body>
           </worldbody>
         </mujoco>
-        """
+        """,
+  }
 
   @parameterized.parameters(
-    (_BOX_PLANE),
-    (_SPHERE_SPHERE),
-    (_CAPSULE_CAPSULE),
+    "box_plane",
+    "plane_sphere",
+    "sphere_sphere",
+    "plane_capsule",
+    "capsule_capsule",
   )
-  def test_convex_collision(self, xml_string):
-    """Tests convex collision with different geometries."""
-    m = mujoco.MjModel.from_xml_string(xml_string)
+  def test_contact(self, name):
+    """Tests contact calculation with different collision functions."""
+    m = mujoco.MjModel.from_xml_string(self._MJCFS[name])
     d = mujoco.MjData(m)
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
     dx = mjx.put_data(m, d)
     mjx.collision(mx, dx)
-    mujoco.mj_step(m, d)
-    actual_dist = dx.contact.dist.numpy()[0]
-    actual_pos = dx.contact.pos.numpy()[0, :]
-    actual_frame = dx.contact.frame.numpy()[0].flatten()
-    np.testing.assert_array_almost_equal(actual_dist, d.contact.dist[0], 4)
-    np.testing.assert_array_almost_equal(actual_pos, d.contact.pos[0], 4)
-    np.testing.assert_array_almost_equal(actual_frame, d.contact.frame[0], 4)
+    mujoco.mj_collision(m, d)
+    self.assertEqual(d.ncon, dx.ncon.numpy()[0])
+    for i in range(d.ncon):
+      actual_dist = dx.contact.dist.numpy()[i]
+      actual_pos = dx.contact.pos.numpy()[i, :]
+      actual_frame = dx.contact.frame.numpy()[i].flatten()
+      np.testing.assert_array_almost_equal(actual_dist, d.contact.dist[i], 4)
+      np.testing.assert_array_almost_equal(actual_pos, d.contact.pos[i], 4)
+      np.testing.assert_array_almost_equal(actual_frame, d.contact.frame[i], 4)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/mujoco/mjx/_src/collision_functions.py b/mujoco/mjx/_src/collision_functions.py
@@ -22,7 +22,6 @@
 from .math import closest_segment_to_segment_points
 from .math import normalize_with_norm
 from .support import group_key
-from .support import mat33_from_cols
 
 
 @wp.struct
@@ -279,7 +278,8 @@ def plane_capsule(
     else:
       b = wp.vec3(0.0, 0.0, 1.0)
 
-  frame = mat33_from_cols(n, b, wp.cross(n, b))
+  c = wp.cross(n, b)
+  frame = wp.mat33(n[0], n[1], n[2], b[0], b[1], b[2], c[0], c[1], c[2])
   segment = axis * cap.halfsize
 
   dist1, pos1 = _plane_sphere(n, plane.pos, cap.pos + segment, cap.radius)
@@ -289,13 +289,6 @@ def plane_capsule(
   write_contact(d, dist2, pos2, frame, margin, geom_indices, worldid)
 
 
-@wp.func
-def distance_point_plane(plane_normal: wp.vec3, plane_pos: wp.vec3, point: wp.vec3):
-  plane_normal = wp.normalize(plane_normal)
-  dist = wp.dot(point - plane_pos, plane_normal)
-  return dist, plane_pos - plane_normal * dist
-
-
 @wp.func
 def plane_box(
   plane: GeomPlane,
@@ -305,29 +298,31 @@ def plane_box(
   margin: float,
   geom_indices: wp.vec2i,
 ):
-  contact_count = int(0)
+  count = int(0)
+  corner = wp.vec3()
+  dist = wp.dot(box.pos - plane.pos, plane.normal)
 
-  # Check all 8 corners of the box
+  # test all corners, pick bottom 4
   for i in range(8):
-    corner = wp.vec3(box.size.x * 0.5, box.size.y * 0.5, box.size.z * 0.5)
-    if i % 2 == 0:
-      corner.x = -corner.x
-    if (i // 2) % 2 == 0:
-      corner.y = -corner.y
-    if i < 4:
-      corner.z = -corner.z
-
-    corner_world = box.rot * (corner) + box.pos
-
-    dist, pos = distance_point_plane(plane.normal, plane.pos, corner_world)
-
-    if dist < 0.0:
-      write_contact(
-        d, dist, pos, make_frame(plane.normal), margin, geom_indices, worldid
-      )
-      contact_count += 1
-
-    if contact_count >= 4:
+    # get corner in local coordinates
+    corner.x = wp.select(i & 1, -box.size.x, box.size.x)
+    corner.y = wp.select(i & 2, -box.size.y, box.size.y)
+    corner.z = wp.select(i & 4, -box.size.z, box.size.z)
+
+    # get corner in global coordinates relative to box center
+    corner = box.rot * corner
+
+    # compute distance to plane, skip if too far or pointing up
+    ldist = wp.dot(plane.normal, corner)
+    if dist + ldist > margin or ldist > 0:
+      continue
+
+    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)
+    count += 1
+    if count >= 4:
       break
 
 
diff --git a/mujoco/mjx/_src/io.py b/mujoco/mjx/_src/io.py
@@ -248,6 +248,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.geom_bodyid = wp.array(mjm.geom_bodyid, dtype=wp.int32, ndim=1)
   m.geom_conaffinity = wp.array(mjm.geom_conaffinity, dtype=wp.int32, ndim=1)
   m.geom_contype = wp.array(mjm.geom_contype, dtype=wp.int32, ndim=1)
+  m.geom_condim = wp.array(mjm.geom_condim, dtype=wp.int32, ndim=1)
   m.geom_pos = wp.array(mjm.geom_pos, dtype=wp.vec3, ndim=1)
   m.geom_quat = wp.array(mjm.geom_quat, dtype=wp.quat, ndim=1)
   m.geom_size = wp.array(mjm.geom_size, dtype=wp.vec3, ndim=1)
diff --git a/mujoco/mjx/_src/types.py b/mujoco/mjx/_src/types.py
@@ -292,6 +292,7 @@ class Model:
   geom_bodyid: wp.array(dtype=wp.int32, ndim=1)
   geom_conaffinity: wp.array(dtype=wp.int32, ndim=1)
   geom_contype: wp.array(dtype=wp.int32, ndim=1)
+  geom_condim: wp.array(dtype=wp.int32, ndim=1)
   geom_pos: wp.array(dtype=wp.vec3, ndim=1)
   geom_quat: wp.array(dtype=wp.quat, ndim=1)
   geom_size: wp.array(dtype=wp.vec3, ndim=1)
