[3.0.0-beta2] Fix rl_games play checkpoint fallback#6133
Conversation
- add an opt-in fallback checkpoint pattern to the shared get_checkpoint_path() helper - use that fallback from rl_games play so short runs can load the latest available checkpoint when the preferred best-checkpoint file has not been written yet - naturally sort numbered checkpoint filenames so epoch 10 is selected after epoch 9 - python3 -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 source/isaaclab_tasks/test/core/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/core/test_checkpoint_path.py - pre-commit run --files scripts/reinforcement_learning/rl_games/play.py scripts/reinforcement_learning/rl_games/play_rl_games.py source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py source/isaaclab_tasks/test/core/test_checkpoint_path.py source/isaaclab_tasks/changelog.d/fix-rlgames-checkpoint-fallback.rst - env PYTHONPATH=source/isaaclab_tasks:source/isaaclab:source/isaaclab_rl /home/zhengyuz/Projects/IsaacLab.wt/feature-unify/env_isaaclab/bin/python scripts/reinforcement_learning/rl_games/play.py --help - env PYTHONPATH=source/isaaclab_tasks:source/isaaclab:source/isaaclab_rl /home/zhengyuz/Projects/IsaacLab.wt/feature-unify/env_isaaclab/bin/python scripts/reinforcement_learning/play.py --rl_library rl_games --help --------- Co-authored-by: jichuanh <jichuanh@nvidia.com> (cherry picked from commit de78c80)
github-actions
Bot
added
bug
Greptile SummaryThis cherry-pick from
Confidence Score: 4/5Safe to merge; the core logic change is small and well-tested, with the only practical concern being the silent removal of --use_last_checkpoint from sb3 scripts The checkpoint fallback logic is correct and covered by six focused tests. The natural sort key change is a clear improvement with no edge-case hazard. The two items worth watching are: re.match leaving patterns unanchored at the end (consistent with the existing checkpoint behaviour, and harmless with real SB3/rl_games filenames), and the removal of --use_last_checkpoint from the sb3 scripts without a transitional deprecation warning, which will hard-fail any existing automation that passes the flag. The sb3 play scripts (play_sb3.py and play.py) warrant a second look given the CLI-breaking removal of --use_last_checkpoint; all other files are straightforward. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[get_checkpoint_path called] --> B{preferred_checkpoint\nprovided?}
B -- No --> E[Match files with checkpoint pattern]
B -- Yes --> C[Match files with preferred_checkpoint]
C --> D{Any matches?}
D -- Yes --> G[Natural sort matches]
D -- No --> E
E --> F{Any matches?}
F -- No --> ERR[Raise ValueError]
F -- Yes --> G
G --> H[Return last sorted match]
Reviews (1): Last reviewed commit: "Fix rl_games play checkpoint fallback (#..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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)] |
There was a problem hiding this comment.
re.match does not anchor at the end of the filename
Both preferred_checkpoint and the existing checkpoint parameter use re.match, which anchors only at the start of the string. A pattern like r"model\.zip" would match model.zip_backup because re.match succeeds as soon as the prefix matches without consuming the rest of the string. Using re.fullmatch (or appending $ to all patterns) would make the behaviour unambiguous and consistent with the intent of selecting an exact named file. This is the same pre-existing behaviour as checkpoint, so widening to re.fullmatch here would need a matching change to line 239 for consistency.
| 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" |
There was a problem hiding this comment.
--use_last_checkpoint removed without a deprecation warning
--use_last_checkpoint is removed from both play_sb3.py and play.py (sb3). Any script, CI job, or user invocation that passes --use_last_checkpoint will now fail immediately with argparse's unrecognized arguments error. The rl_games counterparts retain the flag, so the breaking change is only on the sb3 side. Since the new automatic fallback already covers the "short run / no final model" case, the flag is no longer needed — but a short deprecation window (e.g. accepting the flag with a warnings.warn) would smooth the transition for existing automation pipelines.
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review: Fix rl_games play checkpoint fallback
Clean approach to a real usability pain point. The preferred_checkpoint parameter with automatic fallback is a solid design choice—it eliminates the need for users to know about --use_last_checkpoint in the common case while preserving explicit control for power users.
Summary of findings:
- Natural sort implementation works correctly but has a subtle edge case worth noting
- The
preferred_checkpointregex is not validated, which could produce confusing errors - SB3 raw-string regex patterns are correct and consistent
- Test coverage is good for the happy path and the fallback path; one gap noted below
- Documentation update is correct (removes the now-unnecessary flag from the tutorial)
Overall: well-structured change with good test coverage. A few minor suggestions below.
| 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}.") |
There was a problem hiding this comment.
suggestion (minor): The natural sort via re.split(r"(\d+)", m) produces empty strings at the boundaries (e.g. re.split(r'(\d+)', '10_foo') → ['', '10', '_foo']). This works because empty strings sort before non-empty ones, but it means the sort isn't purely "natural" for filenames that start with digits—in that case the leading empty token becomes the primary sort key and all such files cluster together, which is actually fine for checkpoint naming conventions here.
No action needed—just flagging for awareness since the old zero-padded approach (f"{m:0>15}") would break on filenames longer than 15 chars, so this is strictly an improvement.
|
|
||
| # 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 |
There was a problem hiding this comment.
suggestion (robustness): If preferred_checkpoint is an invalid regex, re.match will raise re.error with a potentially confusing traceback that doesn't mention preferred_checkpoint by name. Consider wrapping this in a try/except that raises a clearer ValueError, e.g.:
try:
model_checkpoints = [f for f in os.listdir(run_path) if re.match(preferred_checkpoint, f)]
except re.error as e:
raise ValueError(f"Invalid regex for preferred_checkpoint: '{preferred_checkpoint}': {e}") from eSame applies to the checkpoint pattern on the next branch. Low priority since callers currently pass known-good patterns, but it's a public API.
| 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( |
There was a problem hiding this comment.
nit: The raw-string r"model_.*\.zip" and r"model\.zip" patterns are correct. Just confirming that the \. is intentional to avoid matching e.g. model_zip (without the dot)—good defensive regex.
| str(tmp_path), ".*", r"model_.*\.zip", sort_alpha=False, preferred_checkpoint=r"model\.zip" | ||
| ) | ||
|
|
||
| assert checkpoint_path == str(expected_checkpoint) |
There was a problem hiding this comment.
gap (test coverage): There's no test for the case where preferred_checkpoint matches multiple files. For example, if the regex is r"model\.zip" and somehow both model.zip and (pathologically) another file matches, the sort order determines which wins. It might be worth adding a test that verifies the "last after natural sort" behavior when preferred_checkpoint matches >1 file, ensuring the contract is clear.
Also, no test exercises sort_alpha=False combined with preferred_checkpoint. The sb3 scripts use this combination—would be good to have coverage.
| 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 |
There was a problem hiding this comment.
question: With this change, when --use_last_checkpoint is passed, best_checkpoint is None, so preferred_checkpoint=None is passed and the function skips the preferred branch entirely, falling through to the ".*" pattern which matches all checkpoints and picks the latest by natural sort. This is correct and equivalent to the old behavior.
However, should --use_last_checkpoint still be documented/advertised for rl_games, given that the default behavior now auto-falls-back anyway? The flag's only effect is to skip preferring the best checkpoint—which only matters when the best checkpoint does exist but you explicitly don't want it. That's a valid but niche use case; might be worth a comment or docstring update.
a855496
into
isaac-sim:release/3.0.0-beta2
1 participant
Cherry-picks #6091 to release/3.0.0-beta2.
Upstream merge commit: de78c80
Cherry-pick commit: e598169
Conflict resolution:
source/isaaclab_tasks/test/test_checkpoint_path.py.Validation: