Refactor train and play scripts for single entry point#5554
Refactor train and play scripts for single entry point#5554StafaH wants to merge 5 commits intoisaac-sim:developfrom
Conversation
StafaH
requested review from
Mayankm96,
hhansen-bdai,
jtigue-bdai,
kellyguo11 and
ooctipus
as code owners
github-actions
Bot
added
documentation
Greptile SummaryThis PR centralizes the previously duplicated RL training/play entrypoints across six libraries (rl_games, rlinf, rsl_rl, sb3, skrl) into unified
Confidence Score: 4/5The refactor is primarily additive — existing per-library scripts remain functional with deprecation warnings and the new unified entry-points replicate original logic faithfully. The core dispatch logic, pretrained-checkpoint constant updates, and CLI alias tests are all consistent. The heavy module-level imports in common.py and the undocumented run_as_script API contract are design-level concerns rather than functional regressions on the standard train/play workflows. scripts/reinforcement_learning/common.py for module-level import overhead and undocumented run_as_script contract; source/isaaclab/isaaclab/cli/utils.py for the Path normalization edge case in the alias lookup. Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant CLI as "isaaclab.sh -p"
participant utils as "cli/utils.py"
participant train as "train.py"
participant common as "common.py"
participant lib as "train_lib.py"
User->>CLI: "./isaaclab.sh -p train.py --library rsl_rl --task X"
CLI->>utils: "run_python_command(train.py, args)"
utils->>utils: "_PYTHON_SCRIPT_ALIASES lookup -> absolute path"
utils->>train: "python .../train.py --library rsl_rl --task X"
train->>common: "dispatch_library_entrypoint(argv, entrypoints)"
common->>common: "parse --library rsl_rl"
common->>lib: "import_local_module(train_rsl_rl, path)"
lib-->>common: "module loaded (top-level code runs)"
common->>lib: "module.run(library_args)"
lib->>lib: "_parse_args + launch_simulation + runner.learn()"
lib-->>User: "training complete"
Note over common,lib: "Play path uses runpy.run_path() instead of import+run()"
Reviews (1): Last reviewed commit: "Refactor play scripts and create alias" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| from __future__ import annotations | ||
|
|
||
| import argparse | ||
| import importlib.util | ||
| import logging | ||
| import os | ||
| import runpy | ||
| import sys | ||
| from pathlib import Path | ||
| from types import ModuleType | ||
| from typing import Any | ||
|
|
||
| import gymnasium as gym | ||
|
|
||
| from isaaclab.envs import DirectMARLEnvCfg, ManagerBasedRLEnvCfg | ||
| from isaaclab.utils.dict import print_dict | ||
| from isaaclab.utils.io import dump_yaml | ||
|
|
There was a problem hiding this comment.
common.py imports gymnasium, isaaclab.envs (two env base classes), isaaclab.utils.dict, isaaclab.utils.io, and isaaclab_tasks.utils unconditionally at import time. Since train.py and play.py both do from common import dispatch_library_entrypoint, these imports fire even when the user just runs train.py --help (no --library selected) or when any test imports the dispatch helper. Only dispatch_library_entrypoint itself is needed for routing; the remaining helpers (create_isaaclab_env, wrap_record_video, etc.) are only called from individual library scripts. Moving those heavyweight imports to be lazy (inside the functions that need them) or splitting dispatch_library_entrypoint into its own thin module would keep the help/routing path fast.
| def dispatch_library_entrypoint( | ||
| argv: list[str] | None, | ||
| entrypoints: dict[str, Path], | ||
| *, | ||
| action: str, | ||
| description: str, | ||
| library_help: str, | ||
| run_as_script: bool = False, | ||
| ) -> int: | ||
| """Dispatch a unified entrypoint to a library-specific implementation. | ||
|
|
||
| Args: | ||
| argv: Command-line arguments, excluding the script path. | ||
| entrypoints: Mapping from library name to implementation path. | ||
| action: Action name used to create a unique module name. | ||
| description: Top-level parser description. | ||
| library_help: Help text for the ``--library`` argument. | ||
| run_as_script: Whether to execute the selected implementation as a script. | ||
|
|
||
| Returns: | ||
| Process exit code. | ||
| """ | ||
| if argv is None: | ||
| argv = sys.argv[1:] | ||
|
|
||
| parser = argparse.ArgumentParser(add_help=False) | ||
| parser.add_argument("--library", choices=sorted(entrypoints)) | ||
| args_cli, library_args = parser.parse_known_args(argv) | ||
|
|
||
| if args_cli.library is None: | ||
| help_parser = argparse.ArgumentParser(description=description) | ||
| help_parser.add_argument("--library", choices=sorted(entrypoints), required=True, help=library_help) | ||
| help_parser.add_argument("args", nargs=argparse.REMAINDER, help="Arguments forwarded to the selected library.") | ||
| help_parser.print_help() | ||
| return 0 if "-h" in argv or "--help" in argv else 2 | ||
|
|
||
| module_path = entrypoints[args_cli.library] | ||
| if run_as_script: | ||
| original_argv = sys.argv | ||
| original_path = list(sys.path) | ||
| try: | ||
| sys.argv = [str(module_path)] + library_args | ||
| sys.path.insert(0, str(module_path.parent)) | ||
| runpy.run_path(str(module_path), run_name="__main__") | ||
| finally: | ||
| sys.argv = original_argv | ||
| sys.path[:] = original_path | ||
| return 0 | ||
|
|
||
| module = import_local_module(f"isaaclab_rl_{action}_{args_cli.library}", module_path) | ||
| module.run(library_args) | ||
| return 0 | ||
|
|
||
|
|
||
| def add_common_train_args( | ||
| parser: argparse.ArgumentParser, | ||
| *, | ||
| agent_default: str | None, | ||
| agent_help: str, | ||
| include_agent: bool = True, | ||
| include_distributed: bool = True, | ||
| ) -> None: |
There was a problem hiding this comment.
dispatch_library_entrypoint supports two incompatible modes keyed on run_as_script:
False(train): loads the target file as a module and callsmodule.run(argv)— the file must expose arun()function and must not parsesys.argvat module level.True(play): usesrunpy.run_path(..., run_name="__main__")and setssys.argvitself — the file must not have arun()function and must parsesys.argvat module level.
Neither mode is validated at call time. If a new contributor creates a play script that exposes run() and uses run_as_script=True, the dispatch silently ignores run(). Conversely, if a new train script parses sys.argv at module level, it will read the wrong argv. Adding an assertion or a clear note in the docstring about the run(argv) requirement for run_as_script=False would prevent future breakage.
| else: | ||
| script_or_module = _PYTHON_SCRIPT_ALIASES.get(str(script_or_module), script_or_module) |
There was a problem hiding this comment.
Path("./train.py") silently matches the alias while "./train.py" string does not
str(Path("./train.py")) normalises to "train.py" on all platforms, so a caller passing Path("./train.py") would be redirected to the global script, whereas the string "./train.py" correctly passes through. The test covers the string form but not the Path form. Guarding on len(Path(_key).parts) == 1 would make the intent explicit and consistent.
| else: | |
| script_or_module = _PYTHON_SCRIPT_ALIASES.get(str(script_or_module), script_or_module) | |
| else: | |
| # Only resolve bare filenames (e.g. "train.py"), not explicit paths such as | |
| # "./train.py" or "/abs/path/train.py". Using Path(x).parts ensures a leading | |
| # "./" or directory component prevents a match. | |
| _key = str(script_or_module) | |
| if len(Path(_key).parts) == 1: | |
| script_or_module = _PYTHON_SCRIPT_ALIASES.get(_key, script_or_module) |
StafaH
changed the title
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review Summary
This is a substantial quality-of-life improvement for the RL workflow. The unified train/play entrypoints with --library dispatch remove significant duplication, and the uv run path for fresh clones is a great developer experience improvement.
Strengths
- Clean dispatch pattern via
dispatch_library_entrypoint— extensible for new libraries common.pyconsolidates duplicated arg definitions and env creation logic well- Deprecation warnings on old scripts with clear migration guidance
_PYTHON_SCRIPT_ALIASESis a nice touch forisaaclab.sh -p train.pyshorthand- Python discovery improvements (active venv → repo-local → kit) are sensible prioritization
- Root
pyproject.tomlwithuvsources/overrides is well-structured
Concerns
See inline comments for specific issues. High-level:
- The
common.pyrelative import (from common import ...) requires the script's parent dir to be onsys.path— this works when run directly but may fail in certain packaging scenarios. - Hard-pinning
torch==2.10.0in the rootpyproject.tomlwill need frequent updates. - The
run_as_scriptpath indispatch_library_entrypointmutatessys.argvandsys.pathwhich could have side effects if the dispatched script doesn't clean up properly.
|
|
||
| from pathlib import Path | ||
|
|
||
| from common import dispatch_library_entrypoint |
There was a problem hiding this comment.
🟡 Relative import fragility: from common import dispatch_library_entrypoint relies on the script being run from its parent directory or with sys.path containing the scripts/reinforcement_learning/ directory. This works for ./isaaclab.sh -p and python scripts/... but would break if someone tries to import this module from elsewhere.
Consider using a relative path-based import pattern:
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import dispatch_library_entrypointOr make common.py a proper package (scripts/reinforcement_learning/__init__.py).
| return 0 if "-h" in argv or "--help" in argv else 2 | ||
|
|
||
| module_path = entrypoints[args_cli.library] | ||
| if run_as_script: |
There was a problem hiding this comment.
🟡 sys.argv / sys.path mutation in run_as_script: Mutating global state (sys.argv, sys.path) and relying on runpy.run_path is fragile. If the dispatched script raises an exception, the finally block restores state, but if it calls sys.exit(), the restore runs but the process is exiting anyway.
More concerning: runpy.run_path executes in the caller's process — any global state changes in the dispatched script persist. This is fine for train scripts that exit, but if used in a library context it could leak state.
Consider documenting this behavior or adding a comment that run_as_script=True is only safe for scripts that are expected to terminate the process.
| """ | ||
| python_path = Path(python_exe) | ||
| return (python_path.parent.parent / "pyvenv.cfg").is_file() | ||
|
|
There was a problem hiding this comment.
🟡 Side effect: os.environ["UV_PYTHON"] set globally: Setting UV_PYTHON as an environment variable affects all subsequent subprocess.run calls that use uv. If get_pip_command is called multiple times with different python_exe values, only the last one wins globally.
Consider returning the env var as part of the command context rather than setting it globally:
# Return (cmd, env_override) tuple, or document the mutationOr at minimum, document that this function has a global side effect.
|
|
||
| def train(args: list[str] | None = None) -> None: | ||
| """Run the unified reinforcement learning training script.""" | ||
| if args is None: |
There was a problem hiding this comment.
💡 CLI dispatch pattern: The if sys.argv[1] == "train" / "play" check before argparse means these sub-commands bypass the normal argparse help system. Running isaaclab.sh --help would show the epilog mentioning train and play, but isaaclab.sh train --help delegates directly to the train script.
This is fine for now, but if more commands are added, consider using argparse subparsers or a proper CLI framework (click, typer) for the top-level dispatch.
|
|
||
| """Script to train RL agent with RL-Games.""" | ||
|
|
||
| import warnings |
There was a problem hiding this comment.
💡 Deprecation warning fires at import time: Because warnings.warn(...) is at module level, it fires even if someone imports a function from this module without intending to run it as a script. This is probably fine since the old scripts are meant to be run, not imported — but worth noting for anything that does from scripts.reinforcement_learning.rl_games.train import main.
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot
Summary
Well-structured refactoring that unifies RL train/play scripts behind a single dispatch entry point (train.py --library <lib> / play.py --library <lib>). The architecture is sound — clean separation between shared utilities (common.py), unified dispatchers, and library-specific implementations. Old scripts are preserved with deprecation warnings for backwards compatibility. A few design concerns around import fragility, namespace collision, and test coverage should be addressed.
Design Assessment
Acceptable, but import resolution relies on an implicit invariant. The unified dispatch pattern is the right approach. The from common import ... in library-specific implementations works because Python adds the script's directory to sys.path[0] when running train.py as a subprocess — but this invariant isn't explicitly enforced by import_local_module() or dispatch_library_entrypoint(). If these modules are ever imported from a different context (tests, REPL, another script), they'll fail with ModuleNotFoundError.
Findings
🟡 Warning: Generic train/play console_script names risk collision — source/isaaclab/setup.py
The entry points "train=isaaclab.cli:train" and "play=isaaclab.cli:play" are extremely generic names that could conflict with other ML packages (fairseq, lightning, etc.). Since isaaclab train / isaaclab play already works via the CLI dispatcher, the bare train/play entry points are redundant. Consider namespacing them to isaaclab-train/isaaclab-play, or removing them entirely.
🟡 Warning: Import resolution relies on implicit sys.path[0] invariant — scripts/reinforcement_learning/common.py
When dispatch_library_entrypoint() calls import_local_module() to load a library module (e.g., train_rsl_rl.py), the from common import ... statement in that module resolves ONLY because scripts/reinforcement_learning/ is already on sys.path[0] (set by Python when train.py runs as a subprocess). This is fragile — if dispatch_library_entrypoint is ever called from a different context (tests, REPL), it fails. Consider adding explicit path management:
rl_scripts_dir = str(Path(__file__).resolve().parent)
if rl_scripts_dir not in sys.path:
sys.path.insert(0, rl_scripts_dir)🟡 Warning: common.py has no unit tests — 12+ shared utility functions (dispatch_library_entrypoint, import_local_module, apply_env_overrides, validate_distributed_device, etc.) have zero direct test coverage. These are the foundation used by all 5 library implementations. A lightweight test_common.py with edge-case tests (missing --library, invalid module path, CPU+distributed rejection) would significantly reduce regression risk.
🔵 Suggestion: env.close() skipped on KeyboardInterrupt — scripts/reinforcement_learning/rsl_rl/train_rsl_rl.py
The except KeyboardInterrupt: pass at the end of training skips env.close(). Consider using finally to ensure cleanup (logging final metrics, releasing GPU resources):
try:
runner.learn(...)
print(f"Training time: ...")
except KeyboardInterrupt:
pass
finally:
env.close()Test Coverage
- CLI wiring (alias resolution, command dispatch): ✅ Well-tested with mocks
common.pyutility functions (12+ functions): ❌ No unit tests- Dispatch mechanism edge cases: ❌ Not directly tested
- Deprecated script backwards compatibility: ❌ Not tested
- Integration tests:
⚠️ Present but currently skipped
CI Status
- Build Wheel: ✅ pass
- Most checks: ⏳ pending
Verdict
Minor fixes needed
The architecture and implementation are solid for all standard execution paths. The dispatch pattern is clean, backwards compatibility is preserved, and the code is well-organized. Address the namespace collision concern (train/play console_scripts) and consider adding explicit sys.path management for robustness. Unit tests for common.py would be a valuable addition but aren't strictly blocking.
1 participant
Description
This PR refactors the reinforcement learning train/play scripts into unified entrypoints while preserving the existing library folder structure and adding a lightweight
uvworkflow for fresh source checkouts.The main changes are:
scripts/reinforcement_learning/train.py --library <library>scripts/reinforcement_learning/play.py --library <library>scripts/reinforcement_learning/rsl_rl/train_rsl_rl.pyscripts/reinforcement_learning/rsl_rl/play_rsl_rl.pytrain.pyandplay.pyscripts intact, with deprecation warnings and migration examples.scripts/reinforcement_learning/common.py../isaaclab.sh train --library <library> ..../isaaclab.sh play --library <library> ..../isaaclab.sh -p train.py ...and./isaaclab.sh -p play.py ....train --library <library> ...play --library <library> ...pyproject.tomlproject so a fresh clone can run kitless Newton training with:uv run train --library rsl_rl --task Isaac-Cartpole-Direct-v0 presets=newton_mjwarp --num_envs 4096uvenvironment to the same PyTorch family used by the Isaac Lab installer to avoid CUDA stack churn when users switch betweenuv runand./isaaclab.sh../isaaclab.shprefers an active or repo-local virtual environment before falling back to system Python../isaaclab.sh --installin uv-created virtual environments that do not include thepipmodule by usinguv pipfor venv-targeted pip operations.Motivation:
The previous RL scripts duplicated substantial train/play setup logic across libraries. This made behavior harder to keep consistent and increased maintenance cost when updating shared functionality. The new structure keeps library-specific logic in each library folder while centralizing shared dispatch and common helpers.
The
uvworkflow gives users a fast path from a fresh clone to kitless Newton training without manually creating an environment first. Isaac Sim / Kit workflows, including PhysX, continue to use the existing full installation path.Dependencies:
No new required runtime dependencies are added to Isaac Lab packages. The root source-checkout
pyproject.tomldescribes the local development environment used byuv run.Fixes # N/A
Type of change
Screenshots
Not applicable.
Validation
Ran:
uv lockuv lock --checkuv run --frozen train --help./isaaclab.sh --install./isaaclab.sh -p -m pytest source/isaaclab/test/cli/test_install.py -q./isaaclab.sh -p -m pytest source/isaaclab/test/cli/test_install.py::TestGetPipCommand source/isaaclab/test/cli/test_install.py::TestExtractPythonExe -q./isaaclab.sh -p train.py --help/./isaaclab.sh -p play.py --helpChecklist
pre-commitchecks with./isaaclab.sh --formatCONTRIBUTORS.mdor my name already exists there