diff --git a/doc/changelog.rst b/doc/changelog.rst index 9fabc3c2..bb66f206 100644 --- a/doc/changelog.rst +++ b/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) ----------------------------------- diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 196b5b78..6095994f 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/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 - # face. + # Clip the segment against side planes and create two contact points against + # 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( - math.closest_segment_to_segment_points, in_axes=[0, 0, None, None] + edge_closest, cap_closest, t_a, t_b = jax.vmap( + 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_axis, edge_dist = math.normalize_with_norm(edge_axis) + edge_dist = math.norm(cap_closest_pt - edge_closest_pt) + 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 diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index b146c837..7eda941d 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/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( + '', '' + ) + 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( + '', '' + ) + 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): - + @@ -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('', '') + + _, 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 = """ @@ -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( + ' 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