Implementation:Google deepmind Mujoco MJWarp Support
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, GPU_Computing, NVIDIA_Warp |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
MJWarp Support provides utility functions for mass matrix multiplication, Jacobian computation, external force accumulation, contact force extraction, and state get/set operations.
Description
This module contains essential support functions used across the MuJoCo Warp pipeline. It implements sparse and dense mass matrix-vector multiplication (mul_m), body Jacobian computation (jac), external force accumulation into generalized forces (xfrc_accumulate), contact force extraction from constraint forces (contact_force), and state serialization/deserialization (get_state/set_state). The module supports both sparse and dense representations and includes cache-decorated kernel builders for performance.
Usage
Used throughout the simulation pipeline for linear algebra operations on the mass matrix, Jacobian computations for sensors and constraints, and state management for integrators and the solver.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
- Lines: 1-824
Key Functions
# Mass matrix multiplication
def mul_m(m: Model, d: Data, res, vec, M=None, ...)
def mul_m_sparse(check_skip: bool) # Cached kernel builder
def mul_m_dense(nv: int, check_skip: bool) # Cached kernel builder
# External forces
def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d)
def apply_ft(m: Model, d: Data, ft, qfrc, flg_add)
# Jacobian computation
def jac(m: Model, d: Data, jacp, jacr, point, bodyid)
def jac_dof(...)
def jac_dot_dof(...)
# Contact forces
def contact_force(m: Model, d: Data, ...)
# State management
def get_state(m: Model, d: Data, state, sig, active=None)
def set_state(m: Model, d: Data, state, sig, active=None)
# Spatial force transform
def transform_force(force, torque, offset) -> wp.spatial_vector
Import
from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
from mujoco.mjx.third_party.mujoco_warp._src.support import jac
from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force
from mujoco.mjx.third_party.mujoco_warp._src.support import get_state
from mujoco.mjx.third_party.mujoco_warp._src.support import set_state
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | Model | Yes | Warp model with kinematic tree and mass matrix structure |
| d | Data | Yes | Warp simulation data with mass matrix and body transforms |
| vec | wp.array2d | No | Input vector for matrix-vector multiplication |
| point | wp.vec3 | No | Point for Jacobian computation |
| bodyid | int | No | Body ID for Jacobian computation |
Outputs
| Name | Type | Description |
|---|---|---|
| res | wp.array2d | Result of mass matrix multiplication or Jacobian computation |
| qfrc | wp.array2d | Generalized forces from external forces or contact |
| state | wp.array2d | Serialized or deserialized simulation state |