Are there any existing suites of MJX based manipulation environments with gym style interfaces?

[StoneT2000]

Intro

Hi!

I am a graduate student at UC San Diego, I am currently investigating and comparing GPU robotics simulators on different manipulation tasks.

My setup

Ubuntu 22.02, latest MJX/Mujoco python API

My question

I was wondering if there are any existing manipulation benchmarks built on top of MJX, similar to IsaacLab/ManiSkill 3. I did a quick search on google and here but couldn't find any already built ones. My goal is to run some RL algorithms on the GPU simulation and wanted to try out MJX instead of just the physx based simulators and learn about the differences.

Or does anyone know if e.g. MetaWorld, DM-Control, Gymnasium Robotics have any progress/plans to support MJX? I think I saw gymnasium maintainers discussing using MJX but I have not seen any updates about it since a few months ago.

Minimal model and/or code that explain my question

No response

Confirmations

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

[erikfrey]

Hi @StoneT2000 ,

We spoke with Gymnasium folks a while back - at the time MJX was still missing some features that are now in or will be in shortly. We'll probably check back in with them now that those are out.

We use MJX for manipulation tasks with the Aloha platform - if it's useful to you to see them, I can get a few of those tasks open sourced. We also are doing some experiments with Madrona for vision-based policies, although that is early work that is not yet ready.

Manipulation is an area of active focus for us, so if there's some way we can support your investigation or comparison, please let us know.

[erikfrey]

I will add one nitpicky addendum: there are scenarios where MJX slows down as more contacts are dynamically encountered during simulation. The constraint solver can early exit, and will do more work during states where there are more active contacts. But this slowdown will be less pronounced than PhysX.

[StoneT2000]

Yeah I find benchmarking simulation speed incredibly tricky to be honest. It is simple when comparing frameworks with already built tasks (e.g. robosuite vs maniskill) since I would generally assume authors tune their simulations well and cannot be at fault for poor tuning of a different framework. But trying to fairly compare the underlying simulators e.g. mujoco vs sapien vs isaac sim is fairly complicated.

While strange, somehow this makes me think running RL on the environment with the same e.g. PPO code and settings is more or less one of the better ways to compare simulation speeds. E.g. IsaacLab typically uses a low simulation to control frequency ratio = higher raw FPS = slower RL learning as each action moves the robot less distance. Harder to manipulate FPS numbers to an extent if comparing RL wall training time.

Has your team tried to benchmark other simulators before and if so, what is your methodology for fair comparisons?

[btaba]

Hi @StoneT2000

We just open sourced a manipulation env (pickup cube) on the Franka Panda in the public MJX tutorial colab. I copied some of your reward terms, but not all. I was satisfied enough with the training time so didn't do further tuning. We train for 20M env steps (at a 50Hz control frequency and 200Hz sim frequency). It trains for around 3 minutes on an A100.

Agreed that benchmarking different sims is difficult. In principle I agree about RL being a good benchmark since it tests a real-world workload. However in practice, you'll have to somehow control for differences in [1] the RL algo/impl and [2] reward terms that both may contribute to different sample efficiencies between the set-ups. For example, one could potentially keep tuning the rewards in the Panda env above to get training time down (less env steps), but it's probably not a good use of time since training time is already 3 minutes. As another example, if one were to copy the same exact RL algo/impl and reward terms from a different set-up, I wouldn't be surprised if they still needed to tune RL hypers and maybe rewards due to differences in physics between simulators.

The way I've typically benchmarked between simulators in the past: [1] load the same scene, [2] check that both sims are stable at the same sim and control frequency with random perturbations, [3] ensure contacts occur in the scene for the benchmark time window, and [4] get raw FPS numbers using provided benchmark scripts on the same device. It's not perfect, but it's a bit easier to run through.

[StoneT2000]

Thanks I just saw the video on your tweet as well, good to see confirmation and prewritten baseline is working. I'll take a look and play around.

Also agree that RL is not a perfect solution. I think maybe at best what we can do is just share baselines for RL+gpu sim for different frameworks, post raw numbers (that maintainers of respective frameworks can tune themselves as much as they want), share video of results, and let users decide. This is partly a suggestion from the maintainer of Simulately, which seems to be a bit more of a mostly unbiased repo/wiki for simulators. Perhaps it can be where we just store all benchmarking scripts for all things sim and robotics?

This also ended up being partly what my lab did with our recent maniskill paper release on arxiv. We compared with Isaac Lab and for e.g. rendering fidelity we don't bother to say which simulator's parallel rendering looks better, we just post results and let user decide

[btaba]

Thanks for the link to Simulately, happy to provide scripts where they're needed. Also 100% agree, users should (and will) use whatever sim meets their use-case best. We try to advertise the sharp bits about MJX as its implemented today. So far MJX has been great for locomotion + sim2real, but not really proven for manipulation workloads. As such, there are bottlenecks in the impl that we're actively working to get rid of!