mjbatch

mjbatch

(opens in a new tab) (opens in a new tab)

mjbatch is a Python library for running thousands of MuJoCo simulations in parallel on CPU.

[!NOTE] This repository is archived. Batched simulation is now part of MuJoCo itself, which this project inspired. Please use the mjbatch API (opens in a new tab) in the main repository instead.

Features include:

• C++ thread pool execution, with the GIL released;
• Live array access to simulation state and controls across the batch, with bind for MjData fields;
• Per-simulation model parameters, with expand for MjModel fields and set_const to 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.

Run with uv run examples/<file>.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.