From fc4f0ddfd8aa6afff71af8594ffc1342c26f6bf8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Thu, 7 May 2026 01:55:10 +0900 Subject: [PATCH] =?UTF-8?q?Deep=20CFR=20resume=20=EB=8F=99=EC=9E=91=20?= =?UTF-8?q?=EB=B3=B4=EA=B0=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../games/classic/deep_cfr/cli.py | 23 ++++- .../games/classic/deep_cfr/config.py | 1 + .../games/classic/deep_cfr/trainer.py | 15 +++- tests/games/classic/test_deep_cfr_trainer.py | 83 +++++++++++++++++++ 4 files changed, 116 insertions(+), 6 deletions(-) diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py index 2c546b5..a585644 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py +++ b/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) diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py index 34aea0f..e74f73e 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/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: diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 9e793f8..6402f2d 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/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) diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index c04ffe5..2439592 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -10,7 +10,10 @@ from coolrl_lost_cities.games.classic.deep_cfr.benchmark import ( benchmark_traversal, benchmark_traversal_modes, ) +from coolrl_lost_cities.games.classic.deep_cfr.checkpoints import load_checkpoint from coolrl_lost_cities.games.classic.deep_cfr.cli import ( + _RESUME_LATEST, + _resolve_resume_path, _train_overrides_from_args, _with_overrides, ) @@ -77,6 +80,7 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None: "eval_every": None, "eval_games": None, "no_save": True, + "exact_resume": False, }, )() config = load_config("configs/deep_cfr/deep_cfr_selfplay_full_depth_slot_playability.yaml") @@ -92,6 +96,27 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None: assert overridden.checkpoint.save_every_iteration is False +def test_deep_cfr_resume_latest_resolution_uses_config_checkpoint_dir(tmp_path) -> None: + config = _deep_cfr_config({"checkpoint": {"directory": str(tmp_path)}}) + latest = tmp_path / "latest.pt" + latest.write_bytes(b"checkpoint") + + assert _resolve_resume_path(config, _RESUME_LATEST) == str(latest) + assert _resolve_resume_path(config, "custom.pt") == "custom.pt" + assert _resolve_resume_path(config, None) is None + + +def test_deep_cfr_resume_latest_resolution_requires_latest(tmp_path) -> None: + config = _deep_cfr_config({"checkpoint": {"directory": str(tmp_path)}}) + + try: + _resolve_resume_path(config, _RESUME_LATEST) + except FileNotFoundError as exc: + assert "latest checkpoint does not exist" in str(exc) + else: # pragma: no cover + raise AssertionError("expected FileNotFoundError") + + def test_deep_cfr_playability_encoding_extends_input_shape() -> None: state = GameState.new_game(LostCitiesConfig(seed=61), seed=61) base_dim = input_dim(state) @@ -332,6 +357,7 @@ def test_deep_cfr_trainer_saves_loads_and_evaluates_checkpoint(tmp_path) -> None restored.load_checkpoint(latest) assert latest.exists() + assert load_checkpoint(latest)["resume_semantics"] == "networks_optimizers_iteration_only" assert (checkpoint_dir / "config.json").exists() assert (checkpoint_dir / "metrics.jsonl").exists() assert (checkpoint_dir / "runtime_progress.json").exists() @@ -339,6 +365,7 @@ def test_deep_cfr_trainer_saves_loads_and_evaluates_checkpoint(tmp_path) -> None train_log = (checkpoint_dir / "train.log").read_text(encoding="utf-8") assert re.search(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}", train_log) assert "Iteration complete:" in train_log + assert "reservoir memories and RNG state are not restored" in train_log assert restored.iteration == 1 assert "eval_random_games" in metrics[0].eval_metrics assert "eval_random_play_action_rate" in metrics[0].eval_metrics @@ -352,6 +379,62 @@ def test_deep_cfr_trainer_saves_loads_and_evaluates_checkpoint(tmp_path) -> None ) +def test_deep_cfr_trainer_always_saves_latest_checkpoint(tmp_path) -> None: + checkpoint_dir = tmp_path / "latest-each-iteration" + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"iterations": 1, "seed": 42}, + "network": {"hidden_size": 16}, + "traversal": {"traversals_per_iteration": 1, "max_depth": 1}, + "checkpoint": { + "directory": str(checkpoint_dir), + "save_every_iteration": False, + "save_iteration_interval": 10, + }, + } + ), + LostCitiesConfig(seed=42), + ) + + trainer.train() + + assert (checkpoint_dir / "latest.pt").exists() + assert not (checkpoint_dir / "iteration_00001.pt").exists() + + +def test_deep_cfr_exact_resume_is_explicitly_not_implemented(tmp_path) -> None: + checkpoint_dir = tmp_path / "exact" + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"iterations": 1, "seed": 44}, + "network": {"hidden_size": 16}, + "traversal": {"traversals_per_iteration": 1, "max_depth": 1}, + "checkpoint": {"directory": str(checkpoint_dir), "save_every_iteration": True}, + } + ), + LostCitiesConfig(seed=44), + ) + trainer.train() + exact = DeepCFRTrainer( + _deep_cfr_config( + { + "network": {"hidden_size": 16}, + "checkpoint": {"directory": str(checkpoint_dir), "exact_resume": True}, + } + ), + LostCitiesConfig(seed=44), + ) + + try: + exact.load_checkpoint(checkpoint_dir / "latest.pt") + except NotImplementedError as exc: + assert "exact_resume" in str(exc) + else: # pragma: no cover + raise AssertionError("expected NotImplementedError") + + def test_deep_cfr_trainer_multiprocessing_smoke_run(tmp_path) -> None: trainer = DeepCFRTrainer( _deep_cfr_config(