Fix curriculum auto-resume stage handoff

experiments/exp_base.py

@@ -9,6 +9,7 @@ from abc import ABC, abstractmethod
 from typing import Optional, Union, Literal, List, Dict
 import pathlib
 import os
+import re
 from datetime import timedelta
 
 import hydra
@@ -72,6 +73,13 @@ _DEMEMWM_FRAME_MEMORY_MARKERS = (
 )
 
 
+def _checkpoint_step_from_path(ckpt_path) -> Optional[int]:
+    if ckpt_path is None:
+        return None
+    step_match = re.search(r"step[=_-]?(\d+)", pathlib.Path(str(ckpt_path)).stem)
+    return int(step_match.group(1)) if step_match else None
+
+
 def _short_list(items, limit=8):
     items = list(items)
     suffix = "" if len(items) <= limit else f", ... (+{len(items) - limit} more)"
@@ -483,9 +491,18 @@ class BaseLightningExperiment(BaseExperiment):
         return None if self.customized_load else self.ckpt_path
 
     def _fit_training_curriculum(self, stages, ckpt_path) -> None:
+        resume_step = _checkpoint_step_from_path(ckpt_path) if self.auto_resuming else None
         for stage_idx, stage in enumerate(stages):
             until_step = self._curriculum_until_step(stage, stage_idx)
             self._apply_curriculum_stage(stage, stage_idx, until_step)
+            if resume_step is not None and resume_step >= until_step:
+                if is_rank_zero:
+                    print(
+                        cyan("Skipping curriculum stage:"),
+                        f"{stage.get('name', stage_idx + 1)} until step {until_step}; "
+                        f"resume checkpoint is already at step {resume_step}",
+                    )
+                continue
             trainer = self._build_training_trainer(max_steps=until_step)
             trainer.fit(
                 self.algo,

tests/test_resume_checkpoint_logic.py

@@ -184,6 +184,46 @@ class ResumeCheckpointLogicTests(unittest.TestCase):
         self.assertEqual(DummyTrainer.save_checkpoint_calls, [Path("/tmp/curriculum_stage.ckpt")])
         self.assertEqual(float(experiment.root_cfg.algorithm.memory_selection.pose_similarity_radius), 8.0)
 
+    def test_auto_resume_curriculum_skips_completed_stages(self):
+        DummyTrainer.init_calls = []
+        DummyTrainer.fit_calls = []
+        DummyTrainer.save_checkpoint_calls = []
+        cfg = _root_cfg(auto_resuming=True)
+        cfg.experiment.training.curriculum = {
+            "enabled": True,
+            "stages": [
+                {"name": "near", "until_step": 3, "dataset": {"wo_updown": True}},
+                {"name": "full", "until_step": 5, "dataset": {"wo_updown": False}},
+            ],
+        }
+        experiment = DummyExperiment(cfg, logger=None, ckpt_path="/tmp/epoch1_step4.ckpt")
+        experiment.algo = object()
+        seen_dataset = []
+
+        def build_training_loader():
+            seen_dataset.append(bool(experiment.root_cfg.dataset.wo_updown))
+            return None
+
+        experiment._build_training_loader = build_training_loader
+
+        with patch.object(exp_base.pl, "Trainer", DummyTrainer), patch.object(
+            DummyExperiment,
+            "_curriculum_checkpoint_path",
+            return_value=Path("/tmp/curriculum_stage.ckpt"),
+        ):
+            experiment.training()
+
+        self.assertEqual(seen_dataset, [False])
+        self.assertEqual(
+            [call["kwargs"]["ckpt_path"] for call in DummyTrainer.fit_calls],
+            ["/tmp/epoch1_step4.ckpt"],
+        )
+        self.assertEqual(
+            [call["kwargs"]["max_steps"] for call in DummyTrainer.init_calls],
+            [5],
+        )
+        self.assertEqual(DummyTrainer.save_checkpoint_calls, [])
+
     def test_validation_load_checkpoint_takes_priority_over_separate_load(self):
         DummyTrainer.validate_calls = []
         experiment = DummyExperiment(_root_cfg(auto_resuming=False), logger=None, ckpt_path="/tmp/trained.ckpt")