[Fix] Publish and fetch the checkpoints a task's components declare - #7485
[Fix] Publish and fetch the checkpoints a task's components declare#7485hujc7 wants to merge 4 commits into
Conversation
github-actions
Bot
added
enhancement
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
play --checkpoint pretrained crashed on the Shadow Hand camera tasks: the vision CNN the policy needs was never published beside it, and nothing in the tooling knew the file existed. A component now declares what it writes with a Checkpoint on its own config. The tooling walks the resolved environment config to find every declaration, so a task declares nothing, and publishes each file beside the policy as <policy stem>_<name><extension>. The download path fetches them with the policy.
8a871d9 to
eda9b46
Compare
hujc7
changed the title
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
"Auxiliary" ranked these files below the policy and did not generalise: any component can declare a run artifact, and the policy is not special among them. The discovery and path helpers, their parameter, and the collect locals now say declared.
17433c7 to
b87603e
Compare
hujc7
changed the title
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
…ned-feature-extractor-ckpt # Conflicts: # source/isaaclab_rl/isaaclab_rl/utils/pretrained_checkpoint.py # source/isaaclab_rl/test/test_pretrained_checkpoint.py
|
run-ci |
isaaclab-bot
Bot
added
ci:run-docker
Greptile SummaryThis PR introduces component-owned checkpoint declarations and updates training, publication, download, playback, and export paths to carry companion weights alongside policy checkpoints.
Confidence Score: 3/5This PR is not yet safe to merge because incomplete bundles remain playable and RL-Games resolves downloaded companions from the wrong directory. Two concrete playback failures remain: RL-Games searches outside the companion's downloaded directory, and missing required companions are only warned about despite causing component initialization to fail. Files Needing Attention: source/isaaclab_rl/isaaclab_rl/utils/pretrained_checkpoint.py, source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rl_games.py, scripts/tools/train_and_publish_checkpoints.py Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
C[Component Checkpoint declaration] --> D[Discover from environment config]
D --> T[Training run writes artifact]
T --> K[Collect policy and companion]
K --> P[Publish beside policy]
P --> F[Fetch into checkpoint cache]
F --> L[Pass policy directory as log_dir]
L --> R[Component resolves companion]
Reviews (1): Last reviewed commit: "Merge remote-tracking branch 'upstream/d..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if args_cli.checkpoint == "pretrained": | ||
| backend_names = get_pretrained_checkpoint_backend_names(env_cfg) | ||
| resume_path = get_published_pretrained_checkpoint("rl_games", train_task_name, *backend_names) | ||
| resume_path = get_published_pretrained_checkpoint("rl_games", train_task_name, env_cfg=env_cfg) |
There was a problem hiding this comment.
When an RL-Games task declares a companion checkpoint, the policy and companion are downloaded beside each other. The RL-Games play path then sets log_dir to the policy's grandparent, so the component searches one directory above the downloaded companion and raises FileNotFoundError. As a result, pretrained playback fails for RL-Games tasks that need companion weights.
| except FileNotFoundError: | ||
| # the asset server reports a checkpoint that was never published and a server it | ||
| # cannot reach the same way, so both are covered by the same message | ||
| backends = ( | ||
| "" | ||
| if physics_backend is None | ||
| else f" with the '{physics_backend}' physics and '{render_backend}' render backends" | ||
| ) | ||
| print( | ||
| "A pre-trained checkpoint is currently unavailable for this task.\n" | ||
| f" The asset server does not provide '{ov_path}'.\n" | ||
| f" Either no checkpoint is published for task '{task_name}'{backends}, or the asset" | ||
| " server could not be reached.\n" | ||
| " Train the task, or pass --checkpoint <path> to use a checkpoint of your own." | ||
| ) | ||
| return None | ||
| print(f"[WARNING]: The asset server does not provide the {checkpoint.name} checkpoint '{declared_path}'.") |
There was a problem hiding this comment.
When a declared companion is missing, this code only warns and still returns the policy. Publication likewise skips an uncollected companion without failing the job. Camera play requires the feature extractor and raises FileNotFoundError when it is absent, so an incomplete publication is reported as successful and --checkpoint pretrained later crashes instead of rejecting the unusable bundle.
| """Published location of pre-existing weights.""" | ||
|
|
||
| @property | ||
| def is_run_artifact(self) -> bool: | ||
| """Whether this run produces the file, as opposed to fetching a published one.""" | ||
| return self.run_glob is not None |
There was a problem hiding this comment.
Checkpoint requires exactly one of run_glob and url, but the configuration class does not validate that invariant. Setting both silently ignores the URL, while setting neither eventually passes None to retrieval or extension handling. This makes malformed component declarations fail later with unrelated errors instead of a direct configuration diagnostic.
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
Isaac Lab Review Bot
The component-owned Checkpoint declaration and environment-config discovery provide a focused mechanism for publishing and fetching companion artifacts. However, the new cache layout does not match the existing log_dir derivation for RL-Games and skrl, so those consumers cannot resolve fetched companion checkpoints.
- Design and architecture: Artifact ownership remains localized to component configs, and discovery avoids task-level duplication. The cache-to-consumer contract needs correction: RL-Games and skrl move two directory levels up from the returned policy path, while companions remain in the policy’s immediate cache directory.
- API: The keyword-only
env_cfgaddition preserves the existing positional backend arguments, and the migrated play, demo, and export callers consistently pass the resolved environment config. The newCheckpointtype is exported through the package typing surface. No compatibility break is established by the accepted finding. - Implementation: Companion discovery, naming, collection, publication, and retrieval are connected across the tooling. Fix the RL-Games and skrl playback/export paths so
env_cfg.log_dirpoints to the directory containing the fetched companion, or arrange the cache layout to match those workflows’ existing two-level path convention.
Minor fixes needed. Posted 1 actionable finding inline.
Automated review; human maintainers own approval decisions.
| print(f"Fetching pre-trained checkpoint : {ov_path}") | ||
| # one cache directory per published checkpoint: play treats it as the run log directory and | ||
| # writes videos, exported policies, and additional checkpoints into it | ||
| download_dir = os.path.join( |
There was a problem hiding this comment.
🟡 Warning · Design Architecture — Companion cache dir misses rl_games/skrl log_dir
Companions are fetched into .pretrained_checkpoints/<workflow>/<stem>/, and the new docstring promises a component reading its log directory finds them. That holds for rsl_rl and sb3, but play_rl_games.py (line 149) and play_skrl.py (line 209), plus their LEAPP exporters, derive log_dir = dirname(dirname(resume_path)) = .pretrained_checkpoints/<workflow>, one level above the companion, so Checkpoint.resolve(log_dir) raises. Mirror each workflow's native nesting or use the download directory.
1 participant
Summary
play --checkpoint pretrainedcrashed on the Shadow Hand camera tasks: the vision CNN the policyneeds was never published beside it, and nothing in the tooling knew the file existed. This PR adds
the smallest mechanism that fixes it — a component declares what it writes, and the tooling
publishes and fetches it with the policy.
The refactor that this mechanism invites is deliberately not here; it is stacked in
#7509 — Replace the pretrained checkpoint functions with CheckpointBundle.
Description
1. Declaring a checkpoint
A component declares what it writes on its own config:
get_declared_checkpointswalks the resolved environment config to find every declaration, so atask declares nothing and the component that writes the file owns its name. Each file is published
beside the policy as
<policy stem>_<name><extension>and fetched with it.Checkpoint.resolvehands the component its local file, so no component needs to know the naming convention.
Checkpointalso covers weights that already exist (url=): a frozen encoder is fetched but neverpublished by the checkpoint tooling.
2. Published checkpoints
Both companion pairs are trained and published, so the fix has something to fetch:
Isaac-Reorient-Cube-Shadow-CameraIsaac-Reorient-Cube-Shadow-Camera-DirectEach is a policy
.ptplus its_feature_extractor.pth.3. Verification
test_pretrained_checkpoint.py,test_checkpoints.py,test_train_and_publish_checkpoints.py,test_shadow_hand_camera_presets.pyandtest_pretrained_checkpoint_lookup.py.isaaclab -fclean.urlweights,and that a published file keeps the extension the component declared.
Type of change
Release backport
developChecklist
pre-commitchecks with./isaaclab.sh --formatsource/<pkg>/changelog.d/for every touched packageCONTRIBUTORS.mdor my name already exists there