DRAFT: POC reset-time variations

AGENTS.md

@@ -55,7 +55,7 @@ Tests require Isaac Sim and run via pytest:
 /isaac-sim/python.sh -m pytest isaaclab_arena/tests/ -m with_cameras
 ```
 
-### Linting, Formatting, and Coding Style
+### Linting, Formatting
 
 Pre-commit hooks enforce the style guide: black (line length 120), flake8, isort, pyupgrade (py310+), and codespell. Run checks **before** committing — not after:
 
@@ -115,3 +115,60 @@ def test_foo():  # pytest-visible outer function
     result = run_simulation_app_function(_test_foo)
     assert result
 ```
+
+## Coding Style
+
+### Checks/Asserts
+
+In Arena we prefer asserts over if...raises in most cases where such an check represents a coding error
+(i.e) something that shouldn't be recovered from. Asserts are briefer than the exception based counterpart.
+The raises-based approach is more verbose (for example it always uses multiple lines), and because
+this type of error is not intended to be recovered from, it offers no advantage.
+
+### Docstrings
+
+Prefer short docstrings. Where a single line is not sufficient, a short (i.e. 2-3 line)
+paragraph can follow. Document args and returns, but not raises.
+We follow google-style for docstrings which is repeated below:
+
+A docstring should give enough information to write a call to the function without reading the function’s code. The docstring should describe the function’s calling syntax and its semantics, but generally not its implementation details, unless those details are relevant to how the function is to be used. For example, a function that mutates one of its arguments as a side effect should note that in its docstring. Otherwise, subtle but important details of a function’s implementation that are not relevant to the caller are better expressed as comments alongside the code than within the function’s docstring.
+
+The docstring may be descriptive-style ("""Fetches rows from a Bigtable.""") or imperative-style ("""Fetch rows from a Bigtable."""), but the style should be consistent within a file. The docstring for a @property data descriptor should use the same style as the docstring for an attribute or a function argument ("""The Bigtable path.""", rather than """Returns the Bigtable path.""").
+
+Certain aspects of a function should be documented in special sections, listed below. Each section begins with a heading line, which ends with a colon. All sections other than the heading should maintain a hanging indent of two or four spaces (be consistent within a file). These sections can be omitted in cases where the function’s name and signature are informative enough that it can be aptly described using a one-line docstring.
+
+#### Args
+List each parameter by name. A description should follow the name, and be separated by a colon followed by either a space or newline. If the description is too long to fit on a single 80-character line, use a hanging indent of 2 or 4 spaces more than the parameter name (be consistent with the rest of the docstrings in the file). The description should include required type(s) if the code does not contain a corresponding type annotation. If a function accepts *foo (variable length argument lists) and/or **bar (arbitrary keyword arguments), they should be listed as *foo and **bar.
+
+#### Returns: (or Yields: for generators)
+Describe the semantics of the return value, including any type information that the type annotation does not provide. If the function only returns None, this section is not required. It may also be omitted if the docstring starts with “Return”, “Returns”, “Yield”, or “Yields” (e.g. """Returns row from Bigtable as a tuple of strings.""") and the opening sentence is sufficient to describe the return value. Do not imitate older ‘NumPy style’ (example), which frequently documented a tuple return value as if it were multiple return values with individual names (never mentioning the tuple). Instead, describe such a return value as: “Returns: A tuple (mat_a, mat_b), where mat_a is …, and …”. The auxiliary names in the docstring need not necessarily correspond to any internal names used in the function body (as those are not part of the API). If the function uses yield (is a generator), the Yields: section should document the object returned by next(), instead of the generator object itself that the call evaluates to.
+
+#### Raises
+We don't document raises in Arena
+
+#### Example
+Here is an example docstring
+
+```python
+def fetch_smalltable_rows(
+    table_handle: smalltable.Table,
+    keys: Sequence[bytes | str],
+    require_all_keys: bool = False,
+) -> Mapping[bytes, tuple[str, ...]]:
+    """Fetches rows from a Smalltable.
+
+    Retrieves rows pertaining to the given keys from the Table instance
+    represented by table_handle.  String keys will be UTF-8 encoded.
+
+    Args:
+        table_handle: An open smalltable.Table instance.
+        keys: A sequence of strings representing the key of each table
+          row to fetch.  String keys will be UTF-8 encoded.
+        require_all_keys: If True only rows with values set for all keys will be
+          returned.
+
+    Returns:
+        A dict mapping keys to the corresponding table row data
+        fetched. Each row is represented as a tuple of strings.
+    """
+```

isaaclab_arena/assets/asset.py

@@ -3,6 +3,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from isaaclab_arena.variations.variation_base import VariationBase
+
 
 class Asset:
     """
@@ -17,3 +24,30 @@ def __init__(self, name: str, tags: list[str] | None = None, **kwargs):
         assert name is not None, "Name is required for all assets"
         self.name = name
         self.tags = tags
+
+    def add_variation(self, variation: VariationBase) -> None:
+        """Attach a variation to this asset under its class-level ``name``.
+
+        Subclasses call this from their ``__init__`` to declare the variations
+        they support. Re-registering the same name overwrites the existing entry.
+        The ``_variations`` dict is created on first attach so subclasses that
+        don't route through ``Asset.__init__`` (e.g. some
+        :class:`~isaaclab_arena.embodiments.embodiment_base.EmbodimentBase`
+        subclasses) still work.
+        """
+        if not hasattr(self, "_variations"):
+            self._variations: dict[str, VariationBase] = {}
+        self._variations[variation.name] = variation
+
+    def get_variation(self, name: str) -> VariationBase:
+        """Return the variation with the given name."""
+        variations = getattr(self, "_variations", {})
+        assert name in variations, (
+            f"Asset '{self.name}' ({type(self).__name__}) does not support variation '{name}'. "
+            f"Supported variations: {sorted(variations)}."
+        )
+        return variations[name]
+
+    def get_variations(self) -> list[VariationBase]:
+        """Return every variation attached to this asset, enabled or not."""
+        return list(getattr(self, "_variations", {}).values())

isaaclab_arena/assets/object.py

@@ -17,6 +17,7 @@
 from isaaclab_arena.utils.pose import Pose
 from isaaclab_arena.utils.usd.rigid_bodies import find_shallowest_rigid_body
 from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd, has_light, open_stage
+from isaaclab_arena.variations.object_color import ObjectColorVariation
 
 
 class Object(ObjectBase):
@@ -61,6 +62,7 @@ def __init__(
         self.bounding_box = None
         self.object_cfg = self._init_object_cfg()
         self.event_cfg = self._init_event_cfg()
+        self.add_variation(ObjectColorVariation(self))
 
     def add_relation(self, relation: RelationBase) -> None:
         """Add a relation to this object."""

isaaclab_arena/assets/object_library.py

@@ -490,11 +490,14 @@ def __init__(
         spawner_cfg: sim_utils.DomeLightCfg = default_spawner_cfg,
         hdr: "HDRImage | None" = None,  # noqa: F821
     ):
+        from isaaclab_arena.variations.hdr_image_variation import HDRImageVariation
+
         super().__init__(
             instance_name=instance_name, prim_path=prim_path, initial_pose=initial_pose, spawner_cfg=spawner_cfg
         )
         if hdr is not None:
             self.add_hdr(hdr)
+        self.add_variation(HDRImageVariation(self))
 
     def add_hdr(self, hdr: "HDRImage") -> None:  # noqa: F821
         """Attach an HDR environment map texture to this dome light.

isaaclab_arena/embodiments/droid/droid.py

@@ -376,15 +376,16 @@ class DroidEventCfg:
             ],
         },
     )
-    randomize_franka_joint_state = EventTerm(
-        func=franka_stack_events.randomize_joint_by_gaussian_offset,
-        mode="reset",
-        params={
-            "mean": 0.0,
-            "std": 0.02,
-            "asset_cfg": SceneEntityCfg("robot"),
-        },
-    )
+    randomize_franka_joint_state = None
+    # EventTerm(
+    #     func=franka_stack_events.randomize_joint_by_gaussian_offset,
+    #     mode="reset",
+    #     params={
+    #         "mean": 0.0,
+    #         "std": 0.02,
+    #         "asset_cfg": SceneEntityCfg("robot"),
+    #     },
+    # )
 
 
 @configclass

isaaclab_arena/environments/arena_env_builder.py

@@ -8,6 +8,7 @@
 import argparse
 import datetime
 import gymnasium as gym
+from typing import Any
 
 from isaaclab.devices.device_base import DeviceCfg, DevicesCfg
 from isaaclab.envs import ManagerBasedRLMimicEnv
@@ -38,6 +39,9 @@
 from isaaclab_arena.utils.isaaclab_utils.simulation_app import reapply_viewer_cfg
 from isaaclab_arena.utils.multiprocess import get_local_rank
 from isaaclab_arena.utils.pose import Pose, PosePerEnv
+from isaaclab_arena.variations import variations_hydra
+from isaaclab_arena.variations.variation_base import BuildTimeVariationBase, RunTimeVariationBase, VariationBase
+from isaaclab_arena.variations.variations_recorder import VariationRecorder
 
 
 class ArenaEnvBuilder:
@@ -173,6 +177,85 @@ def _apply_pool_layouts_to_objects(
             else:
                 obj.set_initial_pose(PosePerEnv(poses=poses))
 
+    def get_all_variations(self) -> dict[str, list[VariationBase]]:
+        """Return ``{asset_name: [variation, ...]}`` for every variation host in the env.
+
+        Combines variations attached to scene assets (via
+        :meth:`~isaaclab_arena.scene.scene.Scene.get_asset_variations`) with
+        any variations attached to the embodiment, keyed by
+        :attr:`~isaaclab_arena.embodiments.embodiment_base.EmbodimentBase.name`.
+        Embodiment-level variations let the embodiment own its own knobs (e.g.
+        camera decalibration on a robot's wrist cam) without needing to be
+        registered as a scene-side ``ObjectBase``.
+        """
+        by_asset: dict[str, list[VariationBase]] = dict(self.arena_env.scene.get_asset_variations())
+        embodiment = self.arena_env.embodiment
+        if embodiment is not None:
+            embodiment_variations = embodiment.get_variations()
+            if embodiment_variations:
+                assert embodiment.name not in by_asset, (
+                    f"Embodiment name '{embodiment.name}' collides with a scene asset that also has variations; "
+                    "rename the scene asset or the embodiment so variation namespaces stay unique."
+                )
+                by_asset[embodiment.name] = list(embodiment_variations)
+        return by_asset
+
+    def _compose_variations_event_cfg(self) -> Any | None:
+        """Build a configclass instance holding an :class:`EventTermCfg` per enabled run-time variation.
+
+        Walks every variation host (scene assets + embodiment), skips disabled
+        ones and any build-time variations (which are applied via
+        :meth:`_apply_build_time_variations`), and asks each remaining
+        variation for its event term. Returns ``None`` when nothing is enabled.
+        """
+        fields: list[tuple[str, type, EventTermCfg]] = []
+        seen: set[str] = set()
+        for asset_variations in self.get_all_variations().values():
+            for variation in asset_variations:
+                if not variation.enabled:
+                    continue
+                if not isinstance(variation, RunTimeVariationBase):
+                    continue
+                event_name, event_cfg = variation.build_event_cfg()
+                assert event_name not in seen, (
+                    f"Duplicate variation event term name '{event_name}'. "
+                    "Each variation must produce a unique name; consider prefixing with the asset name."
+                )
+                seen.add(event_name)
+                fields.append((event_name, EventTermCfg, event_cfg))
+        if not fields:
+            return None
+        VariationsEventCfg = make_configclass("VariationsEventCfg", fields)
+        return VariationsEventCfg()
+
+    def _apply_build_time_variations(self) -> None:
+        """Sample and apply every enabled :class:`BuildTimeVariationBase` from scene + embodiment.
+
+        Build-time variations mutate asset configs in place (e.g. swapping
+        a dome light's spawner texture), so this must run before ``scene_cfg``
+        is materialised. The :class:`VariationRecorder` is attached earlier in
+        ``compose_manager_cfg`` so it captures the samples drawn here.
+        """
+        for asset_variations in self.get_all_variations().values():
+            for variation in asset_variations:
+                if not variation.enabled:
+                    continue
+                if not isinstance(variation, BuildTimeVariationBase):
+                    continue
+                variation.apply()
+
+    def get_variations_schema(self) -> type | None:
+        """Return the dataclass describing every variation in the env, or ``None`` if none."""
+        return variations_hydra.build_schema(self.get_all_variations())
+
+    def load_variations_cfg_from_flags(self, hydra_overrides: list[str]) -> Any | None:
+        """Compose Hydra override strings into a typed ``VariationsCfg`` instance."""
+        return variations_hydra.load_cfg_from_flags(self.get_all_variations(), hydra_overrides)
+
+    def apply_hydra_variation_overrides(self, hydra_overrides: list[str]) -> None:
+        """Apply Hydra-style variation overrides across scene + embodiment variations."""
+        variations_hydra.apply_overrides(self.get_all_variations(), hydra_overrides)
+
     def _modify_recorder_cfg_dataset_filename(self, recorder_cfg: RecorderManagerBaseCfg) -> RecorderManagerBaseCfg:
         """Modify the recorder dataset filename to include the timestamp and rank."""
         base = getattr(recorder_cfg, "dataset_filename", "dataset")
@@ -201,6 +284,17 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg:
         if self.args.solve_relations:
             self._solve_relations()
 
+        # Attach the variation recorder before any sampling happens so it can
+        # observe both build-time variation samples (drawn below) and run-time
+        # variation samples (drawn during simulation). The combined dict covers
+        # both scene-side assets and the embodiment.
+        variations_recorder = VariationRecorder()
+        variations_recorder.attach(self.get_all_variations())
+
+        # Apply build-time variations now: they mutate asset configs (e.g.
+        # DomeLight.spawner_cfg) and must run before scene_cfg is materialised.
+        self._apply_build_time_variations()
+
         # Constructing the environment by combining inputs from the scene, embodiment, and task.
         embodiment = self.arena_env.embodiment or NoEmbodiment()
         task = self.arena_env.task or NoTask()
@@ -224,12 +318,14 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg:
                 [("placement_reset", EventTermCfg, self._placement_event_cfg)],
             )
             placement_event_cfg = PlacementEventCfg()
+        variations_event_cfg = self._compose_variations_event_cfg()
         events_cfg = combine_configclass_instances(
             "EventsCfg",
             embodiment.get_events_cfg(),
             self.arena_env.scene.get_events_cfg(),
             task.get_events_cfg(),
             placement_event_cfg,
+            variations_event_cfg,
         )
         termination_cfg = combine_configclass_instances(
             "TerminationCfg",
@@ -304,6 +400,7 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg:
                 teleop_devices=teleop_devices_cfg,
                 recorders=recorder_manager_cfg,
                 metrics=metrics,
+                variations_recorder=variations_recorder,
                 isaaclab_arena_env=isaaclab_arena_env,
                 viewer=viewer_cfg,
             )
@@ -333,6 +430,7 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg:
                 # I assume that they're not needed for the mimic env.
                 # recorders=recorder_manager_cfg,
                 # metrics=metrics,
+                # variations_recorder=variations_recorder,
                 isaaclab_arena_env=isaaclab_arena_env,
                 viewer=viewer_cfg,
             )

isaaclab_arena/environments/isaaclab_arena_manager_based_env.py

@@ -15,6 +15,7 @@
 
 from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
 from isaaclab_arena.metrics.metric_base import MetricBase
+from isaaclab_arena.variations.variations_recorder import VariationRecorder
 
 
 @configclass
@@ -68,6 +69,10 @@ class IsaacLabArenaManagerBasedRLEnvCfg(ManagerBasedRLEnvCfg):
     # Metrics
     metrics: list[MetricBase] | None = None
 
+    # Variation recorder. Callers can then read ``env.cfg.variations_recorder.records``
+    # after a run to recover the sampled variations.
+    variations_recorder: VariationRecorder | None = None
+
     # Isaaclab Arena Env. Held as a member to allow use of internal functions
     isaaclab_arena_env: IsaacLabArenaEnvironment | None = None
 

isaaclab_arena/evaluation/policy_runner.py

@@ -15,7 +15,11 @@
 from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext
 from isaaclab_arena.utils.multiprocess import get_local_rank, get_world_size
 from isaaclab_arena.utils.random import set_seed
-from isaaclab_arena_environments.cli import get_arena_builder_from_cli, get_isaaclab_arena_environments_cli_parser
+from isaaclab_arena_environments.cli import (
+    get_arena_builder_from_cli,
+    get_isaaclab_arena_environments_cli_parser,
+    split_hydra_overrides,
+)
 from isaaclab_arena_gr00t.utils.groot_path import ensure_groot_deps_in_path
 
 if TYPE_CHECKING:
@@ -161,14 +165,16 @@ def main():
         # Add the example environment arguments + policy-related arguments to the parser
         args_parser = get_isaaclab_arena_environments_cli_parser(args_parser)
         args_parser = policy_cls.add_args_to_parser(args_parser)
-        args_cli = args_parser.parse_args()
+        # Use parse_known_args so positional Hydra variation overrides fall through into unknown.
+        args_cli, unknown = args_parser.parse_known_args()
+        hydra_overrides = split_hydra_overrides(unknown, args_parser)
         # Re-apply per-rank device after parse preventing device got overwritten by the default value
         if is_distributed(args_cli):
             args_cli.distributed = True
             args_cli.device = f"cuda:{local_rank}"
 
         # Build scene
-        arena_builder = get_arena_builder_from_cli(args_cli)
+        arena_builder = get_arena_builder_from_cli(args_cli, hydra_overrides=hydra_overrides)
         env, cfg = arena_builder.make_registered_and_return_cfg()
 
         # Per-rank seed when distributed so each process has a different seed