Restore RL entrypoint backend state by AntoineRichard · Pull Request #6631 · isaac-sim/IsaacLab

AntoineRichard · 2026-07-20T18:00:17Z

Description

The reusable reinforcement learning entrypoints introduced in #6553 restore sys.argv after in-process calls, but backend-specific global settings can still leak into the caller.

This PR scopes those mutations to each invocation:

RSL-RL training parses arguments before temporarily enabling its Torch backend flags and restores all four flags on success or failure.
SKRL training restores skrl.config.jax.backend, including deleting it when the attribute did not exist before the call.
SKRL playback scopes the same state inside main(), covering both unified dispatch and direct imported calls.
Adds failure-path regressions for rejected RSL-RL requests and exceptional SKRL train/play calls.

Relationship

Focused follow-up to Refactor train and play into rl subpackage #6553.
Independent of the benchmark PR stack.

Type of change

Bug fix

Validation

./isaaclab.sh -p -m pytest source/isaaclab_rl/test/test_entrypoints.py -q (11 passed)
./isaaclab.sh -p tools/changelog/cli.py check develop
./isaaclab.sh -f

The full source/isaaclab_rl/test directory was also attempted locally but did not complete or emit results within the five-minute command timeout.

Checklist

I have read and understood the contribution guidelines
I have run the pre-commit checks
I have added regression tests
I have added the required changelog fragment
My name already exists in CONTRIBUTORS.md

Scope Torch and SKRL global backend settings to reusable in-process training and playback calls. Preserve attributes that were initially absent and add failure-path regressions for RSL-RL and SKRL.

greptile-apps · 2026-07-20T18:05:59Z

Greptile Summary

This PR fixes global state leakage from in-process RL entrypoints by scoping backend mutations to each invocation. A new _state.py module provides two context managers — preserve_attribute and scoped_torch_backend_flags — used by the RSL-RL and SKRL backends to save and restore caller state on both success and failure paths.

train_rsl_rl.run() now parses arguments first, then enters scoped_torch_backend_flags so that a --help / early-exit request never mutates the caller's Torch flags at all.
train_skrl.run() conditionally preserves skrl.config.jax.backend only when ml_framework is JAX, while play_skrl.main() unconditionally wraps _main() with the same guard regardless of framework — a minor asymmetry, but both paths are functionally correct.
Four regression tests are added covering the exception, rejection, and JAX-backend-cleanup scenarios.

Confidence Score: 4/5

Safe to merge; the state-restoration logic is correct on all success and failure paths for both RSL-RL and SKRL.

The core fix is well-designed: preserve_attribute and scoped_torch_backend_flags correctly bracket mutations in all backends, and the ExitStack-based approach ensures cleanup on exceptions too. The only gaps are in the tests — one test re-executes module-level script code (including a version-guarded exit()) via runpy.run_module, making it brittle if the installed skrl version changes; another test only exercises the path where flags are never written rather than the path where they are written and then must be restored. These are test-quality issues, not correctness issues in the production code.

The new regression tests in test_entrypoints.py would benefit from a second look, particularly test_skrl_play_main_restores_jax_backend and test_rejected_rsl_training_preserves_torch_backend_state.

Important Files Changed

Filename	Overview
source/isaaclab_rl/isaaclab_rl/entrypoints/_state.py	New module: clean, correct implementations of preserve_attribute and scoped_torch_backend_flags context managers using ExitStack
source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py	run() now parses args before entering scoped_torch_backend_flags, correctly preventing flag mutation on early-exit paths like --help
source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_skrl.py	JAX backend is set in run() before _run() and conditionally preserved; torch path has no preservation context (correct since _run() no longer mutates jax.backend)
source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py	main() unconditionally wraps _main() with preserve_attribute(skrl.config.jax, "backend"), which is safe but inconsistent with train_skrl's conditional approach
source/isaaclab_rl/test/test_entrypoints.py	Adds four regression tests; the play_skrl test uses runpy.run_module which re-executes module-level side effects including a version-check exit(), creating a fragile dependency on the installed skrl version
source/isaaclab_rl/changelog.d/antoiner-restore-entrypoint-state.rst	Changelog fragment correctly describes the fix

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant run as run() / main()
    participant StateCtx as State Context Manager
    participant _run as _run() / _main()

    Note over Caller,_run: RSL-RL training
    Caller->>run: run(argv)
    run->>run: _parse_args(argv) — SystemExit here prevents any mutation
    run->>StateCtx: enter scoped_torch_backend_flags
    StateCtx->>StateCtx: save 4 Torch flags, set RSL-RL values
    run->>_run: _run(args_cli)
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore 4 Torch flags
    run-->>Caller: return

    Note over Caller,_run: SKRL training (JAX path)
    Caller->>run: run(argv)
    run->>run: _parse_args(argv)
    run->>StateCtx: enter preserve_attribute(jax.backend)
    StateCtx->>StateCtx: save jax.backend
    run->>run: set skrl.config.jax.backend
    run->>_run: _run(args_cli)
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore or delete jax.backend
    run-->>Caller: return

    Note over Caller,_run: SKRL playback (JAX path)
    Caller->>run: main()
    run->>StateCtx: enter preserve_attribute(jax.backend)
    StateCtx->>StateCtx: save jax.backend
    run->>_run: _main()
    _run->>_run: set skrl.config.jax.backend
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore or delete jax.backend
    run-->>Caller: return

%%{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 Caller
    participant run as run() / main()
    participant StateCtx as State Context Manager
    participant _run as _run() / _main()

    Note over Caller,_run: RSL-RL training
    Caller->>run: run(argv)
    run->>run: _parse_args(argv) — SystemExit here prevents any mutation
    run->>StateCtx: enter scoped_torch_backend_flags
    StateCtx->>StateCtx: save 4 Torch flags, set RSL-RL values
    run->>_run: _run(args_cli)
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore 4 Torch flags
    run-->>Caller: return

    Note over Caller,_run: SKRL training (JAX path)
    Caller->>run: run(argv)
    run->>run: _parse_args(argv)
    run->>StateCtx: enter preserve_attribute(jax.backend)
    StateCtx->>StateCtx: save jax.backend
    run->>run: set skrl.config.jax.backend
    run->>_run: _run(args_cli)
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore or delete jax.backend
    run-->>Caller: return

    Note over Caller,_run: SKRL playback (JAX path)
    Caller->>run: main()
    run->>StateCtx: enter preserve_attribute(jax.backend)
    StateCtx->>StateCtx: save jax.backend
    run->>_run: _main()
    _run->>_run: set skrl.config.jax.backend
    _run-->>run: return / raise
    StateCtx->>StateCtx: restore or delete jax.backend
    run-->>Caller: return

_{Reviews (1): Last reviewed commit: "Restore RL entrypoint backend state" | Re-trigger Greptile}

greptile-apps · 2026-07-20T18:06:03Z

 def main():
-    """Play with skrl agent."""
+    """Play with SKRL while restoring the caller's global settings."""
+    with preserve_attribute(skrl.config.jax, "backend"):
+        _main()


Unconditional JAX preservation inconsistent with train_skrl

main() always enters preserve_attribute(skrl.config.jax, "backend") regardless of args_cli.ml_framework, whereas train_skrl.run() only enters the context when ml_framework.startswith("jax"). For torch-framework play calls the attribute is never mutated inside _main(), so this save/restore cycle is a no-op — harmless, but the inconsistency can mislead readers into thinking the JAX attribute is always mutated. Since args_cli is accessible at module scope, the guard could mirror the training side: if args_cli.ml_framework.startswith("jax"): ....

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

greptile-apps · 2026-07-20T18:06:04Z

+def test_skrl_play_main_restores_jax_backend(monkeypatch) -> None:
+    """Direct SKRL play calls remove the JAX backend setting they created."""
+    monkeypatch.setattr(sys, "argv", ["play_skrl.py"])
+    namespace = runpy.run_module("isaaclab_rl.entrypoints.backends.play_skrl", run_name="test_play_skrl")
+    skrl = namespace["skrl"]
+    monkeypatch.delattr(skrl.config.jax, "backend", raising=False)
+
+    def fail_after_mutation() -> None:
+        skrl.config.jax.backend = "mutated"
+        raise RuntimeError("failed")
+
+    namespace["main"].__globals__["_main"] = fail_after_mutation
+    with pytest.raises(RuntimeError, match="failed"):
+        namespace["main"]()
+


Module-level side effects executed by runpy.run_module

runpy.run_module("isaaclab_rl.entrypoints.backends.play_skrl", ...) re-runs the script's entire top-level body, including the skrl version check at lines 96–101 of play_skrl.py that calls exit() on failure. If skrl is not installed at the required version in the test environment the test bails out with a SystemExit before any assertion is reached, producing a misleading failure. The other SKRL test (test_skrl_training_restores_jax_backend) avoids this by patching _parse_args and _run on the already-imported module rather than re-executing it. A similar patch-based approach — importing play_skrl once and patching its _main attribute — would remove this dependency.

greptile-apps · 2026-07-20T18:06:05Z

+def test_rejected_rsl_training_preserves_torch_backend_state(monkeypatch) -> None:
+    """A rejected in-process RSL-RL request does not mutate its caller."""
+    import torch
+
+    caller_state = (False, False, True, True)
+    settings = (
+        (torch.backends.cuda.matmul, "allow_tf32"),
+        (torch.backends.cudnn, "allow_tf32"),
+        (torch.backends.cudnn, "deterministic"),
+        (torch.backends.cudnn, "benchmark"),
+    )
+    for (target, name), value in zip(settings, caller_state):
+        monkeypatch.setattr(target, name, value)
+
+    with pytest.raises(SystemExit):
+        dispatch._run_backend("isaaclab_rl.entrypoints.backends.train_rsl_rl", ["--help"], run_as_script=False)
+
+    assert _torch_backend_state() == caller_state


Test exercises the pre-mutation exit path only

["--help"] causes argparse to raise SystemExit inside _parse_args, which is called before scoped_torch_backend_flags is entered. The torch flags are therefore never written, so the assertion trivially passes without exercising the context-manager restore path. A complement that actually enters the context — e.g. by patching _run to raise immediately after the flags have been set — would give stronger confidence that the finally branch in preserve_attribute fires correctly for a mid-training failure. (The unit-level test_scoped_backend_state_restores_values_after_exception covers this for the primitives, but not for train_rsl_rl.run() end-to-end.)

StafaH · 2026-07-20T21:29:20Z

@@ -0,0 +1,46 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).


This type of file is not good practice. Let's just use common.py

StafaH · 2026-07-20T21:31:42Z

+    cuda_matmul = torch_module.backends.cuda.matmul
+    cudnn = torch_module.backends.cudnn
+    settings = (
+        (cuda_matmul, "allow_tf32", True),


Why not just allow all these to be args, then this function can set them, and the function can be used by rsl_rl training.

That way no hardcoded values?

Move shared state guards into the common entrypoint utilities and make RSL-RL's Torch backend policy explicit at the call site. Scope SKRL playback preservation to JAX calls and cover exceptional RSL-RL cleanup.

Restore RL entrypoint backend state

2d3e381

Scope Torch and SKRL global backend settings to reusable in-process training and playback calls. Preserve attributes that were initially absent and add failure-path regressions for RSL-RL and SKRL.

AntoineRichard requested a review from a team July 20, 2026 18:00

AntoineRichard requested review from ClemensSchwarke, Mayankm96, StafaH and ooctipus as code owners July 20, 2026 18:00

github-actions Bot added the isaac-lab Related to Isaac Lab team label Jul 20, 2026

greptile-apps Bot reviewed Jul 20, 2026

View reviewed changes

StafaH approved these changes Jul 20, 2026

View reviewed changes

Refine RL entrypoint state restoration

cee3265

Move shared state guards into the common entrypoint utilities and make RSL-RL's Torch backend policy explicit at the call site. Scope SKRL playback preservation to JAX calls and cover exceptional RSL-RL cleanup.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Restore RL entrypoint backend state#6631

Restore RL entrypoint backend state#6631
AntoineRichard wants to merge 2 commits into
isaac-sim:developfrom
AntoineRichard:antoiner/restore-rl-entrypoint-state

AntoineRichard commented Jul 20, 2026

Uh oh!

greptile-apps Bot commented Jul 20, 2026

Uh oh!

greptile-apps Bot Jul 20, 2026

Uh oh!

greptile-apps Bot Jul 20, 2026

Uh oh!

greptile-apps Bot Jul 20, 2026

Uh oh!

StafaH Jul 20, 2026

Uh oh!

StafaH Jul 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,46 @@
		# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).

Uh oh!

Conversation

AntoineRichard commented Jul 20, 2026

Description

Relationship

Type of change

Validation

Checklist

Uh oh!

greptile-apps Bot commented Jul 20, 2026

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

StafaH Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

StafaH Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants