Skip to content

mujoco warp documentation: per-world meshes #3165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
thowell wants to merge 1 commit into
base: main
Choose a base branch
from
+281 −3
Draft

mujoco warp documentation: per-world meshes #3165

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
284 changes: 281 additions & 3 deletions doc/mjwarp/index.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -460,10 +460,288 @@ Certain fields are safe to modify directly without compilation, enabling on-devi
`GitHub issue 893 <https://github.com/google-deepmind/mujoco_warp/issues/893>`__ tracks adding on-device updates for a
subset of fields.

.. admonition:: Heterogeneous worlds
:class: note
Per-world meshes
----------------

Per-world meshes enable heterogeneous worlds where different worlds simulate different meshes. The workflow
is:

1. Create an :ref:`mjSpec` with **all** mesh assets and the **maximum** number of geom slots needed across variants.
2. Compile each variant by mutating the spec and calling ``spec.compile()``.
3. Compile a **base** model and create :class:`mjw.Model <mujoco_warp.Model>` from it.
4. Override the relevant :class:`mjw.Model <mujoco_warp.Model>` fields with per-world arrays built from the compiled
variants.

**Example 1 — Geom-level** randomization (1 body, 1 geom, 2 mesh assets):

The base scene includes all mesh assets. The geom references one mesh (``mesh_a``); a second mesh
(``mesh_b``) is available for per-world substitution.

.. code-block:: xml

<mujoco>
<asset>
<mesh name="mesh_a" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="mesh_b" vertex="0 0 0 2 0 0 0 2 0 0 0 2"/>
</asset>
<worldbody>
<body pos="0 0 1">
<freejoint/>
<geom name="obj" type="mesh" mesh="mesh_a"/>
</body>
</worldbody>
</mujoco>

.. code-block:: python

nworld = 4

# base spec: 1 body with 1 mesh geom, all mesh assets
spec = mujoco.MjSpec()
mesh_a = spec.add_mesh()
mesh_a.name = "mesh_a"
mesh_a.uservert = [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]

mesh_b = spec.add_mesh()
mesh_b.name = "mesh_b"
mesh_b.uservert = [0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 2]

body = spec.worldbody.add_body()
body.pos = [0, 0, 1]
body.add_freejoint()
geom = body.add_geom()
geom.name = "obj"
geom.type = mujoco.mjtGeom.mjGEOM_MESH
geom.meshname = "mesh_a"

# compile each variant
geom.meshname = "mesh_a"
mjm_a = spec.compile()
geom.meshname = "mesh_b"
mjm_b = spec.compile()

# restore and compile base
geom.meshname = "mesh_a"
mjm = spec.compile()

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)

# build per-world arrays: worlds 0-1 use mesh_a, worlds 2-3 use mesh_b
geom_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj")
variants = [mjm_a, mjm_b]
assignment = [0, 0, 1, 1] # variant index per world

# build per-world arrays
dataid = np.tile(mjm.geom_dataid, (nworld, 1))
geom_size = np.zeros((nworld, mjm.ngeom, 3))
geom_aabb = np.zeros((nworld, mjm.ngeom, 2, 3))
geom_rbound = np.zeros((nworld, mjm.ngeom))
geom_pos = np.zeros((nworld, mjm.ngeom, 3))
body_mass = np.zeros((nworld, mjm.nbody))
body_subtreemass = np.zeros((nworld, mjm.nbody))
body_inertia = np.zeros((nworld, mjm.nbody, 3))
body_invweight0 = np.zeros((nworld, mjm.nbody, 2))
body_ipos = np.zeros((nworld, mjm.nbody, 3))
body_iquat = np.zeros((nworld, mjm.nbody, 4))

for w in range(nworld):
ref = variants[assignment[w]]
dataid[w, geom_id] = ref.geom_dataid[geom_id]
geom_size[w] = ref.geom_size
geom_aabb[w] = ref.geom_aabb.reshape(mjm.ngeom, 2, 3)
geom_rbound[w] = ref.geom_rbound
geom_pos[w] = ref.geom_pos
body_mass[w] = ref.body_mass
body_subtreemass[w] = ref.body_subtreemass
body_inertia[w] = ref.body_inertia
body_invweight0[w] = ref.body_invweight0
body_ipos[w] = ref.body_ipos
body_iquat[w] = ref.body_iquat

m.geom_dataid = wp.array(dataid, dtype=int)
m.geom_size = wp.array(geom_size, dtype=wp.vec3)
m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
m.geom_rbound = wp.array(geom_rbound, dtype=float)
m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
m.body_mass = wp.array(body_mass, dtype=float)
m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
m.body_iquat = wp.array(body_iquat, dtype=wp.quat)

**Example 2 — Body-level** randomization (1 body, 1 or 2 geoms, 3 mesh assets):

.. admonition:: Maximum geom count
:class: important

For body-level randomization, the base ``mjModel`` provided to ``mjw.put_model`` should specify the **maximum number
of geoms** required across all variants. Geom slots that are unused in a particular variant can be disabled
(e.g., ``contype=0``, ``conaffinity=0``, ``dataid=-1``), but they should still be present as part of the body in the
base model.

.. code-block:: xml

<mujoco>
<asset>
<mesh name="mA" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="mB" vertex="0 0 0 2 0 0 0 2 0 0 0 2"/>
<mesh name="mC" vertex="0 0 0 3 0 0 0 3 0 0 0 3"/>
</asset>
<worldbody>
<body name="obj" pos="0 0 1">
<freejoint/>
<geom name="obj_0" type="mesh" mesh="mA"/>
<geom name="obj_1" size=".001" contype="0" conaffinity="0" mass="0"/>
</body>
</worldbody>
</mujoco>

.. code-block:: python

nworld = 6

# base spec: body with 2 geom slots (max across variants), all mesh assets
spec = mujoco.MjSpec()
for name, scale in [("mA", 1), ("mB", 2), ("mC", 3)]:
mesh = spec.add_mesh()
mesh.name = name
mesh.uservert = [0, 0, 0, scale, 0, 0, 0, scale, 0, 0, 0, scale]

body = spec.worldbody.add_body()
body.name = "obj"
body.pos = [0, 0, 1]
body.add_freejoint()

g0 = body.add_geom()
g0.name = "obj_0"
g0.type = mujoco.mjtGeom.mjGEOM_MESH
g0.meshname = "mA"

# null geom slot: disabled collision, no mesh
g1 = body.add_geom()
g1.name = "obj_1"
g1.size = [0.001, 0, 0]
g1.contype = 0
g1.conaffinity = 0
g1.mass = 0

# variant A: 1 geom (mesh mA), g1 stays null
mjm_a = spec.compile()

# variant B: 2 geoms (mesh mB + mC)
g0.meshname = "mB"
g1.type = mujoco.mjtGeom.mjGEOM_MESH
g1.meshname = "mC"
g1.contype = 1
g1.conaffinity = 1
mjm_b = spec.compile()

# restore base and compile
g0.meshname = "mA"
g1.type = mujoco.mjtGeom.mjGEOM_SPHERE
g1.contype = 0
g1.conaffinity = 0
mjm = spec.compile()

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)

# worlds 0-2: variant A (1 active geom), worlds 3-5: variant B (2 active geoms)
variants = [mjm_a, mjm_b]
assignment = [0, 0, 0, 1, 1, 1]

geom0_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj_0")
geom1_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj_1")
body_id = mjm.geom_bodyid[geom0_id]

# build per-world arrays
dataid = np.tile(mjm.geom_dataid, (nworld, 1))
geom_size = np.zeros((nworld, mjm.ngeom, 3))
geom_rbound = np.zeros((nworld, mjm.ngeom))
geom_aabb = np.zeros((nworld, mjm.ngeom, 2, 3))
geom_pos = np.zeros((nworld, mjm.ngeom, 3))
body_mass = np.zeros((nworld, mjm.nbody))
body_subtreemass = np.zeros((nworld, mjm.nbody))
body_inertia = np.zeros((nworld, mjm.nbody, 3))
body_invweight0 = np.zeros((nworld, mjm.nbody, 2))
body_ipos = np.zeros((nworld, mjm.nbody, 3))
body_iquat = np.zeros((nworld, mjm.nbody, 4))

for w in range(nworld):
ref = variants[assignment[w]]
dataid[w] = ref.geom_dataid
# disable unused geom slot for variant A
if assignment[w] == 0:
dataid[w, geom1_id] = -1
geom_size[w] = ref.geom_size
geom_rbound[w] = ref.geom_rbound
geom_aabb[w] = ref.geom_aabb.reshape(mjm.ngeom, 2, 3)
geom_pos[w] = ref.geom_pos
body_mass[w] = ref.body_mass
body_subtreemass[w] = ref.body_subtreemass
body_inertia[w] = ref.body_inertia
body_invweight0[w] = ref.body_invweight0
body_ipos[w] = ref.body_ipos
body_iquat[w] = ref.body_iquat

m.geom_dataid = wp.array(dataid, dtype=int)
m.geom_size = wp.array(geom_size, dtype=wp.vec3)
m.geom_rbound = wp.array(geom_rbound, dtype=float)
m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
m.body_mass = wp.array(body_mass, dtype=float)
m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
m.body_iquat = wp.array(body_iquat, dtype=wp.quat)

**Batched fields** — fields that must be overridden for per-world meshes:

.. list-table::
:width: 90%
:align: left
:widths: 3 2 3
:header-rows: 1

Heterogeneous worlds, for example: per-world meshes or number of degrees of freedom, are not currently available.
* - Field
- dtype
- Shape
* - ``geom_dataid``
- ``int``
- ``(nworld, ngeom)``
* - ``geom_size``
- ``wp.vec3``
- ``(nworld, ngeom)``
* - ``geom_aabb``
- ``wp.vec3``
- ``(nworld, ngeom, 2)``
* - ``geom_rbound``
- ``float``
- ``(nworld, ngeom)``
* - ``geom_pos``
- ``wp.vec3``
- ``(nworld, ngeom)``
* - ``body_mass``
- ``float``
- ``(nworld, nbody)``
* - ``body_subtreemass``
- ``float``
- ``(nworld, nbody)``
* - ``body_inertia``
- ``wp.vec3``
- ``(nworld, nbody)``
* - ``body_invweight0``
- ``wp.vec2``
- ``(nworld, nbody)``
* - ``body_ipos``
- ``wp.vec3``
- ``(nworld, nbody)``
* - ``body_iquat``
- ``wp.quat``
- ``(nworld, nbody)``

Batch Rendering
===============
Expand Down
Loading