Fix manifold point function.

PiperOrigin-RevId: 606366838
Change-Id: I8cf031f6f0fd0173777e0a9bb6d11c0c77c617ec

mjx/mujoco/mjx/_src/collision_convex.py

@@ -173,7 +173,7 @@ def _manifold_points(
   bp = b - poly
   dist_bp = jp.abs(bp.dot(bc)) + dist_mask
   dist_ap = jp.abs(ap.dot(ac)) + dist_mask
-  d_idx = jp.concatenate([dist_bp, dist_ap]).argmax() % poly.shape[0]
+  d_idx = (dist_bp + dist_ap).argmax() % poly.shape[0]
   return jp.array([a_idx, b_idx, c_idx, d_idx])
 
 
@@ -514,6 +514,7 @@ def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
   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
 
 

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -265,7 +265,7 @@ class ConvexTest(absltest.TestCase):
       <worldbody>
         <geom size="40 40 40" type="plane"/>
         <body pos="0 0 0.7" euler="45 0 0">
-          <joint axis="1 0 0" type="free"/>
+          <freejoint/>
           <geom size="0.5 0.5 0.5" type="box"/>
         </body>
       </worldbody>
@@ -281,7 +281,35 @@ class ConvexTest(absltest.TestCase):
     # extract the contact points with penetration
     c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
+      _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-5)
+
+  _FLAT_BOX_PLANE = """
+    <mujoco>
+      <worldbody>
+        <geom size="40 40 40" type="plane"/>
+        <body pos="0 0 0.45">
+          <freejoint/>
+          <geom size="0.5 0.5 0.5" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_flat_box_plane(self):
+    """Tests box collision with a plane."""
+    with jax.disable_jit():
+      d, dx = _collide(self._FLAT_BOX_PLANE)
+
+    np.testing.assert_array_less(dx.contact.dist, 0)
+
+    # sort positions for comparison
+    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
+    d.contact.pos[:] = d.contact.pos[idx]
+
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5)
 
   _BOX_BOX = """
     <mujoco>