Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

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

docs/source/entity/entity_data.rst

@@ -0,0 +1,507 @@
+.. _entity_data:
+
+Entity Data
+===========
+
+This page is the property reference for ``EntityData``. For an overview
+of how ``entity.data`` fits into the broader data access story, see
+:ref:`entity`.
+
+All properties are PyTorch tensors backed by MuJoCo Warp's GPU buffers
+with no copy overhead. The first dimension is always ``num_envs``, the
+number of parallel simulation worlds.
+
+.. warning::
+
+   Read properties reflect the state after ``sim.forward()`` is called.
+   If you write simulation state and then read a derived property in the
+   same event term, call ``sim.forward()`` between the write and the
+   read. The environment step sequence already does this; the warning
+   applies only when writing custom event terms that mix reads and
+   writes. See the :ref:`FAQ <faq-sim-forward>` for a detailed
+   explanation.
+
+
+Reference: root state
+---------------------
+
+Root properties describe the position, orientation, and velocity of the
+entity's root body. Properties ending in ``_w`` are expressed in the world
+frame. Properties ending in ``_b`` are expressed in the entity's base frame.
+See :ref:`frame-conventions` for details.
+
+Each entity has two root reference points: the **link origin** (the body
+frame origin defined in the MJCF) and the **center of mass (COM)**.
+Which one is relevant depends on the task.
+
+.. admonition:: MuJoCo's mixed-frame ``qvel``
+
+   For floating-base entities, the free joint stores 6 DOFs in
+   ``qvel``. MuJoCo expresses the **linear** components in the
+   **world frame** but the **angular** components in the **local body
+   frame**. EntityData avoids this pitfall: all ``_w`` velocity
+   properties are computed from ``cvel`` (see
+   :ref:`cvel-section` below) and are fully world-frame. If you
+   read ``env.sim.data.qvel`` directly, be aware of the mixed
+   convention.
+
+.. rubric:: Root link properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``root_link_pose_w``
+     - ``[num_envs, 7]``
+     - world
+     - Root link position (3) and quaternion (4) concatenated
+   * - ``root_link_pos_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link position
+   * - ``root_link_quat_w``
+     - ``[num_envs, 4]``
+     - world
+     - Root link orientation as quaternion (w, x, y, z)
+   * - ``root_link_vel_w``
+     - ``[num_envs, 6]``
+     - world
+     - Root link linear (3) and angular (3) velocity concatenated
+   * - ``root_link_lin_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link linear velocity
+   * - ``root_link_ang_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root link angular velocity
+   * - ``root_link_lin_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root link linear velocity in base frame
+   * - ``root_link_ang_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root link angular velocity in base frame
+
+.. rubric:: Root COM properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``root_com_pose_w``
+     - ``[num_envs, 7]``
+     - world
+     - Root COM position (3) and quaternion (4) concatenated
+   * - ``root_com_pos_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM position
+   * - ``root_com_quat_w``
+     - ``[num_envs, 4]``
+     - world
+     - Root COM orientation as quaternion (w, x, y, z)
+   * - ``root_com_vel_w``
+     - ``[num_envs, 6]``
+     - world
+     - Root COM linear (3) and angular (3) velocity concatenated
+   * - ``root_com_lin_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM linear velocity
+   * - ``root_com_ang_vel_w``
+     - ``[num_envs, 3]``
+     - world
+     - Root COM angular velocity
+   * - ``root_com_lin_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root COM linear velocity in base frame
+   * - ``root_com_ang_vel_b``
+     - ``[num_envs, 3]``
+     - body
+     - Root COM angular velocity in base frame
+
+.. rubric:: Derived root properties
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 20 15 30
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``projected_gravity_b``
+     - ``[num_envs, 3]``
+     - body
+     - Gravity vector (0, 0, -1) rotated into the base frame. Used to measure
+       tilt: a perfectly upright robot reads ``[0, 0, -1]``.
+   * - ``heading_w``
+     - ``[num_envs]``
+     - world
+     - Heading angle (radians) of the root body's forward axis projected onto
+       the XY plane.
+
+
+Reference: body state
+---------------------
+
+Body properties give per-body kinematic state for all bodies belonging to the
+entity. The second dimension is ``num_bodies``, which counts all non-world
+bodies in the entity's kinematic tree.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 15 25
+
+   * - Property
+     - Shape
+     - Frame
+     - Description
+   * - ``body_link_pose_w``
+     - ``[num_envs, num_bodies, 7]``
+     - world
+     - Per-body link position (3) and quaternion (4)
+   * - ``body_link_pos_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link positions
+   * - ``body_link_quat_w``
+     - ``[num_envs, num_bodies, 4]``
+     - world
+     - Per-body link orientations
+   * - ``body_link_vel_w``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - Per-body link linear (3) and angular (3) velocity
+   * - ``body_link_lin_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link linear velocities
+   * - ``body_link_ang_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body link angular velocities
+   * - ``body_com_pose_w``
+     - ``[num_envs, num_bodies, 7]``
+     - world
+     - Per-body COM position (3) and quaternion (4)
+   * - ``body_com_pos_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM positions
+   * - ``body_com_quat_w``
+     - ``[num_envs, num_bodies, 4]``
+     - world
+     - Per-body COM orientations
+   * - ``body_com_vel_w``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - Per-body COM linear (3) and angular (3) velocity
+   * - ``body_com_lin_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM linear velocities
+   * - ``body_com_ang_vel_w``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - Per-body COM angular velocities
+   * - ``body_external_wrench``
+     - ``[num_envs, num_bodies, 6]``
+     - world
+     - External force (3) and torque (3) applied to each body
+   * - ``body_external_force``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - External forces applied to each body
+   * - ``body_external_torque``
+     - ``[num_envs, num_bodies, 3]``
+     - world
+     - External torques applied to each body
+
+
+Reference: joint state
+----------------------
+
+Joint properties cover 1-DOF revolute and prismatic joints. The free joint
+(root floating-base DOF) is excluded; use root state properties for that.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``joint_pos``
+     - ``[num_envs, num_joints]``
+     - Joint positions in radians (revolute) or metres (prismatic)
+   * - ``joint_pos_biased``
+     - ``[num_envs, num_joints]``
+     - Joint positions with encoder bias added. Used when simulating
+       encoder calibration errors via domain randomization.
+   * - ``joint_vel``
+     - ``[num_envs, num_joints]``
+     - Joint velocities in rad/s or m/s
+   * - ``joint_acc``
+     - ``[num_envs, num_joints]``
+     - Joint accelerations in rad/s² or m/s²
+   * - ``actuator_force``
+     - ``[num_envs, num_actuators]``
+     - Scalar actuator output in actuation space (per actuator). This is
+       the force before projection through the transmission Jacobian. For
+       actuator forces in joint space, use ``qfrc_actuator`` instead.
+
+
+.. _generalized-forces:
+
+Reference: generalized forces
+-----------------------------
+
+These properties expose selected components of MuJoCo's generalized
+force decomposition, sliced to this entity's articulated joint DOFs.
+Free joint DOFs are excluded. All shapes are ``[num_envs, nv]`` where
+``nv`` is the number of articulated DOFs belonging to this entity.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Property
+     - Description
+   * - ``qfrc_actuator``
+     - Forces produced by all actuators, mapped into joint space. For
+       motors this is the commanded torque times the gear ratio. For
+       position and velocity actuators this is the force computed by
+       the internal PD law. When ``actuatorgravcomp`` is enabled on a
+       joint, the gravity compensation force is included here.
+   * - ``qfrc_external``
+     - Forces on joints due to Cartesian wrenches applied to bodies
+       via ``xfrc_applied``. This is the :math:`J^\top F` mapping.
+       MuJoCo does not store this term separately; the property
+       recovers it from other force components after ``forward()``.
+
+Reference: geom and site state
+-------------------------------
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``geom_pose_w``
+     - ``[num_envs, num_geoms, 7]``
+     - Per-geom position (3) and quaternion (4) in world frame
+   * - ``geom_pos_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom positions in world frame
+   * - ``geom_quat_w``
+     - ``[num_envs, num_geoms, 4]``
+     - Per-geom orientations in world frame
+   * - ``geom_vel_w``
+     - ``[num_envs, num_geoms, 6]``
+     - Per-geom linear (3) and angular (3) velocity in world frame
+   * - ``geom_lin_vel_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom linear velocities in world frame
+   * - ``geom_ang_vel_w``
+     - ``[num_envs, num_geoms, 3]``
+     - Per-geom angular velocities in world frame
+   * - ``site_pose_w``
+     - ``[num_envs, num_sites, 7]``
+     - Per-site position (3) and quaternion (4) in world frame
+   * - ``site_pos_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site positions in world frame
+   * - ``site_quat_w``
+     - ``[num_envs, num_sites, 4]``
+     - Per-site orientations in world frame
+   * - ``site_vel_w``
+     - ``[num_envs, num_sites, 6]``
+     - Per-site linear (3) and angular (3) velocity in world frame
+   * - ``site_lin_vel_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site linear velocities in world frame
+   * - ``site_ang_vel_w``
+     - ``[num_envs, num_sites, 3]``
+     - Per-site angular velocities in world frame
+
+
+Reference: tendon state
+-----------------------
+
+Tendon properties are only populated for entities that have tendon-driven
+actuators.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 25 40
+
+   * - Property
+     - Shape
+     - Description
+   * - ``tendon_len``
+     - ``[num_envs, num_tendons]``
+     - Tendon lengths
+   * - ``tendon_vel``
+     - ``[num_envs, num_tendons]``
+     - Tendon velocities
+
+
+.. _frame-conventions:
+
+Frame conventions
+-----------------
+
+Property names encode their reference frame with a suffix.
+
+``_w`` (world frame)
+    A fixed global frame. The origin is typically at the scene origin and
+    its axes are constant throughout the episode. World-frame quantities are
+    useful when you need absolute position, such as checking whether the
+    robot has fallen below a height threshold.
+
+``_b`` (body frame / base frame)
+    The entity's root body frame. It translates and rotates with the robot.
+    Most observation terms use body-frame quantities because they are
+    invariant to the robot's heading direction. A velocity expressed in the
+    body frame reads the same whether the robot faces north or south, which
+    makes it easier for the policy to generalize.
+
+``projected_gravity_b`` is a good example of why the frame suffix
+matters. It takes the world-frame gravity vector ``[0, 0, -1]`` and
+rotates it into the base frame. When the robot is upright the result is
+``[0, 0, -1]``; as the robot tilts, the x and y components grow,
+giving the policy a direct signal for orientation correction.
+
+Quaternion convention
+^^^^^^^^^^^^^^^^^^^^^
+
+All quaternions use the ``(w, x, y, z)`` convention, matching MuJoCo.
+
+Reduced state vs. derived quantities
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+EntityData properties fall into two categories that behave differently
+with respect to ``sim.forward()``:
+
+**Reduced state.** ``joint_pos`` and ``joint_vel`` read directly from
+MuJoCo's ``qpos`` and ``qvel`` arrays. Write methods such as
+``write_joint_state_to_sim()`` modify these arrays directly, so reads
+are always current.
+
+**Derived quantities.** All pose and velocity properties (``*_pose_w``,
+``*_vel_w``, ``*_vel_b``) are computed from MuJoCo's internal arrays
+(``xpos``, ``xquat``, ``cvel``, ``subtree_com``, etc.) which are only
+updated when ``sim.forward()`` runs. If you write to ``qpos``/``qvel``
+and then read a derived property without an intervening ``forward()``,
+the read will return stale values.
+
+The environment step sequence calls ``forward()`` at the right time, so
+this only matters if you write custom event terms that both write and
+read in the same function. See the :ref:`FAQ <faq-sim-forward>` for
+details.
+
+.. _cvel-section:
+
+How velocity properties are computed from ``cvel``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+MuJoCo does not store world-frame linear velocities directly. Instead,
+it stores a 6D spatial velocity per body called ``cvel`` (com-based
+velocity), laid out as ``(angular[3], linear[3])``. This vector is
+expressed in the **c-frame**: a frame centered at ``subtree_com`` (the
+center of mass of the body's kinematic subtree) and oriented like the
+world frame. MuJoCo uses this representation to improve numerical
+precision for mechanisms far from the world origin. See
+`c-frame variables <https://mujoco.readthedocs.io/en/stable/APIreference/APItypes.html#c-frame-variables>`_
+and Featherstone's
+`Spatial Algebra <http://royfeatherstone.org/spatial/>`_ for background.
+
+To recover the world-frame linear velocity at an arbitrary point
+:math:`\mathbf{p}` on a rigid body, we apply the standard rigid-body
+velocity transfer formula. Let :math:`\boldsymbol{\omega}` and
+:math:`\mathbf{v}_c` denote the angular and linear components of
+``cvel``, and let :math:`\mathbf{c}` denote ``subtree_com``. Because
+the c-frame is world-aligned, :math:`\boldsymbol{\omega}` is already in
+the world frame. The linear velocity at :math:`\mathbf{p}` is:
+
+.. math::
+
+   \mathbf{v}_p
+     = \mathbf{v}_c
+       - \boldsymbol{\omega} \times (\mathbf{c} - \mathbf{p})
+
+EntityData applies this formula in ``compute_velocity_from_cvel()``:
+
+.. code-block:: python
+
+   def compute_velocity_from_cvel(pos, subtree_com, cvel):
+       lin_vel_c = cvel[..., 3:6]
+       ang_vel_c = cvel[..., 0:3]
+       offset = subtree_com - pos
+       lin_vel_w = lin_vel_c - torch.cross(ang_vel_c, offset, dim=-1)
+       ang_vel_w = ang_vel_c
+       return torch.cat([lin_vel_w, ang_vel_w], dim=-1)
+
+Every velocity property in EntityData (``root_link_vel_w``,
+``body_link_vel_w``, ``geom_vel_w``, ``site_vel_w``, and their COM
+variants) uses this function, substituting the appropriate point:
+
+- **Link velocities** use ``xpos`` (body frame origin).
+- **COM velocities** use ``xipos`` (body center of mass).
+- **Geom/site velocities** use ``geom_xpos``/``site_xpos``, with
+  ``cvel`` looked up from the parent body.
+
+
+Default pose and relative quantities
+--------------------------------------
+
+``entity.data.default_joint_pos`` holds the joint positions from the entity's
+initial-state configuration (the ``init_state.joint_pos`` field of
+``EntityCfg``). It has shape ``[num_envs, num_joints]`` and is replicated
+across all environments at initialization time.
+
+The relative joint position is the deviation of the current joint position
+from this default:
+
+.. code-block:: python
+
+    joint_pos_rel = joint_pos - default_joint_pos
+
+This is what the ``joint_pos_rel`` observation function computes:
+
+.. code-block:: python
+
+    def joint_pos_rel(env, asset_cfg):
+        asset = env.scene[asset_cfg.name]
+        jnt_ids = asset_cfg.joint_ids
+        return (
+            asset.data.joint_pos[:, jnt_ids]
+            - asset.data.default_joint_pos[:, jnt_ids]
+        )
+
+Relative joint positions give the policy a compact representation of posture
+deviation. When the robot is at its default pose, every element is zero.
+
+Similarly, ``default_joint_vel`` is used by the ``joint_vel_rel`` observation
+function. For most configurations the default velocity is zero, so
+``joint_vel_rel`` is identical to ``joint_vel``. The indirection exists to
+allow non-zero reference velocities in tasks such as motion imitation.
+
+The ``use_default_offset=True`` option in joint position action configs uses
+``default_joint_pos`` as the zero point for the action space, so a network
+output of zero commands the robot to its default pose. This is the standard
+configuration for locomotion tasks.

docs/source/entity/index.rst

@@ -0,0 +1,355 @@
+.. _entity:
+
+Entity
+======
+
+An ``Entity`` represents a physical object in the simulation: a robot, a
+manipulated object, or a fixed fixture like a table. It is the central
+abstraction in mjlab's physics layer.
+
+A single ``Entity`` class covers all variants (contrast Isaac Lab, which
+splits this across ``Articulation``, ``RigidObject``, and several other
+subclasses of ``AssetBase``). Two orthogonal boolean properties classify
+each instance:
+
+**Base type.**
+  A *fixed-base* entity is welded to the world and has no free joint. A
+  *floating-base* entity has a free joint giving it 6-DOF movement.
+
+**Articulation.**
+  An *articulated* entity has internal joints (revolute, prismatic, etc.).
+  A *non-articulated* entity has none beyond a possible free joint.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 25 15 15 15
+
+   * - Type
+     - Example
+     - ``is_fixed_base``
+     - ``is_articulated``
+     - ``is_actuated``
+   * - Fixed non-articulated
+     - Table, wall
+     - True
+     - False
+     - False
+   * - Fixed articulated
+     - Robot arm, door
+     - True
+     - True
+     - True/False
+   * - Floating non-articulated
+     - Box, ball, mug
+     - False
+     - False
+     - False
+   * - Floating articulated
+     - Humanoid, quadruped
+     - False
+     - True
+     - True/False
+
+.. note::
+
+   mjlab automatically wraps every fixed-base entity in a
+   `mocap body <https://mujoco.readthedocs.io/en/stable/modeling.html#mocap-bodies>`_
+   so that each parallel environment can place the entity at a different
+   position. Without this wrapping, all fixed-base entities would be
+   welded to the world origin. The wrapping is transparent, but
+   **positioning only happens when a reset event runs**. You must
+   include a reset event such as ``reset_root_state_uniform`` in your
+   event config; without one, every fixed-base entity will remain at
+   the origin. See the :ref:`FAQ <faq>` for a full example. Mocap
+   entities can also be repositioned at runtime via
+   ``entity.write_mocap_pose_to_sim()``.
+
+
+Configuring an entity
+---------------------
+
+Every entity is described by an ``EntityCfg``. Only ``spec_fn`` is
+required in practice; all other fields have sensible defaults. A passive
+floating object needs nothing more than:
+
+.. code-block:: python
+
+    from mjlab.entity import EntityCfg
+
+    cube_cfg = EntityCfg(spec_fn=get_cube_spec)
+
+An actuated robot uses more of the interface:
+
+.. code-block:: python
+
+    from mjlab.entity import EntityCfg, EntityArticulationInfoCfg
+    from mjlab.actuator import IdealPDActuatorCfg
+
+    robot_cfg = EntityCfg(
+        spec_fn=get_spec,
+        init_state=EntityCfg.InitialStateCfg(
+            pos=(0.0, 0.0, 0.8),
+            joint_pos={".*_hip_.*": 0.5, ".*": 0.0},
+        ),
+        articulation=EntityArticulationInfoCfg(
+            actuators=(
+                IdealPDActuatorCfg(
+                    target_names_expr=(".*",),
+                    stiffness={".*": 50.0},
+                    damping={".*": 5.0},
+                ),
+            ),
+        ),
+        collisions=(my_collision_cfg,),
+    )
+
+The following sections describe each field.
+
+``spec_fn``
+^^^^^^^^^^^
+
+A callable that returns an ``mujoco.MjSpec``. The scene calls it during
+composition, attaches the returned spec with a name prefix, and compiles
+everything into a shared ``MjModel``.
+
+For simple cases a lambda suffices:
+
+.. code-block:: python
+
+    spec_fn = lambda: mujoco.MjSpec.from_file("robot.xml")
+
+For anything more involved, use a regular function. MuJoCo resolves mesh
+assets from disk automatically, so ``get_spec`` only needs to load the
+XML:
+
+.. code-block:: python
+
+    def get_spec() -> mujoco.MjSpec:
+        return mujoco.MjSpec.from_file(str(ROBOT_XML))
+
+Because ``spec_fn`` is an arbitrary callable, you can perform any
+`MjSpec edits <https://mujoco.readthedocs.io/en/stable/python.html#spec>`_
+before returning: add bodies, change joint limits, swap materials,
+or build the entire model programmatically without an XML file at all.
+
+``init_state``
+^^^^^^^^^^^^^^
+
+Default root pose, root velocity, and joint positions/velocities. These
+values are stored as a MuJoCo keyframe and used by reset events to
+return the entity to its initial configuration.
+
+``joint_pos`` and ``joint_vel`` are dicts mapping regex patterns to
+values. Patterns are matched against joint names in order, so later
+entries override earlier ones for any joint that matches both:
+
+.. code-block:: python
+
+    init_state = EntityCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 0.8),       # root position
+        rot=(1.0, 0.0, 0.0, 0.0),  # root quaternion (w, x, y, z)
+        joint_pos={
+            ".*": 0.0,              # all joints to zero
+            ".*_hip_.*": 0.5,       # then override hips to 0.5
+        },
+    )
+
+Set ``joint_pos=None`` to use an existing keyframe from the MJCF model
+instead of defining values here.
+
+``articulation``
+^^^^^^^^^^^^^^^^
+
+Actuator configuration. Only needed for entities that have actuated
+joints. Passive objects (boxes, tables, walls) can omit this field
+entirely. See :ref:`actuators` for details on actuator types.
+
+``soft_joint_pos_limit_factor`` (default 1.0) shrinks the joint range
+used by soft-limit penalty rewards, so the policy is penalized before
+reaching the physical hard stop. This does not modify the actual joint
+limits in the MuJoCo model.
+
+Spec editors
+^^^^^^^^^^^^
+
+The remaining fields are optional tuples of spec editor configs that
+modify the ``MjSpec`` before compilation:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 20 80
+
+   * - Field
+     - Purpose
+   * - ``collisions``
+     - Replace the entity's collision structure: which geoms collide, and
+       with what contact parameters.
+   * - ``lights``
+     - Add lights to specific bodies.
+   * - ``cameras``
+     - Add cameras to specific bodies.
+   * - ``textures``
+     - Add procedural textures (checker, gradient, etc.).
+   * - ``materials``
+     - Add materials and optionally assign them to geoms by regex.
+   * - ``geoms``
+     - Patch attributes of existing geoms (visualization group, collision
+       attributes). Unset attributes are left untouched.
+
+Each editor accepts regex patterns to target specific elements. For
+example, a ``CollisionCfg`` with ``geom_names_expr=(".*_foot.*",)``
+sets contact parameters only on foot geoms. See the asset zoo
+(``mjlab.asset_zoo.robots``) for complete examples.
+
+``geoms`` and ``collisions`` both write geom attributes but with
+different semantics. A ``GeomCfg`` is a sparse *patch*: every attribute
+defaults to ``None``, and only attributes you set are written. A
+``CollisionCfg`` is a *policy*: ``contype``, ``conaffinity``,
+``condim``, and ``priority`` are required and always written to every
+matched geom, and non-matched geoms have collision disabled by default,
+so the entity's contact behavior is fully determined by the config
+regardless of the source XML. Collision configs are applied after geom
+configs; mjlab warns if a ``GeomCfg`` sets a collision attribute that a
+``CollisionCfg`` then overwrites.
+
+Heterogeneous worlds
+^^^^^^^^^^^^^^^^^^^^
+
+For scenes that need different mesh assets in different parallel worlds
+(for example, training a manipulation policy that generalizes across
+object shapes), use ``VariantEntityCfg`` instead of ``EntityCfg``. Each
+world is assigned a variant proportional to a configurable weight, and
+mesh-dependent compiled constants (collision bounds, body inertials,
+subtree mass) are stored as per-world arrays so domain randomization and
+viewers stay consistent. See :ref:`heterogeneous_worlds`.
+
+Subclassing Entity
+^^^^^^^^^^^^^^^^^^
+
+``Entity`` and ``EntityCfg`` can be subclassed for specialized behavior.
+mjlab itself does this for terrain: ``TerrainEntity`` extends ``Entity``
+with procedural terrain generation and per-environment origin
+computation, and ``TerrainEntityCfg`` adds fields like
+``terrain_type``, ``env_spacing``, and ``terrain_generator``. The same
+pattern works for any domain-specific entity that needs logic beyond
+what ``EntityCfg`` and spec editors provide.
+
+Finding elements
+^^^^^^^^^^^^^^^^
+
+Entity provides ``find_*`` methods that accept regex patterns and return
+matched element indices and names:
+
+.. code-block:: python
+
+    ids, names = entity.find_joints((".*_hip_.*", ".*_knee_.*"))
+    ids, names = entity.find_geoms((".*foot.*",))
+    ids, names = entity.find_bodies((".*",))
+
+Available methods: ``find_bodies()``, ``find_joints()``,
+``find_geoms()``, ``find_sites()``, ``find_tendons()``.
+These are used internally during scene construction and manager
+initialization. In reward and observation terms, prefer
+``SceneEntityCfg`` with name patterns as described below.
+
+
+Reading runtime state
+---------------------
+
+Once entities are added to a ``SceneCfg`` and the environment is
+constructed, their state is accessible through three interfaces at
+decreasing levels of abstraction.
+
+EntityData
+^^^^^^^^^^
+
+``entity.data`` is the primary interface for reward, observation, and
+termination functions. It exposes kinematic state (poses, velocities, accelerations), actuator forces,
+generalized forces, and derived body-frame quantities such as projected
+gravity, all as PyTorch tensors with
+shape ``(num_envs, ...)``. See :ref:`entity_data` for the full property
+reference.
+
+``SceneEntityCfg`` selects which entity and which elements within it a
+term operates on. Regex patterns in ``joint_names``, ``body_names``,
+``site_names``, etc. are resolved to integer indices once at manager
+initialization, so there is no regex overhead at runtime:
+
+.. code-block:: python
+
+    from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+    def flat_orientation_l2(
+        env,
+        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+    ) -> torch.Tensor:
+        """Penalize non-flat base orientation using projected gravity."""
+        asset = env.scene[asset_cfg.name]
+        return torch.sum(
+            torch.square(asset.data.projected_gravity_b[:, :2]), dim=1
+        )
+
+``SceneEntityCfg`` also supports regex element selection through
+``joint_names``, ``body_names``, ``site_names``, etc. The resolved
+integer indices (e.g., ``asset_cfg.joint_ids``) make the runtime read a
+single tensor slice with no regex overhead.
+
+Sensors
+^^^^^^^
+
+Sensors are configured on the **scene**, not on individual entities.
+A sensor can reference an entity element (e.g., a contact sensor on the
+robot's feet, an accelerometer attached to a body site), but it can also
+be independent of any entity. This is why sensors live in ``SceneCfg``
+rather than ``EntityCfg``.
+
+At runtime, sensors are accessed by name through ``env.scene``, the same
+way entities are:
+
+.. code-block:: python
+
+    def angular_momentum_penalty(env, sensor_name: str) -> torch.Tensor:
+        sensor = env.scene[sensor_name]
+        return torch.sum(torch.square(sensor.data), dim=-1)
+
+Builtin sensors wrap MuJoCo sensor types (accelerometer, gyro, framepos,
+subtreeangmom, etc.). ``ContactSensor``, ``RayCastSensor``, and
+``CameraSensor`` provide higher-level abstractions for contact detection,
+terrain scanning, and RGB-D rendering. See :ref:`sensors` for details.
+
+Raw simulation data
+^^^^^^^^^^^^^^^^^^^
+
+For anything not covered by ``EntityData`` or sensors, the underlying
+MuJoCo Warp arrays are accessible through ``env.sim.data`` and
+``env.sim.model``. These expose the full ``mjData`` and ``mjModel``
+fields as PyTorch tensors (zero-copy), indexed by global MuJoCo IDs
+rather than per-entity IDs:
+
+.. code-block:: python
+
+    # Global joint positions across all entities.
+    qpos = env.sim.data.qpos          # (num_envs, nq)
+
+    # All body positions.
+    xpos = env.sim.data.xpos          # (num_envs, nbody, 3)
+
+    # Model-level constants.
+    body_mass = env.sim.model.body_mass  # (nbody,)
+
+This is useful for low-level operations or when you need quantities
+that span multiple entities.
+
+.. note::
+
+   The main limitation of raw sim data is that you must manage global
+   MuJoCo indices yourself. In the future, we plan to support MuJoCo's
+   `bind <https://mujoco.readthedocs.io/en/latest/python.html#relationship-to-pymjcf-and-bind>`_
+   functionality, which will allow binding spec elements directly to
+   their corresponding data views without manual index bookkeeping.
+
+.. toctree::
+   :maxdepth: 1
+
+   entity_data
+   per_world_mesh

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.