diff --git a/docs/source/concepts/reinforcement_learning.rst b/docs/source/concepts/reinforcement_learning.rst index 943fbac8675c..8dbe477cca28 100644 --- a/docs/source/concepts/reinforcement_learning.rst +++ b/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 `__ 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 ~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/isaaclab/installation.rst b/docs/source/isaaclab/installation.rst index fa1e432efc32..3d427dd193d6 100644 --- a/docs/source/isaaclab/installation.rst +++ b/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:: diff --git a/docs/source/isaaclab/quickstart.rst b/docs/source/isaaclab/quickstart.rst index 4d47b1449d6d..73d92375ce49 100644 --- a/docs/source/isaaclab/quickstart.rst +++ b/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 ...`` diff --git a/pyproject.toml b/pyproject.toml index b143612a98e7..a9c4a0a25edc 100644 --- a/pyproject.toml +++ b/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 diff --git a/source/isaaclab_rl/changelog.d/toni-skrl_warp.added.rst b/source/isaaclab_rl/changelog.d/toni-skrl_warp.added.rst new file mode 100644 index 000000000000..a7637244c8ab --- /dev/null +++ b/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. diff --git a/source/isaaclab_rl/changelog.d/toni-skrl_warp.changed.rst b/source/isaaclab_rl/changelog.d/toni-skrl_warp.changed.rst new file mode 100644 index 000000000000..0b6f9d279f7e --- /dev/null +++ b/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. diff --git a/source/isaaclab_rl/changelog.d/toni-skrl_warp.minor b/source/isaaclab_rl/changelog.d/toni-skrl_warp.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/cli_args_skrl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/cli_args_skrl.py index 6cc0ddcf6dcd..69361db755af 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/cli_args_skrl.py +++ b/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( diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py index 8feb6df5c65e..dae3ce089f2b 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py +++ b/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) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_skrl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_skrl.py index eed4f14e366b..6646284cf13a 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_skrl.py +++ b/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) diff --git a/source/isaaclab_rl/isaaclab_rl/skrl.py b/source/isaaclab_rl/isaaclab_rl/skrl.py index 954fa58356ba..495cfafe9613 100644 --- a/source/isaaclab_rl/isaaclab_rl/skrl.py +++ b/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 diff --git a/uv.lock b/uv.lock index e55477a12fe1..57007d4b0b31 100644 --- a/uv.lock +++ b/uv.lock @@ -1957,6 +1957,9 @@ skrl-jax = [ { name = "jax", version = "0.11.2", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform != 'win32'" }, { name = "skrl", extra = ["jax"] }, ] +skrl-warp = [ + { name = "skrl", extra = ["warp"] }, +] teleop = [ { name = "dex-retargeting", marker = "platform_machine == 'x86_64' and sys_platform != 'win32'" }, { name = "isaaclab-mimic" }, @@ -2089,8 +2092,9 @@ requires-dist = [ { name = "rsl-rl-lib", specifier = "==5.5.1" }, { name = "rsl-rl-lib", marker = "extra == 'rsl-rl'", specifier = "==5.5.1" }, { name = "scipy", specifier = ">=1.11.0" }, - { name = "skrl", marker = "extra == 'skrl'", specifier = ">=2.1.0" }, - { name = "skrl", extras = ["jax"], marker = "extra == 'skrl-jax'", specifier = ">=2.1.0" }, + { name = "skrl", marker = "extra == 'skrl'", specifier = ">=2.1.1" }, + { name = "skrl", extras = ["jax"], marker = "extra == 'skrl-jax'", specifier = ">=2.1.1" }, + { name = "skrl", extras = ["warp"], marker = "extra == 'skrl-warp'", specifier = ">=2.1.1" }, { name = "sphinx", marker = "extra == 'dev'", specifier = ">=7.0,<9" }, { name = "sphinx-book-theme", marker = "extra == 'dev'", specifier = ">=1.1" }, { name = "sphinx-copybutton", marker = "extra == 'dev'" }, @@ -2122,7 +2126,7 @@ requires-dist = [ { name = "wandb", marker = "extra == 'wandb'", specifier = ">=0.13.11" }, { name = "warp-lang", specifier = "==1.17.0" }, ] -provides-extras = ["tetrahedralization", "video", "importers", "test", "dev", "sb3", "skrl", "skrl-jax", "rl-games", "wandb", "rsl-rl", "torchrl", "viser", "rerun", "isaacsim", "ov", "ovphysx", "ovrtx", "mimic", "teleop", "rlinf", "leapp", "all"] +provides-extras = ["tetrahedralization", "video", "importers", "test", "dev", "sb3", "skrl", "skrl-warp", "skrl-jax", "rl-games", "wandb", "rsl-rl", "torchrl", "viser", "rerun", "isaacsim", "ov", "ovphysx", "ovrtx", "mimic", "teleop", "rlinf", "leapp", "all"] [[package]] name = "isaaclab-experimental" @@ -5275,7 +5279,7 @@ wheels = [ [[package]] name = "skrl" -version = "2.1.0" +version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gymnasium" }, @@ -5283,9 +5287,9 @@ dependencies = [ { name = "tensorboard" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0b/bb/8ac912d477c7e18065281021e3d68db2d91034a09ff1668247733490703c/skrl-2.1.0.tar.gz", hash = "sha256:4a1c925b5e025cda3133023331a46f4197ddc5ce420abc01b95e934235402f08", size = 253522, upload-time = "2026-05-10T20:06:54.3Z" } +sdist = { url = "https://files.pythonhosted.org/packages/41/5a/9c116f0f590a52fd21935cae65473779704e80db89e7f55a1de8323b728c/skrl-2.1.1.tar.gz", hash = "sha256:f8eb4939200c85b37a79a6847f78b7eb9fc32e70fa69201b4716204d536b06e7", size = 252550, upload-time = "2026-10-21T07:39:11.156Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/44/de/eca16c14e27a986ddc0314711a6bc42bb70c247073c7f62c390e4ffaa5ed/skrl-2.1.0-py3-none-any.whl", hash = "sha256:5d0cf61fcf81243018220dacbd4a04c84250c1e67274aaeac1779cd53db44e23", size = 458097, upload-time = "2026-05-10T20:06:52.3Z" }, + { url = "https://files.pythonhosted.org/packages/e5/81/25a1962f3bcc3a8208a239e91e33b9823d89a03fa90dbd2e7e4c9b9f4f0e/skrl-2.1.1-py3-none-any.whl", hash = "sha256:ef39a87ef2f34fc1c4963a0b2b56022add7c6b29937bcbc63d2f4c37a8c46c00", size = 459647, upload-time = "2026-10-21T07:39:09.914Z" }, ] [package.optional-dependencies] @@ -5298,6 +5302,10 @@ jax = [ { name = "jaxlib", version = "0.11.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, { name = "optax" }, ] +warp = [ + { name = "warp-lang" }, + { name = "warp-nn" }, +] [[package]] name = "smart-open" @@ -6129,6 +6137,17 @@ wheels = [ { url = "https://pypi.nvidia.com/warp-lang/warp_lang-1.17.0-py3-none-win_amd64.whl", hash = "sha256:ca35c82242a7553046f09023bbe56750b5c3d9af72625bd1a905a56d3189771d" }, ] +[[package]] +name = "warp-nn" +version = "0.4.1" +source = { registry = "https://pypi.nvidia.com/" } +dependencies = [ + { name = "warp-lang" }, +] +wheels = [ + { url = "https://pypi.nvidia.com/warp-nn/warp_nn-0.4.1-py3-none-any.whl", hash = "sha256:5a598043603f9a7c30ee5f02d832ac3c2065b9621d5c57f3a2dd8b64bb3be387" }, +] + [[package]] name = "watchdog" version = "4.0.0"