Adds pre-trained checkpoints and tools for generating and uploading checkpoints (isaac-sim#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 <[email protected]>
Signed-off-by: Kelly Guo <[email protected]>
Signed-off-by: Kelly Guo <[email protected]>
Co-authored-by: peterd-NV <[email protected]>
Co-authored-by: CY Chen <[email protected]>
Co-authored-by: oahmednv <[email protected]>
Co-authored-by: Toni-SM <[email protected]>
Co-authored-by: Kelly Guo <[email protected]>
Co-authored-by: Kelly Guo <[email protected]>

Author: 6 people

.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}")