Add variation system base. Part 1/3, base classes, compiler, and two concrete variations

[alexmillane]

Summary

Adds the basic machinery for variations.

Detailed description

• Variations are a structured API for introducing randomization into the simulation run.
• This change introduces two types of variation: run-time variation, and build-time variations:
  • run-time variations: Randomize per reset (implemented as a reset term)
  • build-time variations: Randomize during construction of the environment (implemented as a mutation on an Arena object)
• Variations are implemented as a combination of
  • Sampler: Generates the random value
  • Mutator: Applies the random value to affect the change.
  • This design is chosen to allow use to tap in and record the random value, as it is applied.
• Concrete Variations
  • Camera extrinsics: Modifies a camera position by a sampled vector in the camera frame.
  • HDR Image: Selects from a passed list of HDR images.

Not implemented (in follow up MRs)

• Recording the variations
• Control of the variations from the command line.

[isaaclab-review-bot]

Variation System Review — PR #746

Summary

This PR introduces a well-structured variation framework with clean separation between samplers, variation base classes, and two concrete variations (HDR image selection and camera decalibration). The architecture — separating build-time vs run-time variations, composable samplers, and listener-based recording hooks — is sound and extensible. There are a few correctness concerns in the camera decalibration implementation and some API robustness items worth addressing before merge.

---

🔬 Architecture & Design

Overall assessment: The variation system design is clean and fits well into the existing Arena framework. The BuildTimeVariationBase / RunTimeVariationBase split correctly maps to the env lifecycle (pre-build mutation vs event-term application). The sampler → variation → asset pipeline is well-layered. The lazy __init__.py (PEP 562) is a pragmatic solution to the SimulationApp import ordering constraint.

Cross-module impact: Changes touch Asset, Scene, ArenaEnvBuilder, and DomeLight. The framework additions are additive — existing envs without variations see no behavioral change since variations default to enabled=False.

---

Implementation Findings

🔴 Critical: Quaternion convention mismatch in camera_decalibration_variation.py

File: isaaclab_arena/variations/camera_decalibration_variation.py, lines 152–160

quat_apply from isaaclab.utils.math expects quaternions in (w, x, y, z) format (per its docstring: "The quaternion in (w, x, y, z). Shape is (..., 4)"). However, the code explicitly converts the wxyz quaternion from get_local_poses() into xyzw format:

self._q_parent_C_xyzw = torch.roll(q_parent_C_wxyz.detach(), shifts=-1, dims=-1).clone()

…and then passes this xyzw tensor directly to quat_apply:

deltas_parent = quat_apply(self._q_parent_C_xyzw[env_ids], deltas_opengl)

This feeds quat_apply a quaternion where the scalar component w is in position [3] instead of position [0], producing incorrect rotations. The empirical sign-flip hack on the next line ((-1.0, 1.0, -1.0)) may be compensating for this error in the specific test scenario, but the composition of two wrongs will break for cameras with non-trivial parent orientations.

Suggested fix: Store the quaternion in wxyz (as returned by get_local_poses) and pass it directly to quat_apply. Then apply the proper ROS-optical → OpenGL coordinate frame conversion (which should be (+x, -y, -z) per the existing TODO comment) without the empirical sign hack:

self._q_parent_C_wxyz = q_parent_C_wxyz.detach().clone()
# ...
deltas_opengl = deltas_input * deltas_input.new_tensor((1.0, -1.0, -1.0))  # ROS optical → OpenGL
deltas_parent = quat_apply(self._q_parent_C_wxyz[env_ids], deltas_opengl)

---

🟡 Warning: Access to private Camera._view attribute

File: isaaclab_arena/variations/camera_decalibration_variation.py, line 148

view = self._camera._view

_view is an internal implementation detail of Camera / TiledCamera — it's an XformPrimView created during _initialize_impl. Accessing it directly is fragile: upstream changes to the Camera internals (renaming, lazy init order changes) would silently break this code. Additionally, _view is None until the sensor's lifecycle hooks run, so the deferred snapshot on first __call__ is a necessary workaround.

Suggestion: Consider whether Camera.set_world_poses(positions, orientations, env_ids) (which is public) can be used instead, or propose a minimal public method on the Camera class (set_local_poses) upstream to make this pattern stable. At minimum, add a guard:

view = self._camera._view
assert view is not None, (
    f"Camera '{self._camera.cfg.prim_path}' _view is not yet initialized; "
    "decalibration event may be firing before sensor startup."
)

---

🟡 Warning: UniformSampler device mismatch risk

File: isaaclab_arena/variations/uniform_sampler.py, line 64

u = torch.rand(shape)

This always samples on CPU (default device). The consumer in camera_decalibration_variation.py then does .to(device=self._t_parent_C.device), which triggers a CPU→GPU transfer every reset. For large num_envs with a 3D sampler this is cheap, but it's a pattern that scales poorly.

Suggestion: Accept an optional device kwarg in _sample (or on the sampler at construction) so GPU-resident consumers avoid the transfer. Not blocking, but worth noting for future samplers with higher-dimensional outputs.

---

🟡 Warning: DomeLight unconditionally registers HDRImageVariation

File: isaaclab_arena/assets/object_library.py, line 507

self.add_variation(HDRImageVariation(self))

Every DomeLight instance now imports and instantiates HDRImageVariation regardless of whether any HDR variation is configured. The variation starts disabled (enabled=False) so there's no runtime cost, but:

1. It creates a circular-ish dependency path: constructing a DomeLight → imports hdr_image_variation → imports categorical_sampler / variation_base.
2. If a user subclasses DomeLight and doesn't want the variation, they can't easily suppress it without overriding __init__.

This is acceptable for the current scope but worth noting: if more variations are added to library objects, a registry/decorator pattern might be cleaner than hardcoding in __init__.

---

🔵 Improvement: Missing __repr__ on samplers and variations

The sampler and variation classes would benefit from __repr__ implementations for debugging and logging. Currently, inspecting a variation in a debugger or log gives an opaque <CameraDecalibrationVariation object at 0x...>. A simple repr showing name, enabled, and sampler params would be very helpful during development.

---

🔵 Improvement: CategoricalSampler uses torch.randint for CPU-only random selection

File: isaaclab_arena/variations/categorical_sampler.py, line 49

indices = torch.randint(low=0, high=len(choices), size=(num_samples,), dtype=torch.long)

For a categorical sampler that returns Python objects (not tensors), using torch.randint adds a torch dependency where random.choices would suffice and be ~10× faster for small num_samples. This matters less for the HDR variation (which samples once at build time) but could matter if categorical samplers are used at scale.

---

🔍 Error Handling

Overall quality: Good — assertions are used consistently with informative messages. Key validation points:

• ✅ UniformSampler validates low <= high and matching shapes
• ✅ CategoricalSampler rejects empty choices
• ✅ HDRImageVariation.apply() validates HDR names against registry
• ✅ decalibrate_camera_from_sampler.__init__ validates camera type and sampler shape
• ✅ get_variation gives clear error listing supported variations

One gap: BuildTimeVariationBase.apply() is called without checking that a sampler is set. HDRImageVariation.apply() does assert self.sampler is not None, but this check should ideally live in _apply_build_time_variations() so all future build-time variations get the same safety net without needing to remember the assertion pattern.

---

🧪 Test Coverage

What's tested:

• ✅ CategoricalSampler: draws, types, empty rejection, cfg builds
• ✅ UniformSampler: bounds, shape, cfg build, mismatched-bounds rejection
• ✅ HDRImageVariation end-to-end: apply mutates dome light spawner, unknown name asserts

Gaps (should address):

🟡 No tests for CameraDecalibrationVariation — Criticality: 7/10

• The run-time variation path (build_event_cfg → event term → __call__) is untested
• Given the quaternion convention concern above, a test that verifies the delta is applied in the expected direction would catch the bug
• Suggested: A unit test that constructs the variation, mocks a camera with known local pose, calls the event term, and verifies the resulting position is nominal + expected_delta

🟡 No integration test for ArenaEnvBuilder variation composition — Criticality: 6/10

• _compose_variations_event_cfg and _apply_build_time_variations are called in compose_manager_cfg but no test exercises the full pipeline with both build-time and run-time variations enabled simultaneously
• The duplicate event-name assertion path is untested

Verdict:

• Feature PR: Coverage is partial — samplers and build-time path are well-tested; run-time path and integration are not
• Test quality: Good for existing tests (deterministic, isolated, clear assertions)

---

Verdict

Minor fixes needed — The quaternion convention mismatch in the camera decalibration is the primary concern; the empirical sign hack suggests it "works" for the tested mount configuration but is mathematically incorrect and will fail for other camera orientations. The remaining items are robustness and coverage improvements that can be addressed iteratively, especially given this is a draft PR.

---

Incremental Review Updates

Update (commit 8e14ae5): The 9 new commits introduce substantial changes to adjacent subsystems (metrics refactoring with MetricsManager/MetricTermCfg, heterogeneous per-env placement with RigidObjectSet variant assignment, task registry via @register_task decorators, relation registry, ManagerBasedRLEnv file split, LightwheelLazyPath lazy loading). None of these commits modify the variation system files reviewed above.

Update (commit 2dad7e8): Docstring trimming plus _variations → self.variations (public dict), Scene.get_asset_variations() simplified, EmbodimentBase.__init__ now calls super properly.

Update (commit 8e71861): This push removes the listener/observer pattern from samplers and variations (good simplification — samplers are now purely stateless), renames sampler → sampler_cfg in all configs for clarity, renames event_shape → shape_per_sample, removes the lazy PEP 562 __init__.py (users import from submodules directly), and moves docstrings to attribute-level per configclass conventions. The mode field was also removed from CameraDecalibrationVariationCfg (hardcoded to "reset"). File renamed to camera_decalibration_variation.py. All clean simplifications.

Current findings status:

Finding	Status
🔴 Quaternion convention mismatch	Not addressed
🟡 Private Camera._view access	Not addressed
🟡 UniformSampler device mismatch	Not addressed
🟡 DomeLight unconditional variation	Not addressed
🟡 Missing camera decalibration tests	Not addressed

New issues in latest commit: None. The simplifications (listener removal, naming consistency) are well-executed.

---

Update (commit 4b9c4b3): Good API simplification — set_sampler() removed in favor of a single apply_cfg() entry point. The sampler parameter is removed from both variation constructors (CameraDecalibrationVariation, HDRImageVariation), and __init__ now delegates to apply_cfg(cfg) which handles sampler building. New test (test_apply_cfg_keeps_sampler_cfg_and_live_sampler_in_sync) properly validates the round-trip. Clean changes, no new issues introduced.

Previous findings remain unchanged (quaternion convention, private _view access, device mismatch, etc.) — these are orthogonal to this commit.

Update (aa9ad70): Cleanup commit — removes the isaaclab_arena_env back-reference from IsaacLabArenaManagerBasedRLEnvCfg (and its assignment in ArenaEnvBuilder), reorders shape_per_sample in SamplerBase (no functional change), and updates the example notebook to demonstrate HDR variation usage with a dome light. No variation system logic changes. No new issues. Previous findings (quaternion convention, private _view access, device mismatch, missing camera decalibration tests) remain unaddressed — expected for a draft PR iterating on API design.

Update (ba3f315): New commit adds test_run_time_variations.py — an integration test that exercises ArenaEnvBuilder.compose_manager_cfg() with a custom RunTimeVariationBase subclass, verifying both the enabled and disabled paths.

Findings addressed:

• ✅ "No integration test for ArenaEnvBuilder variation composition" — Now covered. The test validates that enabled variations inject their EventTermCfg into env_cfg.events with correct func/mode/params, and disabled variations do not.
• 🟡 "No tests for CameraDecalibrationVariation" — Partially addressed. The run-time variation wiring path is now tested (via a noop event term), but the actual camera decalibration logic (quaternion math, delta application) remains untested.

New issues: None. The test is well-structured — uses __test__ = False to prevent pytest collection of helper classes, properly isolates via subprocess, and validates the full lifecycle.

Remaining unaddressed findings (expected for draft PR): Quaternion convention mismatch, private Camera._view access, UniformSampler device mismatch, missing camera-specific unit tests.

Update (f31abd7): Major improvements in this push — the primary critical finding is now resolved.

Findings addressed:

• ✅ Quaternion convention mismatch — Fixed correctly. The code now treats get_local_poses() output as xyzw (documenting the Isaac Lab docs-vs-implementation discrepancy), uses a proper 180° rotation about X (q_ros_to_opengl_xyzw = (1, 0, 0, 0)) for the ROS→OpenGL frame conversion, and removes the empirical sign hack. The math is now correct for arbitrary camera orientations.
• ✅ Missing camera variation tests — Comprehensive test suite added in test_camera_extrinsics_variation.py covering disabled/enabled event cfg injection AND a full runtime test that verifies the decalibration vector round-trips correctly through the quaternion math.
• ✅ Private _view access guard — assert view is not None guard added as suggested.

Additional good changes:

• Rename from CameraDecalibrationVariation → CameraExtrinsicsVariation is cleaner naming
• New test_isaaclab_bug_get_local_poses.py regression test documents the Isaac Lab quaternion bug — good defensive testing
• test_variation_base.py adds thorough unit tests for the cfg plumbing
• test_hdr_image_variation.py extracted as standalone (cleaner test organization)
• Build-time variation tests refactored to use a generic test fixture instead of being tied to HDR specifics

Minor notes on new changes (non-blocking):

• randomize_franka_joint_state is commented out in FrankaEventCfg — intentional for this branch? Consider removing entirely or adding a TODO if temporary.
• ArenaEnvBuilder.get_all_variations() removed the assert embodiment.name not in by_asset guard — the old assertion caught naming collisions between embodiment and scene assets. The simplified version silently overwrites. Low risk but worth noting.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

---

Update (57e8fcb): Large push adds the graph-spec → live IsaacLabArenaEnvironment conversion pipeline. Key additions:

• arena_env_graph_types.py — Types extracted from arena_env_graph_spec.py to break circular imports between spec and conversion modules. Clean separation.
• arena_env_graph_conversion_utils.py — Materializes nodes into assets, attaches spatial constraints, partitions into embodiment + scene, and builds the final env. Well-structured with clear TODO annotations for future aggregation of state_specs.
• arena_env_graph_task_conversion_utils.py — Clever type-introspection approach: inspects task __init__ annotations to automatically distinguish node-ref params (Asset/AffordanceBase-typed) from plain params (float/str). Correctly handles scalar, list, and Optional union shapes.
• graph_spec_utils.py — Adds assert_spatial_constraint_shapes (validates unary/binary shape per relation class), enforces parent-before-child ordering in assert_references_exist, plus utility functions (iter_nested_leaf_values, map_nested_leaf_values, normalize_identifier, camel_to_snake, strip_suffix).
• relations.py — is_unary() static method added to all relation classes to support shape validation at parse time.
• ArenaEnvGraphSpec.validate() — Validation separated from from_dict() so callers can re-validate after mutation. to_arena_env() added with lazy import to avoid SimulationApp requirement at parse time.
• isaaclab_arena_manager_based_env_cfg.py — Disables RTX ambient light so USD light prims control scene illumination.
• Tests — 3 new test files with strong coverage of the conversion pipeline (task-arg resolution unit tests, end-to-end sim test, spec validation edge cases).
• randomize_franka_joint_state now active (previously commented out — noted in prior update). ✅ Fixed.
• Variations — Docstring-only changes clarifying ROS optical frame convention.

No new issues found. The graph-spec conversion design is clean and well-tested. The type-introspection approach for task-arg resolution is particularly elegant — it avoids hardcoded mappings and stays in sync with task signatures automatically.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

Update (d7641e6): Minimal change — fixes a flaky float comparison in test_build_time_variations.py by using pytest.approx(TEST_APPLIED_RADIUS, abs=1e-6) instead of exact equality. Correct fix for float32 precision mismatch. No new issues.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

---

Update (2636f7c): Lifts task_description to IsaacLabArenaManagerBasedRLEnvCfg as a top-level attribute, populated at build time via task.get_task_description(). policy_runner.py now reads env.unwrapped.cfg.task_description instead of reaching through the removed isaaclab_arena_env back-reference. Clean follow-up to the earlier isaaclab_arena_env removal — no new issues.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

---

Update (e2daf2d): Large push adds the random yaw initialization system (random_yaw_init CLI flag, conservative bbox collision checking via rotated_around_z, orientation composition through the full placement pipeline), the NotNextTo spatial constraint (inverse of NextTo — blocks a half-plane with dual-escape-route loss to avoid solver plateaus), CLI support for building envs from graph-spec YAML (--env_graph_spec_yaml as alternative to subcommand), and reproducibility fixes for PooledObjectPlacer (per-env RNGs replace global random.choice).

All changes are well-implemented with thorough test coverage (bounding box rotation, yaw composition, solver convergence, per-env reproducibility, CLI XOR validation). No new issues found.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

---

Update (3b72bff): Good design improvement — name is promoted from a ClassVar[str] to an instance attribute passed via __init__, enabling multiple variation instances of the same class on a single asset (e.g., per-camera extrinsics variations for robots with multiple cameras).

Key changes:

• VariationBase.__init__ now requires a name: str parameter (no longer ClassVar)
• CameraExtrinsicsVariation auto-derives name from camera_name (defaulting to f"camera_extrinsics_{camera_name}") — clean convention
• HDRImageVariation defaults to "hdr_image" but accepts override
• Asset.add_variation() now asserts uniqueness — prevents silent overwrites
• Tests validate both duplicate-rejection and distinct-name coexistence

This is a well-motivated change: multi-camera robots need independent extrinsics variations, and the old ClassVar pattern forced one-per-class. The implementation is clean — backward-compatible default names preserve existing call sites while enabling the multi-instance use case. No new issues.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

---

Update (d60c553): Good sampler hierarchy refactoring — splits the monolithic SamplerBase into a proper type tree:

• SamplerBase → pure marker class (no methods)
• ContinuousSampler → abstract base for tensor-returning samplers (UniformSampler inherits this)
• ChoiceSampler → replaces CategoricalSampler for discrete choice sampling

CategoricalSampler is removed; HDRImageVariation now uses ChoiceSamplerCfg. CameraExtrinsicsVariation type-annotates with ContinuousSampler instead of the overly broad SamplerBase. UniformSampler.sample() is now a direct public method (no more _sample indirection via base class).

This is a clean improvement — the type hierarchy now correctly encodes the semantic difference between continuous (tensor) and discrete (choice) sampling, enabling better type checking and clearer API contracts. No new issues.

Previous improvement note addressed:

• ✅ CategoricalSampler using torch.randint for CPU-only random selection — Resolved by replacement. ChoiceSampler (the successor) presumably handles this differently.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.---

Update (7b18a0a): Single commit — "Add missing files." Adds choice_sampler.py and continuous_sampler.py which were referenced by the previous sampler hierarchy refactoring (d60c553) but accidentally omitted from that push. Contents match the design described in the prior update: ContinuousSampler is the abstract base for tensor-returning samplers with sample(num_samples) → Tensor and shape_per_sample property; ChoiceSampler draws from a per-call choices sequence using torch.randint. No new issues — these files complete the sampler type split already reviewed.

Remaining unaddressed (acceptable for draft): UniformSampler CPU device default, DomeLight unconditional variation registration.

[alexmillane]

Self-review. First cursory pass.

[greptile-apps]

Greptile Summary

This PR introduces the foundational variation system: a Sampler/Mutator architecture that lets any Asset carry named VariationBase instances, which are either applied once at build time (BuildTimeVariationBase) or wired into the Isaac Lab event system per-reset (RunTimeVariationBase). Two concrete variations ship with it: CameraExtrinsicsVariation and HDRImageVariation.

• New modules (variation_base, sampler_base, continuous_sampler, uniform_sampler, choice_sampler, camera_extrinsics_variation, hdr_image_variation) establish the full type hierarchy; Asset, Scene, and ArenaEnvBuilder are extended to register and apply variations.
• ArenaEnvBuilder gains _apply_build_time_variations() and _compose_variations_event_cfg(), called during compose_manager_cfg(), and the isaaclab_arena_env reference on the env cfg is replaced by the simpler task_description string.
• A suite of unit and integration tests covers the disabled/enabled paths for both variation flavours and verifies the camera-frame round-trip at runtime.

Confidence Score: 4/5

Safe to merge with pre-existing test correctness and naming issues still open; the core variation wiring is sound.

The variation system is well-structured and the integration into ArenaEnvBuilder is clean. The main outstanding concerns flagged in earlier rounds (wrong quaternion direction in the camera-extrinsics runtime test, snake_case class name, and mutable tags aliasing) have not been addressed yet, keeping the score below a clean pass, though none affect production behavior.

isaaclab_arena/tests/test_camera_extrinsics_variation.py and isaaclab_arena/variations/camera_extrinsics_variation.py need attention for the quaternion direction and naming issues.

Important Files Changed

Filename	Overview
isaaclab_arena/variations/variation_base.py	Core abstract base classes for the variation system — clean design, but VariationBaseCfg default sampler_cfg raises NotImplementedError for variations that do not override it.
isaaclab_arena/variations/camera_extrinsics_variation.py	Run-time variation that offsets a camera local position each reset; event-term class uses snake_case instead of PascalCase.
isaaclab_arena/variations/hdr_image_variation.py	Build-time variation that randomly picks an HDR image and attaches it to a DomeLight; logic is sound and guarded with assertions.
isaaclab_arena/environments/arena_env_builder.py	Wires build-time and run-time variations into the env-build pipeline; get_all_variations() has a silent overwrite when the embodiment name matches a scene-asset name.
isaaclab_arena/assets/asset.py	Adds a variations dict to every Asset along with add/get helpers — clean and well-asserted.
isaaclab_arena/tests/test_camera_extrinsics_variation.py	Runtime round-trip test applies quat_apply in the wrong direction when converting parent-frame delta back to camera frame.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class Asset {
        +variations: dict
        +add_variation(variation)
        +get_variation(name)
        +get_variations()
    }
    class VariationBase {
        +name: str
        +cfg: VariationBaseCfg
        +enabled: bool
        +sampler: SamplerBase
        +enable()
        +disable()
        +apply_cfg(cfg)
    }
    class RunTimeVariationBase {
        <<abstract>>
        +build_event_cfg() tuple
    }
    class BuildTimeVariationBase {
        <<abstract>>
        +apply()
    }
    class CameraExtrinsicsVariation {
        +camera_name: str
        +build_event_cfg() tuple
    }
    class HDRImageVariation {
        -_light: DomeLight
        +apply()
    }
    class SamplerBase {
        <<abstract>>
    }
    class UniformSampler {
        +low: Tensor
        +high: Tensor
        +sample(num_samples) Tensor
    }
    class ChoiceSampler {
        +sample(num_samples, choices) list
    }
    class ArenaEnvBuilder {
        +get_all_variations() dict
        +_apply_build_time_variations()
        +_compose_variations_event_cfg()
    }
    Asset "1" --> "*" VariationBase : hosts
    VariationBase <|-- RunTimeVariationBase
    VariationBase <|-- BuildTimeVariationBase
    RunTimeVariationBase <|-- CameraExtrinsicsVariation
    BuildTimeVariationBase <|-- HDRImageVariation
    VariationBase --> SamplerBase : drives
    SamplerBase <|-- UniformSampler
    SamplerBase <|-- ChoiceSampler
    ArenaEnvBuilder --> Asset : reads variations from

Loading

_{Reviews (7): Last reviewed commit: "Add missing files." | Re-trigger Greptile}

[cvolkcvolk]

Looks great!