Implementation:Google deepmind Mujoco MJX Support
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, JAX, Engine_Support |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
Engine support functions providing mass matrix construction, Jacobian computation, force application utilities, name/ID lookup, data binding helpers, contact force decoding, tendon wrapping, and muscle model computations for MJX.
Description
This is a large utility module that provides essential support functions used across the MJX pipeline. is_sparse() determines whether to use sparse or dense mass matrices (TPU vs GPU heuristic). make_m(), full_m(), and mul_m() handle mass matrix construction and multiplication respecting the kinematic tree sparsity pattern. jac() and jac_dot() compute body Jacobians and their time derivatives. BindModel and BindData classes provide named-field access to model and data arrays. contact_force() decodes constraint forces into contact-frame forces. wrap(), wrap_circle(), and wrap_inside() implement tendon wrapping geometry. muscle_gain(), muscle_bias(), and muscle_dynamics() implement the muscle actuator model.
Usage
Called by nearly all other MJX modules. Mass matrix functions are used by smooth.py and solver.py. Jacobian functions are used by constraint.py and sensor.py. Muscle functions are used by the actuator/transmission pipeline.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/_src/support.py
- Lines: 1-1162
Key Functions
def is_sparse(m: Union[mujoco.MjModel, Model]) -> bool
def make_m(m: Model, a: jax.Array, b: jax.Array, d: Optional[jax.Array] = None) -> jax.Array
def full_m(m: Model, d: Data) -> jax.Array
def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array
def jac(m: Model, d: Data, point: jax.Array, body_id: jax.Array) -> Tuple
def jac_dot(m: Model, d: Data, point: jax.Array, body_id: jax.Array) -> Tuple
def apply_ft(m: Model, d: Data, force: jax.Array, torque: jax.Array, ...) -> jax.Array
def xfrc_accumulate(m: Model, d: Data) -> jax.Array
def local_to_global(m: Model, d: Data, ...) -> Tuple
class BindModel(object)
class BindData(object)
def contact_force(m: Model, d: Data, contact_id: int, ...) -> jax.Array
def contact_force_dim(m: Model, d: Data, contact_id: int, ...) -> jax.Array
def wrap_circle(side: jax.Array, ...) -> Tuple
def wrap_inside(pos: jax.Array, ...) -> Tuple
def wrap(m: Model, d: Data) -> Tuple
def muscle_gain(m: Model, d: Data, ...) -> jax.Array
def muscle_bias(m: Model, d: Data, ...) -> jax.Array
def muscle_dynamics(m: Model, d: Data, ...) -> jax.Array
Import
from mujoco.mjx._src.support import full_m
from mujoco.mjx._src.support import mul_m
from mujoco.mjx._src.support import jac
from mujoco.mjx._src.support import is_sparse
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | mjx.Model | Yes | JAX model with body tree, joint, and actuator data |
| d | mjx.Data | Yes | JAX simulation data with positions, velocities, and constraint forces |
| vec | jax.Array | No | Vector to multiply with mass matrix (for mul_m) |
| point | jax.Array | No | Point in world frame for Jacobian computation |
| body_id | jax.Array | No | Body ID for Jacobian computation |
Outputs
| Name | Type | Description |
|---|---|---|
| result | jax.Array | Computed mass matrix, Jacobian, force vector, or muscle output depending on function |