Add ellipsoid plane.

Co-authored-by: Alessio Quaglino <quaglino@google.com>
PiperOrigin-RevId: 621208833
Change-Id: I270c1ae0cecd7dd2d9b41da4e4d0b6378b105aaf

mjx/mujoco/mjx/_src/collision_driver.py

@@ -33,6 +33,7 @@ from mujoco.mjx._src.collision_convex import plane_convex
 from mujoco.mjx._src.collision_convex import sphere_convex
 from mujoco.mjx._src.collision_primitive import capsule_capsule
 from mujoco.mjx._src.collision_primitive import plane_capsule
+from mujoco.mjx._src.collision_primitive import plane_ellipsoid
 from mujoco.mjx._src.collision_primitive import plane_sphere
 from mujoco.mjx._src.collision_primitive import sphere_capsule
 from mujoco.mjx._src.collision_primitive import sphere_sphere
@@ -48,6 +49,7 @@ _COLLISION_FUNC = {
     (GeomType.PLANE, GeomType.SPHERE): plane_sphere,
     (GeomType.PLANE, GeomType.CAPSULE): plane_capsule,
     (GeomType.PLANE, GeomType.BOX): plane_convex,
+    (GeomType.PLANE, GeomType.ELLIPSOID): plane_ellipsoid,
     (GeomType.PLANE, GeomType.MESH): plane_convex,
     (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere,
     (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule,

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -198,6 +198,29 @@ class SphereCollisionTest(parameterized.TestCase):
       _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex_center', 1e-4)
 
 
+class EllipsoidCollisionTest(parameterized.TestCase):
+
+  _ELLIPSOID_PLANE = """
+    <mujoco>
+      <worldbody>
+        <geom name="floor" size="0 0 .05" type="plane"/>
+        <body pos="0 0 0.03" euler="45 0 0">
+          <freejoint/>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_plane_ellipsoid(self):
+    """Tests ellipsoid plane contact."""
+    d, dx = _collide(self._ELLIPSOID_PLANE)
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5)
+
+
 class CapsuleCollisionTest(parameterized.TestCase):
   _CAP_PLANE = """
     <mujoco>

mjx/mujoco/mjx/_src/collision_primitive.py

@@ -65,6 +65,19 @@ def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact:
   return jax.tree_map(lambda *x: jp.concatenate(x), *contacts)
 
 
+def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Contact:
+  """Calculates one contact between an ellipsoid and a plane."""
+  n = plane.mat[:, 2]
+  size = ellipsoid.size
+  sphere_support = -math.normalize((ellipsoid.mat.T @ n) * size)
+  pos = ellipsoid.pos + ellipsoid.mat @ (sphere_support * size)
+  dist = jp.dot(n, pos - plane.pos)
+  pos = pos - n * dist * 0.5
+  return jax.tree_map(
+      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
+  )
+
+
 def _sphere_sphere(
     pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array
 ) -> Contact:
@@ -121,6 +134,7 @@ def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact:
 # store ncon as function attributes
 plane_sphere.ncon = 1
 plane_capsule.ncon = 2
+plane_ellipsoid.ncon = 1
 sphere_sphere.ncon = 1
 sphere_capsule.ncon = 1
 capsule_capsule.ncon = 1