Fast and User-Friendly Environment/Task Interfaces for MuJoCo

[hartikainen]

Hey folks,

I'm training reinforcement learning agents with Jax and using MuJoCo as the physics simulator. I'm trying to understand the possibilities and limitations of current MuJoCo and MJX capabilities, especially in terms of simulating larger scenes with many colliding objects. What I'm basically looking for is a user-friendly environment/task interface for MuJoCo that would allow efficient sampling of environments, making it suitable for RL training. It would be interesting to get community's insights on this.

Existing Environment/Task Interface Options

There are various existing options for environment/task interfaces for MuJoCo. I'm mostly familiar with the following three.

Python-Based Interfaces

Environment interfaces written in Python, like those in dm-control, are flexible and provide really nice user experience for defining environments. However, they tend to be super slow to sample (due to Python's overhead) and require special tricks (like distributing the environment samplers) to achieve higher throughputs. This makes RL training either frustratingly slow or the training stack unnecessarily complicated.

MJX-/Brax-Based Environments

MJX-/Brax-based environments (like the ones in the MJX Colab example), are extremely fast for simple models. While they may not be as user-friendly as pure Python environments, the jit-/vmap-interface is really nice to work with, as I only have to reason about single environment at a time using familiar Jax and MuJoCo syntax. However, MJX's performance degrades rapidly as the number of colliding objects in the scene increases, making it currently unsuitable for larger scenes.

C++-Based Interfaces

I’ve recently come across projects like EnvPool, which achieve pretty impressive throughput on CPUs with C++-based environment interface and customized threading. Writing environments in C++, at least for me, is less intuitive than writing them in Jax, but seems like an acceptable trade-off given that the CPU-based simulation seems to scale to larger number of objects much better than GPU-based MJX. Unfortunately, EnvPool appears to be inactive and I don't know what its state really is.

Questions

Here are some specific questions that I have in my mind:

1. Future of large scene simulation with MJX/MuJoCo: What does MJX's future look like in terms of simulating larger scenes (say, with with O(100) or O(1000) colliding bodies like in the model/replicate/particle.xml demo)? I’ve seen some discussions regarding revamping the collision implementation — for instance the comments here and here, and here here in the docs — which seem to at least partly address the collision-related performance. Could anyone share more about what these improvements would mean for MuJoCo and MJX? Is there a timeline or roadmap for when these changes might be implemented?

2. Other high-throughput, user-friendly solutions: How do people implement MuJoCo-based environments efficiently nowadays? Apart from the three I've listed above, are there any notable environment interfaces that I should take a look at? Or any other solutions that could work well for my setting?

Hopefully my questions make sense! Any advice, pointers, or comments are be much appreciated.

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.