kevinzakka/mjbatch

mjbatch – Parallel MuJoCo simulation on CPU

What it is – mjbatch is a small Python package that lets you run thousands of MuJoCo physics simulations at the same time, using a C++ thread‑pool that releases the Python GIL. It is aimed at researchers and engineers who need massive batched roll‑outs for reinforcement learning, model‑predictive control, system identification, or hardware‑co‑design.

Key ideas

• Thread‑pool execution – The heavy lifting is done in native C++ threads, so each simulation runs on a separate core without being blocked by Python’s global interpreter lock.
• Live array bindings – You can bind directly to MuJoCo data fields (e.g., qpos, ctrl) and treat them as NumPy arrays that are updated in‑place across the whole batch.
• Per‑simulation parameters – Model constants such as geom_friction can be varied per instance via expand, and derived constants can be recomputed with set_const.

Typical workflow

import mujoco, numpy as np
from mjbatch import Batch

# Load a MuJoCo model once
model = mujoco.MjModel.from_xml_path('scene.xml')

# Create a batch of 4096 parallel simulations (one thread per logical CPU by default)
batch = Batch(model, num_sims=4096)

# Bind to state and control arrays – these are NumPy views into the whole batch
qpos, ctrl = batch.bind('qpos'), batch.bind('ctrl')

# Randomise a model parameter for each simulation
batch.expand('geom_friction')[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))

for _ in range(1000):
    ctrl[:] = policy(qpos)   # your controller runs on all sims at once
    batch.step()             # step every simulation in parallel; qpos updates in place

The loop above looks just like a single‑simulation script, but under the hood 4 k environments are stepping together.

What you can do with it

• Reinforcement learning – Generate massive replay buffers or evaluate policies quickly (the README shows a Go1 quadruped learning to walk in under a minute on an older laptop).
• Model‑predictive control / iLQR – Run receding‑horizon optimisers across many initial conditions.
• System identification – Fit inertial parameters by simulating many perturbed models and comparing to synthetic data.
• Co‑design – Jointly optimise robot geometry, gear ratios, and control policies using evolutionary or gradient‑based methods.

Examples shipped with the repo

Example	Goal	Technique
cartpole_swingup.py	Swing‑up a double‑pole cart	iLQR
cartpole_mpc.py	Swing‑up with predictive sampling	MPC
g1_flip.py	Back‑flip a humanoid robot	Receding‑horizon iLQR
go1_joystick.py	Tele‑operate a quadruped	PPO‑trained policy
arm_throw.py	Co‑design arm proportions & control	CEM optimisation
rizon_inertia.py	Identify inertial parameters	Damped Gauss‑Newton

Run any example with uv run examples/<file>.py; some need a display (--headless works without one). The repository also provides a CI badge and a PyPI release, indicating that the package is maintained and installable via pip install mjbatch.

Installation

pip install mjbatch   # pulls the compiled extension and pure‑Python wrapper

The library depends on the official MuJoCo Python bindings, so you need a MuJoCo license (free for academic use) and the mujoco package installed.

License – Apache‑2.0, allowing commercial and open‑source use.

---

In short, mjbatch is a practical tool for anyone who wants to scale up MuJoCo‑based experiments without needing a GPU cluster. It abstracts away the threading details while giving direct NumPy‑style access to the full batch of simulation states.