Add ellipsoid sdf collisions to MJX.
Co-authored-by: Baruch Tabanpour <btaba@google.com> PiperOrigin-RevId: 623415479 Change-Id: Id061c8e083357d063f1ccddf0989f2fafe415eed
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
725630c95d
commit
2b3f336bd3
+7
-5
@@ -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<option-flag-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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = """
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user