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.

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

per-world meshes#1191

per-world meshes#1191
thowell wants to merge 2 commits intogoogle-deepmind:mainfrom
thowell:per_world_meshes

thowell commented Feb 26, 2026 •

edited

Loading

Uh oh!

adenzler-nvidia commented Feb 27, 2026

Uh oh!

erikfrey left a comment

Uh oh!

adenzler-nvidia commented Mar 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

thowell commented Feb 26, 2026 •

edited

Loading