mjbatch-uni

Build PyPI version

mjbatch-uni 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.
  • Same-layout compiler-coherent mesh variants, with VariantPack.from_specs() and Batch.from_variant_pack().
  • Explicit topology-affine groups, with ModelAffineBatch routing global ids to incompatible layouts.
  • 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=...).
  • Opt-in split substeps with fresh position/velocity sensor views for body wrench control.
  • 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

Callbacks can also consume tracked sensors at the exact state they are controlling. With Euler, substep_sensor_copyout=(start, stop) splits each substep at mj_step1 / mj_step2, copies the declared sensor columns into a fourth callback view, and applies ctrl plus bound inputs such as xfrc_applied to that substep:

xfrc = batch.bind("xfrc_applied")

def control(k, state, ctrl, sensor):
  ctrl[:] = impedance(sensor)       # position/velocity sensors match state
  xfrc[:, body_id] = wrench(sensor)

batch.step(callback=control, substep_sensor_copyout=(sensor_start, sensor_stop))

Acceleration-stage sensors and contact forces are not available before mj_step2 and may be stale in that view. The default callback path and all outputs are unchanged unless the optional range is passed.

Model randomization and variants

expand(field) returns a live (num_sims, ...) view of a non-asset MjModel field. Rows are applied before each selected simulation runs. model_field_specs() describes the shape, dtype, writability, and recompute dependency of every field. Declare the fields being written to model_update; like mjlab event terms, mjbatch computes the strongest recompute level and runs one conservative stock-MuJoCo mj_setConst pass.

from mjbatch import RecomputeLevel

assert batch.model_field_specs()["body_mass"].recompute == RecomputeLevel.SET_CONST
with batch.model_update("body_mass", "geom_friction", ids=reset_envs):
  batch.expand("body_mass")[reset_envs, body_id] *= 1.1
  batch.expand("geom_friction")[reset_envs, :, 0] = friction_samples

For mesh-only differences, VariantPack.from_specs() independently compiles each source spec, pools and deduplicates meshes, aligns named geom slots, disables missing optional slots, and scatters compiler-derived geometry and inertia fields. Batch.from_variant_pack() applies its fixed assignment and initial recompute.

from mjbatch import Batch, VariantPack

pack = VariantPack.from_specs([tool_spec_0, tool_spec_1, tool_spec_2])
batch = Batch.from_variant_pack(pack, num_sims, np.arange(num_sims) % 3)

Truly incompatible topologies are not silently padded. Put each one in its own Batch and route fixed global assignments through ModelAffineBatch. State and field views remain on each TopologyGroup; global history is rejected.

from mjbatch import ModelAffineBatch

sharded = ModelAffineBatch([batch_a, batch_b], names=["one_joint", "two_joint"])
sharded.step(np.array([0, 3, 4]))
state_a, state_b = sharded["one_joint"].state, sharded["two_joint"].state

Reproducible cold-start, RSS, stepping, and model-field-update measurements:

uv run python benchmarks/topology_groups.py --num-sims 512 --threads 4

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.

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

Project links

Data verified by PyPI on Sep 14, 2026
Data provided by the project maintainers, verified at the time the release was uploaded to PyPI.

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for Yves_Jia from gravatar.com Yves_Jia

Credits

Author: Kevin Zakka

GitHub Statistics

Data verified by PyPI on Sep 14, 2026
The GitHub source repository was provided by the project maintainers and verified by PyPI at the time of upload. Stars, forks, and open issues/PRs are derived from that repository and have not been independently verified.

License expression

Apache-2.0
View SPDX License List

Requires

Python >=3.10

Tags

mujoco physics simulation batched reinforcement-learning

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mjbatch_uni-0.2.1.tar.gz (40.6 kB view details)

Uploaded Source

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (188.4 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (177.7 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

mjbatch_uni-0.2.1-cp313-cp313-macosx_11_0_arm64.whl (152.6 kB view details)

Uploaded CPython 3.13macOS 11.0+ ARM64

mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (188.5 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (177.8 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

mjbatch_uni-0.2.1-cp312-cp312-macosx_11_0_arm64.whl (152.5 kB view details)

Uploaded CPython 3.12macOS 11.0+ ARM64

mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (189.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (178.8 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

mjbatch_uni-0.2.1-cp311-cp311-macosx_11_0_arm64.whl (153.2 kB view details)

Uploaded CPython 3.11macOS 11.0+ ARM64

mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (189.6 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.27+ x86-64manylinux: glibc 2.28+ x86-64

mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (179.1 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

mjbatch_uni-0.2.1-cp310-cp310-macosx_11_0_arm64.whl (153.5 kB view details)

Uploaded CPython 3.10macOS 11.0+ ARM64

File details

Details for the file mjbatch_uni-0.2.1.tar.gz.

File metadata

  • Download URL: mjbatch_uni-0.2.1.tar.gz
  • Upload date:
  • Size: 40.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.13

File hashes

Hashes for mjbatch_uni-0.2.1.tar.gz
Algorithm Hash digest
SHA256 1e0b1e519ea3663896c2110ca74b18af9893fc956e962c74210ce253c6dddbd1
MD5 2aa2d8f1cc1c29b4e2d0d0fd9038b70b
BLAKE2b-256 34a2c18e9faf301582adbe717bcb6a2c4e3a37cd823cce1e7f7cd09ed6f78054

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1.tar.gz:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 1015c207179989a2b7c5324e1a684b26272dffb1ab056ada03d6ebbb0d2f4208
MD5 8a7016088ecf86161b718bfd5b613416
BLAKE2b-256 a84624ab7269337d59b2a4aa56b0a3a19b9a0705ab3493ca5cf35dbae4177263

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 7d7b86106de25c97d5b564e82a303457065a0d6fde34003eb0ed7a2e986b019a
MD5 13fd7c662d0e9181d1b8dcecda79dd05
BLAKE2b-256 1a84ffb05411354f76e4996385ce14e704ae16f0b2db19393b8339867462e25d

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp313-cp313-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 94f2733e14db8e4a4e950c58946238e1ba9c30e34d4ec66c2ebff2b294ee2446
MD5 7bf2dff3df9643ea52e73d59de82050c
BLAKE2b-256 c6b99cea39c4ba62f2ced3d27792bc6c9ec12025aa16009abc03da541934aeed

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp313-cp313-macosx_11_0_arm64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 d8d1f28f89423ebc7c5866a29eb484b5eb8c51a99d2e9a1f7dbdbccbfec351f1
MD5 e7d26a833330eccc670bb74f029167d0
BLAKE2b-256 6e8c9ef3599097716b91c38267c257e50803e9b97dc8f27de00353a83e7b11cc

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 d4dc7191fab864ca90bf4487e76c3532fa03acd8706b0d68b4982a4701aa8ea5
MD5 2bfc3ec157fa929ee5b257e5508b0f5a
BLAKE2b-256 0b9759028d7cf3013ff3338e532ce8bb24751c2f576b7788694e7d7a305a0b19

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 9b0f1fc17b4b86821c6ad12172c94cac74c66f7037f6cef9c3313537892f3aa7
MD5 cc4650c1e96ed0aa018960bd78fee225
BLAKE2b-256 fd3ad4fcc2c6dbe77266f6f8052027c02afbc119ee4549515e3b604895c241e4

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp312-cp312-macosx_11_0_arm64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 3cf085bf81e9ae6665b9fb09585116c6f6647d17c4942312b4d6bb8d43c3bbce
MD5 d4d776574e665afbc461dbf1433bbe04
BLAKE2b-256 b89b8d2bd544b6f2994b94adcd274df50525589b2fc09a5fa27fe95d9fb75027

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a2995b4b3f28f40a6d51c2c2287709732be7f5b40dab4851bed2eb7e419a78e1
MD5 6173265f291bb96f36d2acaf6d0ec44f
BLAKE2b-256 4fa9379d1eeabf2843a53685dd9b3972e0aabfbddcf55952441412e05fd2908f

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 83d6bef718a7cf900507a5c9eccf8ecd87fba54238530000142947ad85ce1653
MD5 57d4a441c39b106b582b7f2d02533563
BLAKE2b-256 5323f7cc6fbe1cfd38ae2f6187d51944f965390a486647a1a102c7d3e98c2dbd

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp311-cp311-macosx_11_0_arm64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 ce3bb2c9991d86bd63633596f420c68f9cbea254c6d289670442ddf00b4e4b60
MD5 18ba6ec6ada4c007703eee9ffd8b7ca8
BLAKE2b-256 47781fd904476bc7861dba3613b4f352db79ab433e062231aadcd56883444a73

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 856e0b761cf99f62b49cbb92ffe831347532ee89b3fceb2a49150070be87fea4
MD5 72600b5f4dc8c57f016b44ea1193ff94
BLAKE2b-256 8affb1ceda253473c118cd1ff815f9b47945c74b28a2daf27eb424eef8c1cbe3

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file mjbatch_uni-0.2.1-cp310-cp310-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.1-cp310-cp310-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 c97e68b93f586e9659323d4c795f002bd972ae6d2e3f9d09683b693c279a953f
MD5 49fe41ef75bc55aa64b2ea6f153eb03a
BLAKE2b-256 8abd0969c6cd81cb072adc9ce0423696c0f6e41626d2c15f5a5270b26d75ca73

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.1-cp310-cp310-macosx_11_0_arm64.whl:

Publisher: ci.yml on unilabsim/mjbatch_uni

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.2.1 This release

13 files

0.2.0

13 files

0.1.0

13 files

0.0.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page