mjbatch
mjbatch is a Python library for running thousands of MuJoCo simulations in parallel on CPU.
Features include:
- C++ thread pool execution, with the GIL released;
- Live array access to simulation state and controls across the batch, with
bindfor MjData fields; - Per-simulation model parameters, with
expandfor MjModel fields andset_constto recompute derived constants.
For example:
import mujoco, numpy as np
from mjbatch import Batch
model = mujoco.MjModel.from_xml_path("scene.xml")
batch = Batch(model, num_sims=4096) # threads default to every logical CPU
qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl")
batch.expand("geom_friction")[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))
for _ in range(1000):
ctrl[:] = policy(qpos) # your controller, all 4096 at once
batch.step() # step them in parallel; qpos updates in place
Examples
We showcase a range of applications built using mjbatch: RL, MPC, SysID, and hardware
co-design. Each example is a self-contained, performant implementation. For instance, the Go1
RL controller learns to walk in under a minute on a five-year-old M1 laptop.
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/cartpole_swingup.py)
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/cartpole_mpc.py)
A two-pole cart swung upright with iLQR
A cart-pole swing-up controller using predictive sampling
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/g1_flip.py)
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/go1_joystick.py)
A G1 humanoid tracking a reference backflip with receding-horizon iLQR
A Go1 quadruped joystick controller trained with PPO
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/arm_throw.py)
[](https://github.com/kevinzakka/mjbatch/blob/main/examples/rizon_inertia.py)
CEM jointly optimizes a robot arm's proportions, gears, and controls
Damped Gauss–Newton fits a Rizon arm's inertial parameters to synthetic motion data
Run with uv run examples/.py; some need uv sync --group examples. The ones that open
a window need a display; --headless runs the solver without one.
License
Apache-2.0.