Add ellipsoid sdf collisions to MJX.

Co-authored-by: Baruch Tabanpour <btaba@google.com>
PiperOrigin-RevId: 623415479
Change-Id: Id061c8e083357d063f1ccddf0989f2fafe415eed

mjx/mujoco/mjx/_src/collision_driver.py

@@ -37,6 +37,8 @@ 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
+from mujoco.mjx._src.collision_sdf import capsule_ellipsoid
+from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
@@ -57,7 +59,9 @@ _COLLISION_FUNC = {
     (GeomType.SPHERE, GeomType.MESH): sphere_convex,
     (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule,
     (GeomType.CAPSULE, GeomType.BOX): capsule_convex,
+    (GeomType.CAPSULE, GeomType.ELLIPSOID): capsule_ellipsoid,
     (GeomType.CAPSULE, GeomType.MESH): capsule_convex,
+    (GeomType.ELLIPSOID, GeomType.ELLIPSOID): ellipsoid_ellipsoid,
     (GeomType.BOX, GeomType.BOX): convex_convex,
     (GeomType.BOX, GeomType.MESH): convex_convex,
     (GeomType.MESH, GeomType.MESH): convex_convex,

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -220,6 +220,50 @@ class EllipsoidCollisionTest(parameterized.TestCase):
       _assert_attr_eq(
           dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5)
 
+  _ELLIPSOID_ELLIPSOID = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+        <body pos="0 0 0.09">
+          <freejoint/>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_ellipsoid_ellipsoid(self):
+    """Tests ellipsoid ellipsoid contact."""
+    d, dx = _collide(self._ELLIPSOID_ELLIPSOID)
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5)
+
+  _ELLIPSOID_CAPSULE = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+        <body pos="0 0 0.0999">
+          <freejoint/>
+          <geom size=".05" fromto="-.1 0 0 .1 0 0" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_capsule_ellipsoid(self):
+    """Tests ellipsoid capsule contact."""
+    d, dx = _collide(self._ELLIPSOID_CAPSULE)
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-4)
+
 
 class CapsuleCollisionTest(parameterized.TestCase):
   _CAP_PLANE = """

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -0,0 +1,150 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Collision functions for shapes represented as signed distance functions (SDF).
+
+A signed distance function at a given point in space is the shortest distance to
+a surface. This enables to define a geometry implicitly and exactly.
+
+See https://iquilezles.org/articles/distfunctions/ for a list of analytic SDFs.
+"""
+
+import functools
+from typing import Callable
+from typing import Tuple
+
+import jax
+from jax import numpy as jp
+from mujoco.mjx._src import math
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.collision_base import Contact
+from mujoco.mjx._src.collision_base import GeomInfo
+from mujoco.mjx._src.dataclasses import PyTreeNode
+# pylint: enable=g-importing-member
+
+
+# the objective function, inputs: pos (that we optimize for) and size (known)
+SDFFn = Callable[[jax.Array], jax.Array]
+
+
+def _plane(pos: jax.Array, size: jax.Array) -> jax.Array:
+  del size
+  return pos[2]
+
+
+def _sphere(pos: jax.Array, size: jax.Array):
+  return math.norm(pos) - size[0]
+
+
+def _capsule(pos: jax.Array, size: jax.Array):
+  pa = -size[1] * jp.array([0, 0, 1])
+  pb = size[1] * jp.array([0, 0, 1])
+  ab = pb - pa
+  ap = pos - pa
+  denom = ab.dot(ab)
+  denom = jp.where(jp.abs(denom) < 1e-12, 1e-12 * math.sign(denom), denom)
+  t = ab.dot(ap) / denom
+  t = jp.clip(t, 0, 1)
+  c = pa + t * ab
+  return math.norm(pos - c) - size[0]
+
+
+def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array:
+  k0 = math.norm(pos / size)
+  k1 = math.norm(pos / (size*size))
+  return k0 * (k0 - 1.0) / (k1 + (k1 == 0.0) * 1e-12)
+
+
+def _to_local(f: SDFFn, pos: jax.Array, mat: jax.Array)-> SDFFn:
+  return lambda p: f(mat.T @ (p - pos))
+
+
+def _intersect(d1: SDFFn, d2: SDFFn) -> SDFFn:
+  return lambda p: jp.maximum(d1(p), d2(p))
+
+
+def _clearance(d1: SDFFn, d2: SDFFn) -> SDFFn:
+  return lambda p: (d1(p) + d2(p) + jp.abs(_intersect(d1, d2)(p))).squeeze()
+
+
+class GradientState(PyTreeNode):
+  dist: jax.Array
+  x: jax.Array
+
+
+def _gradient_step(objective: SDFFn, state: GradientState) -> GradientState:
+  """Performs a step of gradient descent."""
+  # TODO: find better parameters
+  amin = 1e-4  # minimum value for line search factor scaling the gradient
+  amax = 2.  # maximum value for line search factor scaling the gradient
+  nlinesearch = 10  # line search points
+  grad = jax.grad(objective)(state.x)
+  alpha = jp.geomspace(amin, amax, nlinesearch).reshape(nlinesearch, -1)
+  candidates = state.x - alpha * grad.reshape(-1, 3)
+  values = jax.vmap(objective)(candidates)
+  idx = jp.argmin(values)
+  return state.replace(x=candidates[idx], dist=values[idx])
+
+
+def _gradient_descent(
+    objective: SDFFn,
+    x: jax.Array,
+    niter: int,
+) -> Tuple[jax.Array, jax.Array]:
+  """Performs gradient descent with backtracking line search."""
+  state = GradientState(
+      dist=1e10,
+      x=x,
+  )
+
+  state, _ = jax.lax.scan(
+      lambda s, _: (_gradient_step(objective, s), None), state, (), length=niter
+  )
+  return state.dist, state.x
+
+
+def _optim(
+    d1, d2, info1: GeomInfo, info2: GeomInfo
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+  """Optimizes the clearance function."""
+  d1 = functools.partial(d1, size=info1.size)
+  d1 = _to_local(d1, info1.pos, info1.mat)
+  d2 = functools.partial(d2, size=info2.size)
+  d2 = _to_local(d2, info2.pos, info2.mat)
+  fn = _clearance(d1, d2)
+  _, pos = _gradient_descent(fn, 0.5 * (info1.pos + info2.pos), 10)
+  dist = d1(pos) + d2(pos)
+  n = jax.grad(d1)(pos)
+  return pos, dist, n
+
+
+def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact:
+  """"Calculates contact between a capsule and an ellipsoid."""
+  pos, dist, n = _optim(_capsule, _ellipsoid, c, e)
+  return jax.tree_map(
+      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
+  )
+
+
+def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Contact:
+  """"Calculates contact between two ellipsoids."""
+  pos, dist, n = _optim(_ellipsoid, _ellipsoid, e1, e2)
+  return jax.tree_map(
+      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
+  )
+
+
+# store ncon as function attributes
+capsule_ellipsoid.ncon = 1
+ellipsoid_ellipsoid.ncon = 1

mjx/mujoco/mjx/_src/math.py

@@ -245,6 +245,11 @@ def inert_mul(i: jax.Array, v: jax.Array) -> jax.Array:
   return jp.concatenate((ang, vel))
 
 
+def sign(x: jax.Array) -> jax.Array:
+  """Returns the sign of x in the set {-1, 1}."""
+  return jp.where(x < 0, -1, 1)
+
+
 def transform_motion(vel: jax.Array, offset: jax.Array, rotmat: jax.Array):
   """Transform spatial motion.