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=...).
  • 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

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 13, 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:
Newer release available (0.2.1)

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 13, 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.0.tar.gz (38.2 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.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (184.4 kB view details)

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

mjbatch_uni-0.2.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (173.1 kB view details)

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

mjbatch_uni-0.2.0-cp313-cp313-macosx_11_0_arm64.whl (149.0 kB view details)

Uploaded CPython 3.13macOS 11.0+ ARM64

mjbatch_uni-0.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (184.5 kB view details)

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

mjbatch_uni-0.2.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (173.2 kB view details)

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

mjbatch_uni-0.2.0-cp312-cp312-macosx_11_0_arm64.whl (149.0 kB view details)

Uploaded CPython 3.12macOS 11.0+ ARM64

mjbatch_uni-0.2.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (184.9 kB view details)

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

mjbatch_uni-0.2.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (174.4 kB view details)

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

mjbatch_uni-0.2.0-cp311-cp311-macosx_11_0_arm64.whl (149.8 kB view details)

Uploaded CPython 3.11macOS 11.0+ ARM64

mjbatch_uni-0.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (185.1 kB view details)

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

mjbatch_uni-0.2.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (174.7 kB view details)

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

mjbatch_uni-0.2.0-cp310-cp310-macosx_11_0_arm64.whl (150.2 kB view details)

Uploaded CPython 3.10macOS 11.0+ ARM64

File details

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

File metadata

  • Download URL: mjbatch_uni-0.2.0.tar.gz
  • Upload date:
  • Size: 38.2 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.0.tar.gz
Algorithm Hash digest
SHA256 0e883d307f33636daaa8ed92f729c67c2ca67a5ae54b12114ef40baebf052006
MD5 b92344c127027dc0518878a58d3c238e
BLAKE2b-256 fad70353b4cff9aa3dbc87ba194514f27b11a81155df684c2259adde8b3f6d7b

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0.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.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 f9f40fcae35da5d8a29448bdfa62fac70c9b69440e10056f3d1e3061b400c9c9
MD5 745b60f206f81574882737f3e80c8463
BLAKE2b-256 1848e8d64f0e00c56e36b5ed9351e7e62129c83bacff19bbe55d53844aefa89e

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 6a8a4be2acfd29630f133fa65f742c16b29799a3fd3069133fb54434cf29a7ea
MD5 6435fe4b4c8af0d42bac4aca5d0242db
BLAKE2b-256 a0f28a2871c06b987b349e3e18689ce46ca952538be2ad69d34a2274e74db98b

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp313-cp313-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 5431238c5ca50be35c7efec5b3e342cd8ec450a0855a43d27aec9a277c485be6
MD5 6bd3d18ceec3e8f17122f0f884e053b5
BLAKE2b-256 6ecfb51af8bcd2b50b0c5b3497fac1bb673f3e28bccf7a2808bbb2d55a9ee88f

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 b9b9a564773698faf37ad89560cb071b50cbef4bc81743cca39602dc2e4aa512
MD5 d912f9bab16d33f56d7af25eee07d135
BLAKE2b-256 f61365047e2fd0863133f140ec7f195ec915c88c3502f9ac785ce65ded58951f

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 090ec1d0aa4e6abdc6cb761fe447e0cd31fba62888ed2a6897b2557ff7f69543
MD5 2f17f395bcebd062248fae5037e02145
BLAKE2b-256 56fe09a9cac6b6f91b62c23c943e504ad3d2947d868a7c512bca66d59fdd88d9

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 2ff0c70adf00e19a3878ba4cd7d2549137a60ebe60a6709a92c54ad00fe41cd4
MD5 c86902906a63c66ec7db1ad2428e76cb
BLAKE2b-256 36a415f256c6d3a33811bc4bdbce1649ea2d66bdfcb720f0396a4e83b6cd92ea

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 690db3ca8564c154341306c4167123dff1125c4cd334f37b7ed6e5e321b3540f
MD5 f85cfe8642c07a14e7dc9e0e8d7857f8
BLAKE2b-256 ee6be066bad7afb095051f1dbd9aa500ca62977069ef880109757e9ff12376d5

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 84ca4865dbcfa932cbd4e905d8c15b41fe0a8b8216761a06f9f3668bf0d92256
MD5 2675b1b9881c1793b0ba50fc0712f85a
BLAKE2b-256 f75526240fc42e3fe4f8ffc6ea6605ef72de130c7196bc69bd158527edb1f6de

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 f72e7bccfbd8e0d4f744bfc223265665f32b4e32b3a989f7f03123f811729a18
MD5 3d34f73487d5be0169551213d8647cf8
BLAKE2b-256 fdd95dba00afd7f6784740144de079cfde14ef1731533a09cbafc5bec7f945cc

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 ca1da6c406959c7b3c09f5d30e20edf851f091c1b0f139cb57ab1eb69f8e40bc
MD5 303af88d41aea1cccb89cf7671425ff5
BLAKE2b-256 1984be3d64c4f4e40494e83948ea0a08eb00319da722d01d6bbfe6c770e97588

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3b4175221ccd6c50f98456f81e5f52a7709f280275e160c1341d8a64593dbf37
MD5 b994333f80cc5737905aa52269324c75
BLAKE2b-256 c577f3adc58f3acb7e5369865344e3fd2e165c29d6a4731dcdbb991c89868e2a

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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.0-cp310-cp310-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for mjbatch_uni-0.2.0-cp310-cp310-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 1bac2a8fc41cf61f80fbba5dc4c9690bb78ebc51faafe71f6205b60ebd9eda1f
MD5 d68697eab9b6418e93bed767781a4c3f
BLAKE2b-256 bd9c29ded200f34bf03dbca7302c99832bcff15c2557cd11c5208612f146e4aa

See more details on using hashes here.

Provenance

The following attestation bundles were made for mjbatch_uni-0.2.0-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

0.2.1

13 files

This release

0.2.0 This release

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