Replace the pretrained checkpoint functions with CheckpointBundle

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,
so the policy is one checkpoint among them rather than a special case
with auxiliaries hanging off it. 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.

Author: hujc7

scripts/tools/test/test_train_and_publish_checkpoints.py

@@ -42,7 +42,7 @@ def test_build_core_jobs_skips_unsupported_preset_without_normalizing_default(
         lambda _: {PresetTarget.PHYSICS: ["newton_kamino"]},
     )
     monkeypatch.setattr(
-        "scripts.tools.train_and_publish_checkpoints.get_pretrained_checkpoint_backend_names",
+        "scripts.tools.train_and_publish_checkpoints.CheckpointBundle.backend_names",
         lambda _: pytest.fail("preset-only tasks must not normalize their unsupported default backend"),
     )
     args = Namespace(physics_backends="physx,newtonmjwarp", render_backends="rtx,newton")
@@ -104,11 +104,11 @@ def test_select_physics_variants_does_not_fall_back_to_automatic_physx() -> None
     assert selections == []
 
 
-def test_legacy_job_experiment_name_preserves_task_name() -> None:
+def test_legacy_job_stem_preserves_task_name() -> None:
     """Legacy jobs must keep separate experiment directories for each task."""
     job = CheckpointJob(workflow="rsl_rl", task_name="Isaac-Test")
 
-    assert job.experiment_name == "Isaac-Test"
+    assert job.stem == "Isaac-Test"
 
 
 def test_legacy_collection_preserves_task_directory(tmp_path: Path) -> None:

scripts/tools/train_and_publish_checkpoints.py

@@ -59,7 +59,6 @@
 import fnmatch
 import json
 import os
-import posixpath
 import shutil
 import subprocess
 import sys
@@ -78,23 +77,12 @@
 import gymnasium as gym
 
 from isaaclab.envs import DirectMARLEnvCfg
-from isaaclab.utils import Checkpoint
+from isaaclab.utils.assets import NUCLEUS_ASSET_ROOT_DIR
 
 from isaaclab_rl.utils.pretrained_checkpoint import (
     WORKFLOW_EXPERIMENT_NAME_VARIABLE,
     WORKFLOWS,
-    get_auxiliary_checkpoint_path,
-    get_auxiliary_checkpoints,
-    get_latest_job_run_path,
-    get_pretrained_checkpoint_backend_names,
-    get_pretrained_checkpoint_filename,
-    get_pretrained_checkpoint_path,
-    get_pretrained_checkpoint_publish_path,
-    get_pretrained_checkpoint_review,
-    get_pretrained_checkpoint_review_path,
-    has_pretrained_checkpoint_job_finished,
-    has_pretrained_checkpoint_job_run,
-    has_pretrained_checkpoints_asset_root_dir,
+    CheckpointBundle,
 )
 
 import isaaclab_tasks  # noqa: F401
@@ -108,19 +96,13 @@
 
 
 @dataclass(frozen=True)
-class CheckpointJob:
+class CheckpointJob(CheckpointBundle):
     """One workflow, task, physics, and renderer training combination."""
 
-    workflow: str
-    task_name: str
-    physics_backend: str | None = None
-    render_backend: str | None = None
     physics_selector: str | None = None
     render_selector: str | None = None
     agent: str | None = None
     algorithm: str | None = None
-    auxiliary_checkpoints: tuple[Checkpoint, ...] = ()
-    """Run artifacts the task publishes beside its policy."""
 
     @property
     def job_id(self) -> str:
@@ -129,20 +111,6 @@ def job_id(self) -> str:
             return f"{self.workflow}:{self.task_name}"
         return f"{self.workflow}:{self.task_name}:{self.physics_backend}:{self.render_backend}"
 
-    @property
-    def experiment_name(self) -> str:
-        """Return the experiment directory name used by the RL workflow."""
-        if self.physics_backend is None and self.render_backend is None:
-            return self.task_name
-        filename = get_pretrained_checkpoint_filename(
-            self.workflow,
-            self.task_name,
-            self.physics_backend,
-            self.render_backend,
-        )
-        extension = os.path.splitext(filename)[1]
-        return filename.removesuffix(extension)
-
     @property
     def preset_args(self) -> list[str]:
         """Return typed preset selectors for this job."""
@@ -293,7 +261,7 @@ def _resolve_physics_backend(task_name: str, physics_selector: str | None, defau
     if physics_selector is None:
         return default_backend
     env_cfg, _ = resolve_task_config(task_name, None, overrides=(f"physics={physics_selector}",))
-    physics_backend, _ = get_pretrained_checkpoint_backend_names(env_cfg)
+    physics_backend, _ = CheckpointBundle.backend_names(env_cfg)
     return physics_backend
 
 
@@ -343,7 +311,7 @@ def _build_core_jobs(args: argparse.Namespace) -> list[CheckpointJob]:
         workflow, agent, algorithm = _select_workflow(task_spec, env_cfg)
         default_physics = None
         if not physics_variants:
-            default_physics, _ = get_pretrained_checkpoint_backend_names(env_cfg)
+            default_physics, _ = CheckpointBundle.backend_names(env_cfg)
 
         physics_selections = _select_physics_variants(
             task_spec.id,
@@ -356,16 +324,16 @@ def _build_core_jobs(args: argparse.Namespace) -> list[CheckpointJob]:
             physics_backend = _resolve_physics_backend(task_spec.id, physics_selector, default_physics)
             for render_backend, render_selector in render_selections:
                 jobs.append(
-                    CheckpointJob(
-                        workflow=workflow,
-                        task_name=task_spec.id,
-                        physics_backend=physics_backend,
-                        render_backend=render_backend,
+                    CheckpointJob.from_env_cfg(
+                        workflow,
+                        task_spec.id,
+                        env_cfg,
+                        physics_backend,
+                        render_backend,
                         physics_selector=physics_selector,
                         render_selector=render_selector,
                         agent=agent,
                         algorithm=algorithm,
-                        auxiliary_checkpoints=tuple(get_auxiliary_checkpoints(env_cfg)),
                     )
                 )
     return jobs
@@ -411,7 +379,7 @@ def _training_command(job: CheckpointJob, args: argparse.Namespace, smoke: bool)
     if job.algorithm is not None:
         command.extend(["--algorithm", job.algorithm])
 
-    experiment_name = f"{job.experiment_name}_smoke" if smoke else job.experiment_name
+    experiment_name = f"{job.stem}_smoke" if smoke else job.stem
     experiment_variable = WORKFLOW_EXPERIMENT_NAME_VARIABLE[job.workflow]
     if experiment_variable is not None:
         command.append(f"{experiment_variable}={experiment_name}")
@@ -471,30 +439,15 @@ def _run_command(command: list[str], dry_run: bool) -> int:
 
 def _has_training_job_completed(job: CheckpointJob) -> bool:
     """Return whether the latest run exited successfully with a checkpoint."""
-    run_path = get_latest_job_run_path(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    run_path = job.latest_run
     if run_path is None or not os.path.isfile(os.path.join(run_path, _TRAINING_COMPLETE_FILENAME)):
         return False
-    return has_pretrained_checkpoint_job_finished(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    return job.has_finished
 
 
 def _mark_training_job_completed(job: CheckpointJob) -> None:
     """Record that the latest training subprocess exited successfully."""
-    run_path = get_latest_job_run_path(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    run_path = job.latest_run
     if run_path is None:
         raise RuntimeError(f"Unable to determine the latest run for {job.job_id}")
     marker_path = os.path.join(run_path, _TRAINING_COMPLETE_FILENAME)
@@ -514,12 +467,7 @@ def train_job(job: CheckpointJob, args: argparse.Namespace, smoke: bool = False)
         return False
     if smoke or args.dry_run:
         return True
-    if not has_pretrained_checkpoint_job_finished(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    ):
+    if not job.has_finished:
         print(f"Training did not produce a checkpoint for {job.job_id}", file=sys.stderr)
         return False
     _mark_training_job_completed(job)
@@ -528,71 +476,45 @@ def train_job(job: CheckpointJob, args: argparse.Namespace, smoke: bool = False)
 
 def collect_pretrained_checkpoint(job: CheckpointJob, output_dir: str, dry_run: bool = False) -> str | None:
     """Copy the last or best checkpoint into the structured output directory."""
-    destination = _get_collected_checkpoint_path(job, output_dir)
+    destination = job.collected_path(output_dir)
     if dry_run:
         print(f"Would collect the completed checkpoint -> {destination}")
         return destination
 
-    source_path = get_pretrained_checkpoint_path(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    source_path = job.trained_path()
     if source_path is None or not os.path.isfile(source_path):
         print(f"No completed checkpoint to collect for {job.job_id}")
         return None
 
     print(f"Collecting {source_path} -> {destination}")
     os.makedirs(os.path.dirname(destination), exist_ok=True)
     shutil.copy2(source_path, destination)
-    run_path = get_latest_job_run_path(job.workflow, job.task_name, job.physics_backend, job.render_backend)
-    for checkpoint in job.auxiliary_checkpoints:
-        aux_source = checkpoint.find_in(run_path)
+    for checkpoint in job.checkpoints:
+        aux_source = job.trained_path(checkpoint)
         if aux_source is None:
             print(f"No {checkpoint.name} checkpoint matched {checkpoint.run_glob!r} for {job.job_id}")
             continue
-        aux_destination = get_auxiliary_checkpoint_path(destination, job.workflow, checkpoint)
+        aux_destination = job.collected_path(output_dir, checkpoint)
         print(f"Collecting {aux_source} -> {aux_destination}")
         shutil.copy2(aux_source, aux_destination)
     return destination
 
 
 def review_pretrained_checkpoint(job: CheckpointJob, args: argparse.Namespace) -> bool:
     """Play and interactively review one checkpoint."""
-    if not has_pretrained_checkpoint_job_run(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    ):
+    if not job.has_run:
         print(f"Skipping review of {job.job_id}; it has not been trained")
         return False
-    if not has_pretrained_checkpoint_job_finished(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    ):
+    if not job.has_finished:
         print(f"Skipping review of {job.job_id}; training is incomplete")
         return False
 
-    review = get_pretrained_checkpoint_review(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    review = job.review
     if not args.force_review and review and review.get("reviewed"):
         print(f"Review already complete for {job.job_id}")
         return True
 
-    checkpoint_path = get_pretrained_checkpoint_path(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    checkpoint_path = job.trained_path()
     if checkpoint_path is None:
         print(f"Skipping review of {job.job_id}; no checkpoint was found")
         return False
@@ -610,12 +532,7 @@ def review_pretrained_checkpoint(job: CheckpointJob, args: argparse.Namespace) -
     if notes:
         review_data["notes"] = notes
 
-    review_path = get_pretrained_checkpoint_review_path(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    review_path = job.review_path
     if review_path is None:
         raise RuntimeError(f"Unable to determine review path for {job.job_id}")
     with open(review_path, "w", encoding="utf-8") as review_file:
@@ -625,46 +542,27 @@ def review_pretrained_checkpoint(job: CheckpointJob, args: argparse.Namespace) -
 
 def publish_pretrained_checkpoint(job: CheckpointJob, args: argparse.Namespace) -> bool:
     """Publish an accepted checkpoint to the configured Nucleus asset root."""
-    if args.publish_root is None and not has_pretrained_checkpoints_asset_root_dir():
+    if args.publish_root is None and not NUCLEUS_ASSET_ROOT_DIR:
         raise RuntimeError("A pretrained-checkpoint Nucleus asset root is not configured")
-    local_path = _get_collected_checkpoint_path(job, args.output_dir)
+    local_path = job.collected_path(args.output_dir)
     if not os.path.isfile(local_path):
         print(f"Skipping publish of {job.job_id}; no collected checkpoint was found")
         return False
 
     if not args.force_publish:
-        review = get_pretrained_checkpoint_review(
-            job.workflow,
-            job.task_name,
-            job.physics_backend,
-            job.render_backend,
-        )
+        review = job.review
         if not review or review.get("result") != "accepted":
             print(f"Skipping publish of {job.job_id}; it does not have an accepted review")
             return False
 
-    if args.publish_root is None:
-        publish_path = get_pretrained_checkpoint_publish_path(
-            job.workflow,
-            job.task_name,
-            job.physics_backend,
-            job.render_backend,
-        )
-    else:
-        filename = get_pretrained_checkpoint_filename(
-            job.workflow,
-            job.task_name,
-            job.physics_backend,
-            job.render_backend,
-        )
-        publish_path = posixpath.join(args.publish_root.rstrip("/"), job.workflow, filename)
+    publish_path = job.published_path(root=args.publish_root)
     uploads = [(local_path, publish_path)]
-    for checkpoint in job.auxiliary_checkpoints:
-        local_auxiliary = get_auxiliary_checkpoint_path(local_path, job.workflow, checkpoint)
-        if not os.path.isfile(local_auxiliary):
+    for checkpoint in job.checkpoints:
+        local_declared = job.collected_path(args.output_dir, checkpoint)
+        if not os.path.isfile(local_declared):
             print(f"Skipping the {checkpoint.name} checkpoint for {job.job_id}; it was not collected")
             continue
-        uploads.append((local_auxiliary, get_auxiliary_checkpoint_path(publish_path, job.workflow, checkpoint)))
+        uploads.append((local_declared, job.published_path(checkpoint, root=args.publish_root)))
     for source, destination in uploads:
         print(f"Publishing {source} -> {destination}")
     if args.dry_run:
@@ -683,24 +581,10 @@ def publish_pretrained_checkpoint(job: CheckpointJob, args: argparse.Namespace)
 
 def _summary_row(job: CheckpointJob, output_dir: str) -> list[str | bool]:
     """Return one CSV summary row."""
-    has_run = has_pretrained_checkpoint_job_run(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
-    has_finished = (
-        _has_training_job_completed(job)
-        if job.physics_backend is not None
-        else has_pretrained_checkpoint_job_finished(job.workflow, job.task_name)
-    )
-    collected_path = _get_collected_checkpoint_path(job, output_dir)
-    review = get_pretrained_checkpoint_review(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
+    has_run = job.has_run
+    has_finished = _has_training_job_completed(job) if job.physics_backend is not None else job.has_finished
+    collected_path = job.collected_path(output_dir)
+    review = job.review
     return [
         job.workflow,
         job.task_name,
@@ -715,20 +599,6 @@ def _summary_row(job: CheckpointJob, output_dir: str) -> list[str | bool]:
     ]
 
 
-def _get_collected_checkpoint_path(job: CheckpointJob, output_dir: str) -> str:
-    """Return the absolute path of a checkpoint in the collection directory."""
-    filename = get_pretrained_checkpoint_filename(
-        job.workflow,
-        job.task_name,
-        job.physics_backend,
-        job.render_backend,
-    )
-    path_parts = [output_dir, job.workflow]
-    if job.physics_backend is None:
-        path_parts.append(job.task_name)
-    return os.path.abspath(os.path.join(*path_parts, filename))
-
-
 def main(argv: list[str] | None = None) -> int:
     """Run checkpoint management actions."""
     parser = _create_parser()