Implementation:Google deepmind Mujoco MJX Collision SDF
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, JAX, Collision_Detection |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
Collision detection functions for shapes represented as signed distance functions (SDFs), enabling implicit and exact geometry definitions for ellipsoids and cylinders.
Description
This module uses signed distance functions to compute collisions between geometry pairs that lack simple closed-form intersection formulas (e.g., ellipsoid-ellipsoid, cylinder-cylinder, capsule-cylinder). An SDF returns the shortest distance from a point to a surface, allowing geometry to be defined implicitly. Contact points are found via gradient descent optimization on the SDF intersection/clearance objective using _gradient_descent() and _optim(). Individual SDF primitives (_sphere, _capsule, _ellipsoid, _cylinder) are composed using _intersect() and _clearance() combinators. A custom JVP rule is provided for cylinder SDF to handle gradient discontinuities.
Usage
Registered in the collision function table by collision_driver.py for geom type pairs involving ellipsoids and cylinders. Called during the contact detection phase of mjx.step().
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/_src/collision_sdf.py
- Lines: 1-275
Key Functions
def collider(ncon: int)
def _sphere(pos: jax.Array, size: jax.Array) -> jax.Array
def _capsule(pos: jax.Array, size: jax.Array) -> jax.Array
def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array
def _cylinder(pos: jax.Array, size: jax.Array) -> jax.Array
def _intersect(d1: SDFFn, d2: SDFFn) -> SDFFn
def _clearance(d1: SDFFn, d2: SDFFn) -> SDFFn
class GradientState(PyTreeNode)
def _gradient_step(objective: SDFFn, state: GradientState) -> GradientState
def _gradient_descent(objective: SDFFn, x0: jax.Array, ...) -> jax.Array
def _optim(sdf1: SDFFn, sdf2: SDFFn, ...) -> Collision
def sphere_ellipsoid(s: GeomInfo, e: GeomInfo) -> Collision
def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision
def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision
def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision
def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision
def ellipsoid_cylinder(e: GeomInfo, c: GeomInfo) -> Collision
def cylinder_cylinder(c1: GeomInfo, c2: GeomInfo) -> Collision
Import
from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
from mujoco.mjx._src.collision_sdf import cylinder_cylinder
from mujoco.mjx._src.collision_sdf import capsule_cylinder
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | mjx.Model | Yes | JAX model (used in collider wrapper) |
| d | mjx.Data | Yes | JAX simulation data with geom positions/orientations |
| s / c / e / ca / cy | GeomInfo | Yes | Geom shape info (position, rotation matrix, size) |
| geom | jax.Array | Yes | Array of geom index pairs |
Outputs
| Name | Type | Description |
|---|---|---|
| dist | jax.Array | Penetration distance (negative means overlap) |
| pos | jax.Array | Contact point position |
| frame | jax.Array | Contact frame (3x3); normal in row 0 |