Code Skeleton for multiple parallel simulations with MJX

[Nick20500]

Intro

Hi!

I am a undergrad student at ETH Zurich, I use MuJoCo for my research on Multimaterial deformable objects.

My setup

I use MuJoCo 3.2.3 with python bindings.

My question

I am trying to use MJX to let multiple simulations run in parallel on GPU.

I am able to do this with plain mujoco in serial without any problems. But I can't figure out how to use MJX to parallize this code.
Here is the mujoco code which I want to "refactor" to use MJX:
simulate-Table-With-Increasing-Stiffness.zip

The simulations should all have the same xml scene and object as their basis, which is a bottle droping on to the ground. But each simulation should give the bottle diffrent stability parameters. I use flexcomp bodys which I want to assign diffrent youngs moduli to.
Here is the scene I use:
scene_and_msh_object.zip

Can someone give me a simple code skeleton which uses jax.vmap() to map a array of diffrent youngs moduli to multiple simulations and then runs those simulation with jax.jit() in parallel?

All the ressoruces/documentation/collab notebooks I found, aren't helping me to construct, what I want.

Thank you for your time and effort.

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.