Speed up convex-convex collisions.

PiperOrigin-RevId: 613450874
Change-Id: Ie66e0a078e2ff51a136b949dd9459847e1fe957d

mjx/mujoco/mjx/_src/collision_base.py

@@ -52,6 +52,8 @@ class GeomInfo(PyTreeNode):
   vert: Optional[jax.Array] = None
   edge: Optional[jax.Array] = None
   facenorm: Optional[jax.Array] = None
+  face_edge: Optional[jax.Array] = None
+  face_edge_normal: Optional[jax.Array] = None
 
 
 class SolverParams(PyTreeNode):

mjx/mujoco/mjx/_src/collision_convex.py

@@ -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

mjx/mujoco/mjx/_src/collision_driver.py

@@ -20,8 +20,8 @@ import jax
 from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import collision_base
+from mujoco.mjx._src import mesh
 from mujoco.mjx._src import support
-from mujoco.mjx._src import math
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.collision_base import Candidate
 from mujoco.mjx._src.collision_base import CandidateSet
@@ -43,7 +43,6 @@ from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
-
 # pair-wise collision functions
 _COLLISION_FUNC = {
     (GeomType.PLANE, GeomType.SPHERE): plane_sphere,
@@ -91,7 +90,18 @@ def _add_candidate(
   def mesh_key(i):
     convex_data = [[None] * m.ngeom] * 3
     if isinstance(m, Model):
-      convex_data = [m.geom_convex_face, m.geom_convex_vert, m.geom_convex_edge]
+      convex_data = [
+          m.geom_convex_face,
+          m.geom_convex_vert,
+          m.geom_convex_edge_dir,
+      ]
+    elif isinstance(m, mujoco.MjModel):
+      kwargs = mesh.get(m)
+      convex_data = [
+          kwargs['geom_convex_face'],
+          kwargs['geom_convex_vert'],
+          kwargs['geom_convex_edge_dir'],
+      ]
     key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data)
     return key
 
@@ -181,36 +191,39 @@ def _pair_info(
     m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int]
 ) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]:
   """Returns geom pair info for calculating collision."""
-  g1, g2 = jp.array(geom1), jp.array(geom2)
-  info1 = GeomInfo(
-      g1,
-      d.geom_xpos[g1],
-      d.geom_xmat[g1],
-      m.geom_size[g1],
-  )
-  info2 = GeomInfo(
-      g2,
-      d.geom_xpos[g2],
-      d.geom_xmat[g2],
-      m.geom_size[g2],
-  )
-  in_axes1 = in_axes2 = jax.tree_map(lambda x: 0, info1)
-  if m.geom_convex_face[geom1[0]] is not None:
-    info1 = info1.replace(
-        face=jp.stack([m.geom_convex_face[i] for i in geom1]),
-        vert=jp.stack([m.geom_convex_vert[i] for i in geom1]),
-        edge=jp.stack([m.geom_convex_edge[i] for i in geom1]),
-        facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom1]),
+  def mesh_info(geom):
+    g = jp.array(geom)
+    info = GeomInfo(
+        g,
+        d.geom_xpos[g],
+        d.geom_xmat[g],
+        m.geom_size[g],
     )
-    in_axes1 = in_axes1.replace(face=0, vert=0, edge=0, facenorm=0)
-  if m.geom_convex_face[geom2[0]] is not None:
-    info2 = info2.replace(
-        face=jp.stack([m.geom_convex_face[i] for i in geom2]),
-        vert=jp.stack([m.geom_convex_vert[i] for i in geom2]),
-        edge=jp.stack([m.geom_convex_edge[i] for i in geom2]),
-        facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom2]),
-    )
-    in_axes2 = in_axes2.replace(face=0, vert=0, edge=0, facenorm=0)
+    in_axes = jax.tree_map(lambda x: 0, info)
+    is_mesh = m.geom_convex_face[geom[0]] is not None
+    if is_mesh:
+      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]),
+          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]
+          ),
+          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,
+          facenorm=0,
+          face_edge=0,
+          face_edge_normal=0,
+      )
+    return info, in_axes
+
+  info1, in_axes1 = mesh_info(geom1)
+  info2, in_axes2 = mesh_info(geom2)
   return info1, info2, [in_axes1, in_axes2]
 
 
@@ -289,9 +302,8 @@ def _collide_geoms(
   n_pairs = max_pairs if run_broadphase else len(geom1)
   if run_broadphase:
     # broadphase over geom pairs, using bounding spheres
-    size1 = jp.max(m.geom_size[jp.array(geom1)], axis=-1)
-    size2 = jp.max(m.geom_size[jp.array(geom2)], axis=-1)
-    # TODO(btaba): consider re-using collision info for (sphere, sphere)
+    size1 = jp.max(m.geom_size[g1.geom_id], axis=-1)
+    size2 = jp.max(m.geom_size[g2.geom_id], axis=-1)
     dists = jax.vmap(jp.linalg.norm)(g2.pos - g1.pos) - (size1 + size2)
     _, idx = jax.lax.top_k(-dists, k=n_pairs)
     g1, g2, params = jax.tree_map(

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -371,15 +371,14 @@ class ConvexTest(absltest.TestCase):
   _CONVEX_CONVEX = """
     <mujoco>
       <asset>
-        <mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.1 0.1 0.1" />
         <mesh name="dodecahedron" file="meshes/dodecahedron.stl" scale="0.01 0.01 0.01" />
       </asset>
       <worldbody>
         <body pos="0.0 2.0 0.096">
           <joint axis="1 0 0" type="free"/>
-          <geom size="0.2 0.2 0.2" type="mesh" mesh="tetrahedron"/>
+          <geom size="0.2 0.2 0.2" type="mesh" mesh="dodecahedron"/>
         </body>
-        <body pos="0.0 2.0 0.289" euler="0.1 -0.1 45">
+        <body pos="0.0 2.0 0.281" euler="0.1 -0.1 45">
           <joint axis="1 0 0" type="free"/>
           <geom size="0.1 0.1 0.1" type="mesh" mesh="dodecahedron"/>
         </body>
@@ -388,12 +387,9 @@ class ConvexTest(absltest.TestCase):
   """
 
   def test_convex_convex(self):
-    """Tests generic convex-convex collision."""
+    """Tests generic convex-convex collision via _sat_approx."""
     directory = epath.resource_path('mujoco.mjx')
     assets = {
-        'meshes/tetrahedron.stl': (
-            directory / 'test_data' / 'meshes/tetrahedron.stl'
-        ).read_bytes(),
         'meshes/dodecahedron.stl': (
             directory / 'test_data' / 'meshes/dodecahedron.stl'
         ).read_bytes(),
@@ -492,6 +488,31 @@ class NconTest(parameterized.TestCase):
     ncon = collision_driver.ncon(m)
     self.assertEqual(ncon, 0)
 
+  def test_ncon_meshes(self):
+    m = test_util.load_test_file('shadow_hand/scene_right.xml')
+
+    ncon = collision_driver.ncon(m)
+    self.assertEqual(ncon, 15)
+
+    mx = mjx.put_model(m)
+    ncon = collision_driver.ncon(mx)
+    self.assertEqual(ncon, 15)
+
+    # get rid of max_contact_points, test only max_geom_pairs
+    for i in range(m.nnumeric):
+      name_ = (
+          m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0]
+      )
+      if name_ == 'max_contact_points':
+        m.numeric_data[m.numeric_adr[i]] = -1
+
+    ncon = collision_driver.ncon(m)
+    self.assertEqual(ncon, 307)
+
+    mx = mjx.put_model(m)
+    ncon = collision_driver.ncon(mx)
+    self.assertEqual(ncon, 307)
+
 
 class TopKContactTest(absltest.TestCase):
   """Tests top-k contacts."""

mjx/mujoco/mjx/_src/io_test.py

@@ -110,7 +110,7 @@ class ModelIOTest(parameterized.TestCase):
     np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
     self.assertLen(mx.geom_convex_face, 6)
     self.assertLen(mx.geom_convex_vert, 6)
-    self.assertLen(mx.geom_convex_edge, 6)
+    self.assertLen(mx.geom_convex_edge_dir, 6)
     self.assertLen(mx.geom_convex_facenormal, 6)
 
     np.testing.assert_allclose(mx.jnt_type, m.jnt_type)

mjx/mujoco/mjx/_src/mesh.py

@@ -14,8 +14,11 @@
 # ==============================================================================
 """Mesh processing."""
 
+import collections
+import dataclasses
 import itertools
-from typing import Dict, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
+import warnings
 
 import mujoco
 # pylint: disable=g-importing-member
@@ -44,8 +47,10 @@ _CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {}
 _DERIVED_ARGS = [
     'geom_convex_face',
     'geom_convex_vert',
-    'geom_convex_edge',
+    'geom_convex_edge_dir',
     'geom_convex_facenormal',
+    'geom_convex_face_edge',
+    'geom_convex_face_edge_normal',
 ]
 DERIVED = {(Model, d) for d in _DERIVED_ARGS}
 
@@ -73,8 +78,8 @@ def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
   return face_norm
 
 
-def _get_unique_edges(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
-  """Returns unique edges.
+def _get_unique_edge_dir(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
+  """Returns unique edge directions.
 
   Args:
     vert: (n_vert, 3) vertices
@@ -109,6 +114,52 @@ def _get_unique_edges(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
   return edges[unique_edge_idx]
 
 
+def _get_face_edge_normals(
+    face: np.ndarray, face_norm: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+  """Returns face edges and face edge normals."""
+  # get face edges and scatter the face norms
+  r_face = np.roll(face, 1, axis=1)
+  face_edge = np.array([face, r_face]).transpose((1, 2, 0))
+  face_edge.sort(axis=2)
+  face_edge_flat = np.concatenate(face_edge)
+  edge_face_idx = np.repeat(np.arange(face.shape[0]), face.shape[1])
+  edge_face_norm = face_norm[edge_face_idx]
+
+  # get the edge normals associated with each edge
+  edge_map_list = collections.defaultdict(list)
+  for i in range(face_edge_flat.shape[0]):
+    if face_edge_flat[i][0] == face_edge_flat[i][1]:
+      continue
+    edge_map_list[tuple(face_edge_flat[i])].append(edge_face_norm[i])
+
+  edge_map = {}
+  for k, v in edge_map_list.items():
+    v = np.array(v)
+    if len(v) > 2:
+      # Meshes can be of poor quality and contain edges adjacent to more than
+      # two faces. We take the first two unique face normals.
+      v = np.unique(v, axis=0)[:2]
+    elif len(v) == 1:
+      # Some edges are either degenerate or _MAX_HULL_FACE_VERTICES was hit
+      # and face vertices were down sampled. In either case, we ignore these
+      # edges.
+      continue
+    edge_map[k] = 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
+
+
 def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
   """Calculates the convex hull for a set of points on a plane."""
   # project points onto the closest axis plane
@@ -128,7 +179,16 @@ def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
   return hull_point_idx
 
 
-def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray:
+@dataclasses.dataclass
+class MeshInfo:
+  name: str
+  vert: np.ndarray
+  face: np.ndarray
+  convex_vert: Optional[np.ndarray]
+  convex_face: Optional[np.ndarray]
+
+
+def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
   """Merges coplanar facets."""
   if not tm.facets:
     return tm.faces.copy()  # no facets
@@ -152,6 +212,13 @@ def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray:
     face = point_idx[hull_point_idx]
 
     # resize faces that exceed max polygon vertices
+    if face.shape[0] > _MAX_HULL_FACE_VERTICES:
+      warnings.warn(
+          f'Mesh "{mesh_info.name}" has a coplanar face with more than'
+          f' {_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '
+          'issues and inaccuracies in collision detection. Consider '
+          'decimating the mesh.'
+      )
     every = face.shape[0] // _MAX_HULL_FACE_VERTICES + 1
     face = face[::every]
     facets.append(face)
@@ -173,48 +240,82 @@ def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray:
   return np.concatenate([faces, facets])
 
 
-def _get_faces_verts(
+def _mesh_info(
     m: mujoco.MjModel,
-) -> Tuple[Sequence[np.ndarray], Sequence[np.ndarray]]:
-  """Extracts mesh faces and vertices from MjModel."""
-  verts, faces = [], []
+) -> List[MeshInfo]:
+  """Extracts mesh info from MjModel."""
+  mesh_infos = []
   for i in range(m.nmesh):
+    name = mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_MESH.value, i)
+
     last = (i + 1) >= m.nmesh
     face_start = m.mesh_faceadr[i]
     face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0]
     face = m.mesh_face[face_start:face_end]
-    faces.append(face)
 
     vert_start = m.mesh_vertadr[i]
     vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0]
     vert = m.mesh_vert[vert_start:vert_end]
-    verts.append(vert)
-  return verts, faces
+
+    graphadr = m.mesh_graphadr[i]
+    if graphadr < 0:
+      mesh_infos.append(MeshInfo(name, vert, face, None, None))
+      continue
+
+    graph = m.mesh_graph[graphadr:]
+    numvert, numface = graph[0], graph[1]
+
+    # unused vert_edgeadr
+    # vert_edgeadr = graph[2 : numvert + 2]
+    last_idx = numvert + 2
+
+    vert_globalid = graph[last_idx : last_idx + numvert]
+    last_idx += numvert
+
+    # unused edge_localid
+    # edge_localid = graph[last_idx : last_idx + numvert + 3 * numface]
+    last_idx += numvert + 3 * numface
+
+    face_globalid = graph[last_idx : last_idx + 3 * numface]
+    face_globalid = face_globalid.reshape((numface, 3))
+
+    convex_vert = vert[vert_globalid]
+    vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0])))
+    convex_face = np.vectorize(vertex_map.get)(face_globalid)
+    mesh_infos.append(MeshInfo(name, vert, face, convex_vert, convex_face))
+
+  return mesh_infos
 
 
 def _geom_mesh_kwargs(
-    vert: np.ndarray, face: np.ndarray
+    mesh_info: MeshInfo,
 ) -> Dict[str, np.ndarray]:
   """Generates convex mesh attributes for mjx.Model."""
-  tm = trimesh.Trimesh(vertices=vert, faces=face)
-  tm_convex = trimesh.convex.convex_hull(tm)
+  tm_convex = trimesh.Trimesh(
+      vertices=mesh_info.convex_vert, faces=mesh_info.convex_face
+  )
   vert = np.array(tm_convex.vertices)
-  face = _merge_coplanar(tm_convex)
+  face = _merge_coplanar(tm_convex, mesh_info)
+  facenormal = _get_face_norm(vert, face)
+  face_edge, face_edge_normal = _get_face_edge_normals(face, facenormal)
   return {
       'geom_convex_face': vert[face],
       'geom_convex_face_vert_idx': face,
       'geom_convex_vert': vert,
-      'geom_convex_edge': _get_unique_edges(vert, face),
-      'geom_convex_facenormal': _get_face_norm(vert, face),
+      '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,
   }
 
 
 def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
   """Derives geom mesh attributes for mjx.Model from MjModel."""
   kwargs = {k: [] for k in _DERIVED_ARGS}
-  verts, faces = _get_faces_verts(m)
+  mesh_infos = _mesh_info(m)
   geom_con = m.geom_conaffinity | m.geom_contype
   for geomid in range(m.ngeom):
+    mesh_info = None
     dataid = m.geom_dataid[geomid]
     if not geom_con[geomid]:
       # ignore visual-only meshes
@@ -222,15 +323,22 @@ def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
       continue
     elif m.geom_type[geomid] == GeomType.BOX:
       vert, face = _box(m.geom_size[geomid])
-    elif dataid >= 0:
-      vert, face = verts[dataid], faces[dataid]
-    else:
+      mesh_info = MeshInfo(
+          name='box',
+          vert=vert,
+          face=face,
+          convex_vert=vert,
+          convex_face=face,
+      )
+    elif dataid < 0:
       kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
       continue
 
+    mesh_info = mesh_info or mesh_infos[dataid]
+    vert, face = mesh_info.vert, mesh_info.face
     key = (hash(vert.data.tobytes()), hash(face.data.tobytes()))
     if key not in _CONVEX_CACHE:
-      _CONVEX_CACHE[key] = _geom_mesh_kwargs(vert, face)
+      _CONVEX_CACHE[key] = _geom_mesh_kwargs(mesh_info)
 
     kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS}
 

mjx/mujoco/mjx/_src/mesh_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from mujoco.mjx._src import mesh
 import numpy as np
+import trimesh
 
 
 class GeomMeshKwargsTest(absltest.TestCase):
@@ -33,7 +34,18 @@ class GeomMeshKwargsTest(absltest.TestCase):
     face = np.array(
         [[0, 1, 2], [0, 3, 1], [0, 4, 3], [0, 2, 4], [2, 1, 4], [1, 3, 4]]
     )
-    h = mesh._geom_mesh_kwargs(vert, face)
+    tm = trimesh.Trimesh(vertices=vert, faces=face)
+    tm_convex = trimesh.convex.convex_hull(tm)
+    convex_vert = np.array(tm_convex.vertices)
+    convex_face = np.array(tm_convex.faces)
+    mesh_info = mesh.MeshInfo(
+        name='test',
+        vert=vert,
+        face=face,
+        convex_vert=convex_vert,
+        convex_face=convex_face,
+    )
+    h = mesh._geom_mesh_kwargs(mesh_info)
 
     # get index of vertices in h['geom_convex_vert'] for vertices in vert
     dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile(
@@ -57,7 +69,7 @@ class GeomMeshKwargsTest(absltest.TestCase):
     )
 
     # check edges
-    unique_edge = np.vectorize(map_.get)(h['geom_convex_edge'])
+    unique_edge = np.vectorize(map_.get)(h['geom_convex_edge_dir'])
     unique_edge = np.array(sorted(unique_edge.tolist()))
     np.testing.assert_array_equal(
         unique_edge,
@@ -67,6 +79,51 @@ class GeomMeshKwargsTest(absltest.TestCase):
     # face normals
     self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3))
 
+    # face edges
+    edges = np.concatenate(h['geom_convex_face_edge'])
+    edges = np.vectorize(map_.get)(edges)
+    mask = edges[:, 0] != edges[:, 1]
+    edges = edges[mask]
+    sort_col_idx = np.argsort(edges, axis=1)
+    edges = np.take_along_axis(edges, sort_col_idx, axis=1)
+    sort_row_idx = np.lexsort((edges[:, 1], edges[:, 0]))
+    edges = edges[sort_row_idx]
+    np.testing.assert_array_equal(
+        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 = edge_normal[mask]
+    edge_normal = np.take_along_axis(
+        edge_normal, sort_col_idx[..., None], axis=1
+    )
+    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]),
+    )
+
 
 class ConvexHull2DTest(absltest.TestCase):
 
@@ -109,7 +166,7 @@ class UniqueEdgesTest(absltest.TestCase):
         [[-0.1, 0.0, -0.1], [0.0, 0.1, 0.1], [0.1, 0.0, -0.1], [0.0, -0.1, 0.1]]
     )
     face = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 1], [2, 1, 3]])
-    idx = mesh._get_unique_edges(vert, face)
+    idx = mesh._get_unique_edge_dir(vert, face)
     np.testing.assert_array_equal(
         idx, np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]])
     )

mjx/mujoco/mjx/_src/types.py

@@ -394,8 +394,10 @@ class Model(PyTreeNode):
     mesh_face: vertex face data                               (nmeshface, 3)
     geom_convex_face: vertex face data, MJX only              (ngeom,)
     geom_convex_vert: vertex data, MJX only                   (ngeom,)
-    geom_convex_edge: unique edge data, MJX only              (ngeom,)
+    geom_convex_edge_dir: unique edge direction, MJX only     (ngeom,)
     geom_convex_facenormal: normal face data, MJX only        (ngeom,)
+    geom_convex_face_edge: edges for each face                (ngeom,)
+    geom_convex_face_edge_normal: edge normals for each face  (ngeom,)
     pair_dim: contact dimensionality                          (npair,)
     pair_geom1: id of geom1                                   (npair,)
     pair_geom2: id of geom2                                   (npair,)
@@ -539,8 +541,10 @@ class Model(PyTreeNode):
   pair_geom2: np.ndarray
   geom_convex_face: List[Optional[jax.Array]]
   geom_convex_vert: List[Optional[jax.Array]]
-  geom_convex_edge: 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]]
   pair_solref: jax.Array
   pair_solreffriction: jax.Array
   pair_solimp: jax.Array