Heuristic:Google deepmind Mujoco MJX Feature Compatibility
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, Debugging |
| Last Updated | 2026-02-15 05:00 GMT |
Overview
MJX supports a subset of MuJoCo features: no PGS solver, no Implicit integrator, no flex bodies, limited sensor types, and no margin/gap for mesh collisions. Check compatibility before porting models.
Description
MJX re-implements MuJoCo's physics pipeline in JAX/Warp but does not support every feature available in the native C API. Attempting to use unsupported features raises `NotImplementedError` at model loading time (`put_model`). Understanding these limitations upfront prevents debugging dead ends when porting models from native MuJoCo to MJX.
Key unsupported features include the PGS solver, the implicit integrator (only implicitfast is available), flex bodies, several sensor types, and collision margin/gap for mesh geometries. The Warp backend has additional constraints including mandatory CUDA GPU and noslip solver exclusion.
Usage
Use this heuristic when porting a MuJoCo model to MJX, debugging NotImplementedError exceptions, or choosing between native C API and MJX for a specific model.
The Insight (Rule of Thumb)
- Solver: Only CG and Newton are supported. PGS raises `NotImplementedError`.
- Integrator: Only Euler, RK4, and ImplicitFast are supported. Implicit raises `NotImplementedError`.
- ImplicitFast + Fluid: Not compatible. Fluid drag (density > 0 or viscosity > 0) with implicitfast raises error.
- Flex bodies: Not supported in JAX backend. Raises `NotImplementedError`.
- Mesh collisions: margin/gap must be zero. Non-zero values raise `NotImplementedError`.
- Sensors: Contact sensors with site/body/subtree matching or netforce reduction are not supported in JAX.
- EnableBits: OVERRIDE, ENERGY, FWDINV, ISLAND are unsupported.
- Warp-specific: noslip solver not implemented; flex collisions not implemented.
Reasoning
MJX is designed around JAX's functional programming model with static shapes and no side effects. Features that require:
- Dynamic shapes (flex bodies, variable contact counts) conflict with JAX's compilation model.
- Sequential algorithms (PGS) are not parallelizable for GPU execution.
- Complex state tracking (energy bookkeeping, island detection) add overhead without benefit for typical RL workloads.
The design choice prioritizes throughput and differentiability for the most common robotics use cases (articulated rigid bodies with contact) over completeness.
Code Evidence
Unsupported solver from `mjx/mujoco/mjx/_src/types.py:214-217`:
class SolverType(enum.IntEnum):
# unsupported: PGS
CG = mujoco.mjtSolver.mjSOL_CG
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
Unsupported integrator from `mjx/mujoco/mjx/_src/types.py:131-134`:
class IntegratorType(enum.IntEnum):
EULER = mujoco.mjtIntegrator.mjINT_EULER
RK4 = mujoco.mjtIntegrator.mjINT_RK4
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
# unsupported: IMPLICIT
Flex body check from `mjx/mujoco/mjx/_src/io.py:302-303`:
if m.nflex:
raise NotImplementedError('Flex not implemented for JAX backend.')
ImplicitFast + fluid drag check from `mjx/mujoco/mjx/_src/io.py:263-265`:
if implicitfast and has_fluid_params:
raise NotImplementedError('implicitfast not implemented for fluid drag.')
Warp noslip check from `mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py:128-129`:
if mjm.opt.noslip_iterations > 0:
raise NotImplementedError(f"noslip solver not implemented.")
Unsupported enable bits from `mjx/mujoco/mjx/_src/types.py:93-94`:
INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE
# unsupported: OVERRIDE, ENERGY, FWDINV, ISLAND
Deprecated direct attribute access from `mjx/mujoco/mjx/_src/types.py:1046-1050`:
warnings.warn(
f'Accessing `{name}` directly from `Model` is deprecated. '
f'Access it via `model._impl.{name}` instead.',
DeprecationWarning,
stacklevel=2,
)