Add NVIDIA Warp ML framework support for skrl

docs/source/concepts/reinforcement_learning.rst

@@ -64,9 +64,11 @@ Choose a library for the algorithms and configuration tools your experiment need
        distillation, and recurrent and CNN-based policies
      - Included by default
    * - **skrl**
-     - JAX workflows or explicit multi-agent training
-     - PPO, AMP, IPPO, and MAPPO with PyTorch or JAX
-     - ``--extra skrl`` for PyTorch or ``--extra skrl-jax`` for JAX
+     - PyTorch, JAX, or NVIDIA Warp workflows, explicit multi-agent training,
+       and support for different Gymnasium spaces
+     - PPO, AMP, IPPO, and MAPPO with PyTorch, JAX or NVIDIA Warp. Different Gymnasium spaces.
+     - ``--extra skrl`` for PyTorch, ``--extra skrl-jax`` for JAX,
+       or ``--extra skrl-warp`` for NVIDIA Warp
    * - **RL-Games**
      - Existing RL-Games projects or population-based workflows
      - PPO, SAC, A2C, population-based training, distributed training, and Ray
@@ -93,6 +95,7 @@ Install optional dependencies by selecting the corresponding ``uv`` extra when r
 
    uv run --extra skrl isaaclab train --rl_library skrl --task Isaac-Cartpole
    uv run --extra skrl-jax isaaclab train --rl_library skrl --task Isaac-Cartpole --ml_framework jax
+   uv run --extra skrl-warp isaaclab train --rl_library skrl --task Isaac-Cartpole --ml_framework warp
    uv run --extra rl-games isaaclab train --rl_library rl_games --task Isaac-Cartpole
    uv run --extra sb3 isaaclab train --rl_library sb3 --task Isaac-Cartpole
    uv run --extra torchrl isaaclab train --rl_library torchrl --task Isaac-Cartpole
@@ -441,6 +444,19 @@ learner can share the available memory. See the `skrl installation guide
 installation guide <https://docs.jax.dev/en/latest/installation.html>`__ for
 more information.
 
+skrl training with NVIDIA Warp
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+skrl can also train policies with NVIDIA Warp instead of PyTorch. This is useful when
+integrating with an NVIDIA Warp-based learning stack or experimenting with NVIDIA Warp
+compilation and execution:
+
+.. code-block:: bash
+
+   uv run --extra skrl-warp isaaclab train --rl_library skrl \
+       --task Isaac-Reach-Franka \
+       --ml_framework warp
+
 RLinf VLA post-training
 ~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/isaaclab/installation.rst

@@ -246,7 +246,7 @@ before ``isaaclab``; for example, ``--extra ov`` installs both ovphysx and ovrtx
 backends. Pass a comma-separated list or repeat ``--extra``. No extras conflict, so
 any combination resolves into one environment. The ``--extra all`` shortcut installs the
 curated ``ov``, ``rl-games``, ``sb3``, ``skrl``, ``rsl-rl``, ``rerun``, and ``viser`` extras.
-It does not include Isaac Sim or the specialized ``skrl-jax``, ``rlinf``, ``mimic``, ``teleop``,
+It does not include Isaac Sim or the specialized ``skrl-jax``, ``skrl-warp``, ``rlinf``, ``mimic``, ``teleop``,
 ``tetrahedralization``, ``video``, and ``leapp`` extras; request them by name:
 
 .. code-block:: bash
@@ -436,6 +436,8 @@ have dedicated commands below.
      - The corresponding RL framework.
    * - ``skrl-jax``
      - skrl with CUDA-enabled JAX.
+   * - ``skrl-warp``
+     - skrl with NVIDIA Warp.
    * - ``rerun`` / ``viser``
      - The corresponding visualizer.
    * - ``mimic`` / ``teleop``
@@ -453,7 +455,7 @@ have dedicated commands below.
      - Developer test and documentation tooling.
 
 Use ``all`` for the curated list above. Isaac Sim, standalone importers, specialized extras
-(``skrl-jax``, ``rlinf``, ``mimic``, ``teleop``, ``tetrahedralization``, ``video``, ``leapp``), and
+(``skrl-jax``, ``skrl-warp``, ``rlinf``, ``mimic``, ``teleop``, ``tetrahedralization``, ``video``, ``leapp``), and
 the developer ``test`` tooling remain opt-in.
 
 .. note::

docs/source/isaaclab/quickstart.rst

@@ -66,7 +66,7 @@ Training outputs, including checkpoints, are saved under ``logs/``. Add
    For example, ``--extra ovphysx`` makes the OV PhysX integration available, while
    ``physics=ovphysx`` selects it for the task. You can combine extras as needed. The ``--extra all``
    shortcut installs the curated ``ov``, ``rl-games``, ``sb3``, ``skrl``, ``rsl-rl``, ``rerun``,
-   and ``viser`` extras. Isaac Sim, standalone importers, and specialized extras such as ``skrl-jax``,
+   and ``viser`` extras. Isaac Sim, standalone importers, and specialized extras such as ``skrl-jax``, ``skrl-warp``,
    ``rlinf``, ``mimic``, ``teleop``, ``tetrahedralization``, ``video``, ``leapp``, and ``torchrl``
    are not included; add them explicitly. See
    :ref:`installation-optional-extras` for the complete list.
@@ -90,8 +90,8 @@ and is a good starting point for GPU-based training.
      - PPO, SAC, and A2C workflows
      - ``uv run --extra rl-games isaaclab train --rl_library rl_games ...``
    * - ``skrl``
-     - Broad algorithm support with PyTorch and JAX
-     - ``--extra skrl`` for PyTorch; ``--extra skrl-jax`` for JAX
+     - Broad algorithm support with PyTorch, JAX, and NVIDIA Warp
+     - ``--extra skrl`` for PyTorch; ``--extra skrl-jax`` for JAX; ``--extra skrl-warp`` for NVIDIA Warp
    * - ``sb3``
      - Stable-Baselines3 and CPU-oriented experiments
      - ``uv run --extra sb3 isaaclab train --rl_library sb3 ...``

pyproject.toml

@@ -150,9 +150,9 @@ dev = [
 
 sb3 = ["stable-baselines3>=2.6", "tqdm", "rich"]
 
-skrl = ["skrl>=2.1.0"]
-
-skrl-jax = ["jax[cuda13] ; sys_platform == 'linux'", "skrl[jax]>=2.1.0"]
+skrl = ["skrl>=2.1.1"]
+skrl-warp = ["skrl[warp]>=2.1.1"]
+skrl-jax = ["jax[cuda13] ; sys_platform == 'linux'", "skrl[jax]>=2.1.1"]
 
 rl-games = [
     "aiohttp>=3.14.3",  # security floor; allow Isaac Sim kernels to select their exact version

source/isaaclab_rl/changelog.d/toni-skrl_warp.added.rst

@@ -0,0 +1,3 @@
+* Added ``"warp"`` as an ``ml_framework`` option to :func:`~isaaclab_rl.skrl.import_skrl_runner` and to the
+  ``--ml_framework`` argument of the skrl train/play entrypoints, enabling training with skrl's Warp
+  implementation. Install it with the ``skrl-warp`` extra.

source/isaaclab_rl/changelog.d/toni-skrl_warp.changed.rst

@@ -0,0 +1,2 @@
+* Changed the minimum supported skrl version to 2.1.1, which is required for the Warp ML framework.
+* Changed the skrl train/play run summary to show the selected ML framework.

source/isaaclab_rl/isaaclab_rl/entrypoints/backends/cli_args_skrl.py

@@ -21,7 +21,7 @@ def add_skrl_args(parser: argparse.ArgumentParser) -> None:
         "--ml_framework",
         type=str,
         default="torch",
-        choices=["torch", "jax"],
+        choices=["torch", "jax", "warp"],
         help="The ML framework used for training the skrl agent.",
     )
     arg_group.add_argument(

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

@@ -122,7 +122,7 @@ def _run(args_cli: argparse.Namespace) -> None:
         pre_launch_video_config(env_cfg, args_cli)
         screen.stage("Launching simulation")
         with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
-            show_run_summary(screen, args_cli, env_cfg, library="skrl", action="play")
+            show_run_summary(screen, args_cli, env_cfg, library=f"skrl ({args_cli.ml_framework})", action="play")
             runner_cls = import_skrl_runner(args_cli.ml_framework)
             apply_env_overrides(args_cli, env_cfg)
             args_cli.seed = resolve_seed(args_cli.seed)

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

@@ -97,7 +97,7 @@ def _run(args_cli: argparse.Namespace) -> None:
         pre_launch_video_config(env_cfg, args_cli)
         screen.stage("Launching simulation")
         with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
-            show_run_summary(screen, args_cli, env_cfg, library="skrl", action="train")
+            show_run_summary(screen, args_cli, env_cfg, library=f"skrl ({args_cli.ml_framework})", action="train")
             runner_cls = import_skrl_runner(args_cli.ml_framework)
             apply_env_overrides(args_cli, env_cfg)
 

source/isaaclab_rl/isaaclab_rl/skrl.py

@@ -11,14 +11,15 @@
 
     from isaaclab_rl.skrl import SkrlVecEnvWrapper
 
-    env = SkrlVecEnvWrapper(env, ml_framework="torch")  # or ml_framework="jax"
+    env = SkrlVecEnvWrapper(env, ml_framework="torch")  # or "jax", or "warp"
 
 Or, equivalently, by directly calling the skrl library API as follows:
 
 .. code-block:: python
 
-    from skrl.envs.torch.wrappers import wrap_env  # for PyTorch, or...
-    from skrl.envs.jax.wrappers import wrap_env  # for JAX
+    from skrl.envs.wrappers.torch import wrap_env  # for PyTorch, or...
+    from skrl.envs.wrappers.jax import wrap_env  # for JAX, or...
+    from skrl.envs.wrappers.warp import wrap_env  # for Warp
 
     env = wrap_env(env, wrapper="isaaclab")
 
@@ -36,7 +37,7 @@
 if TYPE_CHECKING:
     from isaaclab.envs import DirectMARLEnv, DirectRLEnv, ManagerBasedRLEnv
 
-SKRL_MIN_VERSION = "2.1.0"
+SKRL_MIN_VERSION = "2.1.1"
 """Oldest skrl release supported by the entrypoints."""
 
 
@@ -57,7 +58,7 @@ def check_skrl_version() -> None:
         raise SystemExit(1)
 
 
-def import_skrl_runner(ml_framework: Literal["torch", "jax"]) -> type:
+def import_skrl_runner(ml_framework: Literal["torch", "jax", "warp"]) -> type:
     """Return the skrl :class:`Runner` implementation of the selected ML framework.
 
     Args:
@@ -70,8 +71,12 @@ def import_skrl_runner(ml_framework: Literal["torch", "jax"]) -> type:
         from skrl.utils.runner.torch import Runner
     elif ml_framework == "jax":
         from skrl.utils.runner.jax import Runner
+    elif ml_framework == "warp":
+        from skrl.utils.runner.warp import Runner
     else:
-        raise ValueError(f"Invalid ML framework for the skrl runner: {ml_framework}. Available options: 'torch', 'jax'")
+        raise ValueError(
+            f"Invalid ML framework for the skrl runner: {ml_framework}. Available options: 'torch', 'jax', 'warp'"
+        )
     return Runner