Variations base 2.5/3.0 - In-memory recording by alexmillane · Pull Request #789 · isaac-sim/IsaacLab-Arena

alexmillane · 2026-06-12T16:11:54Z

Summary

Recording variation samples (in memory).

Detailed description

Recording which values were sampled in variations is critical for downstream analysis
This MR adds listeners to the samplers that record samples taken.
The record is in memory and attached to the env
Note: Implementing this required extending out ArenaEnvBuilder to emit constructor args to gym. I wish that was required, but I can't see another way of passing a run-time object from the compiler to the env.

Not Done

We still are not writing the variations to disk where another process can pick it up (e.g. sensitivity analysis).
This will be added in a follow up MR. (it's non-trivial so I'm breaking this up).

Record every value drawn by an enabled variation's sampler so downstream sensitivity-analysis tooling has the input factors that produced each episode. This adds a sample-observer layer the recorder builds on: SamplerBase gains add_listener/remove_listener and a sample() template method that notifies listeners around the concrete _sample(); VariationBase gains add_sample_listener/remove_sample_listener, re-binding subscriptions onto the sampler rebuilt by apply_cfg so they survive cfg swaps. ArenaEnvBuilder constructs a VariationRecorder, attaches it after Hydra overrides but before any sampling, and exposes it on env.unwrapped. The recorder is stashed on the builder rather than the env cfg because the configclass __post_init__ deep-copies its attributes, which would orphan the listener closures. Signed-off-by: alex <amillane@nvidia.com>

…riations_recording

alexmillane

Self review again.

alexmillane · 2026-06-16T18:55:55Z

+        """Register the Gym env and parse the runtime cfg.
+
+        When ``env_cfg`` is omitted it is compiled/composed in this function from the IsaacLabArenaEnvironment
+        passed to the ArenaEnvBuilder at construction.
+
+        Args:
+            env_cfg: The optional environment cfg to use.
+            env_kwargs: The optional environment kwargs to use.
+
+        Returns:
+            A ``(name, cfg, env_kwargs)`` tuple.
+        """


Can we change the docstring to say

"""
Build env cfg and register the env with gym. Stop short of env.make()

The default operation is to call with no arguments, in which case the env_cfg is built from the Arena description passed to the builder at construction.

"""

alexmillane · 2026-06-16T18:57:06Z

+        render_mode: str | None = None,
    ) -> ManagerBasedEnv:
-        env, _ = self.make_registered_and_return_cfg(env_cfg, render_mode=render_mode)
+        """Build and return the environment from the registered environment configuration.


Can we change the docstring to say

"""
Build env cfg, register the env with gym, and make the env.

The default operation is to call with no arguments, in which case the env_cfg is built from the Arena description passed to the builder at construction.

"""

alexmillane · 2026-06-16T18:57:26Z

    ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
-        name, cfg = self.build_registered(env_cfg)
-        env = gym.make(name, cfg=cfg, render_mode=render_mode)
+        """Build and return the environment from the registered environment configuration and return the configuration.


Can we change the docstring to say

"""
Build env cfg, register the env with gym, and make the env.

The default operation is to call with no arguments, in which case the env_cfg is built from the Arena description passed to the builder at construction.

"""

alexmillane · 2026-06-16T18:58:19Z

+# Print the variations record
+print(env.unwrapped.variations_recorder.summary())
+print("--------------------------------")
+print(env.unwrapped.variations_recorder.details())
+
+
+# %%


alexmillane · 2026-06-16T19:00:07Z


    def sample(self, num_samples: int, choices: Sequence[T]) -> list[T]:
-        """Draw ``num_samples`` items from ``choices``.
+        """Draw ``num_samples`` items from ``choices`` and notify listeners.


remove "and notify listeners"

alexmillane · 2026-06-16T19:00:52Z

+    Observers subscribe via ``add_listener`` to see every value drawn; prefer
+    ``VariationBase.add_sample_listener`` so subscriptions survive sampler swaps. Each sampler
+    family declares its own typed ``sample``; that method draws the value and forwards it to
+    listeners via ``_notify``.


Remove extended docstring.

alexmillane · 2026-06-16T19:01:13Z

+
+        Listeners are invoked synchronously inside ``sample``, in registration order, with
+        the raw sample value (no copy / detach).


remove extended docstring.

alexmillane · 2026-06-16T19:02:00Z

+    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)


Do we use remove_sample_listener anywhere except for the tests? I doubt it. Please remove remove functionality, and in the tests.

alexmillane · 2026-06-16T19:03:05Z

        """Install ``cfg`` as the variation's new source of truth.

-        Replaces :attr:`cfg` and rebuilds :attr:`sampler` from ``cfg.sampler_cfg``.
-        Subclasses with extra derived state should override and call
-        ``super().apply_cfg(cfg)`` first.
+        Replaces ``cfg`` and rebuilds ``sampler`` from ``cfg.sampler_cfg``, re-binding any
+        variation-owned sample listeners onto the new sampler. Subclasses with extra derived
+        state should override and call ``super().apply_cfg(cfg)`` first.



First line of docstring should be: `Apply new ``cfg```

alexmillane · 2026-06-16T19:03:38Z

+
+    def _header_lines(self) -> list[str]:
+        """Return the shared preamble (identity, cfg, sample-call count) for renderers."""
+        #


Complete comment: "Add the title for this variation"

greptile-apps · 2026-06-16T19:20:45Z

Greptile Summary

This PR introduces an in-memory variation recording system: samplers now fire observer callbacks on every draw, variations store those callbacks and re-bind them after cfg/sampler swaps, and a new VariationRecorder aggregates per-variation VariationRecord objects that accumulate every sample taken during both build-time and run-time phases. ArenaEnvBuilder is extended to construct the recorder, attach it before any sampling occurs, and forward it to gym.make via a new env_kwargs dict; IsaacLabArenaManagerBasedRLEnv receives and stores the recorder as self.variations_recorder.

Listener pipeline: SamplerBase gains add_listener/_notify; ContinuousSampler and ChoiceSampler split sample() into a template sample() + abstract _sample(), notifying after each draw. VariationBase holds variation-owned listeners and re-registers them on any new sampler built by apply_cfg, so subscriptions survive cfg swaps.
API change: compose_manager_cfg, build_registered, and make_registered* all return/accept an additional env_kwargs dict, and every caller (tests, scripts, eval runner) has been updated to unpack the 3-tuple and forward kwargs.
In-scope limitation: disk persistence is deferred; the recorder is purely in-memory and will accumulate samples for the entire env lifetime.

Confidence Score: 4/5

The core listener/recorder wiring is correct and the broad API migration across 36 files is mechanically sound; the test suite for the new recorder contains assertions that will fail at runtime due to missing protocol methods on VariationRecorder.

The new test file calls recorder["key"] (subscript) and "key" not in recorder (membership test) on VariationRecorder, but the class defines neither getitem nor contains; both raise TypeError. The same test also asserts recorder.records == [] against a dict, which always evaluates False. These issues were flagged in the previous review round and are still unaddressed, meaning the recorder test suite cannot pass as written. The production code path itself is sound — env construction, listener attachment, and sample capture all work correctly.

isaaclab_arena/tests/test_variations_recorder.py and isaaclab_arena/variations/variations_recorder.py need attention: the recorder's missing getitem/contains methods break multiple test assertions.

Important Files Changed

Filename	Overview
isaaclab_arena/variations/variations_recorder.py	New VariationRecorder and VariationRecord classes; missing getitem/contains on VariationRecorder (breaks test assertions) and cfg stored by reference rather than by snapshot
isaaclab_arena/tests/test_variations_recorder.py	New test suite; multiple assertions will fail at runtime because VariationRecorder lacks getitem and contains, and records is a dict compared to []
isaaclab_arena/environments/arena_env_builder.py	compose_manager_cfg, build_registered, and make_registered return/accept env_kwargs; when an external env_cfg is supplied without env_kwargs the recorder is silently absent (assert fires in IsaacLabArenaManagerBasedRLEnv)
isaaclab_arena/environments/isaaclab_arena_manager_based_env.py	New init receives and stores variations_recorder; hard assert prevents direct instantiation without one, which can produce unhelpful errors for callers that bypass ArenaEnvBuilder
isaaclab_arena/variations/variation_base.py	Adds _sample_listeners list and add_sample_listener; apply_cfg correctly re-binds listeners onto the new sampler after a cfg/sampler swap
isaaclab_arena/variations/sampler_base.py	Adds init with _listeners list and add_listener/_notify helpers to SamplerBase; clean and correct
isaaclab_arena/variations/continuous_sampler.py	sample() becomes a concrete template method that calls _sample() and _notify(); _sample() is the new abstract hook for subclasses
isaaclab_arena/variations/uniform_sampler.py	Adds super().init() call and renames sample() to _sample() to fit the new ContinuousSampler template method pattern
isaaclab_arena/variations/choice_sampler.py	sample() split into sample()/_sample() with _notify() call; listener wiring works through inherited SamplerBase.init
isaaclab_arena/evaluation/eval_runner.py	Mechanical update to unpack 3-tuple from build_registered and forward env_kwargs to make_registered
isaaclab_arena/scripts/imitation_learning/record_demos.py	Updated to thread env_kwargs through create_environment_config and create_environment so gym.make receives the variations_recorder

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Builder as ArenaEnvBuilder
    participant Recorder as VariationRecorder
    participant Variation as VariationBase
    participant Sampler as SamplerBase
    participant Record as VariationRecord
    participant Env as IsaacLabArenaManagerBasedRLEnv

    Builder->>Recorder: VariationRecorder()
    Builder->>Recorder: attach(get_all_variations())
    Recorder->>Variation: add_sample_listener(on_sample)
    Variation->>Sampler: add_listener(on_sample)
    Recorder->>Record: VariationRecord(key, variation.cfg)
    Recorder-->>Builder: recorder ready

    Builder->>Variation: apply() [build-time variations]
    Variation->>Sampler: sample(num_samples)
    Sampler->>Sampler: _sample() → result
    Sampler->>Sampler: _notify(result)
    Sampler->>Record: on_sample(result) → samples.append

    Builder->>Builder: compose env_cfg
    Builder-->>Builder: "return (env_cfg, {variations_recorder: recorder})"

    Builder->>Env: "gym.make(name, cfg=env_cfg, variations_recorder=recorder)"
    Env->>Env: "self.variations_recorder = recorder"
    Env->>Env: super().__init__()

    Note over Env,Sampler: During simulation resets (runtime variations)
    Env->>Variation: event fires → func(env)
    Variation->>Sampler: sample(num_samples)
    Sampler->>Record: on_sample(result) → samples.append

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Builder as ArenaEnvBuilder
    participant Recorder as VariationRecorder
    participant Variation as VariationBase
    participant Sampler as SamplerBase
    participant Record as VariationRecord
    participant Env as IsaacLabArenaManagerBasedRLEnv

    Builder->>Recorder: VariationRecorder()
    Builder->>Recorder: attach(get_all_variations())
    Recorder->>Variation: add_sample_listener(on_sample)
    Variation->>Sampler: add_listener(on_sample)
    Recorder->>Record: VariationRecord(key, variation.cfg)
    Recorder-->>Builder: recorder ready

    Builder->>Variation: apply() [build-time variations]
    Variation->>Sampler: sample(num_samples)
    Sampler->>Sampler: _sample() → result
    Sampler->>Sampler: _notify(result)
    Sampler->>Record: on_sample(result) → samples.append

    Builder->>Builder: compose env_cfg
    Builder-->>Builder: "return (env_cfg, {variations_recorder: recorder})"

    Builder->>Env: "gym.make(name, cfg=env_cfg, variations_recorder=recorder)"
    Env->>Env: "self.variations_recorder = recorder"
    Env->>Env: super().__init__()

    Note over Env,Sampler: During simulation resets (runtime variations)
    Env->>Variation: event fires → func(env)
    Variation->>Sampler: sample(num_samples)
    Sampler->>Record: on_sample(result) → samples.append

_{Reviews (2): Last reviewed commit: "Restore compile_env_notebook.py" | Re-trigger Greptile}

greptile-apps · 2026-06-16T19:20:49Z

        name = self.arena_env.name
-        cfg_entry = env_cfg if env_cfg is not None else self.compose_manager_cfg()
+        if env_cfg is None:
+            env_cfg, env_kwargs = self.compose_manager_cfg()
+        elif env_kwargs is None:
+            env_kwargs = {}


Silent variations_recorder-less path in build_registered

When a caller passes an explicit env_cfg but omits env_kwargs (i.e. build_registered(my_cfg)), the branch at line 380 sets env_kwargs = {} — no variations_recorder. Any subsequent gym.make(**env_kwargs) will reach IsaacLabArenaManagerBasedRLEnv.__init__, which hard-asserts that variations_recorder is not None. The result is an unhelpful AssertionError rather than a clear API-usage error.

Consider either always constructing a fresh VariationRecorder() in that branch, or renaming the parameter to make it explicit that callers must supply one whenever env_cfg is provided externally.

alexmillane added 4 commits June 1, 2026 09:59

Add google style guide docstrings to AGENTS.md

089e7a2

Thread variations recorder through IsaacLabArenaRLEnv

41a8c15

Correct tests that don't need build_registered()

64dce30

alexmillane force-pushed the alex/feature/variations_recording branch from 5cb9285 to 64dce30 Compare June 15, 2026 09:38

alexmillane added 13 commits June 15, 2026 13:13

Added printing of details of the variation samples

7d4e856

lint.

d0f0f20

Merge branch 'main' into alex/feature/agents_md_docstrings

7154b50

Drastically shorten the docstring guidance in AGENTS.md

6353b5f

Merge branch 'alex/feature/agents_md_docstrings' into alex/feature/va…

d988ec5

…riations_recording

Reattack the docstrings.

7c74c8c

WIP

8767c5e

Backticks

8af91de

Merge branch 'alex/feature/agents_md_docstrings' into alex/feature/va…

3f3be33

…riations_recording

Put back in backticks

2006653

Merge branch 'alex/feature/agents_md_docstrings' into alex/feature/va…

5887bbb

…riations_recording

Cleanup

02f32d5

Get rid of base class sample

4afa006

alexmillane commented Jun 16, 2026

View reviewed changes

alexmillane changed the title ~~DRAFT: variation recorder~~ Variations base 2.5/3.0 - In-memory recording Jun 16, 2026

Self review

20950ed

alexmillane marked this pull request as ready for review June 16, 2026 19:13

alexmillane requested review from cvolkcvolk, peterd-NV, qianl-nv, viiik-inside, xyao-nv and zhx06 as code owners June 16, 2026 19:14

Merge branch 'main' into alex/feature/variations_recording

3fb75de

greptile-apps Bot reviewed Jun 16, 2026

View reviewed changes

Restore compile_env_notebook.py

89657f9


		Listeners are invoked synchronously inside ``sample``, in registration order, with
		the raw sample value (no copy / detach).

Conversation

alexmillane commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Detailed description

Not Done

Uh oh!

alexmillane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

alexmillane commented Jun 12, 2026 •

edited

Loading

greptile-apps Bot commented Jun 16, 2026 •

edited

Loading