diff --git a/2026_04_13_sensitivity_analysis.md b/2026_04_13_sensitivity_analysis.md new file mode 100644 index 000000000..6c5d2e1a9 --- /dev/null +++ b/2026_04_13_sensitivity_analysis.md @@ -0,0 +1,611 @@ +# Sensitivity Analysis \- Motivation, Goals, and Design + +This document is intended to lay out: + +* The **goals** of the “Sensitivity Analysis” part of Isaac Lab Arena v0.3 +* What pieces are **missing** that we’re required to build to achieve those goals, and +* What **exists** out there that serves as motivation? + +# Goals and Motivation + +# The main goal of the "Sensitivity Analysis" goal in the v0.3 POR is to enable systematic analysis of how factors (environmental/robot) affect VLA policy performance, by using controlled perturbations within the Isaac Lab Arena. + +# When evaluating a VLA policy across a multi-task benchmark, a binary success rate tells you how well the policy performs, but not why it succeeds or fails. Two policies with the same 30% success rate may fail for entirely different reasons: one may be brittle to camera placement, another to lighting, another to object pose. Without a way to systematically vary these factors and measure their impact, we cannot diagnose failure modes or guide improvements to either the policy or the evaluation setup. + +# The [RoboLab](https://research.nvidia.com/labs/srl/projects/robolab/) project (an internal research benchmark) demonstrated the value of solving this issue. Their sensitivity analysis, published in their paper (NVIDIA, 2026), uses Bayesian posterior estimation (MNPE) over evaluation data to answer questions like "How does policy performance degrade with camera miscalibration?" and "Which lighting conditions are most associated with failure?". Example findings included: + +* **Wrist camera** pose was the single strongest predictor of success across all policies tested, with the posterior sharply concentrated near zero displacement. Policies were far more tolerant to displacement of the external camera. +* **Object distance** from the robot showed a clear peak at \~0.5m, corresponding to the robot's reachable workspace. + +These insights are exactly what Arena users need: understanding not just overall success rates, but which scene parameters drive that performance. **Such insights are *actionable* \- they tell a user what should be changed in order to make the policy better.** + +# What's missing in IsaacLab-Arena today? + +Arena currently has the building blocks for running individual evaluations, but lacks the infrastructure to run controlled perturbation studies and analyze the results. Specifically, we will address two key missing pieces: + +1) **Variation System**: We need a unified system for varying environment parameters across evaluation runs in IsaacLab-Arena + +2) **Analysis Tooling**: Provide dedicated sensitivity analysis tooling (eg. Bayesian posterior estimation \- Take evaluation results and determine which parameters matter.) + +# High-Level API + +This section describes ***what we want*** without considering how we might implement it. + +At a high level, we want a user to be able to: + +* Evaluate their policy in simulated environments with variations selected via CLI arguments/configuration files (rather than having them hard-coded in a script), +* Generate a report that details the sensitivity of output metrics to input variations. + +**Run the policy in an environment with variations** + +Running a single policy on a pick and place task, through our policy\_runner.py script, with variations over the pick\_up\_object, destination object, and lighting could be described as: + +``` +python policy_runner.py \ + --num_envs 10 + franka_srl_pick_and_place \ + pick_up_object.name="Choose(apple, banana, pear)" + pick_up_object.mass="Uniform(0.1, 1.0)" + pick_up_object.color="Choose(red, blue, green)" + destination_object.name="Choose(bowl, pot)" + light.color="Choose(red, blue, green)" +``` + +Running a multi-task, multi-policy test using an experiment configuration file. + +``` +# Run +python isaaclab_arena/evaluation/eval_runner.py \ + --viz kit \ + --eval_jobs_config isaaclab_arena_environments/eval_job_configs/franka_pnp_and_open_door_sensitivity.yaml +``` + +With experiment config file: + +``` +# isaaclab_arena_environments/eval_job_configs/franka_pnp_and_open_door_sensitivity.yaml + +policies: [groot_n16, pi_05] + +pick_and_place_srl_table: + embodiment: franka + pick_up_object: + name: [apple, banana, pear] + mass: Uniform(0.2, 0.5) + light: + color: [red, green, blue] + intensity: Normal(1000, 100) + +open_door_kitchen: + embodiment: franka + pick_up_object: + name: [microwave, toaster_oven] + joint_stiffness: Normal(1.0, 0.1) + light: + hdr: [billiard_hall_robolab, home_office_robolab] +``` + +**Analysing the results (sensitivity analysis)** + +Run an analysis of the output results. This will output a plot detailing the sensitivity of one output metric to one input (variation) factor. + +``` +python analyze_sensitivity.py \ + --results_dir output/franka_pick_and_place_experiment \ + --input_factor pick_up_object.mass \ + --output_metric success_rate \ + --figure_path my_plot.png +``` + +It’s likely that a user will want some interactive way of inspecting the sensitivity of the policy. In this website the user could click on some output metric, and then click on some input factor and get some plot of the sensitivity. + +``` +python sensitivity_website.py \ + --results_dir output/franka_pick_and_place_experiment +``` + +These metrics provide a high-level view of a policy’s performance and its sensitivity to various factors. **This information indicates avenues for improving policy performance \- the whole point of Arena\!** + +# Existing Works **TLDR** + +This short section summarizes the analysis of other frameworks. For the full analysis, including code examples, see: [Background](#existing-works-full). + +**Analysis Tooling** + +* **Sensitivity analysis:** Robolab does a great job of this; we should take this verbatim. +* **Subtasks and Metrics:** Robolab does a great job of this, and we should close the gap between our sub-tasking/metrics system and theirs. + +**Variation System** + +* **Modular Variations:** Colosseum and FactorWorld have a modular system for adding variations that we should take inspiration from. +* **Variation configuration:** FactorWorld allows all parameters of all sources of variation to be configured on the CLI through Hydra. When adding new variations, no additional work is required to have their parameters appear on the command line. Maximum configurability \- we want this. +* **Variation Implementation:** In Colosseum and FactorWorld, each variation is represented as a class that can be registered with a central variations system. We want this extensibility and modularity. +* **Experiment specification:** No reviewed system has a great way of specifying experiments through an experiment specification file. We can make our own improvement here (see above for proposal). + +# Existing Works **FULL** {#existing-works-full} + +This section covers relevant details on how others perform sensitivity analysis. The idea is to take inspiration from the good points and see how we could improve. + +## **RoboLab** + +This section covers how Robolab specifies variations. Sensitivity analysis, i.e. how you make sense of the results *after* the variations have been applied, is done via Simulation-Based Inference (SBI), which is a separate topic. + +Robolab has two mechanisms for specifying variations: + +* **Registration Time:** These bake variations in the env\_cfg by generating several environments by varying factors, say the lights, which are baked in the env\_cfg. + +* **Run-time:** These are randomizations that occur in each environment when it is run. For example, the camera extrinsics randomization, which are implemented as event terms. + +Specifying and applying these variations are handled by top-level experiment scripts (e.g. \`run\_eval\_camera\_pose\_variation.py\` and \`run\_eval\_lighting.py\`). This top-level script takes charge of registering environments (including variations) and inserting randomization events into the environments. + +**Interface** +Variations are implemented through top-level scripts. There is 1 script per type of variation. For example: + +``` +bash + +# Runs the cartesian product of all environmental variations and camrea pose variation types +python robolab/examples/run_eval_camera_pose_variation.py +``` + +**Implementation** +Variations are hardcoded in these top-level scripts: + +``` +robolab/examples/run_eval_camera_pose_variation.py + +# Define perturbations as constants +CAMERA_NAMES_EXTERNAL = ["external_cam"] +CAMERA_POSE_RANGE_EXTERNAL = { + "x": (-0.2, 0.2), + "y": (-0.2, 0.2), + "z": (-0.1, 0.1), + "roll": (-0.2, 0.2), + "pitch": (-0.2, 0.2), + "yaw": (-0.2, 0.2), +} + +... + +# Register environment, with the purturbation in. +env, env_cfg = create_env(base_task_env, + device=args_cli.device, + num_envs=num_envs, + use_fabric=True, + events=camera_pose_event, # Here is the purturbation + policy=args_cli.policy) + +# Run the eval +env_results, msgs, timing = run_episode(env=env, + env_cfg=env_cfg, + episode=run_idx, + save_videos=args_cli.save_videos, + headless=args_cli.headless, + remote_host=args_cli.remote_host, + remote_port=args_cli.remote_port) +``` + +**Supported Variations** +What variations does this framework support? + +* **manipulated object:** \- *(Through USDs not configuration)* +* **receiver object:** \- *(Through USDs not configuration)* +* **background:** background\_hdr, lighting, table material +* **physical:** \- +* **robot:** camera\_pose + +**Pros and Cons** +What can we learn from this framework? + +* **Pros** + * It works and produces amazing results (see paper) + * **Sensitivity Analysis:** They have a mature framework for measuring the sensitivity of the policy to varied factors. + * **Subtasks \+ Metrics:** Xuning has a very mature sub-task/event system that tracks progress towards the ultimate goal (task success). This generates a rich output dataset, against which the inputs (environmental variation) can be correlated. + * **Tasks variance:** Unlike the frameworks below Robolab supports more general tasks than a single table top pick and place with a single focus object. +* **Cons** + * **No variation configuration system:** No unified way of specifying an experiment that includes a list of arbitrary variations. Experiments are effectively captured in the top-level scripts and they test one, hard-coded source of variation. We want a user to be able to specify variations more generally, in some experiment configuration file, without rewriting a top-level script. + * **Output format:** Hard-coded output summary per experiment type. We want some automatic output based on the user specification of the input variations. Basically, without the user having to write some top-level experiment script and specify the data to be saved to the output, all data required for all potential analyses (given the input variations, and output metrics) should be saved. + +![][image1] +Figure 1.0 \- An example of a sensitivity analysis performed in Robolab, of the success rate with respect to the extrinsic camera calibration. + +## **The Colluseum** + +**Interface** +Each environment gets a file describing the possible variations + +``` +colosseum/assets/configs/basketball_in_hoop.py + +env: + task_name: "basketball_in_hoop" + seed: 42 + scene: + factors: + + - variation: object_color + name: manip_obj_color + enabled: False + targets: [ball] + seed: ${env.seed} + + - variation: object_color + name: recv_obj_color + enabled: False + targets: [basket_ball_hoop_visual] + seed: ${env.seed} +``` + +Then we have another file describing experiments to be run + +``` +colosseum/assets/json/basketball_in_hoop.json + +{ + "strategy": [ + { + "spreadsheet_idx": 1, + "variation_name" : "all_mixed", + "enabled": true, + "variations": [ + {"type": "object_color", "name": "manip_obj_color", "enabled": true}, + {"type": "object_color", "name": "recv_obj_color", "enabled": true}, + {"type": "object_texture", "name": "manip_obj_tex", "enabled": true}, + {"type": "object_texture", "name": "recv_obj_tex", "enabled": true}, + {"type": "object_size", "name": "manip_obj_size", "enabled": true}, + {"type": "object_size", "name": "recv_obj_size", "enabled": true}, + {"type": "light_color", "name": "any", "enabled": true}, + {"type": "table_color", "name": "any", "enabled": true}, + {"type": "table_texture", "name": "any", "enabled": true}, + {"type": "distractor_object", "name": "any", "enabled": true}, + {"type": "background_texture", "name": "any", "enabled": true}, + {"type": "camera_pose", "name": "any", "enabled": true}, + {"type": "object_friction", "name": "any", "enabled": true}, + {"type": "object_mass", "name": "any", "enabled": true} + ] + }, + { + "spreadsheet_idx": 2, + "variation_name" : "manip_obj_color", + "enabled": true, + "variations": [ + {"type": "object_color", "name": "manip_obj_color", "enabled": true}, + {"type": "object_color", "name": "recv_obj_color", "enabled": false}, + {"type": "object_texture", "name": "manip_obj_tex", "enabled": false}, + {"type": "object_texture", "name": "recv_obj_tex", "enabled": false}, + {"type": "object_size", "name": "manip_obj_size", "enabled": false}, + {"type": "object_size", "name": "recv_obj_size", "enabled": false}, + {"type": "light_color", "name": "any", "enabled": false}, + {"type": "table_color", "name": "any", "enabled": false}, + {"type": "table_texture", "name": "any", "enabled": false}, + {"type": "distractor_object", "name": "any", "enabled": false}, + {"type": "background_texture", "name": "any", "enabled": false}, + {"type": "camera_pose", "name": "any", "enabled": false}, + {"type": "object_friction", "name": "any", "enabled": false}, + {"type": "object_mass", "name": "any", "enabled": false} + ] + }, +``` + +which sets variations on and off per experiment. + +**Implementation** +The colosseum exposes variations through classes inheriting from \`IVariation\` class for example: + +``` +colosseum/variations/object_color.py + +class ObjectColorVariation(IVariation): + """Object color variation, can change objects' color in the simulation""" + + def __init__( + self, + pyrep: PyRep, + name: Optional[str], + targets_names: List[str] + ): + ... + + def randomize(self) -> None: + """ + Samples a random color and sets it to the objects in the simulation. + Depending on the self._color_scame parameter, all objects will receive + the same color or different colors otherwise if the parameter is false + """ + ... +``` + +Internally, when “randomize” is called this object looks up the objects specified through “target\_names” and changes the colors. + +**Supported Variations** +What variations does this framework support? + +* **manipulated object:** color, texture, size +* **receiver object:** color, texture, size +* **background:** light\_color, table\_texture, distractors, background\_texture +* **physical:** object\_friction, object\_mass +* **robot:** camera\_pose + +**Pros and Cons** +What can we learn from this framework? + +**Pros:** + +* **Config files:** Configuration files describe the possible variations in a scene and what variations are run in an experiment. +* **Variations as objects:** Variations are modular, their code is isolated in their own classes, and they can be registered to tasks/environments (rather than harded coded in some experiment script at the top level). + +**Cons:** + +* **Coupling between task and experiments:** 1:1 linking between a task/environment “basketball\_in\_hoop” and an experiment. *We* want many different types of experiments (potentially) to run on a single environment. +* **Experiment’s file is verbose:** Having to list all possible variations per experiment run results in verbose json files + +**Summary:** I think this is more in the direction of what we want when compared with RoboLab. Variations are modular and can be registered to tasks. On the other hand, there are improvements that we can make regarding the experiment description simplicity (not requiring every variation to be listed), and the experiment flexibility (removing the 1:1 mapping between experiments and tasks). + +## **Robotwin** + +**Similar to frameworks above. Short section.** This section covers how Robotwin specifies variations. + +**Interface** +Variations are specified in a yaml file, which are also overridable on the CLI. + +``` +bash + +python script/eval_policy.py \ + --config ./task_config/demo_clean.yml \ + --overrides --domain_randomization "{'random_background': True, 'cluttered_table': False, 'clean_background_rate': 0.5, 'random_head_camera_dis': 0, 'random_table_height': 0.03, 'random_light': True, 'crazy_random_light_rate': 0.02}" + +``` + +With a configuration file like: + +``` +robotwin/task_config/demo_randomized.yml + +# Define perturbations as constants +render_freq: 0 +episode_num: 50 +use_seed: false +save_freq: 15 +embodiment: [aloha-agilex] +language_num: 100 +domain_randomization: + random_background: true + cluttered_table: true + clean_background_rate: 0.02 + random_head_camera_dis: 0 + random_table_height: 0.03 + random_light: true + crazy_random_light_rate: 0.02 +camera: + head_camera_type: D435 +... +``` + +**Implementation** +Variations take place in the \_base\_task.py + +``` +robotwin/envs/_base_task.py + +# Define perturbations as constants +def setup_scene(self): + + ... + + direction_lights = kwargs.get("direction_lights", [[[0, 0.5, -1], [0.5, 0.5, 0.5]]]) + self.direction_light_lst = [] + for direction_light in direction_lights: + if self.random_light: + direction_light[1] = [ + np.random.rand(), + np.random.rand(), + np.random.rand(), + ] + self.direction_light_lst.append( + self.scene.add_directional_light(direction_light[0], direction_light[1], shadow=shadow)) + + ... +``` + +So the key thing is “if random\_light:” then randomize the lights. So the randomizations are harded in the base class for all tasks. There’s no registration system. + +A lot of this is possible because ALL tasks occur on a predefined table top. So all randomizations, apply to all tasks. + +**Supported Variations** +What variations does this framework support? + +* **manipulated object:** \- +* **receiver object:** \- +* **background:** clutter, background textures, tabletop height, lighting +* **physical:** \- +* **robot:** \- + +**Pros and Cons** +What can we learn from this framework? + +* **Pros** + * **Config files:** The variations are somewhat configurable through a file. +* **Cons** + * **Variations not extensible/module:** All possible variations are hardcoded in the \_base\_task.py. A user would have to make code changes to base\_task to add a variation. + * **Specific to single tabletop:** The whole framework is centered around a single table top. For example a randomization is the global tabletop height. Object placement is tabletop specific. + +## **Factor World** + +**Similar to frameworks above. Short section.** This section covers how Robotwin specifies variations. + +**Interface** +Variations are specified in a yaml file, which are also overridable on the CLI. + +``` +bash + +python -m run_scripted_policy \ + mode=save_video \ + output_dir=/tmp/data \ + num_episodes=10 \ + num_episodes_per_randomize=1 \ + seed=0 \ + factors=[arm_pos,light,object_size] \ + task_name=basketball-v2 + +``` + +With a configuration file looks like: + +``` +factor-world/cfgs/data.yaml + + # Used to sample factor values for data generation. + factors: + # ----- Default factors ----- # + arm_pos: + x_range: [-0.5, 0.5] + y_range: [-0.2, 0.4] + z_range: [-0.15, 0.1] + num_resets_per_randomize: default + seed: ${seed} + object_pos: + x_range: [-0.3, 0.3] + y_range: [-0.1, 0.2] + z_range: [-0, 0] + theta_range: [0, 6.2831853] + num_resets_per_randomize: 1 + seed: ${seed} +... +``` + +Each factor has it’s own parameters. + +**Implementation** +Each factor has its own class. + +``` +robotwin/envs/_base_task.py + +class ObjectPosWrapper(FactorWrapper): + + def __init__(self, + env: gym.Env, + x_range: Tuple[float, float] = (-0.3, 0.3), + y_range: Tuple[float, float] = (-0.1, 0.2), + z_range: Tuple[float, float] = (-0, 0), + theta_range: Tuple[float, float] = (0, 2 * np.pi), + seed: int = None, + **kwargs): + """Creates a new wrapper.""" + super().__init__( + env, + factor_space=spaces.Box( + low=np.array([x_range[0], y_range[0], z_range[0], theta_range[0]]), + high=np.array( + [x_range[1], y_range[1], z_range[1], theta_range[1]]), + dtype=np.float32, + seed=seed), + ) + + ... + + def reset(self): + super().reset() + + # Reset object pos. + self._set_object_pos( + self.object_init_pos, + self.object_init_quat) + + ... +``` + +One cool thing is that all of these constructor arguments are available through the CLI, through hydra. + +**Supported Variations** +What variations does this framework support? + +* **manipulated object:** object\_pos, object\_size, object\_texture, +* **receiver object:** \- +* **background:** floor\_texture, table\_texture, table\_pos, distractor\_xml, distractors +* **physical:** camera\_pos, light +* **robot:** arm pos + +**Pros and Cons** +What can we learn from this framework? + +* **Pros** + * **Extensibility:** Each factor can be added as a new class and it’s arguments are automatically available through the CLI, through hydra. + * **CLI \+ Hydra:** All the variation factor control are available through the CLI automatically. + * **Config files:** The variations are configurable through a file. +* **Cons** + * **Specific to single tabletop:** The whole framework is centered around a single table top. For example the variations rely on their being a single central object for the task. + +# What are we missing to achieve this? + +TBD + +| Description | Risk | Time Required | Priority | +| ----- | ----- | ----- | :---: | +| | | | | +| | | | | + +# Open Questions + +Here are a list of open questions that I can think of that we need to decide on: + +* **Heterogeneous vs. Homogeneous variations** + * **Question:** How do we deal with the fact that some variations can be run in parallel environments (object position, camera extrinsics, etc.) while some variations need to be run sequentially (for example lighting). + * **Approach** + * Experiment to see what, of the common factors of variation, can be run in parallel, and what can’t. + * Need to come up with an orchestration system where the requested variations are broken down into a series of parallal and sequential runs (or in parallel nodes on OSMO). + +* **Sensitivity website/report** + * **Question:** Is the sensitivity analysis fast enough that a meaningful report/website which updates to user requests is possible? + * **Approach:** prototype. + +# Appendix + +``` +# Run +python isaaclab_arena/evaluation/eval_runner.py \ + --viz kit \ + --eval_jobs_config isaaclab_arena_environments/eval_job_configs/franka_pnp_and_open_door_sensitivity.json + + + +# isaaclab_arena_environments/eval_job_configs/franka_pnp_and_open_door_sensitivity.json + + + + +{ + "jobs": [ + { + "name": "franka_pick_and_place", + "arena_env_args": { + "environment": "pick_and_place", + "embodiment": "franka", + "pick_up_object.name": "Choose(apple, banana, pear)", + "pick_up_object.mass": "Uniform(0.1, 0.5)", + "hdr": "Choose(billiard_hall_robolab, home_office_robolab)", + }, + "num_episodes": 100, + }, + { + "name": "franka_open_door", + "arena_env_args": { + "environment": "open_door", + "embodiment": "franka", + "object_with_door.name": "Choose(microwave, toaster_oven)", + "object_with_door.joint_stiffness": "Normal(1.0, 0.1)", + "hdr": "Choose(billiard_hall_robolab, home_office_robolab)", + }, + "num_episodes": 100, + ], + # Run all of the above tests with the following policies. + "policies": [ + "groot", + "pi05" + ] +} +``` + +[image1]: diff --git a/AGENTS.md b/AGENTS.md index 646080c54..346fcaf61 100644 --- a/AGENTS.md +++ b/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. + """ +``` diff --git a/isaaclab_arena/assets/asset.py b/isaaclab_arena/assets/asset.py index 397f5da15..827f84256 100644 --- a/isaaclab_arena/assets/asset.py +++ b/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()) diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index bbc2036b4..b584a8ba9 100644 --- a/isaaclab_arena/assets/object.py +++ b/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.""" diff --git a/isaaclab_arena/assets/object_library.py b/isaaclab_arena/assets/object_library.py index ace4dab07..9e593b02a 100644 --- a/isaaclab_arena/assets/object_library.py +++ b/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. diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 2c8566a48..1c9d641e7 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/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 diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index 3b1e6b5d3..28adb9050 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/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, ) diff --git a/isaaclab_arena/environments/isaaclab_arena_manager_based_env.py b/isaaclab_arena/environments/isaaclab_arena_manager_based_env.py index f6058015c..9d0d6bb5a 100644 --- a/isaaclab_arena/environments/isaaclab_arena_manager_based_env.py +++ b/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 diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py index 21cda1f85..42c13c467 100644 --- a/isaaclab_arena/evaluation/policy_runner.py +++ b/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 diff --git a/isaaclab_arena/examples/compile_env_notebook.py b/isaaclab_arena/examples/compile_env_notebook.py index 109f9f17a..5978f7cb3 100644 --- a/isaaclab_arena/examples/compile_env_notebook.py +++ b/isaaclab_arena/examples/compile_env_notebook.py @@ -11,11 +11,26 @@ import pinocchio # noqa: F401 from isaaclab.app import AppLauncher +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + print("Launching simulation app once in notebook") -simulation_app = AppLauncher() +# headless=False enables the Kit viewer window so we can visually verify per-env +# randomizations (e.g. object tints). Set headless=True for CI / non-GUI runs. +args_cli = get_isaaclab_arena_cli_parser().parse_args([]) +args_cli.headless = False +args_cli.visualizer = "kit" +args_cli.enable_cameras = True +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app # %% +import isaaclab.envs.mdp as mdp # noqa: F401 (kept for the commented-out in-place tint notes below) +from isaaclab.managers import ( # noqa: F401 (kept for the commented-out in-place tint notes below) + EventTermCfg, + SceneEntityCfg, +) + from isaaclab_arena.assets.registries import AssetRegistry from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder @@ -24,18 +39,70 @@ from isaaclab_arena.scene.scene import Scene from isaaclab_arena.utils.pose import Pose +# UniformSampler / UniformSamplerCfg are referenced only by the commented-out +# imperative-path block below; kept around so uncommenting that block is a +# single edit rather than a "and re-add the import" combo. +from isaaclab_arena.variations import ( # noqa: F401 + CameraDecalibrationVariation, + UniformSampler, + UniformSamplerCfg, +) + asset_registry = AssetRegistry() background = asset_registry.get_asset_by_name("kitchen")() -embodiment = asset_registry.get_asset_by_name("franka_ik")() +# embodiment = asset_registry.get_asset_by_name("franka_ik")() +# enable_cameras=True spawns the droid camera rig (incl. ``wrist_camera``) so +# the camera-decalibration variation attached below has something to act on. +embodiment = asset_registry.get_asset_by_name("droid_differential_ik")(enable_cameras=True) cracker_box = asset_registry.get_asset_by_name("cracker_box")() tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")() +dome_light = asset_registry.get_asset_by_name("light")() -cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) +cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) cracker_box.add_relation(IsAnchor()) tomato_soup_can.add_relation(On(cracker_box)) -scene = Scene(assets=[background, cracker_box, tomato_soup_can]) +# --- Embodiment-level variation: wrist-camera decalibration ----------------- +# +# Embodiments are :class:`~isaaclab_arena.assets.asset.Asset` instances, so they +# host variations the same way scene assets do — via ``add_variation``. +# ``ArenaEnvBuilder.get_all_variations`` then merges these into the same +# ``{asset_name: [variation, ...]}`` dict the Hydra schema and the variation +# recorder consume. The override path is keyed on the embodiment's ``name`` +# (``"droid_differential_ik"`` here), shown two cells below. +embodiment.add_variation(CameraDecalibrationVariation("wrist_camera")) + +# --- Variation configuration -------------------------------------------------- +# +# Every ``Object`` ships with a registry of built-in variations (currently just +# ``"color"``), pre-configured with a sensible default sampler. Two surfaces +# can drive a variation: +# +# * **Imperative** (Python): call ``variation.set_sampler(...) / .enable()`` +# directly on the variation object. Kept here as commented-out reference +# code so it's easy to compare against the structured-config path. +# * **Structured / Hydra** (cfg-driven): assemble a list of dotted-path +# override strings that mirror the schema returned by +# ``env_builder.get_variations_schema()`` and hand it to +# ``env_builder.apply_hydra_variation_overrides(...)``. This is the form +# that survives serialisation / CLI overrides and is exercised below. +# +# Both objects randomize along a single RGB axis so the per-env tint is obvious +# at a glance: the cracker box varies red, the tomato soup can varies blue. + +# Imperative path (commented out, replaced by the structured-config overrides +# applied after ``ArenaEnvBuilder`` construction below): +# +# cracker_box_color = cracker_box.get_variation("color") +# cracker_box_color.set_sampler(UniformSampler(low=[0.2, 0.2, 0.0], high=[1.0, 1.0, 0.0])) +# cracker_box_color.enable() +# +# tomato_soup_can_color = tomato_soup_can.get_variation("color") +# tomato_soup_can_color.set_sampler(UniformSamplerCfg(low=[0.0, 0.2, 0.2], high=[0.0, 1.0, 1.0])) +# tomato_soup_can_color.enable() + +scene = Scene(assets=[background, cracker_box, tomato_soup_can, dome_light]) isaaclab_arena_environment = IsaacLabArenaEnvironment( name="reference_object_test", embodiment=embodiment, @@ -44,18 +111,137 @@ args_cli = get_isaaclab_arena_cli_parser().parse_args([]) args_cli.solve_relations = True +# Bump num_envs so we can visually verify per-env color variation. +args_cli.num_envs = 4 +args_cli.visualizer = "kit" +args_cli.enable_cameras = True env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) -env = env_builder.make_registered() + +# %% + +# --- Inspecting the dynamic variations schema -------------------------------- +# +# Before applying any overrides, dump the structured-config schema that +# ``ArenaEnvBuilder.get_variations_schema`` builds from the scene. The schema +# uses each variation's existing ``*Cfg`` directly as a per-variation node — +# ``enabled`` lives on ``VariationBaseCfg`` so every variation cfg already +# carries it. The schema therefore lists every variation knob attached to the +# scene (enabled or not); each entry is pre-populated from the variation's +# current cfg, which at this point is the constructor default (e.g. ``color`` +# disabled, full-RGB-uniform sampler). The override paths printed here line +# up one-to-one with the dotted keys we hand to +# ``apply_hydra_variation_overrides`` in the next cell. +from omegaconf import OmegaConf # noqa: E402 + +variations_schema = env_builder.get_variations_schema() +if variations_schema is None: + print("Scene has no variations attached.") +else: + print(OmegaConf.to_yaml(OmegaConf.structured(variations_schema))) + +# %% + +# --- Structured / Hydra-driven variation overrides --------------------------- +# +# Replaces the imperative ``set_sampler / enable`` calls commented out above. +# Each override string is a dotted path into the schema printed in the +# previous cell: +# +# ..= +# +# Hydra validates the paths against the structured-config schema at compose +# time, so typos / unknown fields (e.g. ``cracker_box.colour.enabled=true``) +# are rejected up front rather than silently ignored. The list below mirrors +# the two color variations the imperative path used to set up: the cracker +# box varies red, the tomato soup can varies blue. +# hydra_variation_overrides = [ +# "cracker_box.color.enabled=true", +# "cracker_box.color.sampler.low=[0.2,0.2,0.0]", +# "cracker_box.color.sampler.high=[1.0,1.0,0.0]", +# "tomato_soup_can.color.enabled=true", +# "tomato_soup_can.color.sampler.low=[0.0,0.2,0.2]", +# "tomato_soup_can.color.sampler.high=[0.0,1.0,1.0]", +# ] +hydra_variation_overrides = [ + "cracker_box.color.enabled=true", + "cracker_box.color.sampler.low=[0.2,0.2,0.2]", + "cracker_box.color.sampler.high=[1.0,1.0,1.0]", + "tomato_soup_can.color.enabled=true", + "tomato_soup_can.color.sampler.low=[0.2,0.2,0.2]", + "tomato_soup_can.color.sampler.high=[1.0,1.0,1.0]", + "light.hdr_image.enabled=true", + # Wrist-camera decalibration: ±5 cm per axis (deliberately exaggerated vs. + # the default ±5 mm so the per-env offsets are obvious in the viewport). + "droid_differential_ik.camera_decalibration.enabled=true", + # "droid_differential_ik.camera_decalibration.sampler.low=[-0.05,-0.05,-0.05]", + # "droid_differential_ik.camera_decalibration.sampler.high=[0.05,0.05,0.05]", + "droid_differential_ik.camera_decalibration.sampler.low=[0.0,0.0,0.0]", + "droid_differential_ik.camera_decalibration.sampler.high=[0.0,0.0,0.30]", +] +env_builder.apply_hydra_variation_overrides(hydra_variation_overrides) + +# Re-dump the schema so we can confirm the overrides landed on the live +# variation cfgs (``enabled: true`` plus the narrowed sampler bounds). +print(OmegaConf.to_yaml(OmegaConf.structured(env_builder.get_variations_schema()))) + +# %% + +# ``compose_manager_cfg`` collects every enabled variation from the scene and +# merges their event terms into ``env_cfg.events`` automatically (see +# ``ArenaEnvBuilder._compose_variations_event_cfg``). No manual plumbing. +env_cfg = env_builder.compose_manager_cfg() +assert env_cfg.scene.replicate_physics is False, "Per-env color variation requires replicate_physics=False; got True." + +# %% + + +# %% + +# --- In-place diffuse tint (still a TODO, see 2026_04_21_color_variation_status.md) -- +# +# The in-place tint path (``randomize_visual_diffuse_tint``) preserves the +# asset's diffuse texture but currently doesn't produce a visible change. +# Left here so it's easy to A/B once the shader-path fix lands. +# env_cfg.events.cracker_box_tint = EventTermCfg( +# func=randomize_visual_diffuse_tint, +# mode="reset", +# params={ +# "asset_cfg": SceneEntityCfg(cracker_box.name), +# "colors": {"r": (0.4, 1.0), "g": (0.4, 1.0), "b": (0.4, 1.0)}, +# }, +# ) + +# %% + +env = env_builder.make_registered(env_cfg) env.reset() # %% +RESET_EVERY_N_STEPS = 10 + # Run some zero actions. -NUM_STEPS = 1000 -for _ in tqdm.tqdm(range(NUM_STEPS)): +NUM_STEPS = 200 +for step in tqdm.tqdm(range(NUM_STEPS)): with torch.inference_mode(): actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) env.step(actions) + if step % RESET_EVERY_N_STEPS == 0: + env.reset() + +# %% + +# --- Variation recorder inspection ------------------------------------------- +# +# ``ArenaEnvBuilder`` attaches a fresh ``VariationRecorder`` to ``env.cfg`` during +# compose (see ``arena_env_builder.compose_manager_cfg``), so every enabled +# variation has a record of the values its sampler actually produced. Each +# record bundles the variation's source id (``.``), the cfg +# that drove it (here, the Hydra-overridden bounds set above), and the ordered +# list of sample tensors — one entry per ``SamplerBase.sample()`` call, shape +# ``(num_envs, *event_shape)``. Useful as a quick sanity check that the +# distribution we asked for is what the policy actually saw. +print(env.unwrapped.cfg.variations_recorder) # %% diff --git a/isaaclab_arena/scene/scene.py b/isaaclab_arena/scene/scene.py index 2faf83060..67b1d29e3 100644 --- a/isaaclab_arena/scene/scene.py +++ b/isaaclab_arena/scene/scene.py @@ -13,12 +13,13 @@ from isaaclab_arena.assets.asset import Asset from isaaclab_arena.assets.object import Object -from isaaclab_arena.assets.object_base import ObjectType +from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.assets.object_set import RigidObjectSet from isaaclab_arena.environments.isaaclab_arena_manager_based_env import IsaacLabArenaManagerBasedRLEnvCfg from isaaclab_arena.utils.configclass import make_configclass from isaaclab_arena.utils.phyx_utils import add_contact_report +from isaaclab_arena.variations.variation_base import VariationBase AssetCfg = Union[AssetBaseCfg, RigidObjectCfg, ArticulationCfg, ContactSensorCfg] @@ -102,6 +103,21 @@ def get_curriculum_cfg(self) -> Any: def get_commands_cfg(self) -> Any: return self.commands_cfg + def get_asset_variations(self) -> dict[str, list[VariationBase]]: + """Return a ``{asset_name: [variation, ...]}`` mapping for every asset that has variations. + + Assets without any attached variation are omitted. Variations are + returned enabled-or-not. + """ + by_asset: dict[str, list[VariationBase]] = {} + for asset in self.assets.values(): + if not isinstance(asset, ObjectBase): + continue + variations = asset.get_variations() + if variations: + by_asset[asset.name] = list(variations) + return by_asset + def get_objects_with_relations(self) -> list[Object | ObjectReference]: """Return all objects in the scene that have at least one relation.""" objects_with_relations: list[Object | ObjectReference] = [] diff --git a/isaaclab_arena/tests/test_build_time_variations.py b/isaaclab_arena/tests/test_build_time_variations.py new file mode 100644 index 000000000..3859e9d88 --- /dev/null +++ b/isaaclab_arena/tests/test_build_time_variations.py @@ -0,0 +1,177 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the build-time variation flavor and the HDR-image variation. + +Sampler-only tests are plain Python. End-to-end variation tests run inside +:func:`~isaaclab_arena.tests.utils.subprocess.run_simulation_app_function` +because constructing assets via the registry pulls in ``isaaclab.sim`` and +the asset library. +""" + +import pytest + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function +from isaaclab_arena.variations.categorical_sampler import CategoricalSampler, CategoricalSamplerCfg + + +def test_categorical_sampler_draws_items_from_choices(): + sampler = CategoricalSampler() + choices = ["a", "b", "c", "d", "e"] + samples = sampler.sample(num_samples=100, choices=choices) + assert isinstance(samples, list) + assert len(samples) == 100 + assert all(s in choices for s in samples) + + +def test_categorical_sampler_returns_actual_item_types(): + """Items in the returned list are the same objects (not copies / indices) as in choices.""" + sampler = CategoricalSampler() + choices = [object(), object(), object()] + [drawn] = sampler.sample(num_samples=1, choices=choices) + assert any(drawn is c for c in choices) + + +def test_categorical_sampler_cfg_builds_live_sampler(): + assert isinstance(CategoricalSamplerCfg().build(), CategoricalSampler) + + +def test_categorical_sampler_rejects_empty_choices(): + sampler = CategoricalSampler() + with pytest.raises(AssertionError): + sampler.sample(num_samples=1, choices=[]) + + +def test_categorical_sampler_notifies_listeners(): + sampler = CategoricalSampler() + seen: list = [] + sampler.add_listener(lambda s: seen.append(s)) + sampler.sample(num_samples=3, choices=["x", "y", "z"]) + assert len(seen) == 1 + assert isinstance(seen[0], list) and len(seen[0]) == 3 + + +def _test_hdr_variation_apply_mutates_dome_light(simulation_app): + from isaaclab_arena.assets.object_library import DomeLight + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.variations.hdr_image_variation import HDRImageVariation, HDRImageVariationCfg + from isaaclab_arena.variations.variation_base import BuildTimeVariationBase + + asset_registry = AssetRegistry() + light = asset_registry.get_asset_by_name("light")() + assert isinstance(light, DomeLight) + + variation = light.get_variation("hdr_image") + assert isinstance(variation, HDRImageVariation) + assert isinstance(variation, BuildTimeVariationBase) + assert not variation.enabled + + variation.apply_cfg(HDRImageVariationCfg(enabled=True, hdr_names=["home_office_robolab"])) + variation.apply() + + texture_file = light.spawner_cfg.texture_file + assert texture_file is not None and texture_file.endswith("home_office.exr"), ( + f"Expected DomeLight.spawner_cfg.texture_file to end with 'home_office.exr', got {texture_file!r}." + ) + # The asset cfg consumed by Isaac Lab's scene builder must reflect the + # newly-bound spawner cfg, since add_hdr re-initialises object_cfg. + assert light.object_cfg.spawn.texture_file == texture_file + return True + + +def test_hdr_variation_apply_mutates_dome_light(): + assert run_simulation_app_function(_test_hdr_variation_apply_mutates_dome_light) + + +def _test_hdr_variation_hydra_override_round_trip(simulation_app): + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.variations import variations_hydra + + asset_registry = AssetRegistry() + light = asset_registry.get_asset_by_name("light")() + scene = Scene(assets=[light]) + + variations_hydra.apply_overrides( + scene.get_asset_variations(), + [ + "light.hdr_image.enabled=true", + "light.hdr_image.hdr_names=[empty_warehouse_robolab]", + ], + ) + + variation = light.get_variation("hdr_image") + assert variation.enabled + assert list(variation.cfg.hdr_names) == ["empty_warehouse_robolab"] + + variation.apply() + assert light.spawner_cfg.texture_file is not None + assert light.spawner_cfg.texture_file.endswith("empty_warehouse.hdr") + return True + + +def test_hdr_variation_hydra_override_round_trip(): + assert run_simulation_app_function(_test_hdr_variation_hydra_override_round_trip) + + +def _test_hdr_variation_unknown_name_asserts(simulation_app): + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.variations.hdr_image_variation import HDRImageVariationCfg + + light = AssetRegistry().get_asset_by_name("light")() + + variation = light.get_variation("hdr_image") + variation.apply_cfg(HDRImageVariationCfg(enabled=True, hdr_names=["does_not_exist_hdr"])) + + try: + variation.apply() + except AssertionError as err: + assert "does_not_exist_hdr" in str(err) + return True + return False + + +def test_hdr_variation_unknown_name_asserts(): + assert run_simulation_app_function(_test_hdr_variation_unknown_name_asserts) + + +def _test_hdr_variation_recorder_captures_chosen_hdr_name(simulation_app): + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.variations import variations_hydra + from isaaclab_arena.variations.variations_recorder import VariationRecorder + + light = AssetRegistry().get_asset_by_name("light")() + scene = Scene(assets=[light]) + + pool = ["home_office_robolab", "empty_warehouse_robolab", "billiard_hall_robolab"] + variations_hydra.apply_overrides( + scene.get_asset_variations(), + [ + "light.hdr_image.enabled=true", + f"light.hdr_image.hdr_names=[{','.join(pool)}]", + ], + ) + + # Attach the recorder *after* Hydra overrides but *before* apply(), which + # mirrors the order ArenaEnvBuilder.compose_manager_cfg uses. + recorder = VariationRecorder() + recorder.attach_to_scene(scene) + + record = recorder["light.hdr_image"] + assert len(record.samples) == 0 + + light.get_variation("hdr_image").apply() + + assert len(record.samples) == 1 + sample = record.samples[0] + # The recorder must capture the chosen HDR *name*, not an index. + assert isinstance(sample, list) and len(sample) == 1 + assert sample[0] in pool + return True + + +def test_hdr_variation_recorder_captures_chosen_hdr_name(): + assert run_simulation_app_function(_test_hdr_variation_recorder_captures_chosen_hdr_name) diff --git a/isaaclab_arena/tests/test_split_hydra_overrides.py b/isaaclab_arena/tests/test_split_hydra_overrides.py new file mode 100644 index 000000000..0c07a8cee --- /dev/null +++ b/isaaclab_arena/tests/test_split_hydra_overrides.py @@ -0,0 +1,92 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for :func:`isaaclab_arena_environments.cli.split_hydra_overrides`. + +Plain-Python unit tests: do not require Isaac Sim. +""" + +import argparse + +import pytest + +from isaaclab_arena_environments.cli import split_hydra_overrides + + +def _parser() -> argparse.ArgumentParser: + """Minimal parser standing in for any of the entry-point parsers.""" + return argparse.ArgumentParser(prog="test_runner") + + +def test_accepts_all_four_hydra_shapes(): + """All four shapes named in the helper docstring pass through unchanged.""" + tokens = [ + "cracker_box.color.enabled=true", + "+cracker_box.color.sampler.low=[0.2,0.2,0.0]", + "++tomato_soup_can.color.sampler.high=[0.0,1.0,1.0]", + "~lighting.intensity", + "~lighting.intensity=0.5", # `~` delete with explicit value is also legal + ] + assert split_hydra_overrides(tokens, _parser()) == tokens + + +def test_empty_rhs_accepted(): + """Empty value (`a.b=`) is a valid Hydra null / empty-string assignment.""" + assert split_hydra_overrides(["a.b="], _parser()) == ["a.b="] + + +def test_preserves_order(): + """Order is preserved so Hydra's later-wins semantics still apply.""" + tokens = [ + "a.b=1", + "a.b=2", + "a.c=3", + ] + assert split_hydra_overrides(tokens, _parser()) == tokens + + +def test_empty_input_returns_empty_list(): + """``parse_known_args`` returning ``[]`` is the common case -- no error.""" + assert split_hydra_overrides([], _parser()) == [] + + +@pytest.mark.parametrize( + "bad_token", + [ + "--object", # typo'd flag (the real one is --object on a subparser, but here it's unknown) + "--unknown_flag", # bare unknown flag + "stray_positional", # bare positional (no '=' so not a Hydra set) + "1.0", # numeric -- not a valid key + "key with space=value", # whitespace not allowed in key + "=value_only", # missing key + "+just_plus", # `+` prefix without `=value` is not a delete (`~` is) + "", # empty token + ], +) +def test_rejects_non_hydra_token(bad_token): + """A non-Hydra leftover must cause the parser to exit with non-zero status.""" + parser = _parser() + with pytest.raises(SystemExit) as exc_info: + split_hydra_overrides([bad_token], parser) + # argparse.error uses exit code 2. + assert exc_info.value.code == 2 + + +def test_rejects_when_mixed_with_valid_tokens(): + """A single bad token poisons the whole batch -- we don't silently drop it.""" + parser = _parser() + with pytest.raises(SystemExit) as exc_info: + split_hydra_overrides(["cracker_box.color.enabled=true", "--object"], parser) + assert exc_info.value.code == 2 + + +def test_error_message_names_bad_tokens(capsys): + """The error message includes the offending token(s) so users can fix typos.""" + parser = _parser() + with pytest.raises(SystemExit): + split_hydra_overrides(["--typo", "stray"], parser) + err = capsys.readouterr().err + assert "--typo" in err + assert "stray" in err diff --git a/isaaclab_arena/variations/__init__.py b/isaaclab_arena/variations/__init__.py new file mode 100644 index 000000000..29dbb8470 --- /dev/null +++ b/isaaclab_arena/variations/__init__.py @@ -0,0 +1,44 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Variation system public API.""" + +from isaaclab_arena.variations.camera_decalibration import ( + CameraDecalibrationVariation, + CameraDecalibrationVariationCfg, +) +from isaaclab_arena.variations.categorical_sampler import CategoricalSampler, CategoricalSamplerCfg +from isaaclab_arena.variations.hdr_image_variation import HDRImageVariation, HDRImageVariationCfg +from isaaclab_arena.variations.object_color import ObjectColorVariation, ObjectColorVariationCfg +from isaaclab_arena.variations.sampler_base import SamplerBase, SamplerBaseCfg +from isaaclab_arena.variations.uniform_sampler import UniformSampler, UniformSamplerCfg +from isaaclab_arena.variations.variation_base import ( + BuildTimeVariationBase, + RunTimeVariationBase, + VariationBase, + VariationBaseCfg, +) +from isaaclab_arena.variations.variations_recorder import VariationRecord, VariationRecorder + +__all__ = [ + "BuildTimeVariationBase", + "CameraDecalibrationVariation", + "CameraDecalibrationVariationCfg", + "CategoricalSampler", + "CategoricalSamplerCfg", + "HDRImageVariation", + "HDRImageVariationCfg", + "ObjectColorVariation", + "ObjectColorVariationCfg", + "RunTimeVariationBase", + "SamplerBase", + "SamplerBaseCfg", + "UniformSampler", + "UniformSamplerCfg", + "VariationBase", + "VariationBaseCfg", + "VariationRecord", + "VariationRecorder", +] diff --git a/isaaclab_arena/variations/camera_decalibration.py b/isaaclab_arena/variations/camera_decalibration.py new file mode 100644 index 000000000..84e6c54bd --- /dev/null +++ b/isaaclab_arena/variations/camera_decalibration.py @@ -0,0 +1,153 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Per-env camera decalibration variation. + +Adds a small sampler-drawn translation on top of a camera's nominal local +position so its observed pose drifts from the calibrated reference. The +nominal (un-decalibrated) local translation is snapshotted on the first event +tick, and each subsequent tick rewrites the local translation to +``nominal + delta`` for the envs being touched, so deltas don't compound +across resets. +""" + +from __future__ import annotations + +import torch +from dataclasses import field +from typing import TYPE_CHECKING + +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab.sensors import Camera, TiledCamera +from isaaclab.utils import configclass + +from isaaclab_arena.variations.sampler_base import SamplerBase +from isaaclab_arena.variations.uniform_sampler import UniformSamplerCfg +from isaaclab_arena.variations.variation_base import RunTimeVariationBase, VariationBaseCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +@configclass +class CameraDecalibrationVariationCfg(VariationBaseCfg): + """Configuration for :class:`CameraDecalibrationVariation`. + + Attributes: + mode: Event mode forwarded to :class:`EventTermCfg`. ``"reset"`` resamples + the decalibration on every episode reset; ``"prestartup"`` picks a + stable per-env offset that persists for the run. + sampler: 3D translation distribution in the camera's parent (USD local) + frame. Defaults to a uniform over ``[-5 mm, +5 mm]`` on every axis. + """ + + mode: str = "reset" + sampler: UniformSamplerCfg = field( + default_factory=lambda: UniformSamplerCfg( + low=[-0.005, -0.005, -0.005], + high=[0.005, 0.005, 0.005], + ) + ) + + +class CameraDecalibrationVariation(RunTimeVariationBase): + """Decalibrate a camera by adding a small offset to its nominal local position. + + The camera's nominal (un-decalibrated) local translation is captured once + on the first event tick, and every subsequent tick rewrites the local + translation to ``nominal + delta`` for the envs being touched. This models + a small mounting / calibration error without compounding across resets, + and keeps wrist-mounted cameras tracking their parent body because only + the camera's local transform (not its world pose) is modified. + + Args: + camera_name: The scene-entity name of the target camera (e.g. + ``"wrist_cam"``). Must resolve to a :class:`Camera` or + :class:`TiledCamera` in ``env.scene``. + cfg: Tunable parameters. Defaults to a + :class:`CameraDecalibrationVariationCfg` with a ``±5 mm`` per-axis + uniform sampler and reset-mode resampling. + sampler: Optional override for the translation distribution. If + ``None``, the sampler in ``cfg`` is used. + """ + + name = "camera_decalibration" + + cfg: CameraDecalibrationVariationCfg + + def __init__( + self, + camera_name: str, + cfg: CameraDecalibrationVariationCfg | None = None, + sampler: SamplerBase | UniformSamplerCfg | None = None, + ): + super().__init__(cfg=cfg if cfg is not None else CameraDecalibrationVariationCfg()) + self.camera_name = camera_name + self.set_sampler(sampler if sampler is not None else self.cfg.sampler) + + def build_event_cfg(self) -> tuple[str, EventTermCfg]: + assert self._sampler is not None, ( + f"CameraDecalibrationVariation on '{self.camera_name}' is enabled but no sampler is set; " + "call .set_sampler(...) before building the env." + ) + event_name = f"{self.camera_name}_decalibration_variation" + event_cfg = EventTermCfg( + func=decalibrate_camera_from_sampler, + mode=self.cfg.mode, + params={ + "asset_cfg": SceneEntityCfg(self.camera_name), + "sampler": self._sampler, + }, + ) + return event_name, event_cfg + + +class decalibrate_camera_from_sampler(ManagerTermBase): + """Add a sampler-drawn translation to a camera's nominal local position. + + Snapshots the camera's factory local translation on first call, then on + every subsequent call writes ``nominal + delta`` per env so decalibrations + don't compound across resets. Operates on the camera's underlying + :class:`~isaaclab.sim.views.XformPrimView` via ``set_local_poses`` so + wrist-mounted cameras keep tracking their parent body. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + sampler: SamplerBase = cfg.params["sampler"] + + camera = env.scene[asset_cfg.name] + assert isinstance(camera, (Camera, TiledCamera)), ( + f"decalibrate_camera_from_sampler expects a Camera or TiledCamera at " + f"scene['{asset_cfg.name}']; got {type(camera).__name__}." + ) + assert tuple(sampler.event_shape) == (3,), ( + "decalibrate_camera_from_sampler expects a sampler with event_shape (3,) over XYZ; " + f"got {tuple(sampler.event_shape)}." + ) + + self._camera = camera + # Snapshot lazily on first __call__ so the camera's view has been + # initialised by the sensor's lifecycle hooks before we read from it. + self._nominal_local_pos: torch.Tensor | None = None + + def __call__( + self, + env: ManagerBasedEnv, # noqa: ARG002 + env_ids: torch.Tensor, + asset_cfg: SceneEntityCfg, # noqa: ARG002 + sampler: SamplerBase, + ): + view = self._camera._view + if self._nominal_local_pos is None: + nominal_pos, _ = view.get_local_poses() + self._nominal_local_pos = nominal_pos.detach().clone() + + sample = sampler.sample(num_samples=len(env_ids)) + deltas = sample.to(device=self._nominal_local_pos.device, dtype=self._nominal_local_pos.dtype) + new_local_pos = self._nominal_local_pos[env_ids] + deltas + view.set_local_poses(translations=new_local_pos, orientations=None, indices=env_ids) diff --git a/isaaclab_arena/variations/categorical_sampler.py b/isaaclab_arena/variations/categorical_sampler.py new file mode 100644 index 000000000..a4e7f365b --- /dev/null +++ b/isaaclab_arena/variations/categorical_sampler.py @@ -0,0 +1,52 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Uniform categorical sampler over an arbitrary ``choices`` sequence. + +The choice domain is not a property of the distribution: it's a per-call +argument passed to :meth:`SamplerBase.sample` (as ``choices``), since +callers typically know the pool only at apply time (e.g. the HDRs to pick +from). The returned samples are the actual items from ``choices``, not +indices, so downstream consumers (e.g. the variation recorder) capture +meaningful values directly. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from isaaclab.utils import configclass + +from isaaclab_arena.variations.sampler_base import SamplerBase, SamplerBaseCfg + + +@configclass +class CategoricalSamplerCfg(SamplerBaseCfg): + """Config for :class:`CategoricalSampler`. + + Intentionally empty: the categorical distribution has no declarative + parameters in this implementation. ``choices`` is supplied per call via + :meth:`SamplerBase.sample`. + """ + + def build(self) -> CategoricalSampler: + return CategoricalSampler() + + +class CategoricalSampler(SamplerBase): + """Uniform categorical sampler returning items drawn from a per-call ``choices`` sequence. + + Each call to :meth:`sample` returns a ``list`` of length ``num_samples``, + where each item is uniformly drawn (with replacement) from ``choices``. + """ + + def _sample(self, num_samples: int, choices: Sequence[Any], **kwargs) -> list[Any]: # noqa: ARG002 + assert num_samples >= 0, f"num_samples must be non-negative; got {num_samples}." + assert len(choices) >= 1, "CategoricalSampler requires a non-empty 'choices' sequence." + indices = torch.randint(low=0, high=len(choices), size=(num_samples,), dtype=torch.long) + return [choices[int(i)] for i in indices] diff --git a/isaaclab_arena/variations/hdr_image_variation.py b/isaaclab_arena/variations/hdr_image_variation.py new file mode 100644 index 000000000..f33806cb8 --- /dev/null +++ b/isaaclab_arena/variations/hdr_image_variation.py @@ -0,0 +1,95 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time variation that picks an HDR environment map for a dome light. + +The HDR is sampled once before env-cfg composition (it can't be swapped on a +running dome light without rebuilding its spawner), and applied via +:meth:`~isaaclab_arena.assets.object_library.DomeLight.add_hdr`. +""" + +from __future__ import annotations + +from dataclasses import field +from typing import TYPE_CHECKING + +from isaaclab.utils import configclass + +from isaaclab_arena.variations.categorical_sampler import CategoricalSamplerCfg +from isaaclab_arena.variations.sampler_base import SamplerBase +from isaaclab_arena.variations.variation_base import BuildTimeVariationBase, VariationBaseCfg + +if TYPE_CHECKING: + from isaaclab_arena.assets.object_library import DomeLight + + +@configclass +class HDRImageVariationCfg(VariationBaseCfg): + """Configuration for :class:`HDRImageVariation`. + + Attributes: + hdr_names: Registered HDR names to sample from (see + :class:`~isaaclab_arena.assets.registries.HDRImageRegistry`). When + empty, the variation samples uniformly across every registered HDR. + sampler: Categorical distribution over indices into the resolved HDR + list. The pool size is passed at :meth:`HDRImageVariation.apply` + time, so the default empty :class:`CategoricalSamplerCfg` works + out of the box. + """ + + hdr_names: list[str] = field(default_factory=list) + sampler: CategoricalSamplerCfg = field(default_factory=CategoricalSamplerCfg) + + +class HDRImageVariation(BuildTimeVariationBase): + """Sample a single HDR and attach it to a :class:`DomeLight` at build time. + + Args: + light: The dome light whose HDR will be set. A reference is captured; + ``apply`` mutates this exact instance. + cfg: Tunable parameters. Defaults to an :class:`HDRImageVariationCfg` + with an empty ``hdr_names`` (i.e. sample over every registered HDR) + and the default categorical sampler. + sampler: Optional override for the categorical distribution. If + ``None``, the sampler in ``cfg`` is used. + """ + + name = "hdr_image" + + cfg: HDRImageVariationCfg + + def __init__( + self, + light: DomeLight, + cfg: HDRImageVariationCfg | None = None, + sampler: SamplerBase | CategoricalSamplerCfg | None = None, + ): + super().__init__(cfg=cfg if cfg is not None else HDRImageVariationCfg()) + self._light = light + self.set_sampler(sampler if sampler is not None else self.cfg.sampler) + + def apply(self) -> None: + from isaaclab_arena.assets.hdr_image import HDRImage # noqa: PLC0415 + from isaaclab_arena.assets.registries import HDRImageRegistry # noqa: PLC0415 + + registry = HDRImageRegistry() + if self.cfg.hdr_names: + for name in self.cfg.hdr_names: + assert registry.is_registered(name), ( + f"HDRImageVariation: HDR name '{name}' is not registered. " + f"Registered HDRs: {sorted(registry.get_all_keys())}." + ) + hdr_names = list(self.cfg.hdr_names) + else: + hdr_names = registry.get_all_keys() + assert hdr_names, "HDRImageVariation: no HDRs are registered; cannot sample." + + assert self.sampler is not None, "HDRImageVariation: sampler not set." + # Pass HDR *names* (not indices) as the categorical sampler's choices, + # so the recorder logs the chosen HDR by name instead of an opaque + # index. + [chosen_name] = self.sampler.sample(num_samples=1, choices=hdr_names) + hdr_cls: type[HDRImage] = registry.get_hdr_by_name(chosen_name) + self._light.add_hdr(hdr_cls()) diff --git a/isaaclab_arena/variations/object_color.py b/isaaclab_arena/variations/object_color.py new file mode 100644 index 000000000..6635a3c64 --- /dev/null +++ b/isaaclab_arena/variations/object_color.py @@ -0,0 +1,178 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Per-asset visual color variation. + +Replaces the asset's bound material with a fresh ``OmniPBR`` instance whose +``diffuse_color_constant`` is sampled per env by an Arena +:class:`~isaaclab_arena.variations.sampler_base.SamplerBase`. The asset's +original diffuse texture is dropped. +""" + +from __future__ import annotations + +import re +import torch +from dataclasses import field +from typing import TYPE_CHECKING + +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab.utils import configclass +from isaaclab.utils.version import compare_versions + +from isaaclab_arena.variations.sampler_base import SamplerBase +from isaaclab_arena.variations.uniform_sampler import UniformSamplerCfg +from isaaclab_arena.variations.variation_base import RunTimeVariationBase, VariationBaseCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from isaaclab_arena.assets.object_base import ObjectBase + + +@configclass +class ObjectColorVariationCfg(VariationBaseCfg): + """Configuration for :class:`ObjectColorVariation`. + + Attributes: + mode: Event mode forwarded to :class:`EventTermCfg`. ``"reset"`` resamples + on every episode reset; ``"prestartup"`` picks a stable color per env. + mesh_name: Sub-mesh selector. Empty string targets all meshes under the + asset's prim. + sampler: RGB distribution. Defaults to a 3D uniform over the full + ``[0, 1]^3`` cube. + """ + + mode: str = "reset" + mesh_name: str = "" + sampler: UniformSamplerCfg = field( + default_factory=lambda: UniformSamplerCfg(low=[0.0, 0.0, 0.0], high=[1.0, 1.0, 1.0]) + ) + + +class ObjectColorVariation(RunTimeVariationBase): + """Randomize an object's visual color per env. + + Requires ``scene.replicate_physics=False`` (the Arena default) so each env + owns its own material prim. + + Args: + asset: The object whose visual color will be varied. Its ``name`` is + used to resolve the scene entity at event time. + cfg: Tunable parameters. Defaults to an :class:`ObjectColorVariationCfg` + with full-RGB-uniform reset-time defaults. + sampler: Optional override for the RGB distribution. If ``None``, the + sampler in ``cfg`` is used. + """ + + name = "color" + + cfg: ObjectColorVariationCfg + + def __init__( + self, + asset: ObjectBase, + cfg: ObjectColorVariationCfg | None = None, + sampler: SamplerBase | UniformSamplerCfg | None = None, + ): + super().__init__(cfg=cfg if cfg is not None else ObjectColorVariationCfg()) + self.asset_name = asset.name + self.set_sampler(sampler if sampler is not None else self.cfg.sampler) + + def build_event_cfg(self) -> tuple[str, EventTermCfg]: + assert self._sampler is not None, ( + f"ObjectColorVariation on '{self.asset_name}' is enabled but no sampler is set; " + "call .set_sampler(...) before building the env." + ) + event_name = f"{self.asset_name}_color_variation" + event_cfg = EventTermCfg( + func=randomize_visual_color_from_sampler, + mode=self.cfg.mode, + params={ + "asset_cfg": SceneEntityCfg(self.asset_name), + "sampler": self._sampler, + "mesh_name": self.cfg.mesh_name, + }, + ) + + return event_name, event_cfg + + +class randomize_visual_color_from_sampler(ManagerTermBase): + """Randomize the visual color of bodies on an asset, sampling via an Arena :class:`SamplerBase`. + + Variant of :class:`isaaclab.envs.mdp.randomize_visual_color` that delegates + RGB sampling to a Python-side :class:`SamplerBase` so the drawn values are + visible to the recording layer. Requires ``omni.replicator.core >= 1.12.4`` + and ``scene.replicate_physics=False``. Like the upstream variant, + randomization is applied to all envs on every call regardless of ``env_ids``. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + # enable replicator extension if not already enabled (local: isaacsim only available with Kit) + from isaacsim.core.utils.extensions import enable_extension # noqa: PLC0415 + + enable_extension("omni.replicator.core") + import omni.replicator.core as rep # noqa: PLC0415 + + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + sampler: SamplerBase = cfg.params["sampler"] + mesh_name: str = cfg.params.get("mesh_name", "") + + assert not env.cfg.scene.replicate_physics, ( + "Cannot randomize visual color with scene replication enabled. " + "Set 'replicate_physics=False' on InteractiveSceneCfg." + ) + + # Match the upstream RGB-only contract: we'll be writing a (num_prims, 3) + # array into 'diffuse_color_constant', so the sampler must produce 3D + # samples. Higher-dim distributions (e.g. with alpha) can be added later + # alongside an extension to the upstream attribute write. + assert tuple(sampler.event_shape) == (3,), ( + "randomize_visual_color_from_sampler expects a sampler with event_shape (3,) over RGB; " + f"got {tuple(sampler.event_shape)}." + ) + + asset = env.scene[asset_cfg.name] + if not mesh_name.startswith("/"): + mesh_name = "/" + mesh_name + mesh_prim_path = f"{asset.cfg.prim_path}{mesh_name}" + # TODO: Need to make it work for multiple meshes (mirrors upstream TODO). + + version = re.match(r"^(\d+\.\d+\.\d+)", rep.__file__.split("/")[-5][21:]).group(1) + assert compare_versions(version, "1.12.4") >= 0, ( + "randomize_visual_color_from_sampler requires omni.replicator.core >= 1.12.4 " + f"(found {version}); the legacy OmniGraph sampling path cannot be driven by an Arena SamplerBase." + ) + + stage = env.sim.stage + prims_group = rep.functional.get.prims(path_pattern=mesh_prim_path, stage=stage) + num_prims = len(prims_group) + + for prim in prims_group: + if prim.IsInstanceable(): + prim.SetInstanceable(False) + + # TODO: Should we specify the value when creating the material? (mirrors upstream TODO). + self.material_prims = rep.functional.create_batch.material( + mdl="OmniPBR.mdl", bind_prims=prims_group, count=num_prims, project_uvw=True + ) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfg: SceneEntityCfg, + sampler: SamplerBase, + mesh_name: str = "", + ): + import omni.replicator.core as rep # noqa: PLC0415 + + num_prims = len(self.material_prims) + sample = sampler.sample(num_samples=num_prims) + random_colors = sample.detach().cpu().numpy() + rep.functional.modify.attribute(self.material_prims, "diffuse_color_constant", random_colors) diff --git a/isaaclab_arena/variations/sampler_base.py b/isaaclab_arena/variations/sampler_base.py new file mode 100644 index 000000000..04071500f --- /dev/null +++ b/isaaclab_arena/variations/sampler_base.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Sampler abstract base for the variation system. + +A :class:`SamplerBase` is a stateless description of how values are drawn. +Each sampler ships a parallel declarative :class:`SamplerBaseCfg` that +``build()`` s into the live sampler, keeping the Hydra-facing surface as +plain data. + +The return type of :meth:`SamplerBase.sample` is intentionally loose +(``Any``): continuous samplers (e.g. :class:`UniformSampler`) return a +``torch.Tensor`` of shape ``(num_samples, *event_shape)``, while categorical +samplers return a ``list`` of items drawn from a per-call ``choices`` +sequence. Listeners (e.g. :class:`~isaaclab_arena.variations.variations_recorder.VariationRecorder`) +should handle either shape. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any + +from isaaclab.utils import configclass + + +@configclass +class SamplerBaseCfg: + """Base configclass for declarative sampler descriptions.""" + + def build(self) -> SamplerBase: + """Return the live :class:`SamplerBase` described by this cfg.""" + raise NotImplementedError( + f"{type(self).__name__}.build() is not implemented; every concrete SamplerBaseCfg " + "subclass must provide a build() that returns its live SamplerBase." + ) + + +class SamplerBase(ABC): + """Abstract distribution over values. + + Samplers are stateless: they hold distribution parameters but no RNG + state. External observers can subscribe via :meth:`add_listener` to see + every value drawn; prefer + :meth:`~isaaclab_arena.variations.variation_base.VariationBase.add_sample_listener` + so subscriptions survive sampler swaps. + """ + + def __init__(self) -> None: + self._listeners: list[Callable[[Any], None]] = [] + + def add_listener(self, listener: Callable[[Any], None]) -> None: + """Register ``listener`` to be called with every sample drawn from this sampler. + + Listeners are invoked synchronously inside :meth:`sample`, in registration + order, with the raw sample value (no copy / detach). + """ + self._listeners.append(listener) + + def remove_listener(self, listener: Callable[[Any], None]) -> None: + """Remove a previously-registered ``listener``.""" + self._listeners.remove(listener) + + def sample(self, num_samples: int, **kwargs) -> Any: + """Draw ``num_samples`` values from this distribution. + + Args: + num_samples: Number of independent samples to draw. + **kwargs: Extra per-call arguments passed through to + :meth:`_sample`. Concrete samplers use this for parameters + that aren't part of the distribution's declarative cfg + (e.g. a list of choices for a categorical sampler). + + Returns: + Whatever the concrete sampler produces. Tensor-based samplers + return a tensor of shape ``(num_samples, *event_shape)``; + categorical samplers return a ``list`` of length ``num_samples``. + """ + sample = self._sample(num_samples, **kwargs) + for listener in self._listeners: + listener(sample) + return sample + + @abstractmethod + def _sample(self, num_samples: int, **kwargs) -> Any: + """Draw ``num_samples`` values from this distribution.""" + ... diff --git a/isaaclab_arena/variations/uniform_sampler.py b/isaaclab_arena/variations/uniform_sampler.py new file mode 100644 index 000000000..d2cb7c3f2 --- /dev/null +++ b/isaaclab_arena/variations/uniform_sampler.py @@ -0,0 +1,66 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Uniform distribution sampler over ``[low, high]``.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from dataclasses import field + +from isaaclab.utils import configclass + +from isaaclab_arena.variations.sampler_base import SamplerBase, SamplerBaseCfg + + +@configclass +class UniformSamplerCfg(SamplerBaseCfg): + """Config for :class:`UniformSampler`. + + Attributes: + low: Lower bound per dimension. Length determines the sampler's + ``event_shape``. + high: Upper bound per dimension. Must have the same length as ``low`` + and be element-wise ``>= low``. + """ + + low: list[float] = field(default_factory=lambda: [0.0]) + high: list[float] = field(default_factory=lambda: [1.0]) + + def build(self) -> UniformSampler: + return UniformSampler(low=self.low, high=self.high) + + +class UniformSampler(SamplerBase): + """Uniform sampler over ``[low, high]``. + + ``low`` and ``high`` may be scalars or broadcast-compatible sequences; + samples are drawn with ``low + (high - low) * U(0, 1)``. + """ + + def __init__(self, low: float | Sequence[float], high: float | Sequence[float]): + super().__init__() + low_t = torch.as_tensor(low, dtype=torch.float32) + high_t = torch.as_tensor(high, dtype=torch.float32) + assert ( + low_t.shape == high_t.shape + ), f"UniformSampler low/high must have matching shape; got {tuple(low_t.shape)} vs {tuple(high_t.shape)}." + assert torch.all( + low_t <= high_t + ), f"UniformSampler requires low <= high elementwise; got low={low_t}, high={high_t}." + self.low = low_t + self.high = high_t + + @property + def event_shape(self) -> torch.Size: + """Shape of a single sample.""" + return self.low.shape + + def _sample(self, num_samples: int, **kwargs) -> torch.Tensor: # noqa: ARG002 + assert num_samples >= 0, f"num_samples must be non-negative; got {num_samples}." + shape = (num_samples, *self.event_shape) + u = torch.rand(shape) + return self.low + (self.high - self.low) * u diff --git a/isaaclab_arena/variations/variation_base.py b/isaaclab_arena/variations/variation_base.py new file mode 100644 index 000000000..2c381bbdd --- /dev/null +++ b/isaaclab_arena/variations/variation_base.py @@ -0,0 +1,193 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Variation abstract base classes. + +A :class:`VariationBase` describes one knob to turn on the scene: a sampler +that drives it together with a hook that realises it. Variations are attached +to any :class:`~isaaclab_arena.assets.asset.Asset` (scene objects or +embodiments) via :meth:`~isaaclab_arena.assets.asset.Asset.add_variation`, +start disabled, and are flipped on either imperatively (:meth:`VariationBase.enable`) +or via Hydra overrides. + +Concrete variations subclass one of two flavors: + +* :class:`RunTimeVariationBase` for variations realised via an event term + that runs during simulation (e.g. per-reset randomization). +* :class:`BuildTimeVariationBase` for variations that sample once and mutate + asset configs before the env cfg is composed (e.g. picking an HDR for a + dome light). +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, ClassVar + +from isaaclab.managers import EventTermCfg +from isaaclab.utils import configclass + +from isaaclab_arena.variations.sampler_base import SamplerBase, SamplerBaseCfg + + +@configclass +class VariationBaseCfg: + """Base configclass for :class:`VariationBase` instances. + + Attributes: + enabled: Whether this variation should be applied. Defaults to ``False`` + so users opt in explicitly via :meth:`VariationBase.enable` or a + Hydra override ``..enabled=true``. + """ + + enabled: bool = False + + +class VariationBase(ABC): + """Abstract variation. + + A variation binds a target asset and a + :class:`~isaaclab_arena.variations.sampler_base.SamplerBase` together with + a hook that realises the variation. Concrete subclasses declare a + class-level :attr:`name`, pair themselves with a :class:`VariationBaseCfg` + subclass, and inherit from one of the two flavored bases + (:class:`RunTimeVariationBase` or :class:`BuildTimeVariationBase`). + """ + + #: Short, unique identifier for this variation kind (e.g. ``"color"``). + name: ClassVar[str] + + #: The configclass instance holding this variation's tunable parameters. + cfg: VariationBaseCfg + + def __init__(self, cfg: VariationBaseCfg): + self.cfg = cfg + self._sampler: SamplerBase | None = None + self._sample_listeners: list[Callable[[Any], None]] = [] + + @property + def enabled(self) -> bool: + """Whether this variation is active and should be built into ``events_cfg``.""" + return self.cfg.enabled + + def enable(self) -> None: + """Mark this variation as active.""" + self.cfg.enabled = True + + def disable(self) -> None: + """Mark this variation as inactive.""" + self.cfg.enabled = False + + @property + def sampler(self) -> SamplerBase | None: + """The sampler driving this variation, or ``None`` if not yet set.""" + return self._sampler + + def set_sampler(self, sampler: SamplerBase | SamplerBaseCfg) -> None: + """Replace this variation's sampler. + + A :class:`SamplerBaseCfg` is built into a live sampler and written back to + ``self.cfg.sampler`` if the cfg has one (the declarative path). A bare + :class:`SamplerBase` is stored directly without touching ``self.cfg`` (the + imperative escape hatch). + """ + assert isinstance( + sampler, (SamplerBase, SamplerBaseCfg) + ), f"set_sampler expects a SamplerBase or SamplerBaseCfg; got {type(sampler).__name__}." + if isinstance(sampler, SamplerBaseCfg): + new_sampler = sampler.build() + new_cfg_sampler: SamplerBaseCfg | None = sampler + else: + new_sampler = sampler + new_cfg_sampler = None + + # Re-bind variation-owned listeners so a sampler swap doesn't drop subscriptions. + if self._sampler is not None: + for listener in self._sample_listeners: + self._sampler.remove_listener(listener) + self._sampler = new_sampler + for listener in self._sample_listeners: + self._sampler.add_listener(listener) + + if new_cfg_sampler is not None and hasattr(self.cfg, "sampler"): + self.cfg.sampler = new_cfg_sampler + + def add_sample_listener(self, listener: Callable[[Any], None]) -> None: + """Subscribe ``listener`` to every sample drawn by this variation's sampler. + + Listeners are stored on the variation, so they survive subsequent + :meth:`set_sampler` calls. + """ + self._sample_listeners.append(listener) + if self._sampler is not None: + self._sampler.add_listener(listener) + + def remove_sample_listener(self, listener: Callable[[Any], None]) -> None: + """Unsubscribe a previously-registered ``listener``.""" + self._sample_listeners.remove(listener) + if self._sampler is not None: + self._sampler.remove_listener(listener) + + def apply_cfg(self, cfg: VariationBaseCfg) -> None: + """Install ``cfg`` as the variation's new source of truth. + + Replaces :attr:`cfg` wholesale and rebuilds the live :class:`SamplerBase` + from the new sampler cfg if the cfg carries one. Subclasses with + additional derived state should override and call ``super().apply_cfg(cfg)`` + first. + + Args: + cfg: The cfg to install. Must be an instance of the same + :class:`VariationBaseCfg` subclass this variation accepts. + """ + self.cfg = cfg + sampler_cfg = getattr(cfg, "sampler", None) + if isinstance(sampler_cfg, SamplerBaseCfg): + self.set_sampler(sampler_cfg) + + +class RunTimeVariationBase(VariationBase): + """Variation that is applied during simulation via an ``EventTermCfg``. + + Concrete subclasses produce an event term that the env's event manager + invokes at ``reset`` / ``prestartup`` / ``interval``. Use this flavor when + the underlying property can be flipped at run time (e.g. visual color, + initial pose, mass). + """ + + @abstractmethod + def build_event_cfg(self) -> tuple[str, EventTermCfg]: + """Return the event term that realises this variation. + + Returns: + A ``(name, cfg)`` pair. ``name`` must be unique across all enabled + variations in the scene. + """ + ... + + +class BuildTimeVariationBase(VariationBase): + """Variation that is sampled once and applied before env build. + + Use this flavor for properties that can't (or shouldn't) be changed + in-flight during simulation: HDR environment maps, USD asset swaps, + spawner parameters baked into a config, etc. + + :meth:`apply` is invoked by + :class:`~isaaclab_arena.environments.arena_env_builder.ArenaEnvBuilder` + after Hydra overrides have been pushed through :meth:`apply_cfg` and + before the scene cfg is materialised, so mutations to asset configs are + visible to env cfg composition. Subclasses are expected to hold direct + references to the asset(s) they mutate (captured at construction time). + """ + + @abstractmethod + def apply(self) -> None: + """Sample and mutate the bound asset(s) in place to realise this variation. + + Called exactly once per env build, while the variation is enabled. + """ + ... diff --git a/isaaclab_arena/variations/variations_hydra.py b/isaaclab_arena/variations/variations_hydra.py new file mode 100644 index 000000000..bc9b3b40c --- /dev/null +++ b/isaaclab_arena/variations/variations_hydra.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Hydra-overridable cfg composition for a scene's variations. + +These helpers build a typed Hydra schema from a ``{asset_name: [variation, ...]}`` +mapping, compose override strings against that schema, and push the +resulting per-variation cfgs back through +:meth:`~isaaclab_arena.variations.variation_base.VariationBase.apply_cfg`. +""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import field, make_dataclass +from typing import TYPE_CHECKING, Any + +from hydra import compose, initialize +from hydra.core.config_store import ConfigStore +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +if TYPE_CHECKING: + from isaaclab_arena.variations.variation_base import VariationBase + + +def _asset_class_name(asset_name: str) -> str: + """Convert ``"cracker_box"`` to ``"CrackerBoxVariationsCfg"``.""" + camel = "".join(part.capitalize() for part in asset_name.split("_")) + return f"{camel}VariationsCfg" + + +def build_schema(variations: dict[str, list[VariationBase]]) -> type | None: + """Return the dataclass describing ``variations``, or ``None`` if the mapping is empty. + + The class has one field per asset; each asset field's type is a + dataclass whose fields are the attached variations' cfgs. Each + per-variation field is typed as the variation's own ``*Cfg`` and + pre-populated by deep-copying its current live cfg, so override paths + line up one-to-one with cfg attribute paths. + + Args: + variations: ``{asset_name: [variation, ...]}`` mapping, typically + from :meth:`~isaaclab_arena.scene.scene.Scene.get_asset_variations`. + + Returns: + The dynamically-built ``VariationsCfg`` dataclass type, or ``None`` + when ``variations`` is empty. + """ + if not variations: + return None + + asset_fields: list[tuple[str, type, Any]] = [] + for asset_name, asset_variations in variations.items(): + variation_fields: list[tuple[str, type, Any]] = [] + for variation in asset_variations: + cfg_cls = type(variation.cfg) + default_cfg = deepcopy(variation.cfg) + variation_fields.append((variation.name, cfg_cls, field(default_factory=lambda d=default_cfg: deepcopy(d)))) + asset_cls = make_dataclass(_asset_class_name(asset_name), variation_fields) + asset_fields.append((asset_name, asset_cls, field(default_factory=asset_cls))) + return make_dataclass("VariationsCfg", asset_fields) + + +def load_cfg_from_flags( + variations: dict[str, list[VariationBase]], + hydra_overrides: list[str], +) -> Any | None: + """Compose Hydra override strings into a typed ``VariationsCfg`` instance. + + Builds the schema from :func:`build_schema`, composes the overrides + against it, and converts the result to typed dataclass form via + :func:`omegaconf.OmegaConf.to_object`. Safe to call repeatedly: + :class:`~hydra.core.global_hydra.GlobalHydra` is cleared on entry. + + Args: + variations: ``{asset_name: [variation, ...]}`` mapping that defines + the schema shape. + hydra_overrides: Hydra override strings. See :func:`apply_overrides` + for examples. + + Returns: + The composed ``VariationsCfg`` instance, or ``None`` when + ``variations`` is empty. + """ + schema_cls = build_schema(variations) + if schema_cls is None: + return None + ConfigStore.instance().store(name="arena_variations_schema", node=schema_cls) + if GlobalHydra.instance().is_initialized(): + GlobalHydra.instance().clear() + with initialize(version_base=None, config_path=None): + composed = compose(config_name="arena_variations_schema", overrides=hydra_overrides) + return OmegaConf.to_object(composed) + + +def apply_overrides( + variations: dict[str, list[VariationBase]], + hydra_overrides: list[str], +) -> None: + """Apply Hydra-style overrides to ``variations`` in-place. + + Composes ``hydra_overrides`` into a typed ``VariationsCfg`` and pushes + each per-variation cfg through + :meth:`~isaaclab_arena.variations.variation_base.VariationBase.apply_cfg`. + + Args: + variations: ``{asset_name: [variation, ...]}`` mapping whose cfgs + will be replaced by the composed values. + hydra_overrides: Hydra override strings, dotted-path syntax + mirroring the schema attribute paths. Example:: + + apply_overrides(scene.get_asset_variations(), [ + "cracker_box.color.enabled=true", + "cracker_box.color.sampler.low=[0.2,0.2,0.0]", + "cracker_box.color.sampler.high=[1.0,1.0,0.0]", + ]) + """ + composed = load_cfg_from_flags(variations, hydra_overrides) + if composed is None: + return + for asset_name, asset_variations in variations.items(): + asset_cfg = getattr(composed, asset_name) + for variation in asset_variations: + variation_cfg = getattr(asset_cfg, variation.name) + variation.apply_cfg(variation_cfg) diff --git a/isaaclab_arena/variations/variations_recorder.py b/isaaclab_arena/variations/variations_recorder.py new file mode 100644 index 000000000..439f6a1c8 --- /dev/null +++ b/isaaclab_arena/variations/variations_recorder.py @@ -0,0 +1,142 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""In-memory recording of variation samples. + +A :class:`VariationRecorder` collects every value drawn by an enabled variation's +sampler so downstream sensitivity-analysis tooling has the input factors that +produced each episode. The recorder is explicitly attached by the caller: there +is no singleton or global lookup. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from omegaconf import OmegaConf + +if TYPE_CHECKING: + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.variations.variation_base import VariationBase, VariationBaseCfg + + +class VariationRecord: + """Per-variation slice of a :class:`VariationRecorder`. + + Bundles a variation's identity (:attr:`source_id`), the cfg driving its + sampler at attach time (:attr:`cfg`), and the ordered sequence of samples + drawn into the record (:attr:`samples`). + """ + + def __init__(self, source_id: str, cfg: VariationBaseCfg) -> None: + self.source_id = source_id + #: Cfg reference captured at :meth:`VariationRecorder.attach` time. The recorder + #: is attached after Hydra overrides have been applied, so the cfg is + #: treated as finalized; deep-copy if a frozen archival snapshot is needed. + self.cfg = cfg + #: One entry per :meth:`~isaaclab_arena.variations.sampler_base.SamplerBase.sample` + #: call. Tensor samples are stored detached on CPU; non-tensor samples + #: (e.g. a list returned by a categorical sampler) are stored as-is. + self.samples: list[Any] = [] + + def summary(self) -> str: + """Return a multi-line human-readable summary of this record.""" + lines = [f"--- {self.source_id} ---", "cfg:"] + lines.append(OmegaConf.to_yaml(OmegaConf.structured(self.cfg)).rstrip()) + lines.append(f"sample calls: {len(self.samples)}") + if self.samples: + first = self.samples[0] + if isinstance(first, torch.Tensor): + stacked_shape = (len(self.samples), *tuple(first.shape)) + lines.append(f"stacked shape: {stacked_shape}") + lines.append(f"first call: {first.tolist()}") + lines.append(f"last call: {self.samples[-1].tolist()}") + else: + lines.append(f"first call: {first!r}") + lines.append(f"last call: {self.samples[-1]!r}") + return "\n".join(lines) + + def __str__(self) -> str: + return self.summary() + + +class VariationRecorder: + """Records every sample drawn by attached variations, grouped per variation. + + Records are keyed by ``source_id`` (by convention ``"{asset}.{variation}"``). + Serialisation is intentionally not handled here. + """ + + def __init__(self) -> None: + self._records: dict[str, VariationRecord] = {} + + def attach(self, variations: dict[str, list[VariationBase]]) -> None: + """Attach every enabled variation in ``variations`` under ``"{asset}.{variation}"``. + + Used by :class:`~isaaclab_arena.environments.arena_env_builder.ArenaEnvBuilder` + to attach variations sourced from both the scene and the embodiment. + Disabled variations are skipped. + """ + for asset_name, asset_variations in variations.items(): + for variation in asset_variations: + if not variation.enabled: + continue + self._attach(f"{asset_name}.{variation.name}", variation) + + def attach_to_scene(self, scene: Scene) -> None: + """Attach every enabled variation in ``scene`` under ``"{asset}.{variation}"``.""" + self.attach(scene.get_asset_variations()) + + def _attach(self, source_id: str, variation: VariationBase) -> None: + """Subscribe this recorder to ``variation`` under ``source_id``. + + The listener is registered via + :meth:`~isaaclab_arena.variations.variation_base.VariationBase.add_sample_listener` + so it survives subsequent sampler swaps. + + Args: + source_id: Identifier the record is stored under. By convention + ``"{asset_name}.{variation.name}"``. + variation: The variation to observe. + """ + assert source_id not in self._records, ( + f"VariationRecorder: source_id '{source_id}' is already attached. " + "Re-attaching the same variation would create two independent records and " + "double-record subsequent samples; detach or use a different source_id instead." + ) + + record = VariationRecord(source_id=source_id, cfg=variation.cfg) + self._records[source_id] = record + + def on_sample(sample: Any) -> None: + if isinstance(sample, torch.Tensor): + record.samples.append(sample.detach().cpu()) + else: + record.samples.append(sample) + + variation.add_sample_listener(on_sample) + + @property + def records(self) -> list[VariationRecord]: + """All per-variation records, in attach order.""" + return list(self._records.values()) + + def summary(self) -> str: + """Return a multi-line human-readable summary of every attached record.""" + parts = [f"VariationRecorder: {len(self._records)} record(s)"] + for record in self._records.values(): + parts.append("") + parts.append(record.summary()) + return "\n".join(parts) + + def __str__(self) -> str: + return self.summary() + + def __getitem__(self, source_id: str) -> VariationRecord: + return self._records[source_id] + + def __contains__(self, source_id: str) -> bool: + return source_id in self._records diff --git a/isaaclab_arena_environments/cli.py b/isaaclab_arena_environments/cli.py index d4ceab8d6..33695a3ac 100644 --- a/isaaclab_arena_environments/cli.py +++ b/isaaclab_arena_environments/cli.py @@ -7,6 +7,7 @@ import argparse import importlib +import re from typing import TYPE_CHECKING from isaaclab_arena.assets.registries import EnvironmentRegistry @@ -17,6 +18,47 @@ from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder +# Hydra override token shapes we accept on the CLI after the env subcommand. See +# https://hydra.cc/docs/advanced/override_grammar/basic/ for the upstream grammar. +# We deliberately match a conservative subset: +# - ``key.path=value`` (set / append-or-set) +# - ``+key.path=value`` (force-add) +# - ``++key.path=value`` (force-set, error if missing) +# - ``~key.path`` or ``~key.path=value`` (delete; value optional) +# The leading ``~`` makes the trailing ``=value`` optional; in every other +# shape ``=`` is mandatory, so bare positionals like ``stray_token`` do *not* +# pass through as overrides -- they get raised as unrecognised by +# :func:`split_hydra_overrides`. +_HYDRA_KEY = r"[A-Za-z_][A-Za-z0-9_.]*" +_HYDRA_OVERRIDE_RE = re.compile(rf"^(?:~{_HYDRA_KEY}(?:=.*)?|(?:\+{{1,2}})?{_HYDRA_KEY}=.*)$") + + +def split_hydra_overrides(unknown: list[str], parser: argparse.ArgumentParser) -> list[str]: + """Pull Hydra-shaped override tokens out of an argparse ``unknown`` list. + + Any leftover that does not match a Hydra override shape (see + :data:`_HYDRA_OVERRIDE_RE`) is rejected via ``parser.error``, exiting the + script with code 2 — the same behaviour strict :meth:`parse_args` had. + + Args: + unknown: Second return value of ``parser.parse_known_args()``. + parser: The parser the unknowns came from; used to format the error. + + Returns: + The Hydra override tokens, in original order. + """ + overrides: list[str] = [] + bad: list[str] = [] + for token in unknown: + if _HYDRA_OVERRIDE_RE.match(token): + overrides.append(token) + else: + bad.append(token) + if bad: + parser.error(f"unrecognized arguments: {' '.join(bad)}") + return overrides + + def ensure_environments_registered(): """Trigger registration of all environments in the ``isaaclab_arena_environments`` package. @@ -94,7 +136,18 @@ def get_isaaclab_arena_environments_cli_parser( return args_parser -def get_arena_builder_from_cli(args_cli: argparse.Namespace) -> ArenaEnvBuilder: +def get_arena_builder_from_cli( + args_cli: argparse.Namespace, + hydra_overrides: list[str] | None = None, +) -> ArenaEnvBuilder: + """Build an :class:`ArenaEnvBuilder` from parsed CLI args. + + Args: + args_cli: Parsed argparse namespace; must carry ``example_environment``. + hydra_overrides: Optional Hydra variation override strings (e.g. + ``"cracker_box.color.enabled=true"``). When non-empty, applied via + :meth:`ArenaEnvBuilder.apply_hydra_variation_overrides`. + """ from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder ensure_environments_registered() @@ -107,4 +160,6 @@ def get_arena_builder_from_cli(args_cli: argparse.Namespace) -> ArenaEnvBuilder: example_env = env_registry.get_component_by_name(args_cli.example_environment)() env_builder = ArenaEnvBuilder(example_env.get_env(args_cli), args_cli) + if hydra_overrides: + env_builder.apply_hydra_variation_overrides(hydra_overrides) return env_builder