Adds pre-trained checkpoints and tools for generating and uploading checkpoints (#151)

Created a standalone script that can train all our environments and
publish their checkpoints to a Nucleus server. The play.py scripts were
modified to add a --use_pretrained_checkpoint flag. This downloads and
caches the pre-trained checkpoint to a .pretrained_checkpoints
directory.

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------
Signed-off-by: peterd-NV <peterd@nvidia.com>
Signed-off-by: Kelly Guo <kellyg@nvidia.com>
Signed-off-by: Kelly Guo <kellyguo123@hotmail.com>
Co-authored-by: peterd-NV <peterd@nvidia.com>
Co-authored-by: CY Chen <cyc@nvidia.com>
Co-authored-by: oahmednv <oahmed@Nvidia.com>
Co-authored-by: Toni-SM <aserranomuno@nvidia.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>
Co-authored-by: Kelly Guo <kellyguo123@hotmail.com>

.gitignore

@@ -57,3 +57,6 @@ _build
 **/runs/*
 **/logs/*
 **/recordings/*
+
+# Pre-Trained Checkpoints
+/.pretrained_checkpoints/

CONTRIBUTORS.md

@@ -39,6 +39,7 @@ Guidelines for modifications:
 * Anton Bjørndahl Mortensen
 * Arjun Bhardwaj
 * Brayden Zhang
+* Cameron Upright
 * Calvin Yu
 * Chenyu Yang
 * CY (Chien-Ying) Chen

docs/source/overview/reinforcement-learning/rl_existing_scripts.rst

@@ -25,6 +25,8 @@ RL-Games
             ./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task Isaac-Ant-v0 --headless
             # run script for playing with 32 environments
             ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth
+            # run script for playing a pre-trained checkpoint with 32 environments
+            ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --num_envs 32 --use_pretrained_checkpoint
             # run script for recording video of a trained agent (requires installing `ffmpeg`)
             ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --headless --video --video_length 200
 
@@ -39,6 +41,8 @@ RL-Games
             isaaclab.bat -p source\standalone\workflows\rl_games\train.py --task Isaac-Ant-v0 --headless
             :: run script for playing with 32 environments
             isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth
+            :: run script for playing a pre-trained checkpoint with 32 environments
+            isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --num_envs 32 --use_pretrained_checkpoint
             :: run script for recording video of a trained agent (requires installing `ffmpeg`)
             isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --headless --video --video_length 200
 
@@ -62,6 +66,8 @@ RSL-RL
             ./isaaclab.sh -p source/standalone/workflows/rsl_rl/train.py --task Isaac-Reach-Franka-v0 --headless
             # run script for playing with 32 environments
             ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt
+            # run script for playing a pre-trained checkpoint with 32 environments
+            ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --use_pretrained_checkpoint
             # run script for recording video of a trained agent (requires installing `ffmpeg`)
             ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200
 
@@ -76,6 +82,8 @@ RSL-RL
             isaaclab.bat -p source\standalone\workflows\rsl_rl\train.py --task Isaac-Reach-Franka-v0 --headless
             :: run script for playing with 32 environments
             isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt
+            :: run script for playing a pre-trained checkpoint with 32 environments
+            isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --use_pretrained_checkpoint
             :: run script for recording video of a trained agent (requires installing `ffmpeg`)
             isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200
 
@@ -103,6 +111,8 @@ SKRL
                      ./isaaclab.sh -p source/standalone/workflows/skrl/train.py --task Isaac-Reach-Franka-v0 --headless
                      # run script for playing with 32 environments
                      ./isaaclab.sh -p source/standalone/workflows/skrl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --checkpoint /PATH/TO/model.pt
+                     # run script for playing a pre-trained checkpoint with 32 environments
+                     ./isaaclab.sh -p source/standalone/workflows/skrl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --use_pretrained_checkpoint
                      # run script for recording video of a trained agent (requires installing `ffmpeg`)
                      ./isaaclab.sh -p source/standalone/workflows/skrl/play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200
 
@@ -117,6 +127,8 @@ SKRL
                      isaaclab.bat -p source\standalone\workflows\skrl\train.py --task Isaac-Reach-Franka-v0 --headless
                      :: run script for playing with 32 environments
                      isaaclab.bat -p source\standalone\workflows\skrl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --checkpoint /PATH/TO/model.pt
+                     :: run script for playing a pre-trained checkpoint with 32 environments
+                     isaaclab.bat -p source\standalone\workflows\skrl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --use_pretrained_checkpoint
                      :: run script for recording video of a trained agent (requires installing `ffmpeg`)
                      isaaclab.bat -p source\standalone\workflows\skrl\play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200
 
@@ -191,6 +203,8 @@ Stable-Baselines3
             ./isaaclab.sh -p source/standalone/workflows/sb3/train.py --task Isaac-Cartpole-v0 --headless --device cpu
             # run script for playing with 32 environments
             ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip
+            # run script for playing a pre-trained checkpoint with 32 environments
+            ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --num_envs 32 --use_pretrained_checkpoint
             # run script for recording video of a trained agent (requires installing `ffmpeg`)
             ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200
 
@@ -206,6 +220,8 @@ Stable-Baselines3
             isaaclab.bat -p source\standalone\workflows\sb3\train.py --task Isaac-Cartpole-v0 --headless --device cpu
             :: run script for playing with 32 environments
             isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip
+            :: run script for playing a pre-trained checkpoint with 32 environments
+            isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --num_envs 32 --use_pretrained_checkpoint
             :: run script for recording video of a trained agent (requires installing `ffmpeg`)
             isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200
 

source/extensions/omni.isaac.lab/docs/CHANGELOG.rst

@@ -163,7 +163,7 @@ Added
 * Added full buffer property to :class:`omni.isaac.lab.utils.buffers.circular_buffer.CircularBuffer`
 
 
-0.27.32 (2024-12-15)
+0.27.33 (2024-12-15)
 ~~~~~~~~~~~~~~~~~~~~
 
 Added
@@ -172,7 +172,7 @@ Added
 * Added action clip to all :class:`omni.isaac.lab.envs.mdp.actions`.
 
 
-0.27.31 (2024-12-14)
+0.27.32 (2024-12-14)
 ~~~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -181,7 +181,7 @@ Changed
 * Added check for error below threshold in state machines to ensure the state has been reached.
 
 
-0.27.30 (2024-12-13)
+0.27.31 (2024-12-13)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -190,7 +190,7 @@ Fixed
 * Fixed the shape of ``quat_w`` in the ``apply_actions`` method of :attr:`~omni.isaac.lab.env.mdp.NonHolonomicAction` (previously (N,B,4), now (N,4) since the number of root bodies B is required to be 1). Previously ``apply_actions`` errored because ``euler_xyz_from_quat`` requires inputs of shape (N,4).
 
 
-0.27.29 (2024-12-11)
+0.27.30 (2024-12-11)
 ~~~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -201,7 +201,7 @@ Changed
 * Improved documentation to clarify the usage of the :meth:`~omni.isaac.lab.envs.mdp.rewards.base_height_l2` function in both flat and rough terrain settings.
 
 
-0.27.28 (2024-12-11)
+0.27.29 (2024-12-11)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -211,7 +211,7 @@ Fixed
   Jacobian computed w.r.t. to the root frame of the robot. This helps ensure that root pose does not affect the tracking.
 
 
-0.27.27 (2024-12-09)
+0.27.28 (2024-12-09)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -221,7 +221,7 @@ Fixed
   return only the states of the specified environment IDs.
 
 
-0.27.26 (2024-12-06)
+0.27.27 (2024-12-06)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -231,7 +231,7 @@ Fixed
   :attr:`~omni.isaac.lab.assets.Articulation.root_physx_view` level.
 
 
-0.27.25 (2024-12-06)
+0.27.26 (2024-12-06)
 ~~~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -242,7 +242,7 @@ Changed
   disabled. Using an articulation root for rigid bodies is not needed and decreases overall performance.
 
 
-0.27.24 (2024-12-06)
+0.27.25 (2024-12-06)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -252,7 +252,7 @@ Fixed
   Earlier, the projection names used snakecase instead of camelcase.
 
 
-0.27.23 (2024-12-06)
+0.27.24 (2024-12-06)
 ~~~~~~~~~~~~~~~~~~~~
 
 Added
@@ -270,7 +270,7 @@ Changed
   :class:`~omni.isaac.lab.sensors.Camera` did not clip them and had a different behavior for both types.
 
 
-0.27.22 (2024-12-05)
+0.27.23 (2024-12-05)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -279,7 +279,7 @@ Fixed
 * Fixed the condition in ``isaaclab.sh`` that checks whether ``pre-commit`` is installed before attempting installation.
 
 
-0.27.21 (2024-12-04)
+0.27.22 (2024-12-04)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -288,7 +288,7 @@ Fixed
 * Fixed the order of the incoming parameters in :class:`omni.isaac.lab.envs.DirectMARLEnv` to correctly use ``NoiseModel`` in marl-envs.
 
 
-0.27.20 (2024-12-04)
+0.27.21 (2024-12-04)
 ~~~~~~~~~~~~~~~~~~~~
 
 Added
@@ -303,7 +303,7 @@ Added
 * Added ``replay_demos.py`` script to replay demos loaded from an HDF5 file.
 
 
-0.27.19 (2024-12-02)
+0.27.20 (2024-12-02)
 ~~~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -312,6 +312,16 @@ Changed
 * Changed :class:`omni.isaac.lab.envs.DirectMARLEnv` to inherit from ``Gymnasium.Env`` due to requirement from Gymnasium v1.0.0 requiring all environments to be a subclass of ``Gymnasium.Env`` when using the ``make`` interface.
 
 
+0.27.19 (2024-12-02)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added ``omni.isaac.lab.utils.pretrained_checkpoints`` containing constants and utility functions used to manipulate
+  paths and load checkpoints from Nucleus.
+
+
 0.27.18 (2024-11-28)
 ~~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/utils/assets.py

@@ -49,7 +49,8 @@ def check_file_path(path: str) -> Literal[0, 1, 2]:
     """
     if os.path.isfile(path):
         return 1
-    elif omni.client.stat(path)[0] == omni.client.Result.OK:
+    # we need to convert backslash to forward slash on Windows for omni.client API
+    elif omni.client.stat(path.replace(os.sep, "/"))[0] == omni.client.Result.OK:
         return 2
     else:
         return 0
@@ -91,12 +92,12 @@ def retrieve_file_path(path: str, download_dir: str | None = None, force_downloa
         if not os.path.exists(download_dir):
             os.makedirs(download_dir)
         # download file in temp directory using os
-        file_name = os.path.basename(omni.client.break_url(path).path)
+        file_name = os.path.basename(omni.client.break_url(path.replace(os.sep, "/")).path)
         target_path = os.path.join(download_dir, file_name)
         # check if file already exists locally
         if not os.path.isfile(target_path) or force_download:
             # copy file to local machine
-            result = omni.client.copy(path, target_path)
+            result = omni.client.copy(path.replace(os.sep, "/"), target_path)
             if result != omni.client.Result.OK and force_download:
                 raise RuntimeError(f"Unable to copy file: '{path}'. Is the Nucleus Server running?")
         return os.path.abspath(target_path)
@@ -122,7 +123,7 @@ def read_file(path: str) -> io.BytesIO:
         with open(path, "rb") as f:
             return io.BytesIO(f.read())
     elif file_status == 2:
-        file_content = omni.client.read_file(path)[2]
+        file_content = omni.client.read_file(path.replace(os.sep, "/"))[2]
         return io.BytesIO(memoryview(file_content).tobytes())
     else:
         raise FileNotFoundError(f"Unable to find the file: {path}")

source/extensions/omni.isaac.lab/omni/isaac/lab/utils/pretrained_checkpoint.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Sub-module for handling various pre-trained checkpoint tasks"""
+
+import glob
+import json
+import os
+
+import carb.settings
+
+from omni.isaac.lab.utils.assets import ISAACLAB_NUCLEUS_DIR
+
+from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry  # noqa: F401
+
+from .assets import retrieve_file_path
+
+PRETRAINED_CHECKPOINTS_ASSET_ROOT_DIR = carb.settings.get_settings().get(
+    "/persistent/isaaclab/asset_root/pretrained_checkpoints"
+)
+"""Path to the root directory on the Nucleus Server."""
+
+WORKFLOWS = ["rl_games", "rsl_rl", "sb3", "skrl"]
+"""The supported workflows for pre-trained checkpoints"""
+
+WORKFLOW_TRAINER = {w: f"source/standalone/workflows/{w}/train.py" for w in WORKFLOWS}
+"""A dict mapping workflow to their training program path"""
+
+WORKFLOW_PLAYER = {w: f"source/standalone/workflows/{w}/play.py" for w in WORKFLOWS}
+"""A dict mapping workflow to their play program path"""
+
+PRETRAINED_CHECKPOINT_PATH = str(PRETRAINED_CHECKPOINTS_ASSET_ROOT_DIR) + "/Isaac/IsaacLab/PretrainedCheckpoints"
+"""URL for where we store all the pre-trained checkpoints"""
+
+"""The filename for checkpoints used by the different workflows"""
+WORKFLOW_PRETRAINED_CHECKPOINT_FILENAMES = {
+    "rl_games": "checkpoint.pth",
+    "rsl_rl": "checkpoint.pt",
+    "sb3": "checkpoint.zip",
+    "skrl": "checkpoint.pt",
+}
+
+"""Maps workflow to the agent variable name that determines the logging directory logs/{workflow}/{variable}"""
+WORKFLOW_EXPERIMENT_NAME_VARIABLE = {
+    "rl_games": "agent.params.config.name",
+    "rsl_rl": "agent.experiment_name",
+    "sb3": None,
+    "skrl": "agent.agent.experiment.directory",
+}
+
+
+def has_pretrained_checkpoints_asset_root_dir() -> bool:
+    """Returns True if and only if /persistent/isaaclab/asset_root/pretrained_checkpoints exists"""
+    return PRETRAINED_CHECKPOINTS_ASSET_ROOT_DIR is not None
+
+
+def get_log_root_path(workflow: str, task_name: str) -> str:
+    """Returns the absolute path where the logs are written for a specific workflow and task_name"""
+    return os.path.abspath(os.path.join("logs", workflow, task_name))
+
+
+def get_latest_job_run_path(workflow: str, task_name: str) -> str:
+    """The local logs path of the most recent run of this workflow and task name"""
+    log_root_path = get_log_root_path(workflow, task_name)
+    return _get_latest_file_or_directory(log_root_path)
+
+
+def get_pretrained_checkpoint_path(workflow: str, task_name: str) -> str:
+    """The local logs path where we get the pre-trained checkpoints from"""
+
+    path = get_latest_job_run_path(workflow, task_name)
+    if not path:
+        return None
+
+    if workflow == "rl_games":
+        return os.path.join(path, "nn", f"{task_name}.pth")
+    elif workflow == "rsl_rl":
+        return _get_latest_file_or_directory(path, "*.pt")
+    elif workflow == "sb3":
+        return os.path.join(path, "model.zip")
+    elif workflow == "skrl":
+        return os.path.join(path, "checkpoints", "best_agent.pt")
+    else:
+        raise Exception(f"Unsupported workflow ({workflow})")
+
+
+def get_pretrained_checkpoint_publish_path(workflow: str, task_name: str) -> str:
+    """The path where pre-trained checkpoints are published to"""
+    return os.path.join(
+        PRETRAINED_CHECKPOINT_PATH, workflow, task_name, WORKFLOW_PRETRAINED_CHECKPOINT_FILENAMES[workflow]
+    )
+
+
+def get_published_pretrained_checkpoint_path(workflow: str, task_name: str) -> str:
+    """The path where pre-trained checkpoints are fetched from"""
+    return os.path.join(
+        ISAACLAB_NUCLEUS_DIR,
+        "PretrainedCheckpoints",
+        workflow,
+        task_name,
+        WORKFLOW_PRETRAINED_CHECKPOINT_FILENAMES[workflow],
+    )
+
+
+def get_published_pretrained_checkpoint(workflow: str, task_name: str) -> str | None:
+    """Gets the path for the pre-trained checkpoint.
+
+    If the checkpoint is not cached locally then the file is downloaded.
+    The cached path is then returned.
+
+    Args:
+        workflow: The workflow.
+        task_name: The task name.
+
+    Returns:
+        The path.
+    """
+    ov_path = get_published_pretrained_checkpoint_path(workflow, task_name)
+    download_dir = os.path.join(".pretrained_checkpoints", workflow, task_name)
+    resume_path = os.path.join(download_dir, WORKFLOW_PRETRAINED_CHECKPOINT_FILENAMES[workflow])
+
+    if not os.path.exists(resume_path):
+        print(f"Fetching pre-trained checkpoint : {ov_path}")
+        try:
+            resume_path = retrieve_file_path(ov_path, download_dir)
+        except Exception:
+            print("A pre-trained checkpoint is currently unavailable for this task.")
+            return None
+    else:
+        print("Using pre-fetched pre-trained checkpoint")
+    return resume_path
+
+
+def has_pretrained_checkpoint_job_run(workflow: str, task_name: str) -> bool:
+    """Returns true if an experiment exists in the logs for the workflow and task"""
+    return os.path.exists(get_log_root_path(workflow, task_name))
+
+
+def has_pretrained_checkpoint_job_finished(workflow: str, task_name: str) -> bool:
+    """Returns true if an experiment has results which may or may not be final depending on workflow"""
+    local_path = get_pretrained_checkpoint_path(workflow, task_name)
+    return local_path is not None and os.path.exists(local_path)
+
+
+def get_pretrained_checkpoint_review_path(workflow: str, task_name: str) -> str | None:
+    """The path of the review JSON file for a workflow and task"""
+    run_path = get_latest_job_run_path(workflow, task_name)
+    if not run_path:
+        return None
+    return os.path.join(run_path, "pretrained_checkpoint_review.json")
+
+
+def get_pretrained_checkpoint_review(workflow: str, task_name: str) -> dict | None:
+    """Returns the review JSON file as a dict if it exists"""
+    review_path = get_pretrained_checkpoint_review_path(workflow, task_name)
+    if not review_path:
+        return None
+
+    if os.path.exists(review_path):
+        with open(review_path) as f:
+            return json.load(f)
+
+    return None
+
+
+def _get_latest_file_or_directory(path: str, pattern: str = "*"):
+    """Returns the path to the most recently modified file or directory at a path matching an optional pattern"""
+    g = glob.glob(f"{path}/{pattern}")
+    if len(g):
+        return max(g, key=os.path.getmtime)
+    return None

source/standalone/tools/pretrained_checkpoint.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Script to manage pretrained checkpoints for our environments.
+"""
+
+import argparse
+
+from omni.isaac.lab.app import AppLauncher
+
+# Initialize the parser
+parser = argparse.ArgumentParser(
+    description="""
+Script used for the training and publishing of pre-trained checkpoints for Isaac Lab.
+
+Examples :
+    # Train an agent using the rl_games workflow on the Isaac-Cartpole-v0 environment.
+    pretrained_checkpoint.py --train rl_games:Isaac-Cartpole-v0
+    # Train and publish the checkpoints for all workflows on only the direct Cartpole environments.
+    pretrained_checkpoint.py -tp "*:Isaac-Cartpole-*Direct-v0" \\
+      --/persistent/isaaclab/asset_root/pretrained_checkpoints="/some/path"
+    # Review all repose cube jobs, excluding the Play tasks and skrl
+    pretrained_checkpoint.py -r "*:*Repose-Cube*" --exclude "*:*Play*" --exclude skrl:*
+    # Publish all results (that have been reviewed and approved).
+    pretrained_checkpoint.py --publish --all \\
+      --/persistent/isaaclab/asset_root/pretrained_checkpoints="/some/path"
+""",
+    formatter_class=argparse.RawTextHelpFormatter,
+)
+
+# Add positional arguments that can accept zero or more values
+parser.add_argument(
+    "jobs",
+    nargs="*",
+    help="""
+A job consists of a workflow and a task name separated by a colon (wildcards optional), for example :
+    rl_games:Isaac-Humanoid-*v0
+    rsl_rl:Isaac-Ant-*-v0
+    *:Isaac-Velocity-Flat-Spot-v0
+""",
+)
+parser.add_argument("-t", "--train", action="store_true", help="Train checkpoints for later publishing.")
+parser.add_argument("-p", "--publish_checkpoint", action="store_true", help="Publish pre-trained checkpoints.")
+parser.add_argument("-r", "--review", action="store_true", help="Review checkpoints.")
+parser.add_argument("-l", "--list", action="store_true", help="List all available environments and workflows.")
+parser.add_argument("-f", "--force", action="store_true", help="Force training when results already exist.")
+parser.add_argument("-a", "--all", action="store_true", help="Run all valid workflow task pairs.")
+parser.add_argument(
+    "-E",
+    "--exclude",
+    action="append",
+    type=str,
+    default=[],
+    help="Excludes jobs matching the argument, with wildcard support.",
+)
+parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
+parser.add_argument("--force_review", action="store_true", help="Forces review when one already exists.")
+parser.add_argument("--force_publish", action="store_true", help="Publish checkpoints without review.")
+parser.add_argument("--headless", action="store_true", help="Run training without the UI.")
+
+args, _ = parser.parse_known_args()
+
+# Need something to do
+if len(args.jobs) == 0 and not args.all:
+    parser.error("Jobs must be provided, or --all.")
+
+# Must train, publish, review or list
+if not (args.train or args.publish_checkpoint or args.review or args.list):
+    parser.error("A train, publish, review or list flag must be given.")
+
+# List excludes train and publish
+if args.list and (args.train or args.publish_checkpoint):
+    parser.error("Can't train or publish when listing.")
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=True)
+simulation_app = app_launcher.app
+
+
+import csv
+
+# Now everything else
+import fnmatch
+import gymnasium as gym
+import json
+import numpy as np
+import os
+import subprocess
+import sys
+
+import omni.client
+from omni.client._omniclient import CopyBehavior
+
+from omni.isaac.lab.utils.pretrained_checkpoint import (
+    WORKFLOW_EXPERIMENT_NAME_VARIABLE,
+    WORKFLOW_PLAYER,
+    WORKFLOW_TRAINER,
+    WORKFLOWS,
+    get_log_root_path,
+    get_pretrained_checkpoint_path,
+    get_pretrained_checkpoint_publish_path,
+    get_pretrained_checkpoint_review,
+    get_pretrained_checkpoint_review_path,
+    has_pretrained_checkpoint_job_finished,
+    has_pretrained_checkpoint_job_run,
+    has_pretrained_checkpoints_asset_root_dir,
+)
+
+# Need somewhere to publish
+if args.publish_checkpoint and not has_pretrained_checkpoints_asset_root_dir():
+    raise Exception("A /persistent/isaaclab/asset_root/pretrained_checkpoints setting is required to publish.")
+
+
+def train_job(workflow, task_name, headless=False, force=False, num_envs=None):
+    """
+    This trains a task using the workflow's train.py script, overriding the experiment name to ensure unique
+    log directories.  By default it will return if an experiment has already been run.
+
+    Args:
+        workflow: The workflow.
+        task_name: The task name.
+        headless: Should the training run without the UI.
+        force: Run training even if previous experiments have been run.
+        num_envs: How many simultaneous environments to simulate, overriding the config.
+    """
+
+    log_root_path = get_log_root_path(workflow, task_name)
+
+    # We already ran this
+    if not force and os.path.exists(log_root_path) and len(os.listdir(log_root_path)) > 0:
+        print(f"Skipping training of {workflow}:{task_name}, already has been run")
+        return
+
+    print(f"Training {workflow}:{task_name}")
+
+    # Construct our command
+    cmd = [
+        sys.executable,
+        WORKFLOW_TRAINER[workflow],
+        "--task",
+        task_name,
+        "--enable_cameras",
+    ]
+
+    # Changes the directory name for logging
+    if WORKFLOW_EXPERIMENT_NAME_VARIABLE[workflow]:
+        cmd.append(f"{WORKFLOW_EXPERIMENT_NAME_VARIABLE[workflow]}={task_name}")
+
+    if headless:
+        cmd.append("--headless")
+    if num_envs:
+        cmd.extend(["--num_envs", str(num_envs)])
+
+    print("Running : " + " ".join(cmd))
+
+    subprocess.run(cmd)
+
+
+def review_pretrained_checkpoint(workflow, task_name, force_review=False, num_envs=None):
+    """
+    This initiates a review of the pretrained checkpoint.  The play.py script for the workflow is run, and the user
+    inspects the results.  When done they close the simulator and will be prompted for their review.
+
+    Args:
+        workflow: The workflow.
+        task_name: The task name.
+        force_review: Performs the review even if a review already exists.
+        num_envs: How many simultaneous environments to simulate, overriding the config.
+    """
+
+    # This workflow task pair hasn't been trained
+    if not has_pretrained_checkpoint_job_run(workflow, task_name):
+        print(f"Skipping review of {workflow}:{task_name}, hasn't been trained yet")
+        return
+
+    # Couldn't find the checkpoint
+    if not has_pretrained_checkpoint_job_finished(workflow, task_name):
+        print(f"Training not complete for {workflow}:{task_name}")
+        return
+
+    review = get_pretrained_checkpoint_review(workflow, task_name)
+
+    if not force_review and review and review["reviewed"]:
+        print(f"Review already complete for {workflow}:{task_name}")
+        return
+
+    print(f"Reviewing {workflow}:{task_name}")
+
+    # Construct our command
+    cmd = [
+        sys.executable,
+        WORKFLOW_PLAYER[workflow],
+        "--task",
+        task_name,
+        "--checkpoint",
+        get_pretrained_checkpoint_path(workflow, task_name),
+        "--enable_cameras",
+    ]
+
+    if num_envs:
+        cmd.extend(["--num_envs", str(num_envs)])
+
+    print("Running : " + " ".join(cmd))
+
+    subprocess.run(cmd)
+
+    # Give user a chance to leave the old review
+    if force_review and review and review["reviewed"]:
+        result = review["result"]
+        notes = review.get("notes")
+        print(f"A review already exists for {workflow}:{task_name}, it was marked as '{result}'.")
+        print(f"  Notes: {notes}")
+        answer = input("Would you like to replace it?  Please answer yes or no (y/n) [n]: ").strip().lower()
+        if answer != "y":
+            return
+
+    # Get the verdict from the user
+    print(f"Do you accept this checkpoint for {workflow}:{task_name}?")
+
+    answer = input("Please answer yes, no or undetermined (y/n/u) [u]: ").strip().lower()
+    if answer not in {"y", "n", "u"}:
+        answer = "u"
+    answer_map = {
+        "y": "accepted",
+        "n": "rejected",
+        "u": "undetermined",
+    }
+
+    # Create the review dict
+    review = {
+        "reviewed": True,
+        "result": answer_map[answer],
+    }
+
+    # Maybe add some notes
+    notes = input("Please add notes or hit enter: ").strip().lower()
+    if notes:
+        review["notes"] = notes
+
+    # Save the review JSON file
+    path = get_pretrained_checkpoint_review_path(workflow, task_name)
+    if not path:
+        raise Exception("This shouldn't be possible, something went very wrong.")
+
+    with open(path, "w") as f:
+        json.dump(review, f, indent=4)
+
+
+def publish_pretrained_checkpoint(workflow, task_name, force_publish=False):
+    """
+    This publishes the pretrained checkpoint to Nucleus using the asset path in the
+    /persistent/isaaclab/asset_root/pretrained_checkpoints Carb variable.
+
+    Args:
+        workflow: The workflow.
+        task_name: The task name.
+        force_publish: Publish without review.
+    """
+
+    # This workflow task pair hasn't been trained
+    if not has_pretrained_checkpoint_job_run(workflow, task_name):
+        print(f"Skipping publishing of {workflow}:{task_name}, hasn't been trained yet")
+        return
+
+    # Couldn't find the checkpoint
+    if not has_pretrained_checkpoint_job_finished(workflow, task_name):
+        print(f"Training not complete for {workflow}:{task_name}")
+        return
+
+    # Get local pretrained checkpoint path
+    local_path = get_pretrained_checkpoint_path(workflow, task_name)
+    if not local_path:
+        raise Exception("This shouldn't be possible, something went very wrong.")
+
+    # Not forcing, need to check review results
+    if not force_publish:
+
+        # Grab the review if it exists
+        review = get_pretrained_checkpoint_review(workflow, task_name)
+
+        if not review or not review["reviewed"]:
+            print(f"Skipping publishing of {workflow}:{task_name}, hasn't been reviewed yet")
+            return
+
+        result = review["result"]
+        if result != "accepted":
+            print(f'Skipping publishing of {workflow}:{task_name}, review result was "{result}"')
+            return
+
+    print(f"Publishing {workflow}:{task_name}")
+
+    # Copy the file
+    publish_path = get_pretrained_checkpoint_publish_path(workflow, task_name)
+    omni.client.copy_file(local_path, publish_path, CopyBehavior.OVERWRITE)
+
+
+def get_job_summary_row(workflow, task_name):
+    """Returns a single row summary of the job"""
+
+    has_run = has_pretrained_checkpoint_job_run(workflow, task_name)
+    has_finished = has_pretrained_checkpoint_job_finished(workflow, task_name)
+    review = get_pretrained_checkpoint_review(workflow, task_name)
+
+    if review:
+        result = review.get("result", "undetermined")
+        notes = review.get("notes", "")
+    else:
+        result = ""
+        notes = ""
+
+    return [workflow, task_name, has_run, has_finished, result, notes]
+
+
+def main():
+
+    # Figure out what workflows and tasks we'll be using
+    if args.all:
+        jobs = ["*:*"]
+    else:
+        jobs = args.jobs
+
+    if args.list:
+        print()
+        print("# Workflow, Task, Ran, Finished, Review, Notes")
+
+    summary_rows = []
+
+    # Could be implemented more efficiently, but the performance gain would be inconsequential
+    for workflow in WORKFLOWS:
+        for task_spec in sorted(gym.registry.values(), key=lambda t: t.id):
+            job_id = f"{workflow}:{task_spec.id}"
+
+            # We've excluded this job
+            if any(fnmatch.fnmatch(job_id, e) for e in args.exclude):
+                continue
+
+            # None of our jobs match this pair
+            if not np.any(np.array([fnmatch.fnmatch(job_id, job) for job in jobs])):
+                continue
+
+            # No config for this workflow
+            if workflow + "_cfg_entry_point" not in task_spec.kwargs:
+                continue
+
+            if args.list:
+                summary_rows.append(get_job_summary_row(workflow, task_spec.id))
+                continue
+
+            # Training reviewing and publishing
+            if args.train:
+                train_job(workflow, task_spec.id, args.headless, args.force, args.num_envs)
+
+            if args.review:
+                review_pretrained_checkpoint(workflow, task_spec.id, args.force_review, args.num_envs)
+
+            if args.publish_checkpoint:
+                publish_pretrained_checkpoint(workflow, task_spec.id, args.force_publish)
+
+    if args.list:
+        writer = csv.writer(sys.stdout, quotechar='"', quoting=csv.QUOTE_MINIMAL)
+        writer.writerows(summary_rows)
+
+
+if __name__ == "__main__":
+
+    try:
+        # Run the main function
+        main()
+    except Exception as e:
+        raise e
+    finally:
+        # Close the app
+        simulation_app.close()

source/standalone/workflows/rl_games/play.py

@@ -21,6 +21,11 @@ parser.add_argument(
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.")
+parser.add_argument(
+    "--use_pretrained_checkpoint",
+    action="store_true",
+    help="Use the pre-trained checkpoint from Nucleus.",
+)
 parser.add_argument(
     "--use_last_checkpoint",
     action="store_true",
@@ -53,6 +58,7 @@ from rl_games.torch_runner import Runner
 from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent
 from omni.isaac.lab.utils.assets import retrieve_file_path
 from omni.isaac.lab.utils.dict import print_dict
+from omni.isaac.lab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 
 import omni.isaac.lab_tasks  # noqa: F401
 from omni.isaac.lab_tasks.utils import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg
@@ -72,7 +78,12 @@ def main():
     log_root_path = os.path.abspath(log_root_path)
     print(f"[INFO] Loading experiment from directory: {log_root_path}")
     # find checkpoint
-    if args_cli.checkpoint is None:
+    if args_cli.use_pretrained_checkpoint:
+        resume_path = get_published_pretrained_checkpoint("rl_games", args_cli.task)
+        if not resume_path:
+            print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
+            return
+    elif args_cli.checkpoint is None:
         # specify directory for logging runs
         run_dir = agent_cfg["params"]["config"].get("full_experiment_name", ".*")
         # specify name of checkpoint

source/standalone/workflows/rsl_rl/play.py

@@ -23,6 +23,11 @@ parser.add_argument(
 )
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
+parser.add_argument(
+    "--use_pretrained_checkpoint",
+    action="store_true",
+    help="Use the pre-trained checkpoint from Nucleus.",
+)
 # append RSL-RL cli arguments
 cli_args.add_rsl_rl_args(parser)
 # append AppLauncher cli args
@@ -45,7 +50,9 @@ import torch
 from rsl_rl.runners import OnPolicyRunner
 
 from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent
+from omni.isaac.lab.utils.assets import retrieve_file_path
 from omni.isaac.lab.utils.dict import print_dict
+from omni.isaac.lab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 
 import omni.isaac.lab_tasks  # noqa: F401
 from omni.isaac.lab_tasks.utils import get_checkpoint_path, parse_env_cfg
@@ -69,7 +76,16 @@ def main():
     log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
     log_root_path = os.path.abspath(log_root_path)
     print(f"[INFO] Loading experiment from directory: {log_root_path}")
-    resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
+    if args_cli.use_pretrained_checkpoint:
+        resume_path = get_published_pretrained_checkpoint("rsl_rl", args_cli.task)
+        if not resume_path:
+            print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
+            return
+    elif args_cli.checkpoint:
+        resume_path = retrieve_file_path(args_cli.checkpoint)
+    else:
+        resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
+
     log_dir = os.path.dirname(resume_path)
 
     # create isaac environment

source/standalone/workflows/sb3/play.py

@@ -21,6 +21,11 @@ parser.add_argument(
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.")
+parser.add_argument(
+    "--use_pretrained_checkpoint",
+    action="store_true",
+    help="Use the pre-trained checkpoint from Nucleus.",
+)
 parser.add_argument(
     "--use_last_checkpoint",
     action="store_true",
@@ -50,6 +55,7 @@ from stable_baselines3.common.vec_env import VecNormalize
 
 from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent
 from omni.isaac.lab.utils.dict import print_dict
+from omni.isaac.lab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 
 import omni.isaac.lab_tasks  # noqa: F401
 from omni.isaac.lab_tasks.utils.parse_cfg import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg
@@ -67,8 +73,13 @@ def main():
     # directory for logging into
     log_root_path = os.path.join("logs", "sb3", args_cli.task)
     log_root_path = os.path.abspath(log_root_path)
-    # check checkpoint is valid
-    if args_cli.checkpoint is None:
+    # checkpoint and log_dir stuff
+    if args_cli.use_pretrained_checkpoint:
+        checkpoint_path = get_published_pretrained_checkpoint("sb3", args_cli.task)
+        if not checkpoint_path:
+            print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
+            return
+    elif args_cli.checkpoint is None:
         if args_cli.use_last_checkpoint:
             checkpoint = "model_.*.zip"
         else:

source/standalone/workflows/skrl/play.py

@@ -26,6 +26,11 @@ parser.add_argument(
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.")
+parser.add_argument(
+    "--use_pretrained_checkpoint",
+    action="store_true",
+    help="Use the pre-trained checkpoint from Nucleus.",
+)
 parser.add_argument(
     "--ml_framework",
     type=str,
@@ -77,6 +82,7 @@ elif args_cli.ml_framework.startswith("jax"):
 
 from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent
 from omni.isaac.lab.utils.dict import print_dict
+from omni.isaac.lab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 
 import omni.isaac.lab_tasks  # noqa: F401
 from omni.isaac.lab_tasks.utils import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg
@@ -106,7 +112,12 @@ def main():
     log_root_path = os.path.abspath(log_root_path)
     print(f"[INFO] Loading experiment from directory: {log_root_path}")
     # get checkpoint path
-    if args_cli.checkpoint:
+    if args_cli.use_pretrained_checkpoint:
+        resume_path = get_published_pretrained_checkpoint("skrl", args_cli.task)
+        if not resume_path:
+            print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
+            return
+    elif args_cli.checkpoint:
         resume_path = os.path.abspath(args_cli.checkpoint)
     else:
         resume_path = get_checkpoint_path(