Skip to content

heterogeneous meshes#1009

Draft
thowell wants to merge 1 commit intogoogle-deepmind:mainfrom
thowell:mesh
Draft

heterogeneous meshes#1009
thowell wants to merge 1 commit intogoogle-deepmind:mainfrom
thowell:mesh

Conversation

@thowell
Copy link
Collaborator

@thowell thowell commented Jan 12, 2026

Heterogeneous Mesh Batching

Enables per-world mesh geometries.

This PR adds support for heterogeneous mesh batching - the ability to simulate different mesh tessellations, sizes, and topologies across parallel worlds. Previously, all worlds shared identical mesh geometry. Now each world can have its own unique mesh configuration while maintaining efficient batched execution.

nmesh, nbody, ngeom, and nv are required to be the same for all worlds

Key capabilities:

  • Variable vertex counts - meshes can have different numbers of vertices per world
  • Different tessellations - coarse vs fine mesh approximations in parallel
  • Per-world physical properties - geom bounds, body mass/inertia derived from each mesh

mesh_batch Utility Function

A utility function mesh_batch is added to construct batched mesh arrays from a collection of mjModel instances. Importantly, dependent fields are also updated.

def mesh_batch(m: types.Model, mjms: Sequence[mujoco.MjModel]) -> types.Model

Takes a Model and a sequence of MuJoCo MjModels (one per world) and creates batched arrays for all mesh-dependent fields:

Category Fields
Mesh data mesh_vertadr, mesh_vertnum, mesh_faceadr, mesh_vert, mesh_face, mesh_graph, mesh_quat, mesh_polynum, mesh_polyadr, mesh_polynormal, ...
Geom bounds geom_size, geom_aabb, geom_rbound
Body mass/inertia body_mass, body_subtreemass, body_inertia, body_invweight0, body_ipos, body_iquat
DOF dof_invweight0

Example Usage

import mujoco
import mujoco_warp as mjw

# Create MuJoCo models with different mesh tessellation levels
mjm0 = mujoco.MjModel.from_xml_string("""
<mujoco>
  <asset>
    <mesh name="mesh1" builtin="sphere" params="0"/>
  </asset>
  <worldbody>
    <body>
      <freejoint/>
      <geom type="mesh" mesh="mesh1"/>
    </body>
  </worldbody>
</mujoco>
""")

mjm1 = mujoco.MjModel.from_xml_string("""
<mujoco>
  <asset>
    <mesh name="mesh1" builtin="sphere" params="1"/>
  </asset>
  <worldbody>
    <body>
      <freejoint/>
      <geom type="mesh" mesh="mesh1"/>
    </body>
  </worldbody>
</mujoco>
""")

mjm2 = mujoco.MjModel.from_xml_string("""
<mujoco>
  <asset>
    <mesh name="mesh1" builtin="sphere" params="2"/>
  </asset>
  <worldbody>
    <body>
      <freejoint/>
      <geom type="mesh" mesh="mesh1"/>
    </body>
  </worldbody>
</mujoco>
""")

# Create mjw.Model and mjw.Data with 3 worlds
_, _, m, d = mjw.put_model(mjm0, nworld=3)

# Apply mesh batching - each world gets unique mesh
m = mjw.mesh_batch(m, [mjm0, mjm1, mjm2])

# Run collision detection - uses per-world mesh geometry
mjw.collision(m, d)

@hartikainen
Copy link
Contributor

Is there a plan to support this in MJX?

@btaba
Copy link
Collaborator

btaba commented Jan 13, 2026

Hi @hartikainen , AFAIK there aren't immediate plans to support this in MJX. I believe you'd have to hook into _put_model_warp to get similar behavior; MJX-Warp is mostly in a KTLO state. We want to add additional graph modes for MJX-Warp, and most likely OSS the render shim, but that's all we have slated right now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Heterogeneous model setups support Simulating different meshes in parallel envs

3 participants