[3.0.0-beta2] Fix rl_games play checkpoint fallback (#6133)

Cherry-picks #6091 to release/3.0.0-beta2.

Upstream merge commit: de78c80
Cherry-pick commit: e598169

Conflict resolution:
- Kept the new checkpoint-path test under the beta2 test layout at
`source/isaaclab_tasks/test/test_checkpoint_path.py`.

Validation:
- git diff --check upstream/release/3.0.0-beta2..HEAD
- env
PYTHONPATH=source/isaaclab_tasks:source/isaaclab:source/isaaclab_rl
/home/zhengyuz/Projects/IsaacLab.wt/baseline-newton-fabric-repro/env_isaaclab/bin/python
-m py_compile source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py
scripts/reinforcement_learning/rl_games/play.py
scripts/reinforcement_learning/rl_games/play_rl_games.py
scripts/reinforcement_learning/sb3/play.py
scripts/reinforcement_learning/sb3/play_sb3.py
source/isaaclab_tasks/test/test_checkpoint_path.py
- env
PYTHONPATH=source/isaaclab_tasks:source/isaaclab:source/isaaclab_rl
/home/zhengyuz/Projects/IsaacLab.wt/newton-fabric-body-binding/env_isaaclab/bin/python
-m pytest -q source/isaaclab_tasks/test/test_checkpoint_path.py
- env
PYTHONPATH=source/isaaclab_tasks_experimental:source/isaaclab_tasks:source/isaaclab_rl:source/isaaclab:source/isaaclab_assets:source/isaaclab_mimic:source/isaaclab_newton:source/isaaclab_physx:source/isaaclab_ov:source/isaaclab_ovphysx:source/isaaclab_visualizers
/home/zhengyuz/Projects/IsaacLab.wt/newton-fabric-body-binding/env_isaaclab/bin/python
scripts/reinforcement_learning/rl_games/play.py --help
- env
PYTHONPATH=source/isaaclab_tasks_experimental:source/isaaclab_tasks:source/isaaclab_rl:source/isaaclab:source/isaaclab_assets:source/isaaclab_mimic:source/isaaclab_newton:source/isaaclab_physx:source/isaaclab_ov:source/isaaclab_ovphysx:source/isaaclab_visualizers
/home/zhengyuz/Projects/IsaacLab.wt/newton-fabric-body-binding/env_isaaclab/bin/python
scripts/reinforcement_learning/play.py --rl_library rl_games --help
- pre-commit run --files
docs/source/tutorials/03_envs/run_rl_training.rst
scripts/reinforcement_learning/rl_games/play.py
scripts/reinforcement_learning/rl_games/play_rl_games.py
scripts/reinforcement_learning/sb3/play.py
scripts/reinforcement_learning/sb3/play_sb3.py
source/isaaclab_tasks/changelog.d/fix-rlgames-checkpoint-fallback.rst
source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py
source/isaaclab_tasks/test/test_checkpoint_path.py

Co-authored-by: jichuanh <jichuanh@nvidia.com>

Author: ooctipus

docs/source/tutorials/03_envs/run_rl_training.rst

@@ -142,7 +142,7 @@ Once the training is complete, you can visualize the trained agent by executing
 .. code:: bash
 
    # execute from the root directory of the repository
-   ./isaaclab.sh play --rl_library sb3 --task Isaac-Cartpole --num_envs 32 --use_last_checkpoint --viz kit
+   ./isaaclab.sh play --rl_library sb3 --task Isaac-Cartpole --num_envs 32 --viz kit
 
 The above command will load the latest checkpoint from the ``logs/sb3/Isaac-Cartpole-v0``
 directory. You can also specify a specific checkpoint by passing the ``--checkpoint`` flag.

scripts/reinforcement_learning/rl_games/play.py

@@ -114,11 +114,12 @@ def main():
                 return
         elif args_cli.checkpoint is None:
             run_dir = agent_cfg["params"]["config"].get("full_experiment_name", ".*")
-            if args_cli.use_last_checkpoint:
-                checkpoint_file = ".*"
-            else:
-                checkpoint_file = f"{agent_cfg['params']['config']['name']}.pth"
-            resume_path = get_checkpoint_path(log_root_path, run_dir, checkpoint_file, other_dirs=["nn"])
+            # prefer the best-reward checkpoint (``<name>.pth``); fall back to the latest checkpoint when it has
+            # not been written yet (e.g. short runs). With --use_last_checkpoint, skip the preference entirely.
+            best_checkpoint = None if args_cli.use_last_checkpoint else f"{agent_cfg['params']['config']['name']}.pth"
+            resume_path = get_checkpoint_path(
+                log_root_path, run_dir, ".*", other_dirs=["nn"], preferred_checkpoint=best_checkpoint
+            )
         else:
             resume_path = retrieve_file_path(args_cli.checkpoint)
         log_dir = os.path.dirname(os.path.dirname(resume_path))

scripts/reinforcement_learning/rl_games/play_rl_games.py

@@ -104,11 +104,12 @@ def main():
                 return
         elif args_cli.checkpoint is None:
             run_dir = agent_cfg["params"]["config"].get("full_experiment_name", ".*")
-            if args_cli.use_last_checkpoint:
-                checkpoint_file = ".*"
-            else:
-                checkpoint_file = f"{agent_cfg['params']['config']['name']}.pth"
-            resume_path = get_checkpoint_path(log_root_path, run_dir, checkpoint_file, other_dirs=["nn"])
+            # prefer the best-reward checkpoint (``<name>.pth``); fall back to the latest checkpoint when it has
+            # not been written yet (e.g. short runs). With --use_last_checkpoint, skip the preference entirely.
+            best_checkpoint = None if args_cli.use_last_checkpoint else f"{agent_cfg['params']['config']['name']}.pth"
+            resume_path = get_checkpoint_path(
+                log_root_path, run_dir, ".*", other_dirs=["nn"], preferred_checkpoint=best_checkpoint
+            )
         else:
             resume_path = retrieve_file_path(args_cli.checkpoint)
         log_dir = os.path.dirname(os.path.dirname(resume_path))

scripts/reinforcement_learning/sb3/play.py

@@ -67,11 +67,6 @@
     action="store_true",
     help="Use the pre-trained checkpoint from Nucleus.",
 )
-parser.add_argument(
-    "--use_last_checkpoint",
-    action="store_true",
-    help="When no checkpoint provided, use the last saved model. Otherwise use the best saved model.",
-)
 parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
 parser.add_argument(
     "--keep_all_info",
@@ -114,11 +109,11 @@ def main():
                 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:
-                checkpoint = "model.zip"
-            checkpoint_path = get_checkpoint_path(log_root_path, ".*", checkpoint, sort_alpha=False)
+            # prefer the final model (``model.zip``); fall back to the latest periodic checkpoint when it has
+            # not been written yet (e.g. short or interrupted runs)
+            checkpoint_path = get_checkpoint_path(
+                log_root_path, ".*", r"model_.*\.zip", sort_alpha=False, preferred_checkpoint=r"model\.zip"
+            )
         else:
             checkpoint_path = args_cli.checkpoint
         log_dir = os.path.dirname(checkpoint_path)

scripts/reinforcement_learning/sb3/play_sb3.py

@@ -56,11 +56,6 @@
     action="store_true",
     help="Use the pre-trained checkpoint from Nucleus.",
 )
-parser.add_argument(
-    "--use_last_checkpoint",
-    action="store_true",
-    help="When no checkpoint provided, use the last saved model. Otherwise use the best saved model.",
-)
 parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
 parser.add_argument(
     "--keep_all_info",
@@ -104,11 +99,11 @@ def main():
                 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:
-                checkpoint = "model.zip"
-            checkpoint_path = get_checkpoint_path(log_root_path, ".*", checkpoint, sort_alpha=False)
+            # prefer the final model (``model.zip``); fall back to the latest periodic checkpoint when it has
+            # not been written yet (e.g. short or interrupted runs)
+            checkpoint_path = get_checkpoint_path(
+                log_root_path, ".*", r"model_.*\.zip", sort_alpha=False, preferred_checkpoint=r"model\.zip"
+            )
         else:
             checkpoint_path = args_cli.checkpoint
         log_dir = os.path.dirname(checkpoint_path)

source/isaaclab_tasks/changelog.d/fix-rlgames-checkpoint-fallback.rst

@@ -0,0 +1,13 @@
+Added
+^^^^^
+
+* Added a ``preferred_checkpoint`` regex argument to :func:`~isaaclab_tasks.utils.parse_cfg.get_checkpoint_path`
+  that is matched before ``checkpoint`` and wins when it matches, otherwise resolution falls back to ``checkpoint``.
+
+Fixed
+^^^^^
+
+* Fixed ``rl_games`` and ``sb3`` play failing to load a checkpoint on short runs where the preferred
+  best/final checkpoint has not been written yet. They now prefer the best (``rl_games``) or final
+  (``sb3``) checkpoint and fall back to the latest available checkpoint when it is missing. Numbered
+  checkpoint filenames are now sorted naturally so epoch 10 is selected after epoch 9.

source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py

@@ -171,7 +171,12 @@ def parse_env_cfg(
 
 
 def get_checkpoint_path(
-    log_path: str, run_dir: str = ".*", checkpoint: str = ".*", other_dirs: list[str] = None, sort_alpha: bool = True
+    log_path: str,
+    run_dir: str = ".*",
+    checkpoint: str = ".*",
+    other_dirs: list[str] = None,
+    sort_alpha: bool = True,
+    preferred_checkpoint: str | None = None,
 ) -> str:
     """Get path to the model checkpoint in input directory.
 
@@ -187,10 +192,15 @@ def get_checkpoint_path(
             recent directory created inside :attr:`log_path`.
         other_dirs: The intermediate directories between the run directory and the checkpoint file. Defaults to
             None, which implies that checkpoint file is directly under the run directory.
-        checkpoint: The regex expression for the model checkpoint file. Defaults to the most recent
-            torch-model saved in the :attr:`run_dir` directory.
+        checkpoint: The regex expression for the model checkpoint file. Defaults to ``".*"``, which matches all
+            checkpoints and selects the most recent one.
         sort_alpha: Whether to sort the runs by alphabetical order. Defaults to True.
             If False, the folders in :attr:`run_dir` are sorted by the last modified time.
+        preferred_checkpoint: An optional regex expression that is matched before :attr:`checkpoint`. When it is
+            provided and matches at least one file, that match is used; otherwise resolution falls back to
+            :attr:`checkpoint`. This allows preferring a specific checkpoint (e.g. the best or final model) while
+            still resolving the latest available checkpoint when the preferred file has not been written yet (e.g.
+            for short runs). Defaults to None, which matches :attr:`checkpoint` directly.
 
     Returns:
         The path to the model checkpoint.
@@ -219,13 +229,22 @@ def get_checkpoint_path(
     except IndexError:
         raise ValueError(f"No runs present in the directory: '{log_path}' match: '{run_dir}'.")
 
-    # list all model checkpoints in the directory
-    model_checkpoints = [f for f in os.listdir(run_path) if re.match(checkpoint, f)]
+    # prefer ``preferred_checkpoint`` when given and it matches; otherwise fall back to the general
+    # ``checkpoint`` pattern (e.g. resolve the latest numbered checkpoint when the preferred best/final
+    # checkpoint file has not been written yet)
+    model_checkpoints = []
+    if preferred_checkpoint is not None:
+        model_checkpoints = [f for f in os.listdir(run_path) if re.match(preferred_checkpoint, f)]
+    if len(model_checkpoints) == 0:
+        model_checkpoints = [f for f in os.listdir(run_path) if re.match(checkpoint, f)]
     # check if any checkpoints are present
     if len(model_checkpoints) == 0:
-        raise ValueError(f"No checkpoints in the directory: '{run_path}' match '{checkpoint}'.")
-    # sort alphabetically while ensuring that *_10 comes after *_9
-    model_checkpoints.sort(key=lambda m: f"{m:0>15}")
+        patterns = f"'{checkpoint}'"
+        if preferred_checkpoint is not None:
+            patterns = f"'{preferred_checkpoint}' nor '{checkpoint}'"
+        raise ValueError(f"No checkpoints in the directory: '{run_path}' match {patterns}.")
+    # sort naturally so numbered checkpoints such as *_10 come after *_9 even with long filename prefixes
+    model_checkpoints.sort(key=lambda m: [int(token) if token.isdigit() else token for token in re.split(r"(\d+)", m)])
     # get latest matched checkpoint file
     checkpoint_file = model_checkpoints[-1]
 

source/isaaclab_tasks/test/test_checkpoint_path.py

@@ -0,0 +1,103 @@
+# 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
+
+import pytest
+
+from isaaclab_tasks.utils import get_checkpoint_path
+
+
+def test_get_checkpoint_path_prefers_requested_checkpoint(tmp_path):
+    checkpoint_dir = tmp_path / "cartpole_camera_direct" / "2026-06-07_02-41-41" / "nn"
+    checkpoint_dir.mkdir(parents=True)
+    expected_checkpoint = checkpoint_dir / "cartpole_camera_direct.pth"
+    expected_checkpoint.touch()
+    (checkpoint_dir / "last_cartpole_camera_direct_ep_10_rew_1.0.pth").touch()
+
+    checkpoint_path = get_checkpoint_path(
+        str(tmp_path / "cartpole_camera_direct"), "2026-06-07_02-41-41", "cartpole_camera_direct.pth", other_dirs=["nn"]
+    )
+
+    assert checkpoint_path == str(expected_checkpoint)
+
+
+def test_get_checkpoint_path_uses_latest_checkpoint_when_checkpoint_is_omitted(tmp_path):
+    checkpoint_dir = tmp_path / "cartpole_camera_direct" / "2026-06-07_02-41-41" / "nn"
+    checkpoint_dir.mkdir(parents=True)
+    (checkpoint_dir / "last_cartpole_camera_direct_ep_5_rew_0.5.pth").touch()
+    expected_checkpoint = checkpoint_dir / "last_cartpole_camera_direct_ep_10_rew_1.0.pth"
+    expected_checkpoint.touch()
+
+    checkpoint_path = get_checkpoint_path(
+        str(tmp_path / "cartpole_camera_direct"), "2026-06-07_02-41-41", other_dirs=["nn"]
+    )
+
+    assert checkpoint_path == str(expected_checkpoint)
+
+
+def test_get_checkpoint_path_requested_checkpoint_stays_strict(tmp_path):
+    checkpoint_dir = tmp_path / "cartpole_camera_direct" / "2026-06-07_02-41-41" / "nn"
+    checkpoint_dir.mkdir(parents=True)
+    (checkpoint_dir / "last_cartpole_camera_direct_ep_10_rew_1.0.pth").touch()
+
+    with pytest.raises(ValueError, match="cartpole_camera_direct.pth"):
+        get_checkpoint_path(
+            str(tmp_path / "cartpole_camera_direct"),
+            "2026-06-07_02-41-41",
+            "cartpole_camera_direct.pth",
+            other_dirs=["nn"],
+        )
+
+
+def test_get_checkpoint_path_prefers_preferred_over_checkpoint(tmp_path):
+    # rl_games: the best checkpoint must win over numbered ones when given as the preferred pattern
+    checkpoint_dir = tmp_path / "cartpole_camera_direct" / "2026-06-07_02-41-41" / "nn"
+    checkpoint_dir.mkdir(parents=True)
+    expected_checkpoint = checkpoint_dir / "cartpole_camera_direct.pth"
+    expected_checkpoint.touch()
+    (checkpoint_dir / "last_cartpole_camera_direct_ep_10_rew_1.0.pth").touch()
+
+    checkpoint_path = get_checkpoint_path(
+        str(tmp_path / "cartpole_camera_direct"),
+        "2026-06-07_02-41-41",
+        ".*",
+        other_dirs=["nn"],
+        preferred_checkpoint="cartpole_camera_direct.pth",
+    )
+
+    assert checkpoint_path == str(expected_checkpoint)
+
+
+def test_get_checkpoint_path_falls_back_when_preferred_missing(tmp_path):
+    # rl_games short run: the best checkpoint is never written, so resolution falls back to the latest numbered one
+    checkpoint_dir = tmp_path / "cartpole_camera_direct" / "2026-06-07_02-41-41" / "nn"
+    checkpoint_dir.mkdir(parents=True)
+    (checkpoint_dir / "last_cartpole_camera_direct_ep_9_rew_0.9.pth").touch()
+    expected_checkpoint = checkpoint_dir / "last_cartpole_camera_direct_ep_10_rew_1.0.pth"
+    expected_checkpoint.touch()
+
+    checkpoint_path = get_checkpoint_path(
+        str(tmp_path / "cartpole_camera_direct"),
+        "2026-06-07_02-41-41",
+        ".*",
+        other_dirs=["nn"],
+        preferred_checkpoint="cartpole_camera_direct.pth",
+    )
+
+    assert checkpoint_path == str(expected_checkpoint)
+
+
+def test_get_checkpoint_path_prefers_final_model_over_steps(tmp_path):
+    # sb3: the final ``model.zip`` must win over numbered ``model_<n>_steps.zip`` snapshots
+    run_dir = tmp_path / "2026-06-07_02-41-41"
+    run_dir.mkdir(parents=True)
+    expected_checkpoint = run_dir / "model.zip"
+    expected_checkpoint.touch()
+    (run_dir / "model_2000_steps.zip").touch()
+
+    checkpoint_path = get_checkpoint_path(
+        str(tmp_path), ".*", r"model_.*\.zip", sort_alpha=False, preferred_checkpoint=r"model\.zip"
+    )
+
+    assert checkpoint_path == str(expected_checkpoint)