diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 8fa56612..f7cec040 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/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, diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index acb52e12..c707a74f 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/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 = """ + + + + + + + + + + """ + + 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 = """ diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py index 8c2dd529..a6fe5132 100644 --- a/mjx/mujoco/mjx/_src/collision_primitive.py +++ b/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