Implementation:Google deepmind Mujoco MJWarp Forward
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, GPU_Computing, NVIDIA_Warp |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
MJWarp Forward implements the forward dynamics pipeline including time integration, actuation, and the full simulation step on the GPU.
Description
This module provides the complete forward dynamics pipeline for MuJoCo Warp. It includes Euler, Runge-Kutta 4, and implicit integration schemes, actuator force computation with gain/bias/dynamics, forward position and velocity computations, and the full simulation step function. The module orchestrates calls to kinematics, collision, constraint, solver, and sensor subsystems to advance the simulation state by one timestep.
Usage
The primary entry point is step(m, d) which performs one complete simulation timestep. Individual stages can also be called independently: fwd_position, fwd_velocity, fwd_actuation, fwd_acceleration, euler, rungekutta4, and implicit.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
- Lines: 1-990
Key Functions
# Full simulation step
def step(m: Model, d: Data)
def step1(m: Model, d: Data)
def step2(m: Model, d: Data)
# Forward dynamics stages
def forward(m: Model, d: Data)
def fwd_position(m: Model, d: Data, factorize: bool = True)
def fwd_velocity(m: Model, d: Data)
def fwd_actuation(m: Model, d: Data)
def fwd_acceleration(m: Model, d: Data, factorize: bool = False)
# Integration schemes
def euler(m: Model, d: Data)
def rungekutta4(m: Model, d: Data)
def implicit(m: Model, d: Data)
# Internal kernels
def _next_position(...) # Integrate positions
def _next_velocity(...) # Integrate velocities
def _next_act(...) # Integrate activations
def _actuator_force(...) # Compute actuator forces
Import
from mujoco.mjx.third_party.mujoco_warp._src.forward import step
from mujoco.mjx.third_party.mujoco_warp._src.forward import forward
from mujoco.mjx.third_party.mujoco_warp._src.forward import euler
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | Model | Yes | Warp model with dynamics parameters, actuator definitions, and integrator settings |
| d | Data | Yes | Warp simulation data with current state (qpos, qvel, act, ctrl) |
Outputs
| Name | Type | Description |
|---|---|---|
| d.qpos | wp.array2d | Updated generalized positions |
| d.qvel | wp.array2d | Updated generalized velocities |
| d.qacc | wp.array2d | Computed generalized accelerations |
| d.act | wp.array2d | Updated actuator activations |
| d.time | wp.array | Updated simulation time |
| d.actuator_force | wp.array2d | Computed actuator forces |