Kevin Zakka's mjbatch Runs 4,096 MuJoCo Worlds on CPU Threads at Once
Kevin Zakka's new library steps thousands of MuJoCo simulations in parallel on a single CPU, unlocking RL and MPC workloads without a GPU.
- Kevin Zakka released mjbatch, a Python library for stepping thousands of MuJoCo simulations in parallel on CPU.
- Uses a C++ thread pool with the GIL released, exposing live numpy views of state and controls.
- Supports per-simulation model parameters for domain randomization, with
set_constrecomputing derived constants. - Trains a Go1 quadruped joystick policy with PPO in under a minute on an M1 MacBook using 1,024 envs.
- Ships examples covering iLQR, predictive-sampling MPC, RL, hardware co-design, and inertial system identification.
- Apache-2.0 licensed, installs via
uv, positioned as a CPU-side complement to GPU stacks like MJX.
mjbatch puts thousands of MuJoCo worlds on CPU threads
Kevin Zakka has released mjbatch, an Apache-2.0 Python library that batches independent MuJoCo simulations across a C++ CPU thread pool. It accepts a standard MjModel, exposes batched state and control arrays as NumPy views, and advances thousands of worlds with one synchronous call. Robotics developers can run reinforcement learning, model predictive control, system identification, and design searches on multicore CPUs while retaining MuJoCo’s native engine.
Many robotics algorithms evaluate the same model from hundreds or thousands of states. A Python loop adds interpreter overhead to every simulation step, while accelerator stacks such as MJX move the workload into JAX and typically involve device arrays, tracing, and compilation. mjbatch keeps the computation on the host CPU and moves the loop into native code.
One call replaces the Python loop
A Batch starts with an MjModel and a simulation count. The library creates the corresponding simulation state, dispatches each step through its thread pool, and releases Python’s global interpreter lock during native execution. The step() call remains synchronous, so the caller resumes after every world has finished; releasing the lock removes interpreter serialization from the native work and allows other Python threads to run.
| API | Purpose | Developer detail |
|---|---|---|
Batch(model, num_sims) |
Creates the simulation batch | Uses a regular MuJoCo model as the template |
bind("qpos") |
Exposes a batched data field | Returns a live NumPy view that can be read or updated in place |
expand("geom_friction") |
Exposes a model field per simulation | Supports domain randomization and physical parameter searches |
set_const() |
Recomputes derived model constants | Required after changing parameters that MuJoCo uses to derive cached values |
step() |
Advances every world | Dispatches the batch through the C++ thread pool |
The live views avoid copying state through a separate Python container. Assigning to ctrl updates the controls consumed by the next step, while reading
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
