Domain randomization for MJX

[alexhansson]

Hi,

I'm a student and I'm trying to use MuJoCo for training RL policies for space systems. According to the google collab notebook linked in the MuJoCo documentation, one can simulate and randomize multiple environments in parallel with this code:

rng = jax.random.PRNGKey(0)
rng = jax.random.split(rng, 4096)
batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)

jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))
batch = jit_step(mjx_model, batch)

However, this only varies the initial position of your object. I'm wondering if you can vary the environments more flexibly, like for example randomizing the location/dimension of collision objects in the scene. I know certain parameters can't be changed once the xml has been generated, but would it in theory be possible to have multiple xml files and "concatenate" them in some way for parallel simulation?

And as a side question, is it possible to visualize the parallel environments similarly to issacgym?

Thanks a lot for your help.