diff --git a/doc/changelog.rst b/doc/changelog.rst index 6e813d81..43e5e34e 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,6 +18,7 @@ MJX ^^^ - Added ``mocap_pos`` and ``mocap_quat`` in kinematics. - Added support for :ref:`spatial tendons ` with pulleys and external sphere and cylinder wrapping. +- Add sphere-cylinder and sphere-ellipsoid collision functions (:github:issue:`2126`). Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 565e586e..796ef3fa 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/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, diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 7cd8f54b..611703f1 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/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 = """ + + + + + + + + + + + + """ + + 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 = """ @@ -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 = """ + + + + + + + + + + + + """ + + 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) diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py index e6f04dce..43bcff21 100644 --- a/mjx/mujoco/mjx/_src/collision_sdf.py +++ b/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.""" diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index aaaa3415..d5905b86 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -193,22 +193,6 @@ class ModelIOTest(parameterized.TestCase): """)) - def test_cylinder_not_implemented(self): - with self.assertRaises(NotImplementedError): - mjx.put_model(mujoco.MjModel.from_xml_string(""" - - - - - - - - - - - - """)) - def test_margin_gap_mesh_not_implemented(self): with self.assertRaises(NotImplementedError): mjx.put_model(mujoco.MjModel.from_xml_string("""