Question: Is there a way to vmap models + data over a simulation loop.

[oliverweissl]

Hi,

I'm a student and I'm trying to use MuJoCo for simulating modular robots.

I'm looking for some help with parallelizing these simulations in MJX.
My aim is to utilize GPU devices for parallelizing simulations, which are composed of n mjx.Model and n mjx.Data.

In the vanilla implementation we use multiprocessing to spawn processes on the CPU, that take one model and data object and use them in a while loop (including the step function) for the main simulation.

The Vanilla implementation:

while cond:
   do something
   mujoco.mj_step(model, data)

My idea to use MJX was to do a pmap over the main simulation loop like this:

def _sim_loop(model, data) :
    while cond:
        do something
        mjx.step(model, data)

def main():
   models = [list of mjx.Model]
   data = [list of mjx.Data]

   jax.pmap(_sim_loop, **kwargs)(models, data, *args)

However this fails because the lists i want to map over are pytrees. I have read that batching could be one way to make this work, but I could not find any documentation on this.

Do you have a tip on how to approach this? Also this is my first time using JAX so im still learning and trying to scan through the docs to find something that fits my issue.

Thanks in advance for any suggestions and help!

[btaba]

Hi @oliverweissl ! I encourage you to look at the colab tutorial. The "Introduction to MJX" section has an example of batching environments (look for batched_mj_data).

For the _sim_loop, you'll want to do a jax.lax.scan or jax.lax.while_loop instead of the Python for or while loop, since the just-in-time compilation will unroll the entire loop into many operations (thus increasing compilation time), rather than lowering the loop into a single loop operation in XLA. Check out this example:

mujoco/mjx/mujoco/mjx/testspeed.py

Lines 94 to 104 in e77c3cb

	@jax.pmap
	def unroll(d):
	@jax.vmap
	def step(d, _):
	d = mjx.step(m, d)
	return d, None
	d, _ = jax.lax.scan(step, d, None, length=FLAGS.nstep, unroll=FLAGS.unroll)
	return d

Another thing to note is that pmap parallelizes over the number of devices, so you'll likely just need a vmap unless you have multiple GPUs.

[oliverweissl]

Hello, thanks for your input!
However I still dont understand how i can get my list of mjx.Model into a format that is useable by jax. In the colab tutorial they generate the batch by changing some numerical values, but my batch is composed of different models and data objects. For my usecase I cant just copy one object and change some parameters.

[btaba]

Ideally you would stack your mjx.Model (i.e. jax.tree_map(lambda *x: jp.stack(x), *list_of_models)), otherwise you won't be able to vmap/pmap over them. Can you give a more complete example of your use-case?

[oliverweissl]

Of course!
Im trying to use mjx for evolutionary computing where I get a new population of robots each iteration. I want to take those robots and evaluate them ideally in parallel, which would require to vmap over mjx.models as far as I understood this. But Im unable to get the mjx.Model objects into a format that allows for vmap. Also the treemap did not work due to a ValueError: Mismatch custom node data.

Now Im wondering if mjx is even usable for parallelizing different models, since all examples i found online use small parametric changes in the models rather than different models.
Thanks again for your help!!

[btaba]

I see thanks @oliverweissl for the explanation. To be able to vmap over independent mjModels, their shapes need to be stackable.

Another way to potentially go about this, is to generate an MjModel that contains all of your robot models in them, and disable collisions between robot models. Then you'd evaluate all models at the same time; the downside is the slowdown/instability you'll likely observe since the solver will merge all constraints into one big call, rather than "islanding" the different robot models and solving them independently