Implementation:Google deepmind Mujoco MJWarp Solver
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, GPU_Computing, NVIDIA_Warp |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
MJWarp Solver implements the constraint solver using conjugate gradient optimization with both parallel and iterative linesearch on the GPU.
Description
This module provides the GPU-accelerated constraint solver for MuJoCo Warp. It implements the conjugate gradient (CG) and Newton solver algorithms with support for both pyramidal and elliptic friction cones. The module includes SolverContext and InverseContext dataclasses for workspace management, gradient computation with sparse and dense mass matrix support via blocked Cholesky factorization, and both parallel (fused log-scale) and iterative linesearch strategies. The solver iteratively minimizes the constraint cost function to find contact forces that satisfy all constraints.
Usage
Called during the forward dynamics pipeline via solve(m, d) after constraints have been constructed. The solver iterates until convergence or maximum iterations, updating qacc and qfrc_constraint in the Data object.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
- Lines: 1-2812
Key Functions
# Workspace context classes
@dataclasses.dataclass
class InverseContext:
Jaref: wp.array2d; search_dot: wp.array; gauss: wp.array; ...
@dataclasses.dataclass
class SolverContext:
Jaref: wp.array2d; grad: wp.array2d; search: wp.array2d; ...
# Context creation
def create_inverse_context(m: Model, d: Data) -> InverseContext
def create_solver_context(m: Model, d: Data) -> SolverContext
# Primary entry point
def solve(m: Model, d: Data)
# Solver iteration
def _solver_iteration(m, d, ctx)
def init_context(m, d, ctx, grad=True)
# Linesearch
def _linesearch(m, d, ctx, cost)
def _linesearch_parallel(m, d, ctx, cost)
def _linesearch_iterative(m, d, ctx, fuse_jv)
# Gradient computation
def _update_gradient(m, d, ctx)
def _update_constraint(m, d, ctx)
Import
from mujoco.mjx.third_party.mujoco_warp._src.solver import solve
from mujoco.mjx.third_party.mujoco_warp._src.solver import create_solver_context
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | Model | Yes | Warp model with solver type, iterations, and friction cone settings |
| d | Data | Yes | Warp simulation data with constraint Jacobians and parameters |
Outputs
| Name | Type | Description |
|---|---|---|
| d.qacc | wp.array2d | Generalized accelerations satisfying constraints |
| d.qfrc_constraint | wp.array2d | Constraint forces in generalized coordinates |
| d.efc.force | wp.array2d | Constraint forces in constraint space |
| d.solver_niter | wp.array | Number of solver iterations performed |