Check all edge pairs in convex-convex.

PiperOrigin-RevId: 620004680
Change-Id: I5d75c0fd21ed3d3e420eed762cfc5a7e04896032
This commit is contained in:
Baruch Tabanpour
2024-03-28 11:32:44 -07:00
committed by Copybara-Service
parent a8db22f0d0
commit d258d5e152
9 changed files with 149 additions and 148 deletions
+3 -2
View File
@@ -299,8 +299,9 @@ Collisions between large meshes
SAT works well for smaller meshes but suffers in both runtime and memory for larger meshes.
For
collisions with convex meshes, the convex decompositon of the mesh should have
roughly **200 vertices or less** for reasonable performance.
collisions with convex meshes and primitives, the convex decompositon of the mesh should have
roughly **200 vertices or less** for reasonable performance. For convex-convex collisions,
the convex mesh should have roughly **fewer than 32 vertices**.
With careful
tuning, MJX can simulate scenes with mesh collisions -- see the MJX
`shadow hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
+3 -3
View File
@@ -50,10 +50,10 @@ class GeomInfo(PyTreeNode):
size: jax.Array
face: Optional[jax.Array] = None
vert: Optional[jax.Array] = None
edge: Optional[jax.Array] = None
edge_dir: Optional[jax.Array] = None
facenorm: Optional[jax.Array] = None
face_edge: Optional[jax.Array] = None
face_edge_normal: Optional[jax.Array] = None
edge: Optional[jax.Array] = None
edge_face_normal: Optional[jax.Array] = None
class SolverParams(PyTreeNode):
+74 -94
View File
@@ -336,8 +336,8 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
cap_closest_pt,
)
edge = jp.take(convex.vert, convex.face_edge, axis=0).reshape(-1, 2, 3)
face_edge_normal = convex.face_edge_normal.reshape((-1, 2, 3)) # pytype: disable=attribute-error
edge = jp.take(convex.vert, convex.edge, axis=0)
edge_face_normal = convex.edge_face_normal # pytype: disable=attribute-error
res = jax.vmap(get_edge_axis)(edge.reshape(-1, 2, 3))
e_idx = jp.abs(res[0]).argmin()
@@ -349,8 +349,8 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
cap_closest_pt,
) = jax.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res)
face_edge_normals = face_edge_normal[e_idx]
edge_voronoi_front = ((face_edge_normals @ edge_axis) < 0).all()
edge_face_normals = edge_face_normal[e_idx]
edge_voronoi_front = ((edge_face_normals @ edge_axis) < 0).all()
shallow = ~degenerate_edge_dir & edge_voronoi_front
edge_penetration = jp.where(shallow, cap.size[0] - edge_dist, -1)
@@ -630,8 +630,9 @@ def _sat_bruteforce(
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.).
This method checks all separating axes via a brute force support function, and
is thus costly to run over large meshes, but is more performant for smaller
meshes (boxes, tetrahedra, etc.).
Args:
faces_a: Faces for hull A.
@@ -731,7 +732,7 @@ def _arcs_intersect(
return (cba * dba < 0) & (adc * bdc < 0) & (cba * bdc > 0)
def _sat_approx(
def _sat_gaussmap(
centroid_a: jax.Array,
faces_a: jax.Array,
faces_b: jax.Array,
@@ -739,23 +740,16 @@ def _sat_approx(
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,
edges_a: jax.Array,
edges_b: jax.Array,
edge_face_normals_a: jax.Array,
edge_face_normals_b: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
"""Runs the Separating Axis Test for a pair of hulls.
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".
Runs the separating axis test for all faces. Tests edge separating axes via
edge intersections on gauss maps for all edge pairs. h/t to Dirk Gregorius
for the implementation details and gauss map trick.
Args:
centroid_a: Centroid of hull A.
@@ -765,10 +759,10 @@ def _sat_approx(
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.
edges_a: Edges for hull A.
edges_b: Edges for hull B.
edge_face_normals_a: Face normals for edges in hull A.
edge_face_normals_b: Face normals for edges in hull B.
Returns:
tuple of dist, pos, and normal
@@ -816,75 +810,61 @@ def _sat_approx(
-best_axis,
)
# 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]
# scatter to all edge pairs
edge_idx_a = jp.repeat(
jp.arange(face_edges_a.shape[0]), face_edges_b.shape[0]
# Handle edge separating axes by checking all edge pairs.
a_idx = jp.tile(jp.arange(edges_a.shape[0]), reps=edges_b.shape[0])
b_idx = jp.repeat(
jp.arange(edges_b.shape[0]), repeats=edges_a.shape[0], axis=0
)
normal_a_1 = edge_face_normals_a[a_idx, 0]
normal_a_2 = edge_face_normals_a[a_idx, 1]
normal_b_1 = edge_face_normals_b[b_idx, 0]
normal_b_2 = edge_face_normals_b[b_idx, 1]
is_minkowski_face = jax.vmap(_arcs_intersect)(
normal_a_1, normal_a_2, -normal_b_1, -normal_b_2
)
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]
@jax.vmap
# get distances
edge_a_dir = jax.vmap(math.normalize)(edges_a[:, 0] - edges_a[:, 1])[a_idx]
edge_b_dir = jax.vmap(math.normalize)(edges_b[:, 0] - edges_b[:, 1])[b_idx]
edges_a, edges_b = edges_a[a_idx], edges_b[b_idx]
edge_a_pt, edge_a_pt_2 = edges_a[:, 0], edges_a[:, 1]
edge_b_pt, edge_b_pt_2 = edges_b[:, 0], edges_b[:, 1]
def get_normals(a_dir, a_pt, b_dir):
edge_axis = math.normalize(jp.cross(a_dir, b_dir))
edge_axis = jp.cross(a_dir, b_dir)
degenerate_edge_axis = jp.sum(edge_axis**2) < 1e-6
edge_axis = math.normalize(edge_axis)
# 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
return edge_axis * sign, degenerate_edge_axis
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)
edge_axes, degenerate_edge_axes = jax.vmap(get_normals)(
edge_a_dir, edge_a_pt, edge_b_dir)
edge_dist = jax.vmap(jp.dot)(edge_axes, edge_b_pt - edge_a_pt)
# handle degenerate axis
edge_dist = jp.where(degenerate_edge_axes, -jp.inf, edge_dist)
# ensure edges create minkowski face
edge_dist = jp.where(is_minkowski_face, edge_dist, -jp.inf)
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)
# 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],
)
is_edge_contact = jp.where(
dist.max() < 0, best_edge_dist > dist.max() - 1e-6,
(best_edge_dist < 0) & ~jp.isinf(best_edge_dist)
)
is_edge_contact = is_edge_contact & ~is_face_separating
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,
)
a_closest, b_closest = math.closest_segment_to_segment_points(
edge_a_pt[best_edge_idx], edge_a_pt_2[best_edge_idx],
edge_b_pt[best_edge_idx], edge_b_pt_2[best_edge_idx])
pos = jp.where(
is_edge_contact,
jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos)
return dist, pos, normal
@@ -924,15 +904,15 @@ 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)
face_edges1 = jp.take(vertices1, c1.face_edge, axis=0)
face_edges2 = jp.take(vertices2, c2.face_edge, axis=0)
edges1 = jp.take(vertices1, c1.edge, axis=0)
edges2 = jp.take(vertices2, c2.edge, axis=0)
face_edge_normals1 = c1.face_edge_normal @ to_local_mat.T
face_edge_normals2 = c2.face_edge_normal
edge_face_normals1 = c1.edge_face_normal @ to_local_mat.T
edge_face_normals2 = c2.edge_face_normal
enable_bruteforce = (
unique_edges1.shape[0] * unique_edges2.shape[0]
< face_edges1[0].shape[0] * face_edges2[0].shape[0]
< edges1[0].shape[0] * edges2[0].shape[0]
)
if enable_bruteforce:
dist, pos, normal = _sat_bruteforce(
@@ -946,7 +926,7 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
unique_edges2,
)
else:
dist, pos, normal = _sat_approx(
dist, pos, normal = _sat_gaussmap(
to_local_pos,
faces1,
faces2,
@@ -954,10 +934,10 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
vertices2,
normals1,
normals2,
face_edges1,
face_edges2,
face_edge_normals1,
face_edge_normals2,
edges1,
edges2,
edge_face_normals1,
edge_face_normals2,
)
# Go back to world frame.
+7 -7
View File
@@ -230,20 +230,20 @@ def _pair_info(
info = info.replace(
face=jp.stack([m.geom_convex_face[i] for i in geom]),
vert=jp.stack([m.geom_convex_vert[i] for i in geom]),
edge=jp.stack([m.geom_convex_edge_dir[i] for i in geom]),
edge_dir=jp.stack([m.geom_convex_edge_dir[i] for i in geom]),
facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom]),
face_edge_normal=jp.stack(
[m.geom_convex_face_edge_normal[i] for i in geom]
edge=jp.stack([m.geom_convex_edge[i] for i in geom]),
edge_face_normal=jp.stack(
[m.geom_convex_edge_face_normal[i] for i in geom]
),
face_edge=jp.stack([m.geom_convex_face_edge[i] for i in geom]),
)
in_axes = in_axes.replace(
face=0,
vert=0,
edge=0,
edge_dir=0,
facenorm=0,
face_edge=0,
face_edge_normal=0,
edge=0,
edge_face_normal=0,
)
return info, in_axes
+41 -3
View File
@@ -527,7 +527,7 @@ class ConvexTest(absltest.TestCase):
"""
def test_convex_convex(self):
"""Tests generic convex-convex collision via _sat_approx."""
"""Tests generic convex-convex collision via _sat_gaussmap."""
directory = epath.resource_path('mujoco.mjx')
assets = {
'meshes/dodecahedron.stl': (
@@ -542,6 +542,44 @@ class ConvexTest(absltest.TestCase):
np.testing.assert_array_less(0, c.dist[1:])
np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1]))
_CONVEX_CONVEX_THIN = """
<mujoco>
<asset>
<mesh name="poly"
vertex="0.3 0 0 0 0.5 0 -0.3 0 0 0 -0.5 0 0 -1 1 0 1 1"
face="0 1 5 0 5 4 0 4 3 3 4 2 2 4 5 1 2 5 0 2 1 0 3 2"/>
</asset>
<worldbody>
<body pos="0.0 2.0 0.35" euler="0 0 90">
<freejoint/>
<geom size="0.2 0.2 0.2" type="mesh" mesh="poly"/>
</body>
<body pos="0.0 2.0 2.281" euler="180 0 0">
<freejoint/>
<geom size="0.2 0.2 0.2" type="mesh" mesh="poly"/>
</body>
</worldbody>
</mujoco>
"""
def test_convex_convex_edge(self):
"""Tests convex-convex collisions with edge contact via _sat_gaussmap."""
_, dx = _collide(self._CONVEX_CONVEX_THIN)
c = dx.contact
# Only one contact point for an edge contact.
self.assertLess(c.dist[0], 0)
np.testing.assert_array_less(0, c.dist[1:])
np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1]))
np.testing.assert_array_almost_equal(
c.pos[0], np.array([0, 2, 1.3155]), decimal=5)
_, dx = _collide(
self._CONVEX_CONVEX_THIN.replace(
'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'))
c = dx.contact
self.assertTrue((c.dist > 0).all())
class BodyPairFilterTest(absltest.TestCase):
"""Tests that certain body pairs get filtered."""
@@ -647,11 +685,11 @@ class NconTest(parameterized.TestCase):
m.numeric_data[m.numeric_adr[i]] = -1
ncon = collision_driver.ncon(m)
self.assertEqual(ncon, 307)
self.assertEqual(ncon, 98)
mx = mjx.put_model(m)
ncon = collision_driver.ncon(mx)
self.assertEqual(ncon, 307)
self.assertEqual(ncon, 98)
class TopKContactTest(absltest.TestCase):
+10 -19
View File
@@ -49,8 +49,8 @@ _DERIVED_ARGS = [
'geom_convex_vert',
'geom_convex_edge_dir',
'geom_convex_facenormal',
'geom_convex_face_edge',
'geom_convex_face_edge_normal',
'geom_convex_edge',
'geom_convex_edge_face_normal',
]
DERIVED = {(Model, d) for d in _DERIVED_ARGS}
@@ -114,7 +114,7 @@ def _get_unique_edge_dir(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
return edges[unique_edge_idx]
def _get_face_edge_normals(
def _get_edge_normals(
face: np.ndarray, face_norm: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""Returns face edges and face edge normals."""
@@ -133,7 +133,7 @@ def _get_face_edge_normals(
continue
edge_map_list[tuple(face_edge_flat[i])].append(edge_face_norm[i])
edge_map = {}
edges, edge_face_normals = [], []
for k, v in edge_map_list.items():
v = np.array(v)
if len(v) > 2:
@@ -145,19 +145,10 @@ def _get_face_edge_normals(
# and face vertices were down sampled. In either case, we ignore these
# edges.
continue
edge_map[k] = v
edges.append(k)
edge_face_normals.append(v)
# for each face, list the edge normals
face_edge_normal = []
for face_idx in range(face_edge.shape[0]):
normals = []
for edge in face_edge[face_idx]:
k = tuple(edge)
normals.append(edge_map.get(k, np.zeros((2, 3))))
face_edge_normal.append(np.array(normals))
face_edge_normal = np.array(face_edge_normal)
return face_edge, face_edge_normal
return np.array(edges), np.array(edge_face_normals)
def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
@@ -297,15 +288,15 @@ def _geom_mesh_kwargs(
vert = np.array(tm_convex.vertices)
face = _merge_coplanar(tm_convex, mesh_info)
facenormal = _get_face_norm(vert, face)
face_edge, face_edge_normal = _get_face_edge_normals(face, facenormal)
edge, edge_face_normal = _get_edge_normals(face, facenormal)
return {
'geom_convex_face': vert[face],
'geom_convex_face_vert_idx': face,
'geom_convex_vert': vert,
'geom_convex_edge_dir': _get_unique_edge_dir(vert, face),
'geom_convex_facenormal': facenormal,
'geom_convex_face_edge': face_edge,
'geom_convex_face_edge_normal': face_edge_normal,
'geom_convex_edge': edge,
'geom_convex_edge_face_normal': edge_face_normal,
}
+4 -13
View File
@@ -80,7 +80,7 @@ class GeomMeshKwargsTest(absltest.TestCase):
self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3))
# face edges
edges = np.concatenate(h['geom_convex_face_edge'])
edges = h['geom_convex_edge']
edges = np.vectorize(map_.get)(edges)
mask = edges[:, 0] != edges[:, 1]
edges = edges[mask]
@@ -92,27 +92,18 @@ class GeomMeshKwargsTest(absltest.TestCase):
edges,
np.array([
[0, 2],
[0, 2],
[0, 3],
[0, 3],
[0, 4],
[0, 4],
[1, 2],
[1, 2],
[1, 3],
[1, 3],
[1, 4],
[1, 4],
[2, 4],
[2, 4],
[3, 4],
[3, 4],
]),
)
# face edge normals
edge_normal = h['geom_convex_face_edge_normal']
edge_normal = np.concatenate(edge_normal)
edge_normal = h['geom_convex_edge_face_normal']
edge_normal = edge_normal[mask]
edge_normal = np.take_along_axis(
edge_normal, sort_col_idx[..., None], axis=1
@@ -120,8 +111,8 @@ class GeomMeshKwargsTest(absltest.TestCase):
edge_normal = edge_normal[sort_row_idx]
edge_normal_02 = np.array([[0.4472136, -0.0, 0.89442719], [-1.0, 0.0, 0.0]])
np.testing.assert_array_almost_equal(
edge_normal[:2],
np.array([edge_normal_02, edge_normal_02]),
edge_normal[:1],
np.array([edge_normal_02]),
)
+2 -2
View File
@@ -543,8 +543,8 @@ class Model(PyTreeNode):
geom_convex_vert: List[Optional[jax.Array]]
geom_convex_edge_dir: List[Optional[jax.Array]]
geom_convex_facenormal: List[Optional[jax.Array]]
geom_convex_face_edge: List[Optional[jax.Array]]
geom_convex_face_edge_normal: List[Optional[jax.Array]]
geom_convex_edge: List[Optional[jax.Array]]
geom_convex_edge_face_normal: List[Optional[jax.Array]]
pair_solref: jax.Array
pair_solreffriction: jax.Array
pair_solimp: jax.Array
@@ -156,7 +156,7 @@
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_FFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214" contype="0" conaffinity="2"/>
</body>
</body>
</body>
@@ -181,7 +181,7 @@
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_MFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214" contype="0" conaffinity="2"/>
</body>
</body>
</body>
@@ -206,7 +206,7 @@
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_RFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214" contype="0" conaffinity="2"/>
</body>
</body>
</body>
@@ -236,7 +236,7 @@
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_LFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst_214" contype="0" conaffinity="2"/>
</body>
</body>
</body>
@@ -266,7 +266,7 @@
diaginertia="2.37794e-06 2.27794e-06 1e-06"/>
<joint name="rh_THJ1" class="thdistal"/>
<geom class="plastic_visual" mesh="th_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="th_distal_pst_190"/>
<geom class="plastic_collision" type="mesh" mesh="th_distal_pst_190" contype="0" conaffinity="2"/>
</body>
</body>
</body>