Restore RL entrypoint backend state

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.

Author: AntoineRichard

source/isaaclab_rl/changelog.d/antoiner-restore-entrypoint-state.rst

@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed reusable reinforcement learning entrypoints to restore caller Torch and
+  SKRL backend settings after in-process training and playback.

source/isaaclab_rl/isaaclab_rl/entrypoints/_state.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Scoped restoration for backend-global reinforcement learning settings."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import ExitStack, contextmanager
+from types import ModuleType
+
+_MISSING = object()
+
+
+@contextmanager
+def preserve_attribute(target: object, name: str) -> Iterator[None]:
+    """Restore an attribute after a scoped operation, including an initially missing attribute."""
+    previous = getattr(target, name, _MISSING)
+    try:
+        yield
+    finally:
+        if previous is _MISSING:
+            if hasattr(target, name):
+                delattr(target, name)
+        else:
+            setattr(target, name, previous)
+
+
+@contextmanager
+def scoped_torch_backend_flags(torch_module: ModuleType) -> Iterator[None]:
+    """Temporarily configure the Torch backend flags used by RSL-RL training."""
+    cuda_matmul = torch_module.backends.cuda.matmul
+    cudnn = torch_module.backends.cudnn
+    settings = (
+        (cuda_matmul, "allow_tf32", True),
+        (cudnn, "allow_tf32", True),
+        (cudnn, "deterministic", False),
+        (cudnn, "benchmark", False),
+    )
+    with ExitStack() as cleanup:
+        for target, name, value in settings:
+            cleanup.enter_context(preserve_attribute(target, name))
+            setattr(target, name, value)
+        yield

source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py

@@ -27,6 +27,7 @@
 from isaaclab.utils.dict import print_dict
 from isaaclab.utils.seed import configure_seed
 
+from isaaclab_rl.entrypoints._state import preserve_attribute
 from isaaclab_rl.entrypoints.common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector
 from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 
@@ -108,7 +109,13 @@
 
 
 def main():
-    """Play with skrl agent."""
+    """Play with SKRL while restoring the caller's global settings."""
+    with preserve_attribute(skrl.config.jax, "backend"):
+        _main()
+
+
+def _main():
+    """Execute SKRL playback."""
     env_cfg, experiment_cfg = resolve_task_config(args_cli.task, agent_cfg_entry_point)
     with launch_simulation(env_cfg, args_cli):
         if args_cli.ml_framework.startswith("torch"):

source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py

@@ -20,6 +20,7 @@
 
 from isaaclab.app import add_launcher_args
 
+from isaaclab_rl.entrypoints._state import scoped_torch_backend_flags
 from isaaclab_rl.entrypoints.backends import cli_args_rsl_rl as cli_args
 from isaaclab_rl.entrypoints.common import (
     CHECKPOINT_SELECTORS,
@@ -96,8 +97,16 @@ def _parse_args(argv: list[str]) -> argparse.Namespace:
 
 
 def run(argv: list[str]) -> None:
-    """Train an RSL-RL agent."""
+    """Train an RSL-RL agent while restoring the caller's Torch backend settings."""
     import torch
+
+    args_cli = _parse_args(argv)
+    with scoped_torch_backend_flags(torch):
+        _run(args_cli)
+
+
+def _run(args_cli: argparse.Namespace) -> None:
+    """Execute RSL-RL training with parsed arguments."""
     from rsl_rl.runners import DistillationRunner, OnPolicyRunner
 
     from isaaclab.app import launch_simulation
@@ -108,12 +117,6 @@ def run(argv: list[str]) -> None:
 
     from isaaclab_tasks.utils import get_checkpoint_path, resolve_task_config
 
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    torch.backends.cudnn.deterministic = False
-    torch.backends.cudnn.benchmark = False
-
-    args_cli = _parse_args(argv)
     installed_version = _check_rsl_rl_version()
     env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
 

source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_skrl.py

@@ -19,6 +19,7 @@
 
 from isaaclab.app import add_launcher_args
 
+from isaaclab_rl.entrypoints._state import preserve_attribute
 from isaaclab_rl.entrypoints.common import (
     CHECKPOINT_SELECTORS,
     add_common_train_args,
@@ -99,7 +100,19 @@ def _get_distributed_rank(args_cli: argparse.Namespace) -> int:
 
 
 def run(argv: list[str]) -> None:
-    """Train a skrl agent."""
+    """Train a skrl agent while restoring the caller's global SKRL settings."""
+    import skrl
+
+    args_cli = _parse_args(argv)
+    with contextlib.ExitStack() as cleanup:
+        if args_cli.ml_framework.startswith("jax"):
+            cleanup.enter_context(preserve_attribute(skrl.config.jax, "backend"))
+            skrl.config.jax.backend = "jax" if args_cli.ml_framework == "jax" else "numpy"
+        _run(args_cli)
+
+
+def _run(args_cli: argparse.Namespace) -> None:
+    """Execute SKRL training with parsed arguments."""
     import skrl
 
     from isaaclab.app import launch_simulation
@@ -111,8 +124,6 @@ def run(argv: list[str]) -> None:
 
     from isaaclab_tasks.utils import resolve_task_config
 
-    args_cli = _parse_args(argv)
-
     if version.parse(skrl.__version__) < version.parse(SKRL_VERSION):
         skrl.logger.error(
             f"Unsupported skrl version: {skrl.__version__}. "
@@ -139,9 +150,6 @@ def run(argv: list[str]) -> None:
             agent_cfg["trainer"]["timesteps"] = args_cli.max_iterations * agent_cfg["agent"]["rollouts"]
         agent_cfg["trainer"]["close_environment_at_exit"] = False
 
-        if args_cli.ml_framework.startswith("jax"):
-            skrl.config.jax.backend = "jax" if args_cli.ml_framework == "jax" else "numpy"
-
         if args_cli.seed == -1:
             args_cli.seed = random.randint(0, 10000)
 

source/isaaclab_rl/test/test_entrypoints.py

@@ -8,9 +8,12 @@
 from __future__ import annotations
 
 import importlib
+import runpy
 import sys
 import types
 
+import pytest
+
 from isaaclab_rl.entrypoints import PlaybackRequest, TrainingRequest, api, dispatch
 
 
@@ -140,3 +143,92 @@ def test_dispatch_fuses_option_like_kit_args(monkeypatch) -> None:
 def test_dispatch_requires_a_backend() -> None:
     """Missing backend selection returns the conventional CLI error status."""
     assert dispatch.run_train_cli(["--task", "Isaac-Cartpole"]) == 2
+
+
+def _torch_backend_state() -> tuple[bool, bool, bool, bool]:
+    import torch
+
+    return (
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.backends.cudnn.deterministic,
+        torch.backends.cudnn.benchmark,
+    )
+
+
+def test_scoped_backend_state_restores_values_after_exception() -> None:
+    """Backend-global settings are restored when a scoped operation fails."""
+    import torch
+
+    from isaaclab_rl.entrypoints._state import preserve_attribute, scoped_torch_backend_flags
+
+    original = _torch_backend_state()
+    holder = types.SimpleNamespace(value="original")
+
+    with pytest.raises(RuntimeError, match="failed"):
+        with scoped_torch_backend_flags(torch), preserve_attribute(holder, "value"):
+            holder.value = "temporary"
+            assert _torch_backend_state() == (True, True, False, False)
+            raise RuntimeError("failed")
+
+    assert _torch_backend_state() == original
+    assert holder.value == "original"
+
+
+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
+
+
+def test_skrl_training_restores_jax_backend(monkeypatch) -> None:
+    """SKRL training removes the JAX backend setting it created after an exception."""
+    import skrl
+
+    from isaaclab_rl.entrypoints.backends import train_skrl
+
+    monkeypatch.delattr(skrl.config.jax, "backend", raising=False)
+    monkeypatch.setattr(train_skrl, "_parse_args", lambda argv: types.SimpleNamespace(ml_framework="jax"))
+
+    def fail_after_mutation(_args_cli) -> None:
+        assert skrl.config.jax.backend == "jax"
+        skrl.config.jax.backend = "mutated"
+        raise RuntimeError("failed")
+
+    monkeypatch.setattr(train_skrl, "_run", fail_after_mutation)
+    with pytest.raises(RuntimeError, match="failed"):
+        train_skrl.run([])
+
+    assert not hasattr(skrl.config.jax, "backend")
+
+
+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"]()
+
+    assert not hasattr(skrl.config.jax, "backend")