diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 6095994f..30e64401 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -302,12 +302,12 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: # the face. edge_p0 = jp.roll(face, 1, axis=0) edge_p1 = face - edge_normals = jax.vmap(jp.cross, in_axes=[0, None])( + side_planes = jax.vmap(jp.cross, in_axes=[0, None])( edge_p1 - edge_p0, normal, ) cap_pts_clipped, mask = _clip_edge_to_planes( - cap_pts[0], cap_pts[1], edge_p0, edge_normals + cap_pts[0], cap_pts[1], edge_p0, side_planes ) cap_pts_clipped = cap_pts_clipped - normal * cap.size[0] face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])( @@ -321,28 +321,34 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: ) # Get a potential edge contact. - edge_closest, cap_closest, t_a, t_b = jax.vmap( - math.closest_segment_to_segment_points_w_barycentric, + edge_closest, cap_closest = jax.vmap( + math.closest_segment_to_segment_points, 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_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) + edge_dir = edge_closest_pt - cap_closest_pt + 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() + edge_axis = math.normalize(-edge_closest_pt) # approximate edge axis + edge_axis = jp.where( + ~degenerate_edge_dir & in_edge_voronoi, + edge_dir, # shallow edge penetration + edge_axis, # deep edge penetration + ) # Determine edge contact position. edge_pos = ( edge_closest_pt + (cap_closest_pt + edge_axis * cap.size[0]) ) * 0.5 edge_penetration = cap.size[0] - edge_dist - has_edge_contact = (edge_penetration > 0) & in_segment & ~degenerate_edge_axis + # prefer face contact if the edge axis is parallel to the face normal + 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 # Get the contact info. pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 7eda941d..c914ad5c 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -248,8 +248,8 @@ class CapsuleCollisionTest(parameterized.TestCase): """ - def test_capsule_convex(self): - """Tests a capsule-convex collision for a face contact.""" + def test_capsule_convex_face(self): + """Tests face contact.""" d, dx = _collide(self._CAP_BOX) # sort positions for comparison @@ -263,7 +263,8 @@ class CapsuleCollisionTest(parameterized.TestCase): for field in dataclasses.fields(Contact): _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4) - # test deep face penetration + def test_capsule_convex_face_deep(self): + """Tests deep face penetration.""" xml = self._CAP_BOX.replace('', '') _, dx = _collide(xml) @@ -294,7 +295,7 @@ class CapsuleCollisionTest(parameterized.TestCase): """ def test_capsule_convex_edge(self): - """Tests a capsule-convex collision for an edge contact.""" + """Tests edge contact.""" d, dx = _collide(self._CAP_EDGE_BOX) c = dx.contact @@ -305,19 +306,38 @@ 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 + def test_capsule_convex_edge_deep(self): + """Tests deep edge penetration.""" xml = self._CAP_EDGE_BOX.replace( ' Tuple[jax.Array, jax.Array, jax.Array, jax.Array]: +) -> Tuple[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,10 +364,6 @@ def closest_segment_to_segment_points_w_barycentric( 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 @@ -381,12 +377,4 @@ def closest_segment_to_segment_points_w_barycentric( 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, 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 + return best_a, best_b