Implementation:Google deepmind Mujoco MJX Warp Forward

Knowledge Sources	Google_deepmind_Mujoco
Domains	Physics Simulation, Forward Dynamics, GPU Computing, JAX-Warp Bridge
Last Updated	2026-02-15 04:00 GMT

Overview

Auto-generated bridge module that dispatches the complete MuJoCo forward dynamics pipeline from JAX to NVIDIA Warp GPU kernels.

Description

forward.py is the largest auto-generated MJX Warp bridge file at 3617 lines. It defines the _forward_shim function decorated with @ffi.format_args_for_warp that accepts the full set of model parameters (actuator properties, body kinematics, joint definitions, constraint solver settings) and simulation data arrays needed for a complete forward dynamics step. The shim initializes placeholder Model, Data, Option, Statistic, Contact, and Constraint dataclass instances and maps the comprehensive set of actuator, body, joint, tendon, and solver parameters to Warp kernel invocations.

Usage

This module is the primary entry point for running forward dynamics in the MJX Warp backend. It is called during each simulation step to compute the full forward pass including actuator forces, constraint solving, and state integration on the GPU.

Code Reference

Source Location

Repository: Google_deepmind_Mujoco
File: mjx/mujoco/mjx/warp/forward.py
Lines: 1-3617

Key Functions

@ffi.format_args_for_warp
def _forward_shim(
    # Model
    nworld: int,
    M_rowadr: wp.array(dtype=int),
    M_rownnz: wp.array(dtype=int),
    actuator_acc0: wp.array(dtype=float),
    actuator_actadr: wp.array(dtype=int),
    actuator_actearly: wp.array(dtype=bool),
    actuator_actlimited: wp.array(dtype=bool),
    actuator_actnum: wp.array(dtype=int),
    actuator_actrange: wp.array2d(dtype=wp.vec2),
    actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f),
    actuator_biastype: wp.array(dtype=int),
    ...
    # Data arrays for state, forces, contacts
)

Import

from mujoco.mjx.warp import forward

I/O Contract

Inputs

Name	Type	Required	Description
nworld	int	Yes	Number of parallel simulation worlds
M_rowadr	wp.array(int)	Yes	Row addresses for sparse mass matrix
M_rownnz	wp.array(int)	Yes	Number of non-zeros per row in mass matrix
actuator_*	various wp.array	Yes	Full actuator model parameters (bias, gain, dynamics, limits, gear ratios)
body_*	various wp.array	Yes	Body kinematic tree parameters (positions, orientations, DOF addresses)
jnt_*	various wp.array	Yes	Joint parameters (axes, limits, types, stiffness)
block_dim	mjwp_types.BlockDim	Yes	Tiled kernel launch block dimension settings
Data arrays	various wp.array	Yes	Simulation state (qpos, qvel, qacc, ctrl, forces)

Outputs

Name	Type	Description
Updated Data fields	various wp.array	Updated simulation state after one forward dynamics step including positions, velocities, accelerations, and contact forces

Related Pages

Google_deepmind_Mujoco_MJX_Warp_Collision_Driver - Collision detection invoked during the forward pass
Google_deepmind_Mujoco_MJX_Warp_Smooth - Smooth dynamics (kinematics) invoked during the forward pass
Google_deepmind_Mujoco_MJX_Warp_FFI - FFI decorator and argument flattening used by this module
Google_deepmind_Mujoco_MJX_Warp_Types - Type definitions for all dataclass structures

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment