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 bind for MjData fields;
• Per-simulation model parameters, with expand for MjModel fields and set_const to recompute derived constants.
• Batched queries beyond stepping: site Jacobians with jac_site and heightfield sampling with sample_hfield (world/yaw grid alignment); query ops return caller-allocated results without refreshing the bound views.
• Per-substep control from Python with step(..., callback=...).
• Optional per-worker CPU pinning on Linux, with cpu_ids binding pool worker i to cpu_ids[i].

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

sample_hfield samples a shared heightfield at XY offsets around a body origin, returning the world z of the sampled surface. alignment="yaw" rotates the sampling grid by the frame body's yaw about world z, so a per-sim geom_quat rotates each simulation's scan pattern.

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.

A two-pole cart swung upright with iLQR	A cart-pole swing-up controller using predictive sampling
A G1 humanoid tracking a reference backflip with receding-horizon iLQR	A Go1 quadruped joystick controller trained with PPO
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/<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.