From 2b3f336bd3bd441b2470a25652966022b2c44388 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 10 Apr 2024 01:39:52 -0700 Subject: [PATCH] Add ellipsoid sdf collisions to MJX. Co-authored-by: Baruch Tabanpour PiperOrigin-RevId: 623415479 Change-Id: Id061c8e083357d063f1ccddf0989f2fafe415eed --- doc/changelog.rst | 12 +- mjx/mujoco/mjx/_src/collision_driver.py | 4 + mjx/mujoco/mjx/_src/collision_driver_test.py | 44 ++++++ mjx/mujoco/mjx/_src/collision_sdf.py | 150 +++++++++++++++++++ mjx/mujoco/mjx/_src/math.py | 5 + 5 files changed, 210 insertions(+), 5 deletions(-) create mode 100644 mjx/mujoco/mjx/_src/collision_sdf.py diff --git a/doc/changelog.rst b/doc/changelog.rst index 8a737fcd..81a20c3f 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -8,6 +8,7 @@ Upcoming version (not yet released) MJX ^^^ 1. Added support for userdata. +2. Added ellipsoid-ellipsoid and ellipsoid-capsule collisions using signed distance functions (SDFs). Version 3.1.4 (April 8th, 2024) ------------------------------- @@ -43,10 +44,11 @@ MJX 5. Fixed bug for sphere/capsule-convex deep penetration. 6. Fixed bug where ``mjx.Data`` produced by ``mjx.put_data`` had different treedef than ``mjx.make_data``. 7. Throw an error for margin/gap for convex mesh collisions, since they are not supported. +8. Added ellipsoid plane collisions. Simulate ^^^^^^^^ -8. Fixed bug in order of enable flag strings. Before this change, using the simulate UI to toggle the +9. Fixed bug in order of enable flag strings. Before this change, using the simulate UI to toggle the :ref:`invdiscrete` or the (now removed) ``sensornoise`` flags would actually toggle the other flag. @@ -57,10 +59,10 @@ Python bindings :align: right :width: 240px -9. Added the ``mujoco.minimize`` Python module for nonlinear least-squares, designed for System Identification (sysID). - The sysID tutorial is work in progress, but a pedagogical colab notebook with examples, including Inverse - Kinematics, is available here: |ls_colab| - |br| The video on the right shows example clips from the tutorial. +10. Added the ``mujoco.minimize`` Python module for nonlinear least-squares, designed for System Identification (sysID). + The sysID tutorial is work in progress, but a pedagogical colab notebook with examples, including Inverse + Kinematics, is available here: |ls_colab| + |br| The video on the right shows example clips from the tutorial. .. |ls_colab| image:: https://colab.research.google.com/assets/colab-badge.svg :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index f7cec040..f4da34ef 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/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, diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index c707a74f..d070ed7b 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/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 = """ + + + + + + + + + + + + """ + + 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 = """ + + + + + + + + + + + + """ + + 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 = """ diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py new file mode 100644 index 00000000..914a8e5b --- /dev/null +++ b/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 diff --git a/mjx/mujoco/mjx/_src/math.py b/mjx/mujoco/mjx/_src/math.py index 631b4425..00c9f7fd 100644 --- a/mjx/mujoco/mjx/_src/math.py +++ b/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.