Implementation:Google deepmind Mujoco MJWarp Benchmark

Knowledge Sources	Google_deepmind_Mujoco MJWarp Documentation
Domains	Physics_Simulation, GPU_Computing, NVIDIA_Warp
Last Updated	2026-02-15 04:00 GMT

Overview

MJWarp Benchmark provides utilities for benchmarking MuJoCo Warp simulation performance with CUDA graph capture and event tracing.

Description

This module contains the benchmark function that measures the execution time of MuJoCo Warp simulation functions. It captures the simulation step as a CUDA graph for optimal replay performance, applies control noise using an Ornstein-Uhlenbeck process with Halton quasi-random sampling, and supports optional event tracing to profile individual sub-routines. The module also tracks JIT compilation time, per-step wall-clock time, contact/constraint allocation counts, solver iteration counts, and simulation stability (NaN detection).

Usage

Used by performance testing scripts and the viewer to measure throughput of the simulation pipeline. The benchmark function returns JIT time, total runtime, trace data, and diagnostics.

Code Reference

Source Location

• Repository: Google_deepmind_Mujoco
• File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
• Lines: 1-171

Key Functions

# Primary benchmarking function
def benchmark(
    fn: Callable[[Model, Data], None],
    m: Model,
    d: Data,
    nstep: int,
    ctrls: np.ndarray | None = None,
    event_trace: bool = False,
    measure_alloc: bool = False,
    measure_solver_niter: bool = False,
    render_context: RenderContext | None = None,
) -> Tuple[float, float, dict, list, list, list, int]

# Control noise kernel (Ornstein-Uhlenbeck process)
@wp.kernel
def ctrl_noise(opt_timestep, actuator_ctrllimited, actuator_ctrlrange,
               ctrl_in, ctrl_center, step, ctrlnoisestd, ctrlnoiserate,
               ctrl_out)

Import

from mujoco.mjx.third_party.mujoco_warp._src.benchmark import benchmark

I/O Contract

Inputs

Name	Type	Required	Description
fn	Callable	Yes	Function to benchmark (e.g., mjw.step)
m	Model	Yes	Warp model
d	Data	Yes	Warp simulation data
nstep	int	Yes	Number of simulation steps to run
ctrls	np.ndarray	No	Optional control sequence to apply
event_trace	bool	No	Enable per-routine event tracing
measure_alloc	bool	No	Record contact/constraint counts
measure_solver_niter	bool	No	Record solver iteration counts

Outputs

Name	Type	Description
jit_duration	float	Time to JIT compile the function
run_duration	float	Total wall-clock time for all steps
trace	dict	Per-routine timing trace (if event_trace enabled)
nacon	list	Contact counts per step (if measure_alloc enabled)
nefc	list	Constraint counts per step (if measure_alloc enabled)
solver_niter	list	Solver iterations per step
nsuccess	int	Number of worlds that did not diverge (no NaN)

Related Pages

• Environment:Google_deepmind_Mujoco_MJX_Warp_CUDA_Environment