Add sphere-ellipsoid, sphere-cylinder. #2126

PiperOrigin-RevId: 683747371
Change-Id: Ic11803bb395944380d25ce848d9fbcf19afaad77

doc/changelog.rst

@@ -18,6 +18,7 @@ MJX
 ^^^
 - Added ``mocap_pos`` and ``mocap_quat`` in kinematics.
 - Added support for :ref:`spatial tendons <tendon-spatial>` with pulleys and external sphere and cylinder wrapping.
+- Add sphere-cylinder and sphere-ellipsoid collision functions (:github:issue:`2126`).
 
 Bug fixes
 ^^^^^^^^^

mjx/mujoco/mjx/_src/collision_driver.py

@@ -66,6 +66,8 @@ from mujoco.mjx._src.collision_sdf import capsule_ellipsoid
 from mujoco.mjx._src.collision_sdf import cylinder_cylinder
 from mujoco.mjx._src.collision_sdf import ellipsoid_cylinder
 from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
+from mujoco.mjx._src.collision_sdf import sphere_cylinder
+from mujoco.mjx._src.collision_sdf import sphere_ellipsoid
 from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
@@ -89,6 +91,8 @@ _COLLISION_FUNC = {
     (GeomType.HFIELD, GeomType.MESH): hfield_convex,
     (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere,
     (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule,
+    (GeomType.SPHERE, GeomType.CYLINDER): sphere_cylinder,
+    (GeomType.SPHERE, GeomType.ELLIPSOID): sphere_ellipsoid,
     (GeomType.SPHERE, GeomType.BOX): sphere_convex,
     (GeomType.SPHERE, GeomType.MESH): sphere_convex,
     (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule,

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -242,6 +242,29 @@ class EllipsoidCollisionTest(parameterized.TestCase):
       _assert_attr_eq(
           dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5)
 
+  _ELLIPSOID_SPHERE = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+        <body pos="0 0 0.08">
+          <freejoint/>
+          <geom size=".05" type="sphere"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_sphere_ellipsoid(self):
+    """Tests ellipsoid capsule contact."""
+    d, dx = _collide(self._ELLIPSOID_SPHERE)
+    d.contact.pos[0][2] = 0.03  # MJX finds the point on the surface
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4)
+
   _ELLIPSOID_CAPSULE = """
     <mujoco>
       <worldbody>
@@ -535,6 +558,29 @@ class CylinderTest(absltest.TestCase):
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(dx.contact, d.contact, field.name, 'cylinder_plane', 1e-5)
 
+  _SPHERE_CYLINDER = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <geom size=".15 .05" type="cylinder"/>
+        </body>
+        <body pos="0 0 0.12">
+          <freejoint/>
+          <geom size=".15" type="sphere"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_sphere_cylinder(self):
+    """Tests sphere cylinder contact."""
+    d, dx = _collide(self._SPHERE_CYLINDER)
+    d.contact.pos[0][2] = 0.05  # MJX finds the deepest point on the surface
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4)
+
 
 class ConvexTest(absltest.TestCase):
   """Tests the convex contact functions."""
@@ -558,7 +604,9 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(dx.contact.dist[:2], 0)
     np.testing.assert_array_less(-dx.contact.dist[2:], 0)
     # extract the contact points with penetration
-    c = jax.tree_util.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
+    c = jax.tree_util.tree_map(
+        lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact
+    )
     c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-5)

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -196,6 +196,21 @@ def _optim(
   return dist, pos, math.make_frame(n)
 
 
+@collider(ncon=1)
+def sphere_ellipsoid(s: GeomInfo, e: GeomInfo) -> Collision:
+  """"Calculates contact between a sphere and an ellipsoid."""
+  x0 = 0.5 * (s.pos + e.pos)
+  return _optim(_sphere, _ellipsoid, s, e, x0)
+
+
+@collider(ncon=1)
+def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision:
+  """"Calculates contact between a sphere and a cylinder."""
+  # TODO: implement analytical version.
+  x0 = 0.5 * (s.pos + c.pos)
+  return _optim(_sphere, _cylinder, s, c, x0)
+
+
 @collider(ncon=1)
 def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
   """"Calculates contact between a capsule and an ellipsoid."""

mjx/mujoco/mjx/_src/io_test.py

@@ -193,22 +193,6 @@ class ModelIOTest(parameterized.TestCase):
           </tendon>
         </mujoco>"""))
 
-  def test_cylinder_not_implemented(self):
-    with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
-        <mujoco>
-          <worldbody>
-            <body>
-              <freejoint/>
-              <geom type="cylinder" size="0.05 0.05"/>
-            </body>
-            <body>
-              <freejoint/>
-              <geom size="0.05"/>
-            </body>
-          </worldbody>
-        </mujoco>"""))
-
   def test_margin_gap_mesh_not_implemented(self):
     with self.assertRaises(NotImplementedError):
       mjx.put_model(mujoco.MjModel.from_xml_string("""