Implementation:Google deepmind Mujoco MJX Warp Testspeed
| Knowledge Sources | |
|---|---|
| Domains | Benchmarking, Performance Testing, GPU Computing, JAX-Warp Bridge |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
Benchmarking script that measures the performance of MJX-Warp simulation functions across JAX, Warp, and JAX-Warp FFI backends.
Description
testspeed.py provides a comprehensive benchmarking harness for comparing simulation throughput across three execution backends: pure JAX (jax), pure Warp (warp), and the JAX-Warp FFI bridge (jax_warp). It supports benchmarking individual pipeline functions (kinematics, collision, forward, smooth, step) with configurable parameters including model file, number of environments (up to 8192 default), step count, unroll depth, maximum contacts/constraints, and CUDA graph capture modes (NONE, WARP, WARP_STAGED, WARP_STAGED_EX). The script also supports rendering benchmarks with configurable resolution and texture settings.
Usage
This script is run from the command line using Abseil flags to benchmark MJX-Warp performance. It is used by developers to profile GPU kernel throughput, compare backend implementations, and validate that the Warp bridge achieves expected speedups over pure JAX execution.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/warp/testspeed.py
- Lines: 1-350
Key Functions
# Command-line flags
_MODELFILE = flags.DEFINE_string('modelfile', 'humanoid/humanoid.xml', 'path to model')
_FUNCTION = flags.DEFINE_string('function', 'kinematics', 'function to benchmark')
_NSTEP = flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
_NENV = flags.DEFINE_integer('nenv', 8192, 'number of environments to simulate')
_UNROLL = flags.DEFINE_integer('unroll', 4, 'number of steps to unroll')
_BENCHMARK = flags.DEFINE_enum('benchmark', 'jax_warp', ['jax_warp', 'jax', 'warp'], 'Which benchmark to run.')
_GRAPH_MODE = flags.DEFINE_enum('graph_mode', 'WARP', ['NONE', 'WARP', 'WARP_STAGED', 'WARP_STAGED_EX'], 'Graph capture mode.')
Import
from mujoco.mjx.warp import testspeed
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| --modelfile | string | No | Path to MuJoCo XML model file (default: humanoid/humanoid.xml) |
| --function | string | No | Pipeline function to benchmark (kinematics, collision, forward, smooth, step) |
| --nstep | int | No | Number of simulation steps per rollout (default: 1000) |
| --nenv | int | No | Number of parallel environments (default: 8192) |
| --unroll | int | No | Number of steps to unroll in the computation graph (default: 4) |
| --benchmark | enum | No | Backend to benchmark: jax_warp, jax, or warp |
| --graph_mode | enum | No | CUDA graph capture mode: NONE, WARP, WARP_STAGED, WARP_STAGED_EX |
Outputs
| Name | Type | Description |
|---|---|---|
| Benchmark results | stdout | Timing statistics including steps per second and wall-clock time per step |
Related Pages
- Google_deepmind_Mujoco_MJX_Warp_Forward - Forward dynamics function benchmarked by this script
- Google_deepmind_Mujoco_MJX_Warp_Smooth - Smooth dynamics function benchmarked by this script
- Google_deepmind_Mujoco_MJX_Warp_Collision_Driver - Collision detection function benchmarked by this script
- Google_deepmind_Mujoco_JAX_Experimental_FFI - FFI layer whose graph modes are tested here