Sharing a data object for multiple models #2922

bd-tablett · 2025-10-31T20:12:55Z

bd-tablett
Oct 31, 2025

Intro

Hi!

My name is Trevor, I'm a researcher at Boston Dynamics, and I use MuJoCo for my research on Atlas.

My setup

MuJoCo version 3.3.2 - 3.3.7
Python
x86_64
Ubuntu 22.04

My question

My real setup uses mujoco.rollout.Rollout, but I'm trying to sample different objects within the sequences of models that I pass to the rollout function. Specifically, each model would contain the robot and one object selected from many, but each model would otherwise have the same number of bodies.

This breaks the assumption that the mujoco data matches the model, although I didn't think that it fully broke the "homogeneous models" requirement of mujoco.rollout.Rollout. I was hoping that I could get around this by defining a "dummy" object that has more geoms than any of the other objects I care about, and base my data on that object. This seems to partially work, but also fails in surprising ways, and also in different ways depending on which version of MuJoCo I'm using.

I've attached a minimal example that fails in two different ways. When not using shared data, this is what I get, which would be the expected behaviour:

Screencast.from.10-31-2025.03.56.00.PM.webm

When using shared data, with version 3.3.2 - 3.3.5, I get this, where the primitive-based objects seem to be okay, but the sphere mesh doesn't work:

Screencast.from.10-31-2025.04.00.21.PM.webm

Where as soon as the sphere is loaded, I get this error:

WARNING: Nan, Inf or huge value in QACC at DOF 0. The simulation is unstable. Time = 1.6020.

When I use 3.3.6 - 3.3.7, I get this, where the second primitive object is clearly not simulating correctly:

Screencast.from.10-31-2025.04.02.13.PM.webm

and no warnings.

My questions are:

Is this approach to sharing data reasonable, and if so, what am I doing wrong?
If it's not reasonable, is there a recommended way to randomize across objects while still using mujoco.rollout.Rollout and shared data?

Minimal model and/or code that explain my question

To run and reproduce my errors, set USE_DUMMY_DATA=True, and to run without sharing data, set USE_DUMMY_DATA=False.

Code:

# Standard Library
import os
import shutil
import tempfile
import time
from importlib.metadata import version

# Third Party
import mujoco
import mujoco.viewer
import trimesh
from packaging.version import Version

DUMMY = """
<mujoco model="dummy_body">
  <worldbody>
    <site name="world_frame" pos="0 0 0" size="0.02" rgba="1 0 0 1"/>
    <light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
    <body name="ground" pos="0 0 0">
      <geom type="plane" size="1 1 0.1" rgba=".9 .9 .9 1"/>
    </body>
    
    <body name="dummy_body" pos="0 0 1.0">
      <joint type="free"/>
      <geom type="sphere" size="1.05" rgba="1 0 0 1"/>
      <geom type="sphere" size="1.05" rgba="1 0 0 1"/>
      <geom type="sphere" size="1.05" rgba="1 0 0 1"/>
      <geom type="sphere" size="1.05" rgba="1 0 0 1"/>
    </body>
  </worldbody>
</mujoco>
"""

MULTI_GEOM_BODY_1 = """
<mujoco model="multi_geom_body_1">
  <worldbody>
    <body name="multi_geom_body_1" pos="0.0 0.0 30.0" quat="1.0 0.0 0.0 0.0">
      <joint type="free"/>
      <geom type="box" size="0.005 0.04 0.08" group="2" rgba="0 1 0 1"/>
      <geom type="box" size="0.04 0.005 0.08" group="2" rgba="0 1 0 1"/>
    </body>
  </worldbody>
</mujoco>
"""

MULTI_GEOM_BODY_2 = """
<mujoco model="multi_geom_body_2">
  <worldbody>
    <body name="multi_geom_body_2" pos="0.0 0.0 30.0" quat="1.0 0.0 0.0 0.0">
      <joint type="free"/>
      <geom type="box" size="0.005 0.04 0.08" group="2" rgba="0 0 1 1"/>
      <geom type="box" size="0.4 0.005 0.08" group="2" rgba="0 0 1 1"/>
      <geom type="box" size="0.04 0.3 0.01" group="2" rgba="0 0 1 1"/>
    </body>
  </worldbody>
</mujoco>
"""

MESH_BODY_TEMPLATE = """
<mujoco model="mesh_body">
  <asset>
    <mesh name="my_mesh" file="{mesh_path}"/>
  </asset>
  <worldbody>
    <body name="mesh_body" pos="0.0 0.0 2.0">
      <joint type="free"/>
      <geom type="mesh" mesh="my_mesh" rgba="0.5 0.5 1 1"/>
    </body>
  </worldbody>
</mujoco>
"""

NUM_RESETS = 5
NUM_ENV_STEPS = 400
USE_DUMMY_DATA = True


def create_mesh_file():
    mesh = trimesh.primitives.Sphere(radius=0.1, subdivisions=3)
    temp_dir = tempfile.mkdtemp()
    mesh_path = os.path.join(temp_dir, 'example_mesh.stl')
    mesh.export(mesh_path)
    print(f"Created temporary mesh at: {mesh_path}")

    return temp_dir, mesh_path


def reset_body(model: mujoco.MjModel, data: mujoco.MjData, body_name: str, pos,
               quat=(1, 0, 0, 0)) -> None:
    """Reset the position of the dummy body to a specified height."""
    body_id = model.body(body_name).id
    joint_id = model.body_jntadr[body_id]
    qpos_addr = model.jnt_qposadr[joint_id]
    data.qpos[qpos_addr:qpos_addr + 3] = pos
    data.qpos[qpos_addr + 3:qpos_addr + 7] = quat
    data.qvel.fill(0.0)
    data.qacc.fill(0.0)
    data.ctrl.fill(0.0)
    data.xpos.fill(0.0)
    mujoco.mj_forward(model, data)


def copy_mj_data(src_data: mujoco.MjData, dest_data: mujoco.MjData) -> None:
    """Copy relevant fields from src_data to dest_data."""
    dest_data.time = src_data.time
    dest_data.qpos[:] = src_data.qpos
    dest_data.qvel[:] = src_data.qvel
    dest_data.qacc[:] = src_data.qacc
    dest_data.ctrl[:] = src_data.ctrl
    dest_data.xpos[:] = src_data.xpos[:dest_data.xpos.shape[0]]
    mujoco.mj_forward(dest_data.model, dest_data)


def demonstrate_mismatch():
    temp_dir = None
    try:
        # Generate mesh file and XML with it
        temp_dir, mesh_path = create_mesh_file()

        # Generate specs, models, data from each XML
        specs = {}
        models = {}
        datas = {}
        for name, xml in [("dummy_body", DUMMY), ("multi_geom_body_1", MULTI_GEOM_BODY_1),
                          ("multi_geom_body_2", MULTI_GEOM_BODY_2),
                          ("mesh_body", MESH_BODY_TEMPLATE.format(mesh_path=mesh_path))]:
            spec = mujoco.MjSpec.from_string(xml)
            model = spec.compile()
            data = mujoco.MjData(model)
            specs[name] = spec
            models[name] = model
            datas[name] = data

        # Detach the dummy body, attach the multi-geom bodies
        new_models = {}
        for name in ["multi_geom_body_1", "multi_geom_body_2", "mesh_body"]:
            new_spec = specs["dummy_body"].copy()
            dummy_body = new_spec.body("dummy_body")
            if Version(version("mujoco")) > Version("3.3.3"):
                new_spec.delete(dummy_body)
            else:
                new_spec.detach_body(dummy_body)
            new_spec.attach(specs[name], site=new_spec.site("world_frame"))
            print(f"After attaching {name}:")
            print(new_spec.to_xml())
            new_models[name] = new_spec.compile()

        # View the simulation with the mismatched bodys and models
        reset_body(models["dummy_body"], datas["dummy_body"], body_name="dummy_body", pos=(0, 0, 1),
                   quat=(0, -.487, 0, -.873))
        with mujoco.viewer.launch_passive(models["dummy_body"], datas["dummy_body"]) as viewer:
            for _ in range(NUM_RESETS):
                for name in ["multi_geom_body_1", "multi_geom_body_2", "mesh_body"]:
                    # Replace the model and data with the ones we're actually testing
                    sim = viewer._get_sim()
                    model = new_models[name]
                    viewer_mj_data = mujoco.MjData(
                        model)  # has to match for viewer, but unused otherwise
                    sim.load(model, viewer_mj_data, "")

                    # Reset the dummy body position
                    if USE_DUMMY_DATA:
                        reset_body(models["dummy_body"], datas["dummy_body"],
                                   body_name="dummy_body", pos=(0, 0, 1), quat=(0, -.487, 0, -.873))
                    else:
                        reset_body(model, viewer_mj_data, body_name=name, pos=(0, 0, 1),
                                   quat=(0, -.487, 0, -.873))

                    # Step the simulation with the new model and the dummy data that doesn't match
                    for _ in range(NUM_ENV_STEPS):
                        if USE_DUMMY_DATA:
                            mujoco.mj_step(model, datas["dummy_body"])
                            copy_mj_data(datas["dummy_body"], viewer_mj_data)
                        else:
                            mujoco.mj_step(model, viewer_mj_data)
                        viewer.sync()
                        time.sleep(0.005)
    finally:
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
            print(f"\nCleaned up temporary directory: {temp_dir}")


if __name__ == "__main__":
    demonstrate_mismatch()

Confirmations

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

yuvaltassa · 2025-11-02T09:11:42Z

yuvaltassa
Nov 2, 2025
Maintainer

mujoco.rollout supports multiple datas and multiple models, which I think should address what you're trying to do, why is that not a good solution?

4 replies

yuvaltassa Nov 2, 2025
Maintainer

In other words, if you make 10 models, each with a different object, (these are indeed homogeneous), and then make 10 datas one from each model, then you should be able to run the 10 rollouts in parallel. Like, I don't understand why you'd want the same mjData.

cc @aftersomemath in case I am missing something.

aftersomemath Nov 2, 2025

Rollout assumes it can use any mjData with any model. The number of mjData passed controls how many threads are used. The pairing is nondeterministic when the number of rollouts exceeds the number of threads.

Based on the description, it sounds like these mjData are not compatible with every mjModel? The only safe workaround is to use several Rollout objects and asynchronously execute them.

An early prototype of the new rollout allowed completely different models/data. However, it requires more memory.

bd-tablett Nov 2, 2025
Author

Thank you both for your quick replies!

Indeed, as @aftersomemath describes, the number of rollouts exceeds the number of threads (generally by very large amounts), which is why we started talking about sharing mjData objects between models, and what is and is not possible or recommended. It sounds like using the same mjData object for different models is possible but behavior is undefined.

The asynchronous use of several Rollout objects is interesting, but might pose a problem if the number of objects we wanted to randomize over was very large. We may have to end up going this route, and I'll take some time to investigate it. If I understand correctly, if we limit the number of rollouts to be the same as the number of threads, the mjData use will be deterministic.

An early prototype of the new rollout allowed completely different models/data. However, it requires more memory.

This would potentially be ideal for us, since memory is less of a concern than speed. Is there any possibility that this version will be made available?

aftersomemath Nov 3, 2025

Here is the old branch. Keep in mind it is nothing more than a prototype. After it was written, rollout's current interface was decided on and so I started over and abandoned this branch.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Sharing a data object for multiple models #2922

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 1 comment 4 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

Sharing a data object for multiple models #2922

Uh oh!

Uh oh!

bd-tablett Oct 31, 2025

Intro

My setup

My question

Minimal model and/or code that explain my question

Confirmations

Replies: 1 comment · 4 replies

Uh oh!

yuvaltassa Nov 2, 2025 Maintainer

Uh oh!

yuvaltassa Nov 2, 2025 Maintainer

Uh oh!

aftersomemath Nov 2, 2025

Uh oh!

bd-tablett Nov 2, 2025 Author

Uh oh!

aftersomemath Nov 3, 2025

bd-tablett
Oct 31, 2025

Replies: 1 comment 4 replies

yuvaltassa
Nov 2, 2025
Maintainer

yuvaltassa Nov 2, 2025
Maintainer

bd-tablett Nov 2, 2025
Author