Deep CFR resume 동작 보강

src/coolrl_lost_cities/games/classic/deep_cfr/cli.py

@@ -24,6 +24,8 @@ from coolrl_lost_cities.games.classic.deep_cfr.policy_gradient import (
 from coolrl_lost_cities.games.classic.deep_cfr.trainer import DeepCFRTrainer
 from coolrl_lost_cities.games.classic.game import classic_config
 
+_RESUME_LATEST = "__latest__"
+
 
 def _load_config(path: str | None) -> DeepCFRConfig:
     if path is None:
@@ -45,6 +47,17 @@ def _with_overrides(config: DeepCFRConfig, overrides: dict[str, Any]) -> DeepCFR
     return DeepCFRConfig.model_validate(data)
 
 
+def _resolve_resume_path(config: DeepCFRConfig, resume: str | None) -> str | None:
+    if resume != _RESUME_LATEST:
+        return resume
+    latest_path = config.checkpoint_path / "latest.pt"
+    if not latest_path.exists():
+        raise FileNotFoundError(
+            f"--resume was used without a path, but latest checkpoint does not exist: {latest_path}"
+        )
+    return str(latest_path)
+
+
 def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]:
     overrides: dict[str, Any] = {}
     run_overrides = overrides.setdefault("run", {})
@@ -73,6 +86,8 @@ def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]:
         overrides.setdefault("evaluation", {})["games"] = args.eval_games
     if args.no_save:
         overrides.setdefault("checkpoint", {})["save_every_iteration"] = False
+    if args.exact_resume:
+        overrides.setdefault("checkpoint", {})["exact_resume"] = True
     return overrides
 
 
@@ -80,13 +95,14 @@ def train_command(args: argparse.Namespace) -> None:
     config = _load_config(args.config)
     overrides = _train_overrides_from_args(args)
     config = _with_overrides(config, overrides)
+    resume_path = _resolve_resume_path(config, args.resume)
     trainer = DeepCFRTrainer(
         config,
         config.rules.to_lost_cities_config(seed=config.run.seed),
         device=args.device or config.run.device,
     )
-    if args.resume:
-        trainer.load_checkpoint(args.resume)
+    if resume_path:
+        trainer.load_checkpoint(resume_path)
     trainer.train()
 
 
@@ -183,7 +199,8 @@ def main(argv: list[str] | None = None) -> None:
     train.add_argument("--traversals-per-iteration", type=int)
     train.add_argument("--num-workers")
     train.add_argument("--checkpoint-dir")
-    train.add_argument("--resume")
+    train.add_argument("--resume", nargs="?", const=_RESUME_LATEST, default=None)
+    train.add_argument("--exact-resume", action="store_true")
     train.add_argument("--device")
     train.add_argument("--eval-every", type=int)
     train.add_argument("--eval-games", type=int)

src/coolrl_lost_cities/games/classic/deep_cfr/config.py

@@ -218,6 +218,7 @@ class CheckpointConfig(StrictModel):
     save_iteration_interval: int = 0
     save_latest_only: bool = False
     progress_interval_seconds: float = 20.0
+    exact_resume: bool = False
 
     @property
     def path(self) -> Path:

src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py

@@ -171,6 +171,7 @@ class DeepCFRTrainer:
         return {
             "config": self.config.to_dict(),
             "game_config": self.game_config.to_snapshot(),
+            "resume_semantics": "networks_optimizers_iteration_only",
             "iteration": self.iteration,
             "input_dim": self.input_dim,
             "action_size": self.action_size,
@@ -188,8 +189,17 @@ class DeepCFRTrainer:
         return save_checkpoint(path, self.checkpoint_payload(metrics))
 
     def load_checkpoint(self, path: str | Path) -> None:
+        if self.config.checkpoint.exact_resume:
+            # TODO: Implement exact resume by checkpointing reservoir memories, RNG state,
+            # and any worker/traversal sampling state needed for deterministic continuation.
+            raise NotImplementedError("checkpoint.exact_resume is not implemented yet")
         payload = load_checkpoint(path, device=self.device)
         self.iteration = int(payload.get("iteration", 0))
+        self.tracker.log_event(
+            f"Resuming from {path} with resume_semantics="
+            f"{payload.get('resume_semantics', 'networks_optimizers_iteration_only')}; "
+            "reservoir memories and RNG state are not restored"
+        )
         for network, state_dict in zip(
             self.advantage_networks, payload["advantage_networks"], strict=True
         ):
@@ -417,8 +427,7 @@ class DeepCFRTrainer:
             metrics.append(item)
             self._append_metrics(item, elapsed)
             self._maybe_record_self_play_snapshot(iteration)
-            if self._should_save_iteration(iteration):
-                self._save_iteration_checkpoints(iteration, item)
+            self._save_iteration_checkpoints(iteration, item)
             if self._time_limit_reached(run_started):
                 break
             iteration += 1
@@ -445,7 +454,7 @@ class DeepCFRTrainer:
 
     def _save_iteration_checkpoints(self, iteration: int, item: IterationMetrics) -> None:
         checkpoint_dir = self.run_dir
-        if not self.config.checkpoint.save_latest_only:
+        if self._should_save_iteration(iteration) and not self.config.checkpoint.save_latest_only:
             self.save_checkpoint(checkpoint_dir / f"iteration_{iteration:05d}.pt", item)
         self.save_checkpoint(checkpoint_dir / "latest.pt", item)