Speed up convex-convex collisions.
PiperOrigin-RevId: 613450874 Change-Id: Ie66e0a078e2ff51a136b949dd9459847e1fe957d
This commit is contained in:
Baruch Tabanpour
committed by
Copybara-Service
Copybara-Service
parent
106453fa2b
commit
22cd96fbbc
@@ -14,6 +14,7 @@
|
||||
# ==============================================================================
|
||||
"""Convex collisions."""
|
||||
|
||||
import functools
|
||||
from typing import Tuple
|
||||
|
||||
import jax
|
||||
@@ -177,6 +178,174 @@ def _manifold_points(
|
||||
return jp.array([a_idx, b_idx, c_idx, d_idx])
|
||||
|
||||
|
||||
def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contacts between a plane and a convex object."""
|
||||
vert = convex.vert
|
||||
|
||||
# get points in the convex frame
|
||||
plane_pos = convex.mat.T @ (plane.pos - convex.pos)
|
||||
n = convex.mat.T @ plane.mat[:, 2]
|
||||
support = (plane_pos - vert) @ n
|
||||
idx = _manifold_points(vert, support > 0, n)
|
||||
pos = vert[idx]
|
||||
|
||||
# convert to world frame
|
||||
pos = convex.pos + pos @ convex.mat.T
|
||||
n = plane.mat[:, 2]
|
||||
|
||||
frame = jp.stack([math.make_frame(n)] * 4, axis=0)
|
||||
unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1
|
||||
dist = jp.where(unique, -support[idx], 1)
|
||||
pos = pos - 0.5 * dist[:, None] * n
|
||||
return dist, pos, frame
|
||||
|
||||
|
||||
def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contact between a sphere and a convex object."""
|
||||
faces = convex.face
|
||||
normals = convex.facenorm
|
||||
|
||||
# Put sphere in convex frame.
|
||||
sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)
|
||||
|
||||
# Get support from face normals.
|
||||
@jax.vmap
|
||||
def get_support(faces, normal):
|
||||
pos = sphere_pos - normal * sphere.size[0]
|
||||
return jp.dot(pos - faces[0], normal)
|
||||
|
||||
support = get_support(faces, normals)
|
||||
|
||||
# Pick the face with minimal penetration as long as it has support.
|
||||
support = jp.where(support >= 0, -1e12, support)
|
||||
best_idx = support.argmax()
|
||||
face = faces[best_idx]
|
||||
normal = normals[best_idx]
|
||||
|
||||
# Get closest point between the polygon face and the sphere center point.
|
||||
# Project the sphere center point onto poly plane. If it's inside polygon
|
||||
# edge normals, then we're done.
|
||||
pt = _project_pt_onto_plane(sphere_pos, face[0], normal)
|
||||
edge_p0 = jp.roll(face, 1, axis=0)
|
||||
edge_p1 = face
|
||||
edge_normals = jax.vmap(jp.cross, in_axes=[0, None])(
|
||||
edge_p1 - edge_p0,
|
||||
normal,
|
||||
)
|
||||
edge_dist = jax.vmap(
|
||||
lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm)
|
||||
)(edge_p0, edge_normals)
|
||||
inside = jp.all(edge_dist <= 0) # lte to handle degenerate edges
|
||||
|
||||
# If the point is outside edge normals, project onto the closest edge plane
|
||||
# that the point is in front of.
|
||||
degenerate_edge = jp.all(edge_normals == 0, axis=1)
|
||||
behind = edge_dist < 0.0
|
||||
edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist)
|
||||
idx = edge_dist.argmin()
|
||||
edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt)
|
||||
|
||||
pt = jp.where(inside, pt, edge_pt)
|
||||
|
||||
# Get the normal, dist, and contact position.
|
||||
n, d = math.normalize_with_norm(pt - sphere_pos)
|
||||
spt = sphere_pos + n * sphere.size[0]
|
||||
dist = d - sphere.size[0]
|
||||
pos = (pt + spt) * 0.5
|
||||
|
||||
# Go back to world frame.
|
||||
n = convex.mat @ n
|
||||
pos = convex.mat @ pos + convex.pos
|
||||
|
||||
return jax.tree_map(
|
||||
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
|
||||
)
|
||||
|
||||
|
||||
def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contacts between a capsule and a convex object."""
|
||||
# Get convex transformed normals, faces, and vertices.
|
||||
faces = convex.face
|
||||
normals = convex.facenorm
|
||||
|
||||
# Put capsule in convex frame.
|
||||
cap_pos = convex.mat.T @ (cap.pos - convex.pos)
|
||||
axis, length = cap.mat[:, 2], cap.size[1]
|
||||
axis = convex.mat.T @ axis
|
||||
seg = axis * length
|
||||
cap_pts = jp.array([
|
||||
cap_pos - seg,
|
||||
cap_pos + seg,
|
||||
])
|
||||
|
||||
# Get support from face normals.
|
||||
@jax.vmap
|
||||
def get_support(face, normal):
|
||||
pts = cap_pts - normal * cap.size[0]
|
||||
sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts)
|
||||
return sup.min()
|
||||
|
||||
support = get_support(faces, normals)
|
||||
has_support = jp.all(support < 0)
|
||||
|
||||
# Pick the face with minimal penetration as long as it has support.
|
||||
support = jp.where(support >= 0, -1e12, support)
|
||||
best_idx = support.argmax()
|
||||
face = faces[best_idx]
|
||||
normal = normals[best_idx]
|
||||
|
||||
# Clip the edge 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])(
|
||||
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_clipped = cap_pts_clipped - normal * cap.size[0]
|
||||
face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
|
||||
cap_pts_clipped, face[0], normal
|
||||
)
|
||||
# Create variables for the face contact.
|
||||
pos = (cap_pts_clipped + face_pts) * 0.5
|
||||
norm = 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_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_pos = (
|
||||
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
|
||||
|
||||
# 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)
|
||||
|
||||
# Go back to world frame.
|
||||
pos = convex.pos + pos @ convex.mat.T
|
||||
n = n @ convex.mat.T
|
||||
|
||||
dist = -jp.where(
|
||||
has_edge_contact, penetration.at[0].set(edge_penetration), penetration
|
||||
)
|
||||
frame = jax.vmap(math.make_frame)(n)
|
||||
return dist, pos, frame
|
||||
|
||||
|
||||
def _project_pt_onto_plane(
|
||||
pt: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array
|
||||
) -> jax.Array:
|
||||
@@ -396,7 +565,7 @@ def _create_contact_manifold(
|
||||
return dist, pos, normal
|
||||
|
||||
|
||||
def _sat_hull_hull(
|
||||
def _sat_bruteforce(
|
||||
faces_a: jax.Array,
|
||||
faces_b: jax.Array,
|
||||
vertices_a: jax.Array,
|
||||
@@ -414,13 +583,16 @@ def _sat_hull_hull(
|
||||
We return both the edge and face contacts. Valid contacts can be checked with
|
||||
dist < 0. Resulting edge contacts should be preferred over face contacts.
|
||||
|
||||
This method checks all unique edge-pairs and is thus costly to run over large
|
||||
meshes, but is more performant for smaller meshes (boxes, tetrahedra, etc.).
|
||||
|
||||
Args:
|
||||
faces_a: An ndarray of hull A's polygon faces.
|
||||
faces_b: An ndarray of hull B's polygon faces.
|
||||
faces_a: Faces for hull A.
|
||||
faces_b: Faces for hull B.
|
||||
vertices_a: Vertices for hull A.
|
||||
vertices_b: Vertices for hull B.
|
||||
normals_a: Normal vectors for hull A's polygon faces.
|
||||
normals_b: Normal vectors for hull B's polygon faces.
|
||||
normals_a: Normal vectors for hull A faces.
|
||||
normals_b: Normal vectors for hull B faces.
|
||||
unique_edges_a: Unique edges for hull A.
|
||||
unique_edges_b: Unique edges for hull B.
|
||||
|
||||
@@ -428,30 +600,36 @@ def _sat_hull_hull(
|
||||
tuple of dist, pos, and normal
|
||||
"""
|
||||
# get the separating axes
|
||||
edge_dir_a = unique_edges_a[:, 0] - unique_edges_a[:, 1]
|
||||
edge_dir_b = unique_edges_b[:, 0] - unique_edges_b[:, 1]
|
||||
v_norm = jax.vmap(math.normalize)
|
||||
edge_dir_a = v_norm(unique_edges_a[:, 0] - unique_edges_a[:, 1])
|
||||
edge_dir_b = v_norm(unique_edges_b[:, 0] - unique_edges_b[:, 1])
|
||||
edge_dir_a_r = jp.tile(edge_dir_a, reps=(unique_edges_b.shape[0], 1))
|
||||
edge_dir_b_r = jp.repeat(edge_dir_b, repeats=unique_edges_a.shape[0], axis=0)
|
||||
edge_edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r)
|
||||
edge_edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))(
|
||||
edge_edge_axes
|
||||
edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r)
|
||||
degenerate_edge_axes = (edge_axes**2).sum(axis=1) < 1e-6
|
||||
edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))(edge_axes)
|
||||
n_norm = normals_a.shape[0] + normals_b.shape[0]
|
||||
degenerate_axes = jp.concatenate(
|
||||
[jp.array([False] * n_norm), degenerate_edge_axes]
|
||||
)
|
||||
|
||||
axes = jp.concatenate([normals_a, normals_b, edge_edge_axes])
|
||||
axes = jp.concatenate([normals_a, normals_b, edge_axes])
|
||||
|
||||
# for each separating axis, get the support
|
||||
@jax.vmap
|
||||
def get_support(axis):
|
||||
support_a = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_a)
|
||||
support_b = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_b)
|
||||
def get_support(axis, is_degenerate):
|
||||
# the matmul here is more performant with vmap(dot)
|
||||
dot = functools.partial(jp.dot, precision=jax.lax.Precision.HIGH)
|
||||
support_a = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_a)
|
||||
support_b = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_b)
|
||||
dist1 = support_a.max() - support_b.min()
|
||||
dist2 = support_b.max() - support_a.min()
|
||||
sign = jp.where(dist1 > dist2, -1, 1)
|
||||
dist = jp.minimum(dist1, dist2)
|
||||
dist = jp.where(~jp.all(axis == 0.0), dist, 1e6) # degenerate axis
|
||||
dist = jp.where(~is_degenerate, dist, 1e6) # degenerate axis
|
||||
return dist, sign
|
||||
|
||||
support, sign = get_support(axes)
|
||||
support, sign = get_support(axes, degenerate_axes)
|
||||
|
||||
# choose the best separating axis
|
||||
best_idx = jp.argmin(support)
|
||||
@@ -460,8 +638,8 @@ def _sat_hull_hull(
|
||||
is_edge_contact = best_idx >= (normals_a.shape[0] + normals_b.shape[0])
|
||||
|
||||
# get the (reference) face most aligned with the separating axis
|
||||
dist_a = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_a)
|
||||
dist_b = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_b)
|
||||
dist_a = normals_a @ best_axis
|
||||
dist_b = normals_b @ best_axis
|
||||
a_max = dist_a.argmax()
|
||||
b_max = dist_b.argmax()
|
||||
a_min = dist_a.argmin()
|
||||
@@ -496,172 +674,172 @@ def _sat_hull_hull(
|
||||
return dist, pos, normal
|
||||
|
||||
|
||||
def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contacts between a plane and a convex object."""
|
||||
vert = convex.vert
|
||||
|
||||
# get points in the convex frame
|
||||
plane_pos = convex.mat.T @ (plane.pos - convex.pos)
|
||||
n = convex.mat.T @ plane.mat[:, 2]
|
||||
support = (plane_pos - vert) @ n
|
||||
idx = _manifold_points(vert, support > 0, n)
|
||||
pos = vert[idx]
|
||||
|
||||
# convert to world frame
|
||||
pos = convex.pos + pos @ convex.mat.T
|
||||
n = plane.mat[:, 2]
|
||||
|
||||
frame = jp.stack([math.make_frame(n)] * 4, axis=0)
|
||||
unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1
|
||||
dist = jp.where(unique, -support[idx], 1)
|
||||
pos = pos - 0.5 * dist[:, None] * n
|
||||
return dist, pos, frame
|
||||
def _arcs_intersect(
|
||||
a: jax.Array, b: jax.Array, c: jax.Array, d: jax.Array
|
||||
) -> jax.Array:
|
||||
"""Tests if arcs AB and CD on the unit sphere intersect."""
|
||||
ba, dc = jp.cross(b, a), jp.cross(d, c)
|
||||
cba, dba = jp.dot(c, ba), jp.dot(d, ba)
|
||||
adc, bdc = jp.dot(a, dc), jp.dot(b, dc)
|
||||
return (cba * dba < 0) & (adc * bdc < 0) & (cba * bdc > 0)
|
||||
|
||||
|
||||
def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contact between a sphere and a convex object."""
|
||||
faces = convex.face
|
||||
normals = convex.facenorm
|
||||
def _sat_approx(
|
||||
centroid_a: jax.Array,
|
||||
faces_a: jax.Array,
|
||||
faces_b: jax.Array,
|
||||
vertices_a: jax.Array,
|
||||
vertices_b: jax.Array,
|
||||
normals_a: jax.Array,
|
||||
normals_b: jax.Array,
|
||||
face_edges_a: jax.Array,
|
||||
face_edges_b: jax.Array,
|
||||
face_edge_normals_a: jax.Array,
|
||||
face_edge_normals_b: jax.Array,
|
||||
) -> Tuple[jax.Array, jax.Array, jax.Array]:
|
||||
"""Runs the Separating Axis Test for a pair of hulls.
|
||||
|
||||
# Put sphere in convex frame.
|
||||
sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)
|
||||
Runs the separating axis test for all faces. After obtaining a reference
|
||||
and incident face, tests edge separating axes via edge intersections
|
||||
on gauss maps.
|
||||
|
||||
Certain meshes can have nearly parallel coplanar faces that are not merged.
|
||||
Thus reference/incident faces can be non-overlapping, and face contacts will
|
||||
not get generated. Since we only check edge separating axes on
|
||||
reference/incident faces, the correct edge separating axes may also be
|
||||
missing. We mitigate this issue by checking nearly coplanar anti-parallel
|
||||
reference/incident faces for support, hence why this method is "approximate".
|
||||
|
||||
Args:
|
||||
centroid_a: Centroid of hull A.
|
||||
faces_a: Faces for hull A.
|
||||
faces_b: Faces for hull B.
|
||||
vertices_a: Vertices for hull A.
|
||||
vertices_b: Vertices for hull B.
|
||||
normals_a: Normal vectors for hull A faces.
|
||||
normals_b: Normal vectors for hull B faces.
|
||||
face_edges_a: Edges for faces in hull A.
|
||||
face_edges_b: Edges for faces in hull B.
|
||||
face_edge_normals_a: Edge normals for faces in hull A.
|
||||
face_edge_normals_b: Edge normals for faces in hull B.
|
||||
|
||||
Returns:
|
||||
tuple of dist, pos, and normal
|
||||
"""
|
||||
# Handle face separating axes.
|
||||
axes = jp.concatenate([normals_a, -normals_b])
|
||||
|
||||
# Get support from face normals.
|
||||
@jax.vmap
|
||||
def get_support(faces, normal):
|
||||
pos = sphere_pos - normal * sphere.size[0]
|
||||
return jp.dot(pos - faces[0], normal)
|
||||
def get_support(axis):
|
||||
# the matmul here is more performant with vmap(dot)
|
||||
dot = functools.partial(jp.dot, precision=jax.lax.Precision.HIGH)
|
||||
support_a = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_a)
|
||||
support_b = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_b)
|
||||
dist = support_a.max() - support_b.min()
|
||||
separating = dist < 0
|
||||
dist = jp.where(dist < 0, 1e6, dist)
|
||||
return dist, separating
|
||||
|
||||
support = get_support(faces, normals)
|
||||
support, separating = get_support(axes)
|
||||
is_face_separating = separating.any()
|
||||
|
||||
# Pick the face with minimal penetration as long as it has support.
|
||||
support = jp.where(support >= 0, -1e12, support)
|
||||
best_idx = support.argmax()
|
||||
face = faces[best_idx]
|
||||
normal = normals[best_idx]
|
||||
# choose the best separating axis
|
||||
best_idx = jp.argmin(support)
|
||||
best_axis = axes[best_idx]
|
||||
|
||||
# Get closest point between the polygon face and the sphere center point.
|
||||
# Project the sphere center point onto poly plane. If it's inside polygon
|
||||
# edge normals, then we're done.
|
||||
pt = _project_pt_onto_plane(sphere_pos, face[0], normal)
|
||||
edge_p0 = jp.roll(face, 1, axis=0)
|
||||
edge_p1 = face
|
||||
edge_normals = jax.vmap(jp.cross, in_axes=[0, None])(
|
||||
edge_p1 - edge_p0,
|
||||
normal,
|
||||
)
|
||||
edge_dist = jax.vmap(
|
||||
lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm)
|
||||
)(edge_p0, edge_normals)
|
||||
inside = jp.all(edge_dist <= 0) # lte to handle degenerate edges
|
||||
# get the (reference) face most aligned with the separating axis
|
||||
dist_a = normals_a @ best_axis
|
||||
dist_b = normals_b @ -best_axis
|
||||
face_a_idx = dist_a.argmax()
|
||||
face_b_idx = dist_b.argmax()
|
||||
|
||||
# If the point is outside edge normals, project onto the closest edge plane
|
||||
# that the point is in front of.
|
||||
degenerate_edge = jp.all(edge_normals == 0, axis=1)
|
||||
behind = edge_dist < 0.0
|
||||
edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist)
|
||||
idx = edge_dist.argmin()
|
||||
edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt)
|
||||
|
||||
pt = jp.where(inside, pt, edge_pt)
|
||||
|
||||
# Get the normal, dist, and contact position.
|
||||
n, d = math.normalize_with_norm(pt - sphere_pos)
|
||||
spt = sphere_pos + n * sphere.size[0]
|
||||
dist = d - sphere.size[0]
|
||||
pos = (pt + spt) * 0.5
|
||||
|
||||
# Go back to world frame.
|
||||
n = convex.mat @ n
|
||||
pos = convex.mat @ pos + convex.pos
|
||||
|
||||
return jax.tree_map(
|
||||
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
|
||||
cond = best_idx < normals_a.shape[0]
|
||||
ref_face = jp.where(cond, faces_a[face_a_idx], faces_b[face_b_idx])
|
||||
incident_face = jp.where(cond, faces_b[face_b_idx], faces_a[face_a_idx])
|
||||
ref_face_norm = jp.where(cond, normals_a[face_a_idx], normals_b[face_b_idx])
|
||||
incident_face_norm = jp.where(
|
||||
cond, normals_b[face_b_idx], normals_a[face_a_idx]
|
||||
)
|
||||
|
||||
dist, pos, normal = _create_contact_manifold(
|
||||
ref_face,
|
||||
incident_face,
|
||||
ref_face_norm,
|
||||
incident_face_norm,
|
||||
-best_axis,
|
||||
)
|
||||
|
||||
def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
|
||||
"""Calculates contacts between a capsule and a convex object."""
|
||||
# Get convex transformed normals, faces, and vertices.
|
||||
faces = convex.face
|
||||
normals = convex.facenorm
|
||||
# Handle edge separating axes by checking edge pairs on the reference and
|
||||
# incident faces. In principle, the correct edge-edge separating axis can be
|
||||
# created from edge pairs on the reference and incident faces.
|
||||
# first get the edge directions and face edge normals
|
||||
face_edges_a = face_edges_a[face_a_idx]
|
||||
v_norm = jax.vmap(math.normalize)
|
||||
face_edges_dir_a = v_norm(face_edges_a[:, 0] - face_edges_a[:, 1])
|
||||
face_edges_b = face_edges_b[face_b_idx]
|
||||
face_edges_dir_b = v_norm(face_edges_b[:, 0] - face_edges_b[:, 1])
|
||||
face_edge_normals_a = face_edge_normals_a[face_a_idx]
|
||||
face_edge_normals_b = face_edge_normals_b[face_b_idx]
|
||||
|
||||
# Put capsule in convex frame.
|
||||
cap_pos = convex.mat.T @ (cap.pos - convex.pos)
|
||||
axis, length = cap.mat[:, 2], cap.size[1]
|
||||
axis = convex.mat.T @ axis
|
||||
seg = axis * length
|
||||
cap_pts = jp.array([
|
||||
cap_pos - seg,
|
||||
cap_pos + seg,
|
||||
])
|
||||
# scatter to all edge pairs
|
||||
edge_idx_a = jp.repeat(
|
||||
jp.arange(face_edges_a.shape[0]), face_edges_b.shape[0]
|
||||
)
|
||||
edge_idx_b = jp.tile(jp.arange(face_edges_b.shape[0]), face_edges_a.shape[0])
|
||||
face_edges_dir_a = face_edges_dir_a[edge_idx_a]
|
||||
face_edges_dir_b = face_edges_dir_b[edge_idx_b]
|
||||
face_edges_pt_a = face_edges_a[:, 0][edge_idx_a]
|
||||
face_edges_pt_b = face_edges_b[:, 0][edge_idx_b]
|
||||
face_edge_normals_a = face_edge_normals_a[edge_idx_a]
|
||||
face_edge_normals_b = face_edge_normals_b[edge_idx_b]
|
||||
|
||||
# Get support from face normals.
|
||||
@jax.vmap
|
||||
def get_support(face, normal):
|
||||
pts = cap_pts - normal * cap.size[0]
|
||||
sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts)
|
||||
return sup.min()
|
||||
def get_normals(a_dir, a_pt, b_dir):
|
||||
edge_axis = math.normalize(jp.cross(a_dir, b_dir))
|
||||
# correct normal to point from a to b, object b is at the origin
|
||||
sign = jp.where(jp.dot(edge_axis, a_pt - centroid_a) > 0.0, 1.0, -1.0)
|
||||
return edge_axis * sign
|
||||
|
||||
support = get_support(faces, normals)
|
||||
has_support = jp.all(support < 0)
|
||||
edge_axes = get_normals(face_edges_dir_a, face_edges_pt_a, face_edges_dir_b)
|
||||
edge_dist = jax.vmap(jp.dot)(edge_axes, face_edges_pt_b - face_edges_pt_a)
|
||||
# handle degenerate axes
|
||||
edge_dist = jp.where((edge_axes**2).sum(axis=1) < 1e-6, 1e6, edge_dist)
|
||||
# ensure edges create a minkowski face by testing intersection on gauss maps
|
||||
is_minkowski_face = jax.vmap(_arcs_intersect)(
|
||||
face_edge_normals_a[:, 0],
|
||||
face_edge_normals_a[:, 1],
|
||||
-face_edge_normals_b[:, 0],
|
||||
-face_edge_normals_b[:, 1],
|
||||
)
|
||||
edge_dist = jp.where(is_minkowski_face, edge_dist, 1e6)
|
||||
edge_dist = jp.where(edge_dist > 0, -1e6, edge_dist)
|
||||
|
||||
# Pick the face with minimal penetration as long as it has support.
|
||||
support = jp.where(support >= 0, -1e12, support)
|
||||
best_idx = support.argmax()
|
||||
face = faces[best_idx]
|
||||
normal = normals[best_idx]
|
||||
best_edge_idx = edge_dist.argmax()
|
||||
best_edge_dist = edge_dist[best_edge_idx]
|
||||
# prefer edge over face contacts as long as we have a valid edge contact
|
||||
is_edge_contact = (best_edge_dist > dist.min() + 1e-6) & (best_edge_dist < 0)
|
||||
normal = jp.where(is_edge_contact, edge_axes[best_edge_idx], normal)
|
||||
dist = jp.where(is_edge_contact, jp.array([best_edge_dist, 1, 1, 1]), dist)
|
||||
|
||||
# Clip the edge 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])(
|
||||
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_clipped = cap_pts_clipped - normal * cap.size[0]
|
||||
face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
|
||||
cap_pts_clipped, face[0], normal
|
||||
)
|
||||
# Create variables for the face contact.
|
||||
pos = (cap_pts_clipped + face_pts) * 0.5
|
||||
norm = jp.stack([normal] * 2, 0)
|
||||
penetration = jp.where(
|
||||
mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1
|
||||
# A failure mode occurs if faces are very narrow and nearly parallel
|
||||
# (i.e. faces did not get merged properly as coplanar faces). The face
|
||||
# contacts will be empty since the reference/incident faces may not overlap.
|
||||
# An edge-edge separating axis will not be found, since the reference and
|
||||
# incident faces will also not overlap or necessarily create a minkowski face.
|
||||
# Thus, we approximate the contact for anti-parallel faces with support.
|
||||
anti_parallel = incident_face_norm.dot(ref_face_norm) < -0.97
|
||||
dist = dist.at[0].set(
|
||||
jp.where(
|
||||
anti_parallel
|
||||
& ~is_face_separating
|
||||
& (dist > 0).all()
|
||||
& ~is_edge_contact,
|
||||
-support[best_idx],
|
||||
dist[0],
|
||||
)
|
||||
)
|
||||
|
||||
# Get a potential edge contact.
|
||||
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_axis = cap_closest_pt - edge_closest_pt
|
||||
edge_axis, edge_dist = math.normalize_with_norm(edge_axis)
|
||||
edge_pos = (
|
||||
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
|
||||
|
||||
# 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)
|
||||
|
||||
# Go back to world frame.
|
||||
pos = convex.pos + pos @ convex.mat.T
|
||||
n = n @ convex.mat.T
|
||||
|
||||
dist = -jp.where(
|
||||
has_edge_contact, penetration.at[0].set(edge_penetration), penetration
|
||||
)
|
||||
frame = jax.vmap(math.make_frame)(n)
|
||||
return dist, pos, frame
|
||||
return dist, pos, normal
|
||||
|
||||
|
||||
def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
|
||||
@@ -699,16 +877,41 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
|
||||
unique_edges1 = jp.take(vertices1, c1.edge, axis=0)
|
||||
unique_edges2 = jp.take(vertices2, c2.edge, axis=0)
|
||||
|
||||
dist, pos, normal = _sat_hull_hull(
|
||||
faces1,
|
||||
faces2,
|
||||
vertices1,
|
||||
vertices2,
|
||||
normals1,
|
||||
normals2,
|
||||
unique_edges1,
|
||||
unique_edges2,
|
||||
face_edges1 = jp.take(vertices1, c1.face_edge, axis=0)
|
||||
face_edges2 = jp.take(vertices2, c2.face_edge, axis=0)
|
||||
|
||||
face_edge_normals1 = c1.face_edge_normal @ to_local_mat.T
|
||||
face_edge_normals2 = c2.face_edge_normal
|
||||
|
||||
enable_bruteforce = (
|
||||
unique_edges1.shape[0] * unique_edges2.shape[0]
|
||||
< face_edges1[0].shape[0] * face_edges2[0].shape[0]
|
||||
)
|
||||
if enable_bruteforce:
|
||||
dist, pos, normal = _sat_bruteforce(
|
||||
faces1,
|
||||
faces2,
|
||||
vertices1,
|
||||
vertices2,
|
||||
normals1,
|
||||
normals2,
|
||||
unique_edges1,
|
||||
unique_edges2,
|
||||
)
|
||||
else:
|
||||
dist, pos, normal = _sat_approx(
|
||||
to_local_pos,
|
||||
faces1,
|
||||
faces2,
|
||||
vertices1,
|
||||
vertices2,
|
||||
normals1,
|
||||
normals2,
|
||||
face_edges1,
|
||||
face_edges2,
|
||||
face_edge_normals1,
|
||||
face_edge_normals2,
|
||||
)
|
||||
|
||||
# Go back to world frame.
|
||||
pos = c2.pos + pos @ c2.mat.T
|
||||
|
||||
Reference in New Issue
Block a user