From 00707ddb89392211a9128ca483c70c985eedccca Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 11 Mar 2026 13:15:08 +0000 Subject: [PATCH] mujoco warp documentation: per-world meshes --- doc/mjwarp/index.rst | 284 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 281 insertions(+), 3 deletions(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 85ab28ce..53386a11 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -460,10 +460,288 @@ Certain fields are safe to modify directly without compilation, enabling on-devi `GitHub issue 893 `__ tracks adding on-device updates for a subset of fields. -.. admonition:: Heterogeneous worlds - :class: note +Per-world meshes +---------------- - Heterogeneous worlds, for example: per-world meshes or number of degrees of freedom, are not currently available. +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 ` from it. +4. Override the relevant :class:`mjw.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 + + + + + + + + + + + + + + +.. 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 + + + + + + + + + + + + + + + + +.. 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 + + * - 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 ===============