Implementation:Google deepmind Mujoco MJX Solver

Knowledge Sources	Google_deepmind_Mujoco MJX Documentation
Domains	Physics_Simulation, JAX, Optimization
Last Updated	2026-02-15 04:00 GMT

Overview

Constraint solver implementation for MJX that resolves contact, equality, limit, and friction constraints using projected gradient descent with linesearch.

Description

This module implements the MuJoCo constraint solver in JAX. The Context PyTreeNode holds all per-iteration solver state including accelerations, constraint forces, gradients, linesearch vectors, and friction cone data. Context.create() initializes solver state from model and data. The solver loop alternates between _update_constraint() (computing constraint violations, costs, and forces) and _update_gradient() (computing the master cost gradient). A linesearch (_linesearch()) finds the optimal step size along the search direction. The solver supports both pyramidal and elliptic friction cones and uses _while_loop_scan() for a JIT-compatible fixed-iteration loop.

Usage

Called during mjx.step() after constraint construction. solve() is the main entry point that iterates the solver and writes back qacc, qfrc_constraint, and efc_force into Data.

Code Reference

Source Location

• Repository: Google_deepmind_Mujoco
• File: mjx/mujoco/mjx/_src/solver.py
• Lines: 1-610

Key Functions

class Context(PyTreeNode):
    """Data updated during each solver iteration."""
    @classmethod
    def create(cls, m: Model, d: Data, grad: bool = True) -> 'Context'

class _LSPoint(PyTreeNode):
    """Linesearch point."""

class _LSContext(PyTreeNode):
    """Linesearch context."""

def _while_loop_scan(cond_fun, body_fun, init_val, max_iter)
def _update_constraint(m: Model, d: Data, ctx: Context) -> Context
def _update_gradient(m: Model, d: Data, ctx: Context) -> Context
def _linesearch(m: Model, d: Data, ctx: Context) -> Context
def solve(m: Model, d: Data) -> Data

Import

from mujoco.mjx._src.solver import solve
from mujoco.mjx._src.solver import Context

I/O Contract

Inputs

Name	Type	Required	Description
m	mjx.Model	Yes	JAX model with solver options (iterations, tolerance, cone type)
d	mjx.Data	Yes	JAX simulation data with constraint matrices (efc_J, efc_D, efc_aref)

Outputs

Name	Type	Description
d	mjx.Data	Updated data with solved qacc, qfrc_constraint, efc_force, and solver_niter

Related Pages

• Environment:Google_deepmind_Mujoco_MJX_JAX_Environment