Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

docs/source/entity/per_world_mesh.rst

@@ -0,0 +1,490 @@
+.. _heterogeneous_worlds:
+
+Heterogeneous Worlds
+====================
+
+mjlab can run a single batched simulation in which different parallel
+worlds use different mesh assets for the same logical entity. World 0
+may simulate a cube, world 1 a sphere, world 2 a bowl. All worlds
+share the same compiled scene and the same body and joint structure;
+only the meshes and the per-geom attributes that travel with them
+(friction, contact bits, mass, density, and a few more) differ across
+worlds. Articulated props work too (you can have a hinge or slide
+below the variant's root), as long as the joint topology matches
+across variants. The feature is exposed through ``VariantEntityCfg``.
+The full breakdown of what can and cannot vary across variants is in
+the next section.
+
+
+Quickstart
+----------
+
+Say you want some parallel worlds to hold a sphere and others to hold
+a cone, with a single shared scene running both at once. Define each
+variant as a function that returns an ``MjSpec``, then group them
+under one ``VariantEntityCfg``:
+
+.. code-block:: python
+
+    import mujoco
+
+    from mjlab.entity import EntityCfg, VariantEntityCfg
+
+
+    def make_sphere_spec() -> mujoco.MjSpec:
+        spec = mujoco.MjSpec()
+        mesh = spec.add_mesh(name="visual")
+        mesh.make_sphere(subdivision=3)
+        mesh.scale[:] = (0.05,) * 3
+        body = spec.worldbody.add_body(name="prop")
+        body.add_freejoint()
+        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
+        return spec
+
+
+    def make_cone_spec() -> mujoco.MjSpec:
+        spec = mujoco.MjSpec()
+        mesh = spec.add_mesh(name="visual")
+        mesh.make_cone(nedge=16, radius=0.04)
+        body = spec.worldbody.add_body(name="prop")
+        body.add_freejoint()
+        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
+        return spec
+
+
+    object_cfg = VariantEntityCfg(
+        variants={
+            "sphere": make_sphere_spec,
+            "cone":   make_cone_spec,
+        },
+        assignment={"cone": 2.0},  # twice as many cones as spheres
+        init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
+    )
+
+Plug the variant entity into a :ref:`scene` exactly like a regular
+``EntityCfg``:
+
+.. code-block:: python
+
+    from mjlab.scene import SceneCfg
+
+    scene_cfg = SceneCfg(
+        num_envs=4096,
+        entities={"object": object_cfg},
+    )
+
+Twice as many worlds will hold a cone as a sphere. Variants not listed
+in the ``assignment`` dict default to weight 1.0; omit ``assignment``
+entirely for uniform allocation across all variants.
+
+
+What variants can differ in
+---------------------------
+
+**Free to vary across variants:** the mesh asset assigned to each
+slot, the number of mesh geoms per ``(body, role)`` bucket on the
+variant body (one variant can have more collision meshes than
+another), the per-mesh-geom attributes that travel with the mesh
+(friction, contact bits, mass, density, ``condim``, and a handful of
+others), and explicit body inertial values within whichever single
+inertial mode the variants agree on per body.
+
+**Must match across variants:** the body tree, joint topology,
+primitive (non-mesh) geoms, and any actuators / sensors / tendons /
+equalities. Variants must also agree on the inertial representation
+per body (mesh-derived, diagonal, or fullinertia), and may not use the
+reserved ``mjlab/pad/`` name prefix on any element. Variant entities
+must also be floating-base: the root body declares a freejoint.
+
+The validator runs at entity build time and raises ``ValueError``
+naming the offending variant and the exact mismatch.
+
+
+How variants are assembled
+--------------------------
+
+mjlab merges every variant's mesh assets into a single ``MjSpec`` and
+gives the variant body enough mesh-geom *slots* to cover the maximum
+mesh count any variant uses for each ``(body, role)`` bucket. A slot
+is identified by ``(body_path, role, ordinal)``. ``role`` is "visual"
+or "collision", derived from ``contype``/``conaffinity``;
+mujoco_warp's ``geom_contype``/``geom_conaffinity`` are 1D shared
+(not per-world), so a slot's role is fixed across worlds by
+construction.
+
+A worked example
+~~~~~~~~~~~~~~~~
+
+Say variant ``sphere`` has 1 visual mesh geom and 2 collision mesh
+geoms on the prop body, and variant ``cone`` has 1 visual mesh geom
+and 4 collision mesh geoms on the same body.
+
+.. code-block:: text
+
+    sphere variant body              cone variant body
+    -------------------              -------------------
+    prop body                        prop body
+      [visual] sphere_vis              [visual] cone_vis
+      [coll]   sphere_col_0            [coll]   cone_col_0
+      [coll]   sphere_col_1            [coll]   cone_col_1
+                                       [coll]   cone_col_2
+                                       [coll]   cone_col_3
+
+mjlab walks each variant's body tree, buckets mesh geoms by
+``(body_path, role)``, and lays the union out as slots:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 8 18 8 12 27 27
+
+   * - Slot
+     - body_path
+     - role
+     - ordinal
+     - sphere fills with
+     - cone fills with
+   * - 0
+     - /prop
+     - visual
+     - 0
+     - sphere_vis
+     - cone_vis
+   * - 1
+     - /prop
+     - collision
+     - 0
+     - sphere_col_0
+     - cone_col_0
+   * - 2
+     - /prop
+     - collision
+     - 1
+     - sphere_col_1
+     - cone_col_1
+   * - 3
+     - /prop
+     - collision
+     - 2
+     - *(unfilled)*
+     - cone_col_2
+   * - 4
+     - /prop
+     - collision
+     - 3
+     - *(unfilled)*
+     - cone_col_3
+
+Five slots total. The merged scene's prop body has five mesh geoms:
+slot 0 plus four collision slots (the union of sphere's two and
+cone's four). At merge time, every variant's mesh asset is added to
+the merged spec under a unique name (e.g.
+``sphere/sphere_vis``, ``cone/cone_col_2``).
+
+The merged scene compiles once into a single canonical ``MjModel``
+that every world in the batch agrees on layout-wise: same nbody,
+ngeom, same body and geom IDs. mjlab's per-world overrides on top of
+that one model are what make worlds heterogeneous.
+
+What each world sees at runtime
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Worlds where ``sphere`` is active see only its three meshes; the two
+extra collision slots are disabled via per-world ``geom_dataid = -1``,
+and mujoco_warp skips them. Worlds where ``cone`` is active see all
+five meshes wired up.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 14 14 12 12 12 12 12
+
+   * - World
+     - variant
+     - slot 0
+     - slot 1
+     - slot 2
+     - slot 3
+     - slot 4
+   * - 0
+     - sphere
+     - sphere_vis
+     - sphere_col_0
+     - sphere_col_1
+     - **off (-1)**
+     - **off (-1)**
+   * - 1
+     - cone
+     - cone_vis
+     - cone_col_0
+     - cone_col_1
+     - cone_col_2
+     - cone_col_3
+
+Three categories of per-world override carry the variation:
+
+* **geom_dataid** is a ``(num_envs, ngeom)`` table. Its row for
+  world W picks which compiled mesh each slot points at. ``-1`` is
+  the "skip me" sentinel mujoco_warp already understands.
+* **Mesh-derived fields** (``geom_size``, ``geom_rbound``,
+  ``geom_aabb``, ``geom_pos``, ``geom_quat``, ``body_mass``,
+  ``body_subtreemass``, ``body_inertia``, ``body_invweight0``,
+  ``body_ipos``, ``body_iquat``) are stored as ``(num_envs, ...)``
+  arrays. The values for sphere worlds reflect a sphere-shaped
+  inertia tensor and sphere-sized AABBs; the values for cone worlds
+  reflect the cone. The full list is in
+  ``mjlab.entity.variants.VARIANT_DEPENDENT_FIELDS``.
+* **Per-mesh-geom attributes** (contact bits, friction, mass,
+  density, condim, group, priority, rgba, solref, solimp, margin,
+  gap) are captured per variant in ``VariantGeomSpec`` at merge time
+  and restored verbatim on the slot geom during the per-variant
+  reference compile. So if sphere's collision geoms have
+  ``friction=0.5`` and cone's have ``friction=1.2``, world W's
+  per-step friction reflects the assigned variant's source value.
+  The one exception is ``material``, which is not propagated across
+  variants; if you need per-world appearance variation use DR on
+  ``geom_rgba`` / ``mat_rgba``.
+
+If ``sphere`` adds a body that ``cone`` lacks (or vice versa), the
+validator rejects the configuration before any of the merge logic
+runs. The slot mechanism only flexes mesh geom counts within
+matching bodies; everything structural above the geom level must
+agree.
+
+.. note::
+
+   **Doesn't compiling the merged scene ruin the prop body's
+   inertia?**
+
+   No, but it's worth understanding why, because the naive intuition
+   says it should. If you stuck every variant's mesh geoms on the
+   prop body and called ``spec.compile()``, MuJoCo would sum each
+   geom's inertial contribution, and you would get a body whose mass
+   and inertia tensor are a meaningless mix of every variant's shape.
+
+   mjlab avoids this in two layers:
+
+   * **The merged scene does not stick every variant's geoms on the
+     body.** The prop body in the merged spec carries variant 0's
+     mesh geoms (with their original mass and density) plus, for any
+     slot variant 0 doesn't fill, a synthesized padding geom that has
+     ``mass = 0`` and ``density = 0``. Padding contributes nothing to
+     body inertia. Other variants' meshes are present in the merged
+     spec only as **mesh assets** (in the assets section, not as geoms
+     on any body). They get wired in at runtime via per-world
+     ``geom_dataid`` and never affect the host compile's inertial
+     sums.
+   * **Per-world overrides come from per-variant source compiles.**
+     Even with the above, the merged-scene compile's prop body inertia
+     is only correct for variant 0. For every other variant, mjlab
+     compiles that variant's original source spec in isolation (one
+     body, one variant's worth of meshes), reads the resulting
+     ``body_mass``, ``body_inertia``, ``body_ipos``, ``body_iquat``,
+     ``body_invweight0``, and ``body_subtreemass``, and writes them
+     into the per-world arrays at the prop body's index.
+
+   Net result: world W's prop body inertia is byte-equal to what you
+   would get by compiling variant W's source spec on its own. There
+   is a regression test
+   (``test_visual_collision_split_inertia_matches_independent_compile``
+   in ``tests/test_variants.py``) that asserts exactly this against
+   independent per-variant compiles.
+
+
+World assignment
+----------------
+
+How worlds get mapped to variants is controlled by the ``assignment``
+field on ``VariantEntityCfg``. It accepts three shapes:
+
+* ``None`` (default): uniform allocation across variants.
+* ``dict[str, float]``: per-variant weights. Variants not listed
+  default to weight 1.0.
+* ``Callable[[int], Sequence[int]]``: an explicit assignment function
+  called with ``num_envs`` at simulation init.
+
+Both the ``None`` and dict cases use the
+`largest remainder method
+<https://en.wikipedia.org/wiki/Largest_remainder_method>`_. Each
+variant's quota is ``q_i = (w_i / sum(w)) * num_envs``; each variant
+first receives ``floor(q_i)`` worlds, and the remaining
+``num_envs - sum(floors)`` worlds go to the variants with the largest
+fractional remainders, with ties broken by declaration order. For
+``num_envs = 10`` and weights ``(1.0, 2.0, 1.0)`` this gives
+``(3, 5, 2)`` worlds per variant. Weights are normalized internally,
+so ``{"a": 1, "b": 2, "c": 1}`` and ``{"a": 0.25, "b": 0.5, "c": 0.25}``
+produce identical assignments. A weight of zero is allowed and
+produces zero worlds for that variant; at least one variant must end
+up with positive weight.
+
+The default and dict paths are purely deterministic given
+``(assignment, num_envs)``. With ``assignment={"a": 1, "b": 1}`` and
+``num_envs = 8`` you always get ``[0, 0, 0, 0, 1, 1, 1, 1]``. There is
+no seed involved; rerunning the same config produces the same
+partition every time. Note that the partition's *boundaries* depend
+on ``num_envs``, so world W's variant is not necessarily stable when
+you change ``num_envs``. If you need explicit per-world stability
+across batch sizes (e.g. "world 0 is always variant 0, world 1 is
+always variant 1, regardless of how many envs I launch"), use a
+callable assignment as below.
+
+Variant assignment is fixed at ``Simulation`` initialization and does
+not resample on episode reset. The intended use is heterogeneous
+training across the batch, not per-episode mesh randomization.
+
+Read the resolved assignment from user code via
+``env.sim.world_to_variant``:
+
+.. code-block:: python
+
+    >>> env.sim.world_to_variant["object"]
+    tensor([0, 0, 0, 1, 1, 1, 1, 1, 1, 1])
+
+The mapping is keyed by entity name (without trailing slash) and
+returns a ``(num_envs,)`` tensor of variant indices in the order
+variants were declared in ``VariantEntityCfg.variants``. The dict is
+empty for non-variant scenes.
+
+
+Custom assignment with a callable
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When the weighted default is not what you want, pass a callable to
+``assignment``. The callable receives ``num_envs`` and must return a
+length-``num_envs`` sequence of variant indices in
+``[0, len(variants))``. The returned sequence's length and bounds are
+validated at sim init; mismatches raise a ``ValueError`` naming the
+offending entity.
+
+A few patterns:
+
+**Round-robin** - cycle through variants by world index.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"a": make_a, "b": make_b, "c": make_c},
+        assignment=lambda n: [w % 3 for w in range(n)],
+    )
+
+**Stratified halves** - first half is variant 0, second half is
+variant 1.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"easy": make_easy, "hard": make_hard},
+        assignment=lambda n: [0] * (n // 2) + [1] * (n - n // 2),
+    )
+
+Domain randomization
+--------------------
+
+Domain randomization on variant scenes preserves per-variant baselines
+automatically. When the simulation initializes, mjlab snapshots the
+variant-dependent fields as ``(num_envs, ...)`` tensors and registers
+them in ``sim.per_world_default_fields``. DR operations that read
+defaults (scale, additive offsets) detect this registration and index
+the per-world default array by environment, so a 10% mass scale
+applied across a batch containing a 100 g sphere variant and a 1 kg
+cube variant produces 10% perturbations *around each variant's own
+mass*, not 10% of a shared template mass.
+
+Fields that are not variant-dependent (``geom_friction``,
+``dof_armature``, ``dof_damping``, and so on) behave identically on
+variant and non-variant scenes.
+
+For inertial randomization the recommended path is
+``dr.pseudo_inertia``, which jointly randomizes mass, COM offset,
+principal moments of inertia, and principal frame orientation through
+the pseudo-inertia matrix factorization of `Rucker and Wensing (2022)
+<https://par.nsf.gov/servlets/purl/10347458>`_. It is exact for any
+perturbation magnitude and remains physically consistent across
+variants of different scale. ``dr.body_mass`` modifies ``body_mass``
+without touching the inertia tensor and emits a ``UserWarning`` when
+called; it is appropriate only for modeling a point mass added at the
+COM, not for density-like randomization. The distinction matters more
+on variant scenes than on single-asset scenes because variants often
+differ in mass by an order of magnitude.
+
+
+Viewers
+-------
+
+The native viewer, offscreen renderer, and Viser viewer all sync the
+selected environment's per-world fields into the host ``MjModel``
+before rendering, so the rendered geometry matches the variant
+assigned to the viewed environment. Switching environments in the
+native viewer (the ``,`` and ``.`` keys) updates the displayed mesh
+accordingly.
+
+Viser bakes mesh data into batched handles and cannot rely on a live
+view of ``geom_dataid``. It groups worlds by visual fingerprint (mesh
+selection, local geom frames, baked appearance) and builds one batched
+handle per group, with each environment assigned to its handle. A
+scene with N variants typically produces up to N handles per body.
+Convex hull visualization is computed per variant from the variant's
+mesh vertices.
+
+
+Performance
+-----------
+
+**Per-step cost is unaffected by variant count.** Variant-dependent
+fields are stored as per-world arrays accessed by world index in the
+existing kernels, with no branching or dispatch on variant.
+
+**Construction cost is linear in the total variant count.** mjlab
+compiles the merged scene once to produce the canonical ``MjModel``,
+then compiles each variant's original (un-merged) source spec in
+isolation to recover that variant's per-body and per-geom mesh-derived
+fields. Each per-variant compile sees only that variant's single body
+and mesh, so its cost is independent of the total number of variants
+in the scene.
+
+For a scene with one variant entity declaring k variants, construction
+runs ``1 + k`` compiles. With multiple variant entities, compiles
+decouple across entities: two variant entities of 5 variants each cost
+``1 + 5 + 5 = 11`` compiles, not ``1 + 5 * 5 = 26``. As an order of
+magnitude on CPU with typical procedural meshes, each per-variant
+compile takes around 1-2 ms, so a scene with 100 variants pays a few
+hundred milliseconds at startup and a scene with 1000 variants pays
+roughly two seconds.
+
+The merged spec contains every variant's mesh assets simultaneously,
+so memory at scene-build time scales with the total mesh vertex /
+face count across all variants. This is paid once at startup and does
+not affect training throughput.
+
+
+Limitations
+-----------
+
+**Floating-base only.** Each variant's root body must declare a free
+joint. Fixed-base variants are rejected; mocap auto-wrapping that
+applies to non-variant entities is not applied here.
+
+**Material assets are not propagated.** Each variant's ``contype``,
+``conaffinity``, ``condim``, ``friction``, ``mass``, ``density``,
+``group``, ``priority``, ``rgba``, ``solref``, ``solimp``, ``margin``,
+and ``gap`` are restored per-world during compile, but the
+``material`` reference on slot geoms inherits whichever material the
+template variant set. Use DR on ``geom_rgba`` / ``mat_rgba`` for
+per-world appearance variation.
+
+**Assignment is fixed at sim init.** There is no API to swap a world
+to a different variant on episode reset. World W's mesh asset is
+whatever it was assigned at init for the lifetime of the simulation.
+Per-episode mesh randomization is not supported today; DR can vary
+scalar properties (mass, friction, color, scale) on a fixed variant
+but cannot swap one mesh for another.
+
+**No support for per-world differing kinematic topology.** Variants
+must share the same body tree, joints, and actuator/sensor counts,
+so you cannot configure things like:
+
+* a different number of objects per world (world 0 has two props on
+  the table, world 1 has three);
+* different articulation per world (world 0's prop is an articulated
+  drawer with a slider joint, world 1's prop is a rigid block).
+
+True heterogeneous topology requires upstream support in mujoco_warp
+that does not currently exist.