Implementation:Google deepmind Mujoco MJX Collision Driver
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, JAX, Collision_Detection |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
Orchestrates collision checking for all geom pairs in a MuJoCo model by building and dispatching a collision function table keyed on geom type, mesh ID, and condim.
Description
The collision driver organizes geom pairs into groups based on their collision function type (e.g., sphere-sphere, sphere-capsule) and runs the appropriate collision function for each group. It builds a function table that maps FunctionKey entries (comprising geom types, mesh data IDs, and condim) to batched collision calls. This grouping ensures static JAX array shapes for contacts and Jacobians, which is essential for JIT compilation. The module imports collision functions from collision_primitive, collision_convex, and collision_sdf sub-modules.
Usage
Called during the mjx.step() pipeline to produce contact data. The main entry point collision() takes a model and data, runs broadphase geom pair filtering, and returns updated data with populated contact arrays.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/_src/collision_driver.py
- Lines: 1-464
Key Functions
def has_collision_fn(t1: GeomType, t2: GeomType) -> bool
def geom_pairs(m: Model, ...) -> Iterator[Tuple[FunctionKey, ...]]
def make_condim(m: Model, ...) -> jax.Array
def collision(m: Model, d: Data) -> Data
Import
from mujoco.mjx._src.collision_driver import collision
from mujoco.mjx._src.collision_driver import has_collision_fn
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | mjx.Model | Yes | JAX model containing geom definitions, sizes, and types |
| d | mjx.Data | Yes | JAX simulation data with current geom positions and orientations |
Outputs
| Name | Type | Description |
|---|---|---|
| d | mjx.Data | Updated simulation data with populated Contact fields (dist, pos, frame, geom pairs, etc.) |