MuJoCo 3

[saran-t]

MuJoCo 3

Bringing accelerator support to the MuJoCo ecosystem.

Two years ago, DeepMind took over the development of MuJoCo and made it freely available. Since then, the MuJoCo team made numerous versioned releases accompanied by a long list of new features, improved performance and memory footprint, and made it much easier for new users to get started. We also created the MuJoCo Menagerie, a curated collection of high-quality, ready-to-simulate models, from robot arms, to dogs and mobile manipulators. We have been energised by our steadily growing community on GitHub, and were blown away by the level of interest and engagement at our social event in London during ICRA earlier this year.

Building on this progress, the MuJoCo team recently joined forces with colleagues from Brax and Everyday Robots, who bring expertise in hardware acceleration and broader simulation infrastructure, to form an unified Robotics Simulation team within Google DeepMind. Today, as a result of our combined effort, we are announcing MuJoCo 3: a major new release of MuJoCo, which adds support for accelerator hardware, improves scalability on CPU, and introduces more versatile collision primitives to widen the range of simulatable objects.

The full changelog is available on the release page and in our documentation.

MuJoCo XLA: Accelerated physics with JAX

MuJoCo 3 brings support for accelerated simulation via the new MuJoCo XLA (MJX) module. You can run MuJoCo simulations at millions of steps per second on Google Cloud TPU or your own accelerator hardware.

You can install MJX via pip install mujoco-mjx, and documentation is available here.

Seamless transition

Users of MuJoCo in Python can transition seamlessly between running MuJoCo on CPU, GPU, or TPU. The MJX API is nearly identical to MuJoCo, employing the same data model and simulation algorithms. Simulations in MJX will generally run in the same way as they do in MuJoCo, and machine learning models trained with MJX will operate the same in MuJoCo.

Blazingly fast simulation

MJX increases MuJoCo simulation throughput on a single machine, making it useful for data-hungry learning regimes such as reinforcement learning, evolution strategies, optimization for model predictive control, and more.

(CPU benchmarking was performed on an n2d-standard-128 GCE instance)

In the benchmark above, although classic MuJoCo can provide high throughput on machines with many CPU cores, the Python bindings are subject to the Global Interpreter Lock (GIL). As such, pure Python code cannot benefit from using threads to scale over available CPU cores. Python users upgrading from classic MuJoCo to MJX can expect to see a big increase in throughput through batching.

Because MJX runs on specialised hardware, it has unique performance characteristics that may require some model tuning, see documentation for details.

Speedy machine learning

MJX makes it easy to simulate big batches of environments to power reinforcement learning. This Colab notebook provides an example that teaches Google’s Barkour robot to run using MJX and Brax’s PPO trainer. We train a policy that can be deployed on a real robot and compare the training time to the Barkour publication, which used Nvidia Isaac Gym.

Platform	Time to train PPO at 60M steps
MJX + Brax on TPU v5e-8	249 seconds
MJX + Brax on Nvidia A100 GPU	718 seconds
Nvidia Isaac Gym on A100 GPU	600 seconds

Support for major hardware vendors through an open source ecosystem

MJX is built on XLA, which allows it to run on Nvidia, AMD, and Apple Silicon GPUs, as well as Google Cloud TPUs.

Scalability and parallelism

MuJoCo offers unparalleled speed for small and medium-sized scenes, but large, building-scale scenes with hundreds of robots or free bodies are not handled well. For CPU-based MuJoCo, we have made significant progress in this direction, with ongoing work to enable full support of very large scenes.

Memory sparsification

MuJoCo already uses sparse algebra for large models, but several key matrices use dense memory allocation, while touching only non-zero entries. As a result, very large models simply do not fit in available memory. We have completed the memory-compression refactor for the Conjugate Gradient solver, which reduced the memory footprint of a scene with a humanoid and 100 free objects by a factor of 80, and are almost done sparsifying the remaining algorithms.

Contact islands and multi-threading

Until today, when we spoke about multithreading in MuJoCo, we talked about running many instances of the same scene in parallel. This threading model is useful for optimization and reinforcement learning, but does not speed up cases where a single large scene is being simulated, for example in gaming.

In order to speed up a single, large scene, MuJoCo 3 detects “contact islands”: sub-components of the scene which do not interact. For example in a scene with two robots, they are assigned to separate islands as long as the robots don’t touch each other. Islands can be solved-for independently, for example on separate threads. Using a new thread-pool API, designed to exploit such parallelization opportunities, a model with 22 humanoids runs 3x faster. Threading over contact islands is currently only implemented for the Conjugate Gradient solver, implementation for all other solvers is in progress.

Other such intra-scene parallelization opportunities, for example collision detection, are also forthcoming.

Non-convex geometries

MuJoCo 3 adds support for collision geometries defined via signed distance functions (SDFs), allowing users to create new primitives by specifying the distance from any given location to the closest point on a surface. These geometries are not constrained to be convex. We provide examples of what is possible including tori, gears, nuts and bolts. Additionally, users can import any mesh and generate a multiresolution voxelized SDF at model compilation time.

In contrast to meshes, SDFs allow the number of contacts to be independent from the mesh resolution, making the cost of collision detection cheaper and more predictable. MuJoCo 3 SDFs can collide with all existing MuJoCo primitives and meshes.

Deformable bodies

MuJoCo 3 introduces a new kind of deformable body called flex. These are collections of segments (1D), triangles (2D) and tetrahedra (3D), corresponding to rope, cloth and deformable volumetric shapes like biological tissue. Flex bodies are not defined in a hierarchical kinematic tree, and can therefore simulate closed loop structures of any topological genus like rubber bands and T-shirts.

Deformation can be controlled either with equality constraints (for relatively stiff bodies) or with passive forces (for relatively soft bodies).

---

The Robotics Simulation Team

(in alphabetical order) Kyle Bayes (@kbayes), Matthew Bennice (@crewmatt), Tom Erez (@erez-tom), Erik Frey (@erikfrey), Chuyuan Kelly Fu, Nimrod Gileadi (@nimrod-gileadi), Alessio Quaglino (@quagla), Baruch Tabanpour (@btaba), Yuval Tassa (@yuvaltassa), Saran Tunyasuvunakool (@saran-t)

[JeanElsner]

Looks amazing! I knew some of these features were coming, but the extent of this release still surprises me! At our institute there is an ongoing discussion about whether to go with Isaac Sim or MuJoCo, but this settles it for me :) Of course at the end of the day, everyone should be free to choose the tools they feel comfortable with. But for me this release checks off a lot of requirements, I can't wait to integrate and play around with it.

[davidireland3]

Will DM control suite introduce functionality to run simulations on the GPU in the future?

[saran-t]

It's highly unlikely that we'll be retrofitting MJX into dm_control, however Brax will continue to be maintained to provide an interface between MJX and RL.

[RayYoh]

A bad news for dm control user. sad

[PeterMitrano]

In my opinion, dm_control offers very little over just using mujoco directly. Unless you're interested in their prebuilt environments, I don't see the advantage of using dm_control.

[kevinzakka]

@PeterMitrano Hard disagree, MJCF and composer are super useful features in dm_control.

[erikfrey]

Re-writing environments to be JAX-friendly is pretty easy, especially since dm_control already uses numpy! This could be a really nice community project for someone to take on.

Brax does have the Gym MuJoCo envs which are not as comprehensive, but could be a useful reference if someone wants to give this a shot.

[BolunDai0216]

Are RGBD and Lidar sensors in MuJoCo also accelerated?

[BolunDai0216]

Is that in your plan for future releases?

[aftersomemath]

Do you mean is there such sensors available? Or have they been accelerated by MuJoCo XLA? RGBD is available through the rendering API. For an example see renderer.py.

[BolunDai0216]

I mean if it is on GPU or CPU. For the camera, I am interested in what kind of resolution and speed can I expect. For the lidar, I am interested in the number of rays and also speed.

[aftersomemath]

RGBD can be done via OpenGL which typically runs on the GPU. I recently tested performance in the first post of #978. It's not as nicely supported right now as other MuJoCo sensors.

For discussion about Lidar support see #688.

[Andrew-Luo1]

MJX has an accelerated mjx.ray function which I have seen render 3.3e10 rays per second in a very simple scene. It can probably be used to implement a performant lidar without too much difficulty.

[bayerj]

Does it feature jax-transformable functions, just as brax?

[erikfrey]

Hi - Brax/MJX author here. Yes, MJX is pure JAX and is designed to work with vmap, scan and friends.

We haven't yet studied how MJX behaves under grad but it is a goal of the project to be usefully differentiable at some point in the future.

[kevinzakka]

Amazing work, excited to dive into the new features!!

[areiner222]

Very exciting!

How easy is it to batch the environment? In the tutorial notebook, I notice that the Model is batched over friction/gain/bias as an example, but can it enable batches of geometries? i.e., I'm simulating a humanoid like system and have various data based initial configurations for meshes that are all the same mesh topologically but are different shapes / sizes.

Would it be possible to augment the geometries (and update corresponding inertias / physical properties) across a batch dimension once I have instantiated a blueprint mjModel/mjx.Model ?

[btaba]

Hi @areiner222 , yes it's possible to batch by the scale of a geom (mjx.Model.geom_size, mjx.Model.geom_convex*), but not necessarily the shape (due to jax JIT). This is also true in Brax and is a feature/limitation of jax rather than the library itself. As noted you'll also have to update mass/inertia.

[areiner222]

Got it. Appreciate the insight. So it would be modifying the geometries around here? (assuming I'm modifying mesh geometries)?

For the sake of clarity, when I refer to shape, I'm talking about batching meshes that have the same number of vertices and number of faces but different vertex 3d positions. It sounds like this should work given I don't have changing mesh dimension sizes (in num vertices or num faces) across the batches?

[btaba]

Yeah that's right, as long as the jax array shape stays the same, you should be able to batch over the array and randomize the entries

[traversaro]

Thanks for the release! From what I understand (but I may be wrong), mujoco and mujoco-mjx are completely different implementations. Is any form of "consistency" of results between the two enforced in any way, for example via consistency tests or similar checks?

[erikfrey]

MJX is a direct reimplementation of MuJoCo in JAX - we use the same algorithms and data structures, just carefully architected to be performant on accelerators.

And yep this is tested! See for example:

https://github.com/google-deepmind/mujoco/blob/main/mjx/mujoco/mjx/_src/smooth_test.py#L66

In practice we have found that policies trained in MJX transfer seamlessly to MuJoCo.

[traversaro]

MJX is a direct reimplementation of MuJoCo in JAX

Yes, sorry for the bad wording: with "different implementation" I mean that they do not directly share code.

And yep this is tested! See for example:
https://github.com/google-deepmind/mujoco/blob/main/mjx/mujoco/mjx/_src/smooth_test.py#L66

Cool, thanks!

[snasiriany]

Will contact islands be supported for solvers other than CG? In particular for the currently default solver, Newton?

[yuvaltassa]

Not yet, but this is a high priority.

[vidurvij-Unity]

Are the binaries for MuJoCo 3 available for Unity?

[DMackRus]

Wow this seems really cool!

Is there any data available in terms of the million steps per second for the humanoid model for the previous version of mujoco using a similar CPU so we can get a sense of scale of the speed improvement?

[Nate711]

Will MJX support html rendering in colab like https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb#scrollTo=kF5fS-yo35sI ? I teach RL labs and it's amazing we could potentially move away from isaac gym and just have students use colab. However, as far as I know, mujoco doesn't have an interactive renderer that is compatible with colab like with brax.

[areiner222]

I'd find an interactive viewer really helpful for development in colab/jupyter. It would be nice to be able to render directly from the mjx objects instead of mapping back to mujoco mj* structures. One thing I'm currently trying to handle is how to map batched mjx.Model back into a list of mujoco.MjModel(s). I create a modified Model with varying geometries across batch dimension and so the templated MjModel I begin with is not representative of the system I want to visualize. While device_get_into seems to work for data / batched data, it's not working for a mjx.Model. It would be great to either a) be able to map between corresponding structures (mujoco<>mjx) for the Model (ideally without a result placeholder for ease) and/or b) render directly with mjx structures (interactively or not).

[erikfrey]

@areiner222 this should be pretty doable - I'll add the following to our list of fast follow features: support for mjx.Model in device_get_into.

Rendering with mjx structures is a bit more work but certainly something we could consider as part of this web viewer component that @saran-t was mentioning.

[saran-t]

I honestly think rendering directly from mjx structures is a bad idea since that'll create two forks of the rendering logic. Is there a strong reason why device_get_into doesn't work for you?

[areiner222]

It probably can work for me. There's currently an error when mapping the mjx.Model back to mujoco.mjModel that I think has to do with the fact that you can't set attributes on the MjModel?

e.g.,

model_mj = mujoco.MjModel.from_xml_path('/path/to/xml')

# build mjx.Model with batched attributes
def build_model(params, model_mj):
    model = mjx.device_put(model_mj)
    # Update model using params
    # .....
    # model now has batched:
    # - geom_convex*, body_pos, body_quat, body_mass, body_ipos, body_iquat, body_inertia, geom_pos, geom_quat, body_subtreemass, body_invweight0, dof_invweight0
    return model

model = build_model(params, model_mj)

# simulate
data = init(model, qpos_init)
for i in range(num_sim_steps):
    data = step(model, data, ctrl=...)

# map to mj structures
model_mj_recon = copy.deepcopy(model_mj)
mjx.device_get_into(model_mj_recon, model)  # fails (even if model has no batched attributes)
... # move data off of device, etc.

[Nate711]

Thanks for discussing this! I just sent Baruch an email (I'm Nathan) but hadn't seen these replies at the time. Are there any updates since the last post? My understanding is that the brax html visualizer uses a json-ified brax system and trajectory and that to have similar functionality with MJX there would need to be either 1) a way to convert MJX system to brax system 1.5) a way to convert mjx system to brax json format 2) a separate visualization application that uses MJX systems. And there is also a web-based app similar to Mujoco's "simulate" binary?

[DimaRubyAI]

Great work!
Can I find somewhere the code of the YouTube SDFs examples?

[yuvaltassa]

plugin/sdf/

[DimaRubyAI]

Yes, thank you! I already found it, very nice examples

[Yasu31]

The options to set collisions to all, predefined, and dynamic seems to have been removed in the MJCF modeling format 3.0.0 (https://mujoco.readthedocs.io/en/2.3.7/XMLreference.html#option-collision ), is there an alternative option we could set to enable / disable all collisions in MJCF?

[yuvaltassa]

Yes. The flag for toggling all collisions is here

[Yasu31]

Thank you so much!