per-world meshes by thowell · Pull Request #1191 · google-deepmind/mujoco_warp

thowell · 2026-02-26T18:39:32Z

update: the pr is now primarily about enabling batched geom_dataid. see here for documentation (draft) containing examples utilizing this feature.

per-world meshes

define mesh distributions at the geom or body level using custom elements in xml or mjspec
- body-level randomization supports different numbers of geoms/meshes (eg, aloha pot scene with 2 different objects represented as convex decomposition 1 with 10 geoms and convex decomposition 2 with 20 geoms)

scene with 2 bodies

body 1
- body-level randomization
  - randomization 1 (default) has 2 geoms/meshes, 60% of worlds
  - randomization 2 has 3 geoms/meshes, 40% of worlds
body 2
- geom-level randomization
- 4 geoms/meshes, each for 25% of worlds

XML

<mujoco>
  <asset>
    <!-- body-level: object decomposition variants -->
    <!-- object A -->
    <mesh name="object_A_0" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
    <mesh name="object_A_1" vertex="1 0 0  2 0 0  1 1 0  1 0 1"/>
    <mesh name="object_A_2" vertex="0 1 0  1 1 0  0 2 0  0 1 1"/>
    <!-- object B -->
    <mesh name="object_B_0" vertex="0 0 0  3 0 0  0 3 0  0 0 3"/>
    <mesh name="object_B_1" vertex="3 0 0  6 0 0  3 3 0  3 0 3"/>

    <!-- geom-level: cube mesh candidates -->
    <mesh name="cube_s" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
    <mesh name="cube_m" vertex="0 0 0  2 0 0  0 2 0  0 0 2"/>
    <mesh name="cube_l" vertex="0 0 0  3 0 0  0 3 0  0 0 3"/>
    <mesh name="cube_xl" vertex="0 0 0  4 0 0  0 4 0  0 0 4"/>
  </asset>

  <worldbody>
    <!-- body-level randomization: convex decomposition -->
    <body name="object" pos="0 0 1">
      <freejoint/>
      <geom name="object_col_0" type="mesh" mesh="object_B_0"/>
      <geom name="object_col_1" type="mesh" mesh="object_B_1"/>
    </body>

    <!-- geom-level randomization: single geom, 4 mesh candidates -->
    <body pos="2 0 1">
      <freejoint/>
      <geom name="cube" type="mesh" mesh="cube_s"/>
    </body>
  </worldbody>

  <custom>
    <!-- body-level: object decomposition variants -->
    <tuple name="object_A">
      <element objtype="mesh" objname="object_A_0" prm="0"/>
      <element objtype="mesh" objname="object_A_1" prm="0"/>
      <element objtype="mesh" objname="object_A_2" prm="0"/>
    </tuple>
    <tuple name="object_B">
      <element objtype="mesh" objname="object_B_0" prm="0"/>
      <element objtype="mesh" objname="object_B_1" prm="0"/>
    </tuple>
    <tuple name="object">
      <element objtype="tuple" objname="object_A" prm="0.6"/>
      <element objtype="tuple" objname="object_B" prm="0.4"/>
    </tuple>

    <!-- geom-level: cube mesh candidates -->
    <tuple name="cube">
      <element objtype="mesh" objname="cube_s" prm="0.25"/>
      <element objtype="mesh" objname="cube_m" prm="0.25"/>
      <element objtype="mesh" objname="cube_l" prm="0.25"/>
      <element objtype="mesh" objname="cube_xl" prm="0.25"/>
    </tuple>
  </custom>
</mujoco>

Define via mjSpec

spec2 = mujoco.MjSpec()

# define all meshes
all_meshes = {
  "object_A_0": np.array([0,0,0, 1,0,0, 0,1,0, 0,0,1], dtype=np.float32),
  "object_A_1": np.array([1,0,0, 2,0,0, 1,1,0, 1,0,1], dtype=np.float32),
  "object_A_2": np.array([0,1,0, 1,1,0, 0,2,0, 0,1,1], dtype=np.float32),
  "object_B_0": np.array([0,0,0, 3,0,0, 0,3,0, 0,0,3], dtype=np.float32),
  "object_B_1": np.array([3,0,0, 6,0,0, 3,3,0, 3,0,3], dtype=np.float32),
  "cube_s":  np.array([0,0,0, 1,0,0, 0,1,0, 0,0,1], dtype=np.float32),
  "cube_m":  np.array([0,0,0, 2,0,0, 0,2,0, 0,0,2], dtype=np.float32),
  "cube_l":  np.array([0,0,0, 3,0,0, 0,3,0, 0,0,3], dtype=np.float32),
  "cube_xl": np.array([0,0,0, 4,0,0, 0,4,0, 0,0,4], dtype=np.float32),
}
for name, verts in all_meshes.items():
  m = spec2.add_mesh()
  m.name = name
  m.uservert = verts

# body-level: object with default 2-piece decomposition
body = spec2.worldbody.add_body()
body.name = "object"
body.pos = [0, 0, 1]
body.add_freejoint()
for i in range(2):
  g = body.add_geom()
  g.name = f"object_col_{i}"
  g.type = mujoco.mjtGeom.mjGEOM_MESH
  g.meshname = f"object_B_{i}"

# geom-level: cube with 1 geom
body2 = spec2.worldbody.add_body()
body2.pos = [2, 0, 1]
body2.add_freejoint()
g_cube = body2.add_geom()
g_cube.name = "cube"
g_cube.type = mujoco.mjtGeom.mjGEOM_MESH
g_cube.meshname = "cube_s"

# body-level tuples
t_a = spec2.add_tuple()
t_a.name = "object_A"
t_a.objtype = [MESH, MESH, MESH]
t_a.objname = ["object_A_0", "object_A_1", "object_A_2"]
t_a.objprm = [0, 0, 0]

t_b = spec2.add_tuple()
t_b.name = "object_B"
t_b.objtype = [MESH, MESH]
t_b.objname = ["object_B_0", "object_B_1"]
t_b.objprm = [0, 0]

t_body = spec2.add_tuple()
t_body.name = "object"
t_body.objtype = [TUPLE, TUPLE]
t_body.objname = ["object_A", "object_B"]
t_body.objprm = [0.6, 0.4]

# geom-level tuple
t_cube = spec2.add_tuple()
t_cube.name = "cube"
t_cube.objtype = [MESH, MESH, MESH, MESH]
t_cube.objname = ["cube_s", "cube_m", "cube_l", "cube_xl"]
t_cube.objprm = [0.25, 0.25, 0.25, 0.25]

nworld = 10
spec = mujoco.MjSpec.from_string(XML)
mjm = spec.compile()

m = mjwarp.put_model(mjm)
m = mjwarp.per_world_mesh(m, spec, nworld)

adenzler-nvidia · 2026-02-27T08:11:51Z

Thanks for sharing - will try to dry-run this in newton. I think the direction looks good, we are using MjSpec.

One thing that I'm failing to see on the first try is whether it's possible to control which world gets which randomization. I think that level of control would be useful, also for debugging.

erikfrey

high level design: i really like this approach. it may be worth syncing with @adenzler-nvidia on whether we would want to specify a mapping per world explicitly, vs what you're doing here - i assume it's like a cartesian product or something?

my only suggestion is let's find some kind of naming convention in the custom tuples so it's obviously namespaced for this purpose. we may want to use custom tuples for other purposes too.

adenzler-nvidia · 2026-03-10T12:33:13Z

Unless it's technically impossible, I would vote for having the ability to specify the mapping explicitly. You never know what people want to be doing with this, and it seems like it makes sense to control this if needed.

I didn't get to dry-run this yet because of release crunch but will do ASAP, hopefully still this week.

jvonmuralt · 2026-03-12T12:51:06Z

@thowell if I am not mistaken, I think some properties e.g. masses are not updated correctly in some scenarios, here is an example:

  def test_mesh_randomize_body_level_mass(self):
    """Body-level variant with fewer geoms must not include phantom mass.

    vA: 2 unit cubes (±1), mass 8000 each, total 16000.
    vB: 1 tiny cube (±0.1), mass 8.
    Default body uses vA (max variant). vB disables 1 geom slot.
    Disabled geom slots must not contribute to body_mass.
    """
    _F = "0 2 1  0 3 2  4 5 6  4 6 7  0 1 5  0 5 4  2 3 7  2 7 6  0 4 7  0 7 3  1 2 6  1 6 5"
    xml = f"""
    <mujoco>
      <asset>
        <mesh name="a0" vertex="-1 -1 -1  1 -1 -1  1 1 -1  -1 1 -1  -1 -1 1  1 -1 1  1 1 1  -1 1 1" face="{_F}"/>
        <mesh name="a1" vertex="-1 -1 -1  1 -1 -1  1 1 -1  -1 1 -1  -1 -1 1  1 -1 1  1 1 1  -1 1 1" face="{_F}"/>
        <mesh name="b0" vertex="-.1 -.1 -.1  .1 -.1 -.1  .1 .1 -.1  -.1 .1 -.1  -.1 -.1 .1  .1 -.1 .1  .1 .1 .1  -.1 .1 .1" face="{_F}"/>
      </asset>
      <worldbody>
        <body name="obj" pos="0 0 1">
          <freejoint/>
          <geom name="g0" type="mesh" mesh="a0"/>
          <geom name="g1" type="mesh" mesh="a1"/>
        </body>
      </worldbody>
      <custom>
        <tuple name="vA">
          <element objtype="mesh" objname="a0"/>
          <element objtype="mesh" objname="a1"/>
        </tuple>
        <tuple name="vB">
          <element objtype="mesh" objname="b0"/>
        </tuple>
        <tuple name="obj">
          <element objtype="tuple" objname="vA" prm="0.5"/>
          <element objtype="tuple" objname="vB" prm="0.5"/>
        </tuple>
      </custom>
    </mujoco>
    """
    xml_b = f"""
    <mujoco>
      <asset>
        <mesh name="b0" vertex="-.1 -.1 -.1  .1 -.1 -.1  .1 .1 -.1  -.1 .1 -.1  -.1 -.1 .1  .1 -.1 .1  .1 .1 .1  -.1 .1 .1" face="{_F}"/>
      </asset>
      <worldbody>
        <body name="obj" pos="0 0 1">
          <freejoint/>
          <geom type="mesh" mesh="b0"/>
        </body>
      </worldbody>
    </mujoco>
    """
    expected_B = mujoco.MjSpec.from_string(xml_b).compile().body_mass[1]

    nworld = 2
    spec = mujoco.MjSpec.from_string(xml)
    mjm = spec.compile()
    m = mjwarp.put_model(mjm)
    m = per_world_mesh(m, spec, nworld)

    body_mass = m.body_mass.numpy()
    dataid = m.geom_dataid.numpy()

    for w in range(nworld):
      if -1 in dataid[w]:
        np.testing.assert_allclose(
            body_mass[w, 1], expected_B, rtol=1e-4,
        )

thowell · 2026-03-12T13:15:49Z

@jvonmuralt thank you for reporting this issue and providing a repro!

the implementation was incorrect and should be fixed now. please let us know if the update does not resolve the reported issue, thanks!

omarrayyann · 2026-03-13T02:46:54Z

hello @thowell! I tried this somewhere and had these issues:

opt params are lost in per_world_mesh as it recompiles from spec. I had to adjust it to copy opt and reinstate it after. reprodce:

spec = mujoco.MjSpec.from_string(x)
mjm = spec.compile()

mjm.opt.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST      
m = mjwarp.put_model(mjm)
m = per_world_mesh(m, spec, nworld=1)

self.assertEqual(m.opt.integrator, mujoco.mjtIntegrator.mjINT_IMPLICITFAST)

also the fixes for preserving geom params wouldnt work if one is not unnamed per the g.name logic in _populate_dependent_fields i think. from @jvonmuralt's repro with a name removal from the g1 geom adjustment:

def test_mesh_randomize_body_level_mass(self):
  """Body-level variant with fewer geoms must not include phantom mass.

  vA: 2 unit cubes (±1), mass 8000 each, total 16000.
  vB: 1 tiny cube (±0.1), mass 8.
  Default body uses vA (max variant). vB disables 1 geom slot.
  Disabled geom slots must not contribute to body_mass.
  """
  _F = "0 2 1  0 3 2  4 5 6  4 6 7  0 1 5  0 5 4  2 3 7  2 7 6  0 4 7  0 7 3  1 2 6  1 6 5"
  xml = f"""
  <mujoco>
    <asset>
      <mesh name="a0" vertex="-1 -1 -1  1 -1 -1  1 1 -1  -1 1 -1  -1 -1 1  1 -1 1  1 1 1  -1 1 1" face="{_F}"/>
      <mesh name="a1" vertex="-1 -1 -1  1 -1 -1  1 1 -1  -1 1 -1  -1 -1 1  1 -1 1  1 1 1  -1 1 1" face="{_F}"/>
      <mesh name="b0" vertex="-.1 -.1 -.1  .1 -.1 -.1  .1 .1 -.1  -.1 .1 -.1  -.1 -.1 .1  .1 -.1 .1  .1 .1 .1  -.1 .1 .1" face="{_F}"/>
    </asset>
    <worldbody>
      <body name="obj" pos="0 0 1">
        <freejoint/>
        <geom type="mesh" mesh="a0"/>
        <geom name="g1" type="mesh" mesh="a1"/>
      </body>
    </worldbody>
    <custom>
      <tuple name="vA">
        <element objtype="mesh" objname="a0"/>
        <element objtype="mesh" objname="a1"/>
      </tuple>
      <tuple name="vB">
        <element objtype="mesh" objname="b0"/>
      </tuple>
      <tuple name="obj">
        <element objtype="tuple" objname="vA" prm="0.5"/>
        <element objtype="tuple" objname="vB" prm="0.5"/>
      </tuple>
    </custom>
  </mujoco>
  """
  xml_b = f"""
  <mujoco>
    <asset>
      <mesh name="b0" vertex="-.1 -.1 -.1  .1 -.1 -.1  .1 .1 -.1  -.1 .1 -.1  -.1 -.1 .1  .1 -.1 .1  .1 .1 .1  -.1 .1 .1" face="{_F}"/>
    </asset>
    <worldbody>
      <body name="obj" pos="0 0 1">
        <freejoint/>
        <geom type="mesh" mesh="b0"/>
      </body>
    </worldbody>
  </mujoco>
  """
  expected_B = mujoco.MjSpec.from_string(xml_b).compile().body_mass[1]

  nworld = 2
  spec = mujoco.MjSpec.from_string(xml)
  mjm = spec.compile()
  m = mjwarp.put_model(mjm)
  m = per_world_mesh(m, spec, nworld)

  body_mass = m.body_mass.numpy()
  dataid = m.geom_dataid.numpy()

  for w in range(nworld):
    if -1 in dataid[w]:
      np.testing.assert_allclose(
          body_mass[w, 1], expected_B, rtol=1e-4,
      )

thanks!

thowell · 2026-03-13T13:58:16Z

@omarrayyann thanks for reporting these issues! the per_world_mesh implementation has been updated:

the function now accepts only spec and nworld arguments
the implementation is updated to account for geoms that do not have names

spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST can be used to override the integrator option with the implictfast integrator prior to calling per_world_mesh.

please let us know if the updated implementation does not resolve the reported issues. thanks!

thowell · 2026-03-17T11:36:58Z

update io_test.per_world_mesh to create a copy of spec
updates for the renderer, @StafaH please take a look at this commit 279935d

StafaH · 2026-03-27T20:45:17Z

LGTM. I did a quick check and the renderer works for the multimesh example (50/50 split between two meshes).

StafaH · 2026-03-27T20:52:39Z

  geom_enabled_idx = np.nonzero(geom_enabled_mask)[0]

  mesh_registry = {}
  mesh_bvh_id = [wp.uint64(0) for _ in range(nmesh)]


This was originally placed here to avoid building the mesh BVHs for collision meshes when rendering visual geoms, and vice versa. I think that is still relevant here.

thowell · 2026-03-31T08:39:57Z

@adenzler-nvidia are we good to move forward with this pr? thanks!

adenzler-nvidia · 2026-04-01T06:37:36Z

sorry for the late reply - there is nothing from our side blocking this, we can use this well. That being said, for us it would make things much simpler to just get some of the mesh preprocessing functions exposed such that we could build the mesh array ourselves instead of going through put_model, and then assigning the dataids directly.

This can/should be a follow-up though.

thowell · 2026-04-01T11:33:05Z

@adenzler-nvidia the mesh preprocessing should all be available via mjSpec. mesh assets can be added and then after calling compile the processed mesh fields are available in the returned mjModel instance.

thowell force-pushed the per_world_meshes branch from e6475dd to a2c539c Compare February 27, 2026 10:31

thowell marked this pull request as ready for review March 9, 2026 18:28

thowell force-pushed the per_world_meshes branch from a2c539c to 790b310 Compare March 9, 2026 19:44

erikfrey approved these changes Mar 10, 2026

View reviewed changes

thowell force-pushed the per_world_meshes branch from 790b310 to 41cd042 Compare March 11, 2026 14:26

per-world meshes

25de185

thowell force-pushed the per_world_meshes branch from 41cd042 to 25de185 Compare March 11, 2026 18:49

Merge remote-tracking branch 'upstream/main' into per_world_meshes

f243ee7

geom mass fixes

e1b9d26

thowell force-pushed the per_world_meshes branch from 61d3988 to e1b9d26 Compare March 13, 2026 13:53

andrewrosemberg mentioned this pull request Mar 14, 2026

[WIP] Integrate multi-env mujoco-warp per-world geom_dataid feature mujocolab/mjlab#799

Closed

thowell added 2 commits March 17, 2026 11:02

create a copy of spec in per_world_mesh

54fb4a9

updates for renderer

279935d

StafaH reviewed Mar 27, 2026

View reviewed changes

kevinzakka mentioned this pull request Mar 29, 2026

Add per-world mesh variants mujocolab/mjlab#860

Merged

Merge remote-tracking branch 'upstream/main' into per_world_meshes

234726f

update per_world_mesh

d9385ab

thowell merged commit 368bdca into google-deepmind:main Apr 1, 2026
10 checks passed

thowell linked an issue Apr 1, 2026 that may be closed by this pull request

Simulating different meshes in parallel envs #738

Closed

thowell mentioned this pull request Apr 1, 2026

mujoco warp documentation: per-world meshes google-deepmind/mujoco#3165

Merged

julien-blanchon mentioned this pull request Apr 17, 2026

Simulating different meshes in parallel envs #738

Closed

Conversation

thowell commented Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

XML

Define via mjSpec

Uh oh!

adenzler-nvidia commented Feb 27, 2026

Uh oh!

erikfrey left a comment

Choose a reason for hiding this comment

Uh oh!

adenzler-nvidia commented Mar 10, 2026

Uh oh!

jvonmuralt commented Mar 12, 2026

Uh oh!

thowell commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

omarrayyann commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

thowell commented Mar 13, 2026

Uh oh!

thowell commented Mar 17, 2026

Uh oh!

StafaH commented Mar 27, 2026

Uh oh!

StafaH Mar 27, 2026

Choose a reason for hiding this comment

Uh oh!

thowell commented Mar 31, 2026

Uh oh!

adenzler-nvidia commented Apr 1, 2026

Uh oh!

thowell commented Apr 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

thowell commented Feb 26, 2026 •

edited

Loading

thowell commented Mar 12, 2026 •

edited

Loading

omarrayyann commented Mar 13, 2026 •

edited

Loading