Skip to content

Replace the pretrained checkpoint functions with CheckpointBundle - #7509

Draft
hujc7 wants to merge 6 commits into
isaac-sim:developfrom
hujc7:jichuanh/checkpoint-bundle
Draft

Replace the pretrained checkpoint functions with CheckpointBundle#7509
hujc7 wants to merge 6 commits into
isaac-sim:developfrom
hujc7:jichuanh/checkpoint-bundle

Conversation

@hujc7

@hujc7 hujc7 commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Summary

Stacked on #7485 — [Fix] Publish and fetch the checkpoints a task's components declare, which is the
minimal fix. This PR is the refactor that PR deliberately kept out of scope.

  • One CheckpointBundle replaces 13 free functions that each took the same
    (workflow, task_name, physics_backend, render_backend) tuple; the publish script unpacked it at 20 call sites.
  • One Workflow per RL library replaces four parallel tables and the 13 resolve_checkpoint_selector blocks
    that each spelled that library's file patterns. Two copies had drifted: the benchmark sb3 pattern excluded
    model.zip, so its own preference never matched; skrl preferred best_agent.pt only when publishing.
  • Checkpoint.local_path replaces the cache-dir-as-log-dir hand-off, which put rl_games and skrl companions
    one directory off.

Description

1. Change

  • isaaclab.utils: Checkpoint.local_path, set by the fetch and returned first by resolve(); find_cfgs
    and latest_file become public utilities.
  • isaaclab_rl.utils.pretrained_checkpoint: Workflow (policy glob, preferred file, extension, experiment
    key; selector_args() derives the selector patterns) in WORKFLOWS; CheckpointBundle owns identity,
    companions and the published, cached and collected paths, and fetch() records each companion's
    local_path. get_published_pretrained_checkpoint(…, env_cfg=) is unchanged, so its ten call sites are
    untouched.
  • Entrypoints (4 play, 4 train, 4 benchmark, leapp sb3 export): **WORKFLOWS[lib].selector_args(…).
  • train_and_publish_checkpoints.py: CheckpointJob owns training-run state and the train/play commands;
    from_task resolves the preset-selected config so the backend names and declared checkpoints describe
    what the presets train; the trained policy is selected through the same manifest-based lookup as
    --checkpoint best.

2. Migration

Full retired→new table in source/isaaclab_rl/changelog.d/checkpoint-bundle.major.rst. Behaviour
changes: skrl --checkpoint best prefers best_agent.pt; the publish script only collects runs carrying
a run.json manifest; --publish_root keeps the legacy per-task sub-directory.

3. Verification

200 pass across test_checkpoints.py, test_pretrained_checkpoint.py, test_train_and_publish_checkpoints.py,
test_shadow_hand_camera_presets.py, test_pretrained_checkpoint_lookup.py, test_standalone_scripts.py,
test_entrypoints.py, test_checkpoint_manifest.py and benchmark test_api.py. New tests pin each library's
selector patterns against the files it writes, the local_path write-back, manifest-required selection and
the preset-config plumbing of from_task. pre-commit clean.

Type of change

  • Breaking change (existing functionality will not work without user modification)

Release backport

  • Backport this pull request to the active release branch after it merges into develop

Checklist

  • I have read and understood the contribution guidelines
  • I have run the pre-commit checks with ./isaaclab.sh --format
  • Documentation needs no change: the module has no API page; the changelog fragment carries the migration table
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • I have added a changelog fragment under source/<pkg>/changelog.d/ for every touched package
  • I have added my name to the CONTRIBUTORS.md or my name already exists there

@github-actions github-actions Bot added the isaac-lab Related to Isaac Lab team label Sep 3, 2026
@hujc7

hujc7 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

run-ci

@isaaclab-bot isaaclab-bot Bot added ci:run-docker Trigger the on-demand Docker and GPU CI workflow and removed ci:run-docker Trigger the on-demand Docker and GPU CI workflow labels Sep 4, 2026
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.
"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.
The module answered every path question through 13 free functions that
each took the same (workflow, task_name, physics_backend, render_backend)
tuple, and the publish script unpacked it at 20 call sites.

CheckpointBundle owns that identity and the checkpoints a task declares.
find_cfgs and latest_file move to isaaclab.utils as public utilities.
WORKFLOW_TRAINER and WORKFLOW_PLAYER are deleted: every workflow mapped
to the same entrypoint and nothing read them.
get_published_pretrained_checkpoint is unchanged, so its ten call sites
are untouched.

CheckpointJob gains is_trained and mark_trained: the completion marker is
job state, and the summary no longer re-derives the legacy rule from
physics_backend.
@hujc7
hujc7 force-pushed the jichuanh/checkpoint-bundle branch from e800247 to fc8754e Compare September 4, 2026 07:59
@hujc7

hujc7 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

run-ci

@isaaclab-bot isaaclab-bot Bot added ci:run-docker Trigger the on-demand Docker and GPU CI workflow and removed ci:run-docker Trigger the on-demand Docker and GPU CI workflow labels Sep 4, 2026
The play, train and benchmark entrypoints spelled each library's policy
file pattern, its subdirectory and its preferred file in thirteen
resolve_checkpoint_selector calls, while the publish module kept the same
facts as glob and extension tables. The copies had drifted: the benchmark
sb3 pattern excluded model.zip, so its own preference never matched, and
skrl preferred best_agent.pt only when publishing.

Workflow holds those conventions per library; WORKFLOWS maps the name to
it. selector_args derives the selector patterns from the policy glob, so
every entrypoint asks for them instead of restating them. The publish
script selects the trained policy through the same manifest-based lookup
as --checkpoint best.

Checkpoint.local_path lets the fetch record where a companion landed, so a
component loads the published copy through its own declaration instead of
globbing a naming convention out of a directory the entrypoint pointed it
at. That directory hand-off broke for rl_games and skrl, whose run layout
puts the policy one level deeper than the flat cache.

CheckpointBundle keeps only what a consumer needs; training-run state and
the train and play commands move onto the publish script's CheckpointJob,
which now resolves the preset-selected config to read both the backend
names and the declared checkpoints from what the presets actually train.
…int-bundle

# Conflicts:
#	source/isaaclab_rl/isaaclab_rl/utils/pretrained_checkpoint.py
#	source/isaaclab_rl/test/test_pretrained_checkpoint.py
@hujc7

hujc7 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

run-ci

@isaaclab-bot isaaclab-bot Bot added ci:run-docker Trigger the on-demand Docker and GPU CI workflow and removed ci:run-docker Trigger the on-demand Docker and GPU CI workflow labels Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

isaac-lab Related to Isaac Lab team

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant