There are three ways to get more simulation out of a machine. One is to make each step cheaper, which starts with knowing where its time goes. Another is to spend several cores on one step. The third is to run many steps at once: many trajectories on CPU threads, or many environments as one batched computation on an accelerator. Which one pays depends on the model, and the answers are easy to measure and easy to get wrong without measuring. This lesson profiles six models with MuJoCo’s own timers, tries the engine’s thread pool and mujoco.rollout on up to 32 threads, and runs MJX, MuJoCo’s JAX implementation, against the C engine on this machine’s CPU. The Performance lab page measures throughput from 1 to 1000 environments; this lesson goes underneath it.
[!established] The tools
mjData.timerholds, per stage (mjtTimer), the accumulateddurationand the call countnumber; the Python bindings fill it without further setup.mju_threadpool(d, n)gives onemjDataa pool ofnworker threads, which the engine uses for work that splits into independent parts, such as solving separate constraint islands.mujoco.rollout.rollout(model, datas, states, controls)runs many open-loop trajectories in C, one thread perMjDatagiven. MJX (mujoco-mjx) reimplements the pipeline in JAX, sojax.vmapbatches it andjax.jitcompiles it for a CPU or an accelerator; MuJoCo Warp (mujoco-warp) does the same in NVIDIA Warp. Both are published at 3.14.0.
The benchmark below steps 1 to 1000 copies of a model in this browser, single-threaded, and reports steps per second. It is the cheapest experiment in this lesson. Run it on both models and note how much a contact-rich model costs per step compared with the cart-pole.
```lab perf {“dock”: true}
## Mathematics
### Where a step's time goes
Over 500 steps of each model:
| Model | nv | ncon | step | kinematics | inertia | collision | make | velocity | constraint | advance |
|---|---|---|---|---|---|---|---|---|---|---|
| pendulum | 1 | 0 | 1.2 µs | 15.5% | 5.4% | 16.8% | 8.2% | 10.2% | 1.8% | 3.6% |
| cart-pole | 2 | 0 | 1.3 µs | 16.0% | 6.1% | 14.6% | 7.9% | 11.8% | 1.6% | 6.0% |
| arm7 | 7 | 0 | 2.4 µs | 20.4% | 8.2% | 19.0% | 5.9% | 14.1% | 0.9% | 7.7% |
| pick_place | 27 | 12 | 13.5 µs | 7.0% | 2.4% | 16.7% | 21.8% | 7.8% | 30.8% | 8.2% |
| heap of 50 boxes | 300 | 296 | 268.2 µs | 0.7% | 0.3% | 15.0% | 20.3% | 3.6% | 58.9% | 0.6% |
| heap of 200 boxes | 1200 | 1286 | 2547.2 µs | 0.4% | 0.1% | 9.4% | 11.1% | 1.9% | 76.6% | 0.3% |
Without contact, a step costs one or two microseconds and no stage dominates; for the pendulum and the cart-pole the timed stages also account for only about two thirds of it, the rest being checks, sensors and the timers themselves (inference: the remainder is not broken down further). With contact, the step costs ten to two thousand times more and the time moves into building the constraint problem ("make") and solving it, 77% for the 200-box heap. For contact-rich scenes the levers are the ones of Lesson 20.2 (solver, iterations) and of Level 9 (fewer, simpler contacts), not faster kinematics.
### More cores on one step
`mju_threadpool` parallelizes across independent parts of one step. With 16 separate piles of 12 boxes, which form 16 constraint islands, the step time falls from 1.37 ms on one thread to 0.98 ms on four and 0.90 ms on eight. With one heap of 192 boxes of the same size (nv 1152), which forms 6 islands dominated by one large one, threads do not help: 4.56 ms on one thread, 4.53 ms on four, and 5.14 and 6.75 ms on eight and sixteen. The pool pays only when the work splits into comparable independent pieces.
### Many trajectories on many cores
`mujoco.rollout` of 1000 `pick_place` trajectories of 200 steps each:
| Threads | 1 | 2 | 4 | 8 | 16 | 32 |
|---|---|---|---|---|---|---|
| steps per second | 70,194 | 108,929 | 185,406 | 364,153 | 564,282 | 822,729 |
| speed-up | 1.00 | 1.55 | 2.64 | 5.19 | 8.04 | 11.72 |
| efficiency | 100% | 78% | 66% | 65% | 50% | 37% |
Throughput keeps rising to 32 threads, but each added thread contributes less. On this shared machine, other users' jobs compete for cores and memory bandwidth, so the efficiencies are a lower bound for a dedicated one (inference). Absolute rates vary between runs here by a factor of up to two (see the Performance lab page), while the shape of the table does not.
### MJX on a CPU
MJX accepts six of the course's eight models. It rejects `push` and `pick_place` with "(mjtGeom.mjGEOM_CYLINDER, mjtGeom.mjGEOM_BOX) collisions not implemented": their cylindrical fingertips and pusher meet boxes. Where it runs, on this CPU:
| Model | C engine, one thread | C engine, rollout on 8 threads | MJX batch 1 | MJX 64 | MJX 1024 | MJX 4096 | MJX first call |
|---|---|---|---|---|---|---|---|
| cart-pole | 669,784 | 2,864,893 | 2,385 | 134,373 | 278,589 | 356,838 | 1.5 to 2.2 s |
| cube_table | 199,212 | 928,223 | 1,496 | 3,652 | 10,007 | 18,625 | 2.4 to 3.9 s |
(steps per second). Batching is what makes MJX fast: one environment runs at a few thousand steps per second, and 4096 at once raise the throughput 150 times for the cart-pole and 12 times for `cube_table`. On a CPU it still loses to the C engine: by a factor of 8 for the cart-pole and 50 for the model with contacts, against 8 threads. Every new batch size or model also costs a compilation of a few seconds.
> [!unverified] On a GPU
> This machine has no GPU free for this measurement, so the course quotes no GPU numbers. Whether MJX or MuJoCo Warp beats the C engine on your accelerator is a measurement: run `examples/l20_3_mjx_cpu.py` with a CUDA build of JAX and compare with the rollout rows above on all your CPU cores. Also compare a few hundred steps of the two engines from the same state, since the reimplementations do not match the C engine bit for bit (inference; check before training on one and evaluating on the other).
> [!recommendation] Choosing
> Profile first; for contact-rich scenes, tune the contacts and the solver before adding hardware. Use `mujoco.rollout` threads for many open-loop trajectories (planning, evaluation from saved states); they scaled to 11.7 times on 32 threads here. Use `mju_threadpool` only for a single large scene made of several comparable islands. Consider MJX or MuJoCo Warp when you need thousands of identical environments, your model's features are supported, and your policy already runs on the same accelerator. Measure against the CPU baseline above before committing.
## Implementation
```io
INPUT: six models from a pendulum to a 200-box heap; scenes with 16 islands and with one; 1000 pick_place trajectories; MJX 3.14.0 with JAX on a CPU
PROCESS: mjData.timer per stage; mju_threadpool with 1 to 16 threads; mujoco.rollout with 1 to 32 threads; mjx.put_model on every course model, and batched mjx.step against the C engine
OUTPUT: printed tables
```python file=examples/l20_3_performance.py “"”Lesson 20.3: where a step’s time goes, and how to spend more cores on it (Level 20).
INPUT pendulum, cart-pole, arm7, pick_place (after 0.5 s), and heaps of 50 and 200 boxes in a bin; a scene of 16 separate piles of 12 boxes (16 constraint islands) and one heap of 192 boxes (one island) PROCESS (1) MuJoCo’s timers (mjData.timer) over 500 steps of each model: microseconds per stage and their share; (2) mju_threadpool on one mjData: time per step with 1 to 16 threads, for 16 islands and for one; (3) mujoco.rollout of 1000 pick_place trajectories (200 steps) on 1 to 32 threads: steps per second and the efficiency against one thread OUTPUT printed tables (absolute times depend on the machine and its load; this one is shared)
Run: python examples/l20_3_performance.py (about 2 minutes) “””
import os import time
import mujoco import mujoco.rollout import numpy as np
from mjcourse import model_path
FAST = os.environ.get(“MJC_FAST”) == “1” SIDE = 0.025 STAGES = [“POS_KINEMATICS”, “POS_INERTIA”, “POS_COLLISION”, “POS_MAKE”, “POS_PROJECT”, “VELOCITY”, “ACTUATION”, “CONSTRAINT”, “ADVANCE”]
def box(x: float, y: float, z: float, yaw: float = 0.0) -> str:
return (f’<body pos="{x:.3f} {y:.3f} {z:.3f}" euler="0 0 {yaw:.0f}">
def heap(n: int) -> mujoco.MjModel:
rng, side = np.random.default_rng(0), int(np.ceil(np.sqrt(n / 4)))
boxes = ““.join(box((i % side - side / 2) * 0.06, ((i // side) % side - side / 2) * 0.06, 0.1 + 0.06 * (i // side**2),
rng.uniform(0, 90)) for i in range(n))
wall = side * 0.035 + 0.05
walls = ““.join(f’
def piles(n_piles: int, per_pile: int) -> mujoco.MjModel:
“"”Separate piles far apart: no contact between piles, so each is its own constraint island.”””
grid = int(np.ceil(np.sqrt(n_piles)))
boxes = ““.join(box((p % grid) * 0.5 + 0.004 * (k % 2), (p // grid) * 0.5, SIDE + 2 * SIDE * k, 10 * k)
for p in range(n_piles) for k in range(per_pile))
return mujoco.MjModel.from_xml_string(f’
def library(name: str, key: int | None = None, settle: float = 0.0) -> tuple[mujoco.MjModel, mujoco.MjData]: model = mujoco.MjModel.from_xml_path(str(model_path(name))) data = mujoco.MjData(model) if key is not None: mujoco.mj_resetDataKeyframe(model, data, key) while data.time < settle: mujoco.mj_step(model, data) return model, data
def settled(model: mujoco.MjModel, seconds: float) -> tuple[mujoco.MjModel, mujoco.MjData]: data = mujoco.MjData(model) while data.time < seconds: mujoco.mj_step(model, data) return model, data
def profile(model, data, steps: int = 500) -> tuple[float, dict]: for t in data.timer: t.duration, t.number = 0.0, 0 for _ in range(steps): mujoco.mj_step(model, data) step = data.timer[mujoco.mjtTimer.mjTIMER_STEP] per_step = 1e6 * step.duration / max(step.number, 1) stages = {s: 1e6 * data.timer[getattr(mujoco.mjtTimer, f”mjTIMER_{s}”)].duration / max(step.number, 1) for s in STAGES} return per_step, stages
if name == “main”: print(f”MuJoCo {mujoco.version}; {os.cpu_count()} logical CPUs (shared machine)”) print(“(1) mjData.timer, microseconds per step and share of the step”) scenes = [(“pendulum”, *library(“pendulum”)), (“cart-pole”, *library(“cartpole”)), (“arm7”, *library(“arm7”, key=0)), (“pick_place”, *library(“pick_place”, key=0, settle=0.5)), (“heap of 50”, *settled(heap(50), 1.5))] if not FAST: scenes.append((“heap of 200”, *settled(heap(200), 1.5))) short = {“POS_KINEMATICS”: “kin”, “POS_INERTIA”: “inertia”, “POS_COLLISION”: “collision”, “POS_MAKE”: “make”, “POS_PROJECT”: “project”, “VELOCITY”: “velocity”, “ACTUATION”: “actuation”, “CONSTRAINT”: “constraint”, “ADVANCE”: “advance”} print(f” {‘model’:<12}{‘nv’:>5}{‘ncon’:>6}{‘step us’:>9} “ + ““.join(f”{short[s]:>11}” for s in STAGES)) for label, model, data in scenes: per_step, stages = profile(model, data) cells = ““.join(f”{100 * v / per_step:>10.1f}%” for v in stages.values()) print(f” {label:<12}{model.nv:>5}{data.ncon:>6}{per_step:>9.1f} {cells}”) print(“(2) mju_threadpool on one mjData: milliseconds per step”) for label, model in ((“16 piles of 12 boxes”, piles(16, 12)), (“one heap of 192 boxes”, heap(192))): _, base = settled(model, 1.0) cells = [] for threads in ((1, 4) if FAST else (1, 2, 4, 8, 16)): data = mujoco.MjData(model) data.qpos[:], data.qvel[:] = base.qpos, base.qvel if threads > 1: mujoco.mju_threadpool(data, threads) mujoco.mj_forward(model, data) t0, steps = time.perf_counter(), 200 for _ in range(steps): mujoco.mj_step(model, data) cells.append(f”{threads} threads {1000 * (time.perf_counter() - t0) / steps:.2f}”) print(f” {label:<22} (nv {model.nv}, islands {base.nisland}): “ + “, “.join(cells)) print(“(3) mujoco.rollout, 1000 pick_place trajectories of 200 steps”) model, data = library(“pick_place”, key=0) state = np.empty(mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)) mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) n, nstep = (100, 50) if FAST else (1000, 200) states = np.tile(state, (n, 1)) controls = np.tile(data.ctrl, (n, nstep, 1)) base_rate = None for threads in ((1, 8) if FAST else (1, 2, 4, 8, 16, 32)): datas = [mujoco.MjData(model) for _ in range(threads)] # one MjData per thread t0 = time.perf_counter() mujoco.rollout.rollout(model, datas, states, controls) rate = n * nstep / (time.perf_counter() - t0) base_rate = base_rate or rate print(f” {threads:>3} threads: {rate:>11,.0f} steps/s, {rate / base_rate:5.2f}x one thread, “ f”efficiency {rate / base_rate / threads:5.0%}”)
Output:
```text
MuJoCo 3.14.0; 48 logical CPUs (shared machine)
(1) mjData.timer, microseconds per step and share of the step
model nv ncon step us kin inertia collision make project velocity actuation constraint advance
pendulum 1 0 1.2 15.5% 5.4% 16.8% 8.2% 1.7% 10.2% 4.2% 1.8% 3.6%
cart-pole 2 0 1.3 16.0% 6.1% 14.6% 7.9% 1.5% 11.8% 3.7% 1.6% 6.0%
arm7 7 0 2.4 20.4% 8.2% 19.0% 5.9% 0.8% 14.1% 5.0% 0.9% 7.7%
pick_place 27 12 13.5 7.0% 2.4% 16.7% 21.8% 0.1% 7.8% 1.1% 30.8% 8.2%
heap of 50 300 296 268.2 0.7% 0.3% 15.0% 20.3% 0.0% 3.6% 0.0% 58.9% 0.6%
heap of 200 1200 1286 2547.2 0.4% 0.1% 9.4% 11.1% 0.0% 1.9% 0.0% 76.6% 0.3%
(2) mju_threadpool on one mjData: milliseconds per step
16 piles of 12 boxes (nv 1152, islands 16): 1 threads 1.37, 2 threads 1.08, 4 threads 0.98, 8 threads 0.90, 16 threads 0.90
one heap of 192 boxes (nv 1152, islands 6): 1 threads 4.56, 2 threads 4.55, 4 threads 4.53, 8 threads 5.14, 16 threads 6.75
(3) mujoco.rollout, 1000 pick_place trajectories of 200 steps
1 threads: 70,194 steps/s, 1.00x one thread, efficiency 100%
2 threads: 108,929 steps/s, 1.55x one thread, efficiency 78%
4 threads: 185,406 steps/s, 2.64x one thread, efficiency 66%
8 threads: 364,153 steps/s, 5.19x one thread, efficiency 65%
16 threads: 564,282 steps/s, 8.04x one thread, efficiency 50%
32 threads: 822,729 steps/s, 11.72x one thread, efficiency 37%
The MJX comparison needs pip install "mujoco-mjx==3.14.0" "jax[cpu]" (the course’s pip install -e ".[mjx]"). Importing MJX 3.14.0 without MuJoCo Warp installed prints two “Failed to import warp” lines; they are harmless.
```python file=examples/l20_3_mjx_cpu.py “"”Lesson 20.3: MJX, MuJoCo’s JAX reimplementation, on a CPU, against the C engine (Level 20).
INPUT the course models; cart-pole (2 dof, no contacts) and cube_table (two cubes on a table, contacts) PROCESS (1) which models mjx.put_model accepts, and why it rejects the others; (2) for cart-pole and cube_table: the C engine (a Python loop, and mujoco.rollout on 8 threads) against jax.jit(jax.vmap(mjx.step)) at batch sizes 1 to 4096 on the CPU: the first call’s time (tracing and compilation) and steady steps per second OUTPUT printed tables. This machine has no GPU; nothing here says how MJX performs on one.
Run: pip install “mujoco-mjx==3.14.0” “jax[cpu]”; python examples/l20_3_mjx_cpu.py (about 2 minutes) “””
import os import time
import jax import jax.numpy as jnp import mujoco import mujoco.rollout import numpy as np from mujoco import mjx
from mjcourse import model_path
FAST = os.environ.get(“MJC_FAST”) == “1” MODELS = (“pendulum”, “cartpole”, “double_pendulum”, “arm7”, “cube_table”, “ball_drop”, “push”, “pick_place”)
def c_engine(model: mujoco.MjModel, steps: int = 2000) -> tuple[float, float]: “"”Steps per second of the C engine: one Python loop, and rollout of 512 trajectories on 8 threads.””” data = mujoco.MjData(model) t0 = time.perf_counter() for _ in range(steps): mujoco.mj_step(model, data) loop = steps / (time.perf_counter() - t0) state = np.empty(mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)) mujoco.mj_getState(model, mujoco.MjData(model), state, mujoco.mjtState.mjSTATE_FULLPHYSICS) n, nstep = 512, 200 datas = [mujoco.MjData(model) for _ in range(8)] t0 = time.perf_counter() mujoco.rollout.rollout(model, datas, np.tile(state, (n, 1)), np.zeros((n, nstep, model.nu))) return loop, n * nstep / (time.perf_counter() - t0)
def mjx_rate(model: mujoco.MjModel, batch: int, steps: int) -> tuple[float, float]: “”“(seconds for the first call, steady steps per second) of a jitted, vmapped mjx.step.””” mx = mjx.put_model(model) dx = mjx.make_data(mx) batched = jax.vmap(lambda _: dx)(jnp.arange(batch)) step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0))) t0 = time.perf_counter() out = jax.block_until_ready(step(mx, batched)) first = time.perf_counter() - t0 t0 = time.perf_counter() for _ in range(steps): out = step(mx, out) jax.block_until_ready(out) return first, batch * steps / (time.perf_counter() - t0)
if name == “main”: print(f”MuJoCo {mujoco.version}, JAX {jax.version} on {jax.devices()[0].platform}; {os.cpu_count()} logical CPUs”) print(“(1) mjx.put_model on the course models”) for name in MODELS: model = mujoco.MjModel.from_xml_path(str(model_path(name))) try: mjx.put_model(model) print(f” {name:<16} accepted (nv {model.nv})”) except Exception as err: # MJX names the unsupported feature print(f” {name:<16} rejected: {str(err).splitlines()[0]}”) print(“(2) steps per second on this CPU”) for name in (“cartpole”, “cube_table”): model = mujoco.MjModel.from_xml_path(str(model_path(name))) loop, rollout = c_engine(model) print(f” {name}: C engine, Python loop {loop:,.0f}; rollout on 8 threads {rollout:,.0f}”) for batch in ((1, 256) if FAST else (1, 64, 1024, 4096)): first, rate = mjx_rate(model, batch, 20 if FAST else max(20, 20000 // batch)) print(f” {name}: MJX batch {batch:>5}: first call {first:5.2f} s, then {rate:>11,.0f} steps/s”)
Output:
```text
Failed to import warp: No module named 'warp'
Failed to import mujoco_warp: No module named 'warp'
MuJoCo 3.14.0, JAX 0.11.2 on cpu; 48 logical CPUs
(1) mjx.put_model on the course models
pendulum accepted (nv 1)
cartpole accepted (nv 2)
double_pendulum accepted (nv 2)
arm7 accepted (nv 7)
cube_table accepted (nv 12)
ball_drop accepted (nv 12)
push rejected: (mjtGeom.mjGEOM_CYLINDER, mjtGeom.mjGEOM_BOX) collisions not implemented.
pick_place rejected: (mjtGeom.mjGEOM_CYLINDER, mjtGeom.mjGEOM_BOX) collisions not implemented.
(2) steps per second on this CPU
cartpole: C engine, Python loop 669,784; rollout on 8 threads 2,864,893
cartpole: MJX batch 1: first call 2.17 s, then 2,385 steps/s
cartpole: MJX batch 64: first call 1.46 s, then 134,373 steps/s
cartpole: MJX batch 1024: first call 1.49 s, then 278,589 steps/s
cartpole: MJX batch 4096: first call 1.49 s, then 356,838 steps/s
cube_table: C engine, Python loop 199,212; rollout on 8 threads 928,223
cube_table: MJX batch 1: first call 2.71 s, then 1,496 steps/s
cube_table: MJX batch 64: first call 2.35 s, then 3,652 steps/s
cube_table: MJX batch 1024: first call 3.65 s, then 10,007 steps/s
cube_table: MJX batch 4096: first call 3.88 s, then 18,625 steps/s
Adding threads made it slower. The work did not split, or the pieces were too small to cover the cost of coordinating threads. Measure with 1, 2, 4 and 8 threads before choosing.
MJX rejects a model. It names the missing feature. Change the geometry (here, a box or capsule fingertip instead of a cylinder) or stay on the C engine.
The first MJX call takes seconds. That is tracing and compilation. Keep batch sizes fixed and reuse the jitted function; a new shape compiles again.
Benchmarks disagree between runs. The machine is shared or busy. Repeat, report the spread, and compare ratios measured in the same run.
Profile cube_table and the 50-box heap with opt.cone="elliptic" and with opt.solver="CG". How do the shares of “make” and “constraint” change, and what does the step cost?
If you have a GPU, run examples/l20_3_mjx_cpu.py on it with a CUDA build of JAX. Find the batch size at which MJX matches the C engine’s rollout on all your CPU cores, for the cart-pole and for cube_table, and report the compilation time. Then do the same with MuJoCo Warp.
[!research] Throughput claims in papers Simulation throughput in learning papers is often reported as steps per second at a batch size, on one accelerator, for one model. Before comparing such numbers, check the model (contacts or not), the batch size, whether compilation is counted, and the CPU baseline it is compared against, which is often a single thread rather than all cores. This lesson’s measurements show how much each choice moves the answer: a factor of 4.3 between one C thread and eight on the cart-pole, and a factor of 150 between MJX batch sizes 1 and 4096 on the same CPU.
{"id": "20.3-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "Which stage dominated the step for the 200-box heap?",
"options": ["kinematics", "collision", "constraint (the solver)", "advance"],
"answer": 2,
"explain": "<p>76.6% of 2.5 ms; building the constraint problem took another 11%.</p>"},
{"kind": "mcq", "q": "When did mju_threadpool shorten the step?",
"options": ["Always", "For 16 separate piles, which form 16 islands", "For one large heap", "Only with Newton"],
"answer": 1,
"explain": "<p>1.37 to 0.90 ms with 8 threads for the piles; no gain, then a slowdown, for the heap.</p>"},
{"kind": "mcq", "q": "Why did MJX reject pick_place?",
"options": ["Too many dofs", "Cylinder-box collisions are not implemented in MJX", "It has sensors", "It has a tendon"],
"answer": 1,
"explain": "<p>mjx.put_model raised \"(mjtGeom.mjGEOM_CYLINDER, mjtGeom.mjGEOM_BOX) collisions not implemented\".</p>"},
{"kind": "numeric", "q": "On this CPU, how many times faster was the C engine on 8 threads than MJX at batch 4096 for cube_table? (nearest integer)",
"answer": 50, "tol": 2,
"explain": "<p>928,223 against 18,625 steps per second.</p>"}
]}
Lesson 20.4 uses the engine from outside: derivatives, callbacks, plugins and the C API.