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.

[StoneT2000]

What do you mean always calculated in MJX? If the gripper is not near the cube collisions are still being calculated? So whether there is actual contact or not does not change FPS? I find this surprising as this is not the case with PhysX. Contacts are not computed/considered in the simulation pipeline in gpu sim if two objects are too far away (default is like 2 centimeters in maniskill). Often during RL FPS of the environment goes down as the agent learns to actually grab the intended object more often in e.g. a peg insert task or cube picking task.

My understanding is brax is the RL interfaces for mjx now? What is the difference between the mjx backend and what brax had before?

[btaba]

If the gripper is not near the cube collisions are still being calculated?

Yes this is currently the case. As mentioned, XLA requires static shapes known at compile time, and the MJCF provided always contains 16 contact points. We have a custom parameter to do broadphase, but I didn't add that to the MJCF provided. This is why I mentioned that collisions are always being calculated between the box and the gripper, and that testspeed should be a decent representation of the env performance. You can check mjx_data.contact.geom on the first step.

My understanding is brax is the RL interfaces for mjx now? What is the difference between the mjx backend and what brax had before?

Yes we typically use brax as the RL interface for MJX. MJX is MuJoCo implemented in JAX. Other backends in brax implement different physics pipelines (mass-spring, PBD, etc.).

[StoneT2000]

Thanks for clarifications! I understand how mjx works much better now, I will try and share benchmarking results when I can.

[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?