Further improve edge contact for capsule-convex.

PiperOrigin-RevId: 615900645
Change-Id: I7fd495b10b7b3121da7aa5ed15cde413d916bc84

mjx/mujoco/mjx/_src/collision_convex.py

@@ -331,10 +331,12 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
   degenerate_edge_dir = jp.sum(jp.square(edge_dir)) < 1e-6
   edge_dir, edge_dist = math.normalize_with_norm(edge_dir)
   face_edge_normals = convex.face_edge_normal[best_idx][e_idx]
-  in_edge_voronoi = ((face_edge_normals @ edge_dir) < 0).all()
+  inside_edge_voronoi_front = ((face_edge_normals @ edge_dir) < 0).all()
+  inside_edge_voronoi_back = ((face_edge_normals @ edge_dir) > 0).all()
+  outside_edge_voronoi = ~inside_edge_voronoi_front & ~inside_edge_voronoi_back
   edge_axis = math.normalize(-edge_closest_pt)  # approximate edge axis
   edge_axis = jp.where(
-      ~degenerate_edge_dir & in_edge_voronoi,
+      ~degenerate_edge_dir & inside_edge_voronoi_front,
       edge_dir,  # shallow edge penetration
       edge_axis,  # deep edge penetration
   )
@@ -348,7 +350,11 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
   edge_dir_parallel_to_face = (
       jp.abs(edge_dir.dot(normal)) > 0.99
   ) & ~degenerate_edge_dir
-  has_edge_contact = (edge_penetration > 0) & ~edge_dir_parallel_to_face
+  has_edge_contact = (
+      (edge_penetration > 0)
+      & ~edge_dir_parallel_to_face
+      & ~outside_edge_voronoi
+  )
 
   # Get the contact info.
   pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos)

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -323,6 +323,8 @@ class CapsuleCollisionTest(parameterized.TestCase):
 
   def test_capsule_convex_edge_shallow_tip(self):
     """Tests shallow edge penetration on the tip of the capsule."""
+    # the capsule sphere is inside the edge voronoi region, so there is an
+    # edge contact
     xml = self._CAP_EDGE_BOX.replace(
         '<geom fromto="-0.6 0 0 0.6 0 0" size="0.05"',
         '<geom fromto="0.6 0 0.6 -0.05 0 0" size="0.1"',
@@ -336,7 +338,30 @@ class CapsuleCollisionTest(parameterized.TestCase):
     # extract the contact point with penetration
     c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip', 1e-4)
+      _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4)
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
+    )
+
+    # the capsule sphere is outside the edge voronoi region, so there is a
+    # face contact
+    xml = self._CAP_EDGE_BOX.replace(
+        '<geom fromto="-0.6 0 0 0.6 0 0" size="0.05"',
+        '<geom fromto="-0.6 0 0.6 -0.05 0 0" size="0.1"',
+    )
+    xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.5 0 0.52"')
+    d, dx = _collide(xml)
+
+    c = dx.contact
+    self.assertEqual(c.pos.shape[0], 2)
+    self.assertGreater(c.dist[1], 0)
+    # extract the contact point with penetration
+    c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4)
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
+    )
 
 
 class ConvexTest(absltest.TestCase):