Fix deep penetration for sphere/capsule-convex.

PiperOrigin-RevId: 615539264
Change-Id: I6a5aac0574ddcb39bae8aafaf63f9561ca69288b

doc/changelog.rst

@@ -9,6 +9,7 @@ MJX
 ^^^
 
 1. Improved performance of SAT for convex collisions.
+2. Fixed bug for sphere/capsule-convex deep penetration.
 
 Version 3.1.3 (March 5th, 2024)
 -----------------------------------

mjx/mujoco/mjx/_src/collision_convex.py

@@ -249,6 +249,10 @@ def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
 
   # Get the normal, dist, and contact position.
   n, d = math.normalize_with_norm(pt - sphere_pos)
+  n = jp.where(inside | (d < 1e-6), normal, n)
+  # Ensure normal points towards convex centroid.
+  n *= jp.where(jp.dot(pt, n) < 0, 1, -1)
+
   spt = sphere_pos + n * sphere.size[0]
   dist = d - sphere.size[0]
   pos = (pt + spt) * 0.5
@@ -294,8 +298,8 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
   face = faces[best_idx]
   normal = normals[best_idx]
 
-  # Clip the edge against side planes and create two contact points against the
+  # Clip the segment against side planes and create two contact points against
-  # face.
+  # the face.
   edge_p0 = jp.roll(face, 1, axis=0)
   edge_p1 = face
   edge_normals = jax.vmap(jp.cross, in_axes=[0, None])(
@@ -311,29 +315,40 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
   )
   # Create variables for the face contact.
   pos = (cap_pts_clipped + face_pts) * 0.5
-  norm = jp.stack([normal] * 2, 0)
+  contact_normal = -jp.stack([normal] * 2, 0)
   penetration = jp.where(
       mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1
   )
 
   # Get a potential edge contact.
-  edge_closest, cap_closest = jax.vmap(
+  edge_closest, cap_closest, t_a, t_b = jax.vmap(
-      math.closest_segment_to_segment_points, in_axes=[0, 0, None, None]
+      math.closest_segment_to_segment_points_w_barycentric,
+      in_axes=[0, 0, None, None],
   )(edge_p0, edge_p1, cap_pts[0], cap_pts[1])
   e_idx = ((edge_closest - cap_closest) ** 2).sum(axis=1).argmin()
   cap_closest_pt, edge_closest_pt = cap_closest[e_idx], edge_closest[e_idx]
-  edge_axis = cap_closest_pt - edge_closest_pt
+  edge_dist = math.norm(cap_closest_pt - edge_closest_pt)
-  edge_axis, edge_dist = math.normalize_with_norm(edge_axis)
+  t_a, t_b = t_a[e_idx], t_b[e_idx]
+  in_segment = ((t_a > 0) & (t_a < 1)) | ((t_b > 0) & (t_b < 1))
+  # Ensure edge_axis points towards the convex centroid.
+  edge_axis = jp.cross(edge_p0[e_idx] - edge_p1[e_idx], axis)
+  sign = jp.where(jp.dot(edge_closest_pt, edge_axis) < 0, 1, -1)
+  edge_axis *= sign
+  degenerate_edge_axis = (edge_axis**2).sum() < 1e-6
+  edge_axis = math.normalize(edge_axis)
+
+  # Determine edge contact position.
   edge_pos = (
-      edge_closest_pt + (cap_closest_pt - edge_axis * cap.size[0])
+      edge_closest_pt + (cap_closest_pt + edge_axis * cap.size[0])
   ) * 0.5
-  edge_norm = edge_axis
   edge_penetration = cap.size[0] - edge_dist
-  has_edge_contact = edge_penetration > 0
+  has_edge_contact = (edge_penetration > 0) & in_segment & ~degenerate_edge_axis
 
   # Get the contact info.
   pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos)
-  n = -jp.where(has_edge_contact, norm.at[0].set(edge_norm), norm)
+  n = jp.where(
+      has_edge_contact, contact_normal.at[0].set(edge_axis), contact_normal
+  )
 
   # Go back to world frame.
   pos = convex.pos + pos @ convex.mat.T

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -138,6 +138,32 @@ class SphereCollisionTest(parameterized.TestCase):
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(dx.contact, d.contact, field.name, 'sphere_convex', 1e-4)
 
+    # test deep penetration
+    xml = self._SPHERE_CONVEX.replace(
+        '<body pos="0.52 0 0.52">', '<body pos="0.49 0 0.49">'
+    )
+    d, dx = _collide(xml)
+
+    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((d.contact.dist < 0).all())
+    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+    )
+
+    # test sphere center on vertex
+    xml = self._SPHERE_CONVEX.replace(
+        '<body pos="0.52 0 0.52">', '<body pos="0.5 0 0.5">'
+    )
+    d, dx = _collide(xml)
+
+    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((d.contact.dist < 0).all())
+    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+    )
+
 
 class CapsuleCollisionTest(parameterized.TestCase):
   _CAP_PLANE = """
@@ -212,7 +238,7 @@ class CapsuleCollisionTest(parameterized.TestCase):
       <worldbody>
         <body pos="0 0 0.54">
           <joint axis="1 0 0" type="free"/>
-          <geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
+          <geom fromto="-0.4 0 0 1.0 0 0" size="0.05" type="capsule"/>
         </body>
         <body>
           <joint axis="1 0 0" type="free"/>
@@ -226,9 +252,32 @@ class CapsuleCollisionTest(parameterized.TestCase):
     """Tests a capsule-convex collision for a face contact."""
     d, dx = _collide(self._CAP_BOX)
 
+    # sort positions for comparison
+    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
+    d.contact.pos[:] = d.contact.pos[idx]
+    d.contact.frame[:] = d.contact.frame[idx]
+    d.contact.dist[:] = d.contact.dist[idx]
+
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4)
 
+    # test deep face penetration
+    xml = self._CAP_BOX.replace('<body pos="0 0 0.54">', '<body pos="0 0 0.4">')
+
+    _, dx = _collide(xml)
+    self.assertTrue((dx.contact.dist < 0).all())
+    np.testing.assert_array_almost_equal(
+        dx.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
+    )
+    np.testing.assert_array_almost_equal(
+        dx.contact.dist, np.array([-0.15, -0.15])
+    )
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
+    )
+
   _CAP_EDGE_BOX = """
     <mujoco>
       <worldbody>
@@ -256,6 +305,19 @@ class CapsuleCollisionTest(parameterized.TestCase):
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
 
+    # test deep edge penetration
+    xml = self._CAP_EDGE_BOX.replace(
+        '<body pos="0.5 0 0.55"', '<body pos="0.5 0 0.48"'
+    )
+    _, dx = _collide(xml)
+    np.testing.assert_array_equal(dx.contact.dist < 0, np.array([True, False]))
+    np.testing.assert_array_almost_equal(
+        dx.contact.pos[0], np.array([0.483, 0, 0.471]), decimal=3
+    )
+    np.testing.assert_array_almost_equal(
+        dx.contact.frame[0, 0], np.array([-0.5, 0, -0.866]), decimal=3
+    )
+
 
 class ConvexTest(absltest.TestCase):
   """Tests the convex contact functions."""

mjx/mujoco/mjx/_src/math.py

@@ -331,9 +331,9 @@ def closest_segment_point_and_dist(
   return closest, dist
 
 
-def closest_segment_to_segment_points(
+def closest_segment_to_segment_points_w_barycentric(
     a0: jax.Array, a1: jax.Array, b0: jax.Array, b1: jax.Array
-) -> Tuple[jax.Array, jax.Array]:
+) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
   """Returns closest points between two line segments."""
   # Gets the closest segment points by first finding the closest points
   # between two lines. Points are then clipped to be on the line segments
@@ -364,6 +364,10 @@ def closest_segment_to_segment_points(
   t_a = jp.clip(orig_t_a, -half_len_a, half_len_a)
   t_b = jp.clip(orig_t_b, -half_len_b, half_len_b)
 
+  # reparametrize t_a, t_b
+  t_a_01 = (orig_t_a + half_len_a) / jp.maximum(len_a, 1e-6)
+  t_b_01 = (orig_t_b + half_len_b) / jp.maximum(len_b, 1e-6)
+
   best_a = a_mid + dir_a * t_a
   best_b = b_mid + dir_b * t_b
 
@@ -377,4 +381,12 @@ def closest_segment_to_segment_points(
   best_a = jp.where(d1 < d2, new_a, best_a)
   best_b = jp.where(d1 < d2, best_b, new_b)
 
-  return best_a, best_b
+  return best_a, best_b, t_a_01, t_b_01
+
+
+def closest_segment_to_segment_points(
+    a0: jax.Array, a1: jax.Array, b0: jax.Array, b1: jax.Array
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns closest points between two line segments."""
+  a, b, *_ = closest_segment_to_segment_points_w_barycentric(a0, a1, b0, b1)
+  return a, b