Brax ❤️ MuJoCo = MJX

[erikfrey]

Hi all,

🚨 Major announcement 🚨

Brax has joined forces with the MuJoCo authors, to develop a jax version of the MuJoCo simulator called MuJoCo XLA (MJX). MJX is a strictly more comprehensive and performant version of the Brax generalized pipeline. We've just released MuJoCo XLA in MuJoCo 3.0.

Please check it out!

A bit of history

MJX sprang from a dawning realization that the algorithms employed in our generalized physics pipeline were more scalable than we had originally anticipated. Fun fact: we originally developed generalized as a baseline for a paper, and we did not expect to actively develop it past publication. We were quite surprised when our metrics for generalized suggested that, with careful algorithm design, accelerators could significantly speed up even complex/sophisticated physics algorithms such as those used in MuJoCo!

We still don't think Brax generalized or MJX are suitable for all rigid body simulation workloads, and we expound upon this a bit in 🔪 MJX - The Sharp Bits 🔪, but we are convinced at this point that it's useful for a significant core of real-world robotics tasks, so we're happy to share our tools with the outside world.

It's been a real joy to work with the MuJoCo developers and we've learned a lot! MJX features many improvements over the generalized pipeline - it not only supports more of the MuJoCo feature set than generalized (see this feature parity table) but it scales better across more scenes than Brax generalized. At this point we consider it more or less a strictly better implementation than generalized.

What this means for Brax

Going forward we plan to support both Brax and MJX, and see the world as follows:

• Brax will continue to remain a home for useful reinforcement learning algorithms and reference environments
• Brax also has a nifty implementation of position based dynamics that scales better than MJX in many circumstances, but is more challenging to transfer to real-world robots
• We will work to port MJX into Brax as another physics pipeline
• MJX will be the new generalized and will continue to be developed, while generalized will soon be deprecated

In the coming months, we will make Brax import MJX directly. We will make it easy to select MJX as a pipeline, just like spring, pbd, and generalized. Once that's done, we will mark generalized as deprecated and encourage our users to use the MJX pipeline instead.

Thanks to all of you for your helpful PRs, bug reports, and general comments - it makes this project that much more fun to work on! Onwards and upwards.

• Brax authors

[lianghongzhuo]

What's the relationship between Brax and MuJoCo XLA (MJX)?

[btaba]

MJX is the next iteration of the Brax generalized pipeline. Brax has 2 other physics pipelines spring and positional. We'll port MJX in to Brax in the near future as another physics pipeline and deprecate generalized

[lianghongzhuo]

Thanks for the clarification! The future question is, as a user who wants to do robot simulation, should I choose Brax, or use MuJoCo directly? What are the key benefits that make users choose to use Brax?

[cemlyn007]

Well done guys 🔥

[btaba]

Thanks @cemlyn007 !

[lianghongzhuo]

Thanks for the clarification! The future question is, as a user who wants to do robot simulation, should I choose Brax, or use MuJoCo directly? What are the key benefits that make users choose to use Brax?

[erikfrey]

MJX is a physics engine, Brax is a reinforcement learning suite!

In the near future, Brax will import MJX. Use Brax if you are an RL researcher and you want to develop novel RL algorithms or use predefined Gym benchmark environments.

Use MJX if you're into physics simulation specifically and want to directly hack on low level physics.

[btaba]

Hi @lianghongzhuo , here are some additional suggestions: If you want to do massively parallel simulation for RL or similar, we think you should use Brax. If you require very speedy simulation but not necessarily good sim2real transfer, we think you should use the positional backend in Brax. If you want to shoot for sim2real transfer, you should use either generalized or Brax + MJX.

Here's an example doing sim2real on a quadruped robot in Brax:
https://github.com/google/brax/tree/main/brax/experimental/barkour_v0

Here's a tutorial doing a similar thing with Brax+MJX:
https://github.com/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb

[lianghongzhuo]

Thank you!

[mimoralea]

How is MJX different from MuJoCo? Is MJX using Brax environments with an optimized MuJoCo backend?

I have a custom MuJoCo model from https://github.com/google-deepmind/mujoco_menagerie/. Can one import and use it in Brax with the MJX backend?

Finally, is dm_control then irrelevant with this new release?

[btaba]

Hi @mimoralea , some answers below:

"Is MJX different from MuJoCo?"
MJX is a jax implementation of MuJoCo, see https://mujoco.readthedocs.io/en/stable/mjx.html

"Is MJX is using Brax environments with an optimized MuJoCo backend?"
MJX isn't ported as a backend into Brax yet. MJX lives here https://github.com/google-deepmind/mujoco

"Can one import and use it in Brax with the MJX backend?"
MJX isn't ported into Brax yet, as mentioned above. You can import the mjcf into MuJoCo/MJX or Brax; take note of feature parity and performance tuning for MJX. See also this colab tutorial which uses MJX for the physics and Brax for the RL training.

"is dm_control then irrelevant with this new release?"
There isn't a dm_control impl using MJX/Brax as far as I know

[mimoralea]

Thank you, @btaba.

I understand a lot more after reading the docs and re-reading the announcement. MJX is not meant to replace MuJoCo for RL training, although it can help in many cases. Also, Brax (and Brax+MJX, more specifically) is not meant to replace dm_control, but they are complementary in some ways. I can see how training in Brax+MJX, using low-fidelity or maybe a few models, can cut down training time on tasks that can later be transferred and fine-tuned on more elaborate scenarios using dm_control+MuJoCo.

I'll run the Colab notebooks you reference, including the others on your README.md. Thanks for taking the time--this is exciting!