📣 MuJoCo Playground works with MuJoCo Warp! 📣 (in Beta)

[btaba]

Hello Playground users!

As of 229fb1f, MuJoCo Warp is available in MuJoCo Playground and MJX for all environments. MuJoCo Warp is still in beta in Playground (until the 1.0.0 release), but we have good parity in terms of reward behavior compared to the MuJoCo JAX implementation. We encourage users to start tinkering with Warp in your environments!

Try it out now with:

git clone https://github.com/google-deepmind/mujoco_playground.git
cd mujoco_playground
uv pip install -e . 
python learning/train_jax_ppo.py --env_name Go1JoystickFlatTerrain --impl=warp

What is MuJoCo Warp?

MuJoCo Warp is an implementation of MuJoCo implemented in Warp. Similarly to JAX, Warp allows one to write GPU code in Python, and compile just-in-time to run on a GPU. Warp is targeted specifically to NVIDIA GPUs, which enables thread divergent code (SIMT). That means we can now generate a dynamic number of contacts and constraints per physics step. Previously in JAX, we were forced to generate a fixed number of contacts/constraints per step (SIMD).

tl;dr MuJoCo Warp allows us to scale MuJoCo GPU simulation to much larger scenes!

What is MJX-Warp?

MJX-Warp is simply MJX hooked up to MuJoCo Warp. You can think of MJX-Warp as the JAX frontend API to MuJoCo Warp.

How do I migrate from MJX to MJX-Warp?

If you are familiar with MJX, you've probably generated an MJX model and data via:

mj_model = mujoco.MjModel.from_xml_path(...)
model = mjx.put_model(mj_model)
data = mjx.make_data(mj_model)

To start using MJX-Warp, simply add the impl argument:

mj_model = mujoco.MjModel.from_xml_path(...)
model = mjx.put_model(mj_model, impl='warp')
data = mjx.make_data(mj_model, impl='warp', nconmax=nconmax, njmax=njmax)

In MuJoCo Playground, impl was added to all environment configs (example), which you can simply flip to warp.

nconmax/njmax

You may have noticed that there are two additional parameters that we pass to make_data.

• nconmax defines the maximum number of contacts for all worlds combined.
• njmax defines the maximum number of constraints per world.

If you are developing a new scene, these parameters should be tuned by loading them in the viewer and increasing the values accordingly as overflows occur. Don't forget to then scale nconmax by the number of environments you'll wind up using during training!

Exposing Contacts

As of MuJoCo 3.3.5, contacts were moved from mjx.Data.contact to private mjx.Data._impl. Since JAX and Warp diverge in their implementations of contact buffers, we encourage users to read out contacts solely through contact sensors. You can find many examples in Playground (e.g. here and here). In the environment, you can read sensor values as such.

Why should I use MJX-Warp?

If you are simulating locomotion, loco-manipulation, or manipulation environments with a medium to large number mesh collisions, we encourage you to use the Warp implementation. JAX is performant for a small number of fixed collisions and primitive collisions, but Warp scales much better for medium to large scenes. As an example, many of the Playground manipulation environments have 1.5-2x higher throughput with Warp! If you are already using MJX, the switch to Warp was designed to be as seamless as possible.

Caveats

JAX support

• There may be general hiccups with the JAX<>Warp interop, but we are actively working to resolve issues. Please report any issues in MuJoCo.
• pmap support is not yet available with MJX-Warp.

Rewards

• PandPick* and AlohaSinglePeg environments have slight regressions with MuJoCo Warp. We are working to fix these issues in MuJoCo Warp.
• Heightfield environments may exhibit NaNs in training, but they work with legacy_gjk. We are working to fix this issue as well in MuJoCo Warp.

[btaba]

For posterity, here are the learning curves as of 229fb1f. The orange curve is Warp, and the green one is JAX.

DM Control

Locomotion

Manipulation

[btaba]

We train for longer for some envs for now. Remaining physics issues are being tracked here:

google-deepmind/mujoco_warp#445 (comment)

[kassasin]

How much training speed improvement compare to mjx?

[rademacher-p]

@oursland's question still seems relevant. Wall clock improvements aside, the environments must be doing something substantially different if the same agents are reporting different performance versus total env steps. I think?

[oursland]

That's my thoughts as well, but I believe the explanation is that now it is possible to run these tests, not that they're 1:1 or even more efficient. The DM Control tests are very close, but some of the Locomotion and Manipulation tests need some TLC.

[btaba]

@oursland is correct, most envs match 1:1. Some do not, but we have identified why and are fixing the physics to match what we did in JAX. @thowell @kbayes

We posted because the integration is in a sufficiently good state for general usage

[kevinzakka]

@btaba you’re a legend

[omarrayyann]

Very cool! I was hitting lots of nan errors while using JAX in a custom environment and they went away after switching to --impl=warp

[btaba]

Also posting training sps (orange is Warp, green is JAX). We see speed-up using Warp on several manipulation environments, and we expect perf to keep improving for Warp as we are heavily optimizing the library.

Manipulation

Locomotion

[kassasin]

More efficient in contact rich environment?