Jit-compatible MJX model randomization

[jonasweihing]

Intro

Hi!

I am a masters' student at the University of Tübingen, Germany, and I use MuJoCo/MJX for my research on robotic manipulation for obstacle avoidance in a RL reaching task.

I have implemented my own environment, similar to what was done in the MuJoCo MJX Playground, and I train the agent with Brax's PPO implementation.

My setup

MuJoCo/MJX 3.2.7, Brax 0.12.1, Python 3.12, Ubuntu 24.04

My question

Goal

Do domain randomization on the initial state of the environment, namely the robot position and obstacle properties (position, rotation, size, shape).

What is working

The randomization of the initial robot joint positions and obstacle position + orientation is perfectly working for me. It does not change the model definition and I can set those on episode reset with qpos, qvel, mocap_pos and mocap_quat.

What I struggle with and need help

Randomizing the obstacle size (and possibly shape - but I'll leave that out for this question) requires changing the model definition, which means I have to recompile the model. Sampling the size with jax.numpy and modifying the mjspec with it works fine, as long as I don't jit-compile the function with. However, since all this would be done in the reset function of my environment, which is jit-compiled, means, it will be jit-compiled. That makes sampling the size with jax.numpy and modifying the mjspec not possible.

Ideas I came up with:

• Sampling the size and compiling the model before the training:
  This means that hundreds of thousands of precompiled models must be stored in memory, which - without investigating it further - likely is a far too large memory footprint.

• Using a jax's pure_callback function to compile the model:
  Firstly I coulnt quickly figure out how to specify the outcoming shape as it is a mjx.Model and secondly, this would most likely make an immense performance penalty hit
  This would most likely make an enormous performance penalty hit and I couldn't quickly figure out how to specify the outgoing shape of the callback function, as it is a mjx.Model and not just a jax.Array.

• Using numpy to sample the size:
  This makes building the model possible, but it also means I have to pass through a numpy generator that can be split into children like with jax to mitigate the side effects and allow parallel execution. I haven't gone very far down this path yet, but I have doubts that it's possible because the reset function in the training script is vmapped, and I don't think it's at all possible to jax vmap a function across a list/array of numpy generators.

• Using the randomization_fn parameter of the train function:
  As this function is used when wrapping the reset function of the environment, it is also jit-compiled.

• Modify the mjx.Model directly:
  I know this is a really bad idea, but it seems to be the only way that somehow makes it possible. I tried to go this way and reimplement the calculation of body_quat, body_mass, body_subtreemass, body_inertia, geom_size, geom_rbound, bvh_aabb, geom_aabb and geom_rbound_hfield as well as stat in jax. Those fields changed, when I changed the size of the mocap body/geom. This worked surprisingly well. However, some properties are numpy ndarray's, which makes it impossible to update them in jit-compiled code, especially when it is based on a sampled property.

How would you approach this randomization, and do you have any idea on how I can do this? Or is it even possible in any way and I have to pre-sample and just live with the memory footprint?

What would be my dream solution?

A jit-compatible model compilation function to dynamically change the model between episodes and use it accelerated on GPU/TPU.

Minimal model and/or code that explain my question

Here is a stripped down example of what I am trying to achieve.

Code:

import jax
import mujoco
from mujoco import mjx

rng1, rng2 = jax.random.split(jax.random.key(0))


def create_model_with_sampled_size(rng):
    size = jax.random.uniform(rng, shape=(3,), minval=0.01, maxval=0.1)

    spec = mujoco.MjSpec()
    spec.worldbody.add_body(
        mocap=True,
        name="obstacle_body",
    ).add_geom(
        name="obstacle_geom",
        size=size,
    )
    model = mjx.put_model(spec.compile())
    return model


# This is possible
model = create_model_with_sampled_size(rng1)

# This is not possible
fn = jax.jit(create_model_with_sampled_size)
model = fn(rng2)

Confirmations

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

[Balint-H]

Hello!

One approach you could also consider is to simply have multiple objects in your scene already, and teleport only the currently relevant one in and leave the other far away (to make collision checks cheaper with it). Of course, this had many drawbacks. For one, it forces you to heavily discretize your distributions.

[jonasweihing]

Hello!

Thanks for the quick reply. Yes, I was thinking about that for randomizing the shapes, since it is a limited number of possible randomizations. I could also do this for the size of the obstacle - true.
A question I have regarding this: What do you think is the maximum number of objects I can have in MuJoCo without it having too much of an impact on performance. Is it 10, 100, 1000, etc.? And does it speed up the simulation noticeably if I place each obstacle at a large distance of f.e. 100 meters when it is not used, as you suggested?

As I was digging through the mujoco playground code to see how domain randomization was done there, I came across another great idea:
Don't randomize the domain for each episode within the reset, but "create" different models within a vmap wrapper. So, for example, if I have 2048 parallel environments (parameter num_envs), I at least get 2048 different environments, which is definitely better than one. And since the randomization is done during initialization and therefore not jit-compiled, I could merge multiple models compiled from mjspecs into one vectorized model. That sounds very promising.

[Balint-H]

Yes that sounds like a good way to go about it, being able to do this seems like one of the strengths of parallel environments.

[jonasweihing]

I finally implemented a randomization function, similar to how it is done in mujoco playground, which uses the parallelization of envs with mjx/brax to randomize the model.

This solution has some limitations. The fields bvh_aabb, geom_aabb and geom_rbound_hfield are not batchable, but change on recompile when I change the size of the object. However, bvh_aabb is restricted to mujoco and therefore not used in MJX and geom_rbound_hfield is only accessed if a mjGEOM_HFIELD geometry is used, which I do not intend to do. Therefore, I am fine with them not being batchable. For geom_aabb I'm not sure because I couldn't find any usage within MJX, but I guess I am fine with it as well.

Here's a snippet of the code I used in case anyone stumbles across this discussion and is interested. Be careful because it contains a sampling function that I don't provide, and I'm accessing some variables in my environment.

in_axes = jax.tree.map(lambda _: None, env.model)
in_axes = in_axes.tree_replace(
    {
        "stat.meaninertia": 0,
        "stat.meansize": 0,
        "stat.extent": 0,
        "body_inertia": 0,
        "dof_M0": 0,
        "geom_size": 0,
        "geom_rbound": 0,
        # Theoretically, the following fields should be replaced as well as they can
        # change, but it is not possible to vmap over them.
        # "bvh_aabb": 0,
        # "geom_aabb": 0,
        # "geom_rbound_hfield": 0,
    }
)

model_randomized: mjx.Model = env.model.tree_replace(
    {
        "stat.meaninertia": jnp.repeat(
            jnp.expand_dims(env.model.stat.meaninertia, 0), num_envs, axis=0
        ),
        "stat.meansize": jnp.repeat(
            jnp.expand_dims(env.model.stat.meansize, 0), num_envs, axis=0
        ),
        "stat.extent": jnp.repeat(
            jnp.expand_dims(env.model.stat.extent, 0), num_envs, axis=0
        ),
        "body_inertia": jnp.repeat(
            jnp.expand_dims(env.model.body_inertia, 0), num_envs, axis=0
        ),
        "dof_M0": jnp.repeat(
            jnp.expand_dims(env.model.dof_M0, 0), num_envs, axis=0
        ),
        "geom_size": jnp.repeat(
            jnp.expand_dims(env.model.geom_size, 0), num_envs, axis=0
        ),
        "geom_rbound": jnp.repeat(
            jnp.expand_dims(env.model.geom_rbound, 0), num_envs, axis=0
        ),
    }
)

spec = env.spec
geom = env.spec.geoms[env.obstacle_geom_ids[0]]

for idx in range(num_envs):
    rng, rng1 = jax.random.split(rng, num=2)

    size = obstacles.sample_size(
        rng1, mjx.GeomType(geom.type), env.obstacle_size_limits
    )
    geom.size = size
    mj_model = spec.compile()

    model_randomized = model_randomized.tree_replace(
        {
            "stat.meaninertia": model_randomized.stat.meaninertia.at[idx].set(
                mj_model.stat.meaninertia
            ),
            "stat.meansize": model_randomized.stat.meansize.at[idx].set(
                mj_model.stat.meansize
            ),
            "stat.extent": model_randomized.stat.extent.at[idx].set(
                mj_model.stat.extent
            ),
            "body_inertia": model_randomized.body_inertia.at[idx].set(
                mj_model.body_inertia
            ),
            "dof_M0": model_randomized.dof_M0.at[idx].set(mj_model.dof_M0),
            "geom_size": model_randomized.geom_size.at[idx].set(mj_model.geom_size),
            "geom_rbound": model_randomized.geom_rbound.at[idx].set(
                mj_model.geom_rbound
            ),
        }
    )

Related issue: #1607