Deep CFR resume 동작 보강
This commit is contained in:
@@ -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.deep_cfr.trainer import DeepCFRTrainer
|
||||||
from coolrl_lost_cities.games.classic.game import classic_config
|
from coolrl_lost_cities.games.classic.game import classic_config
|
||||||
|
|
||||||
|
_RESUME_LATEST = "__latest__"
|
||||||
|
|
||||||
|
|
||||||
def _load_config(path: str | None) -> DeepCFRConfig:
|
def _load_config(path: str | None) -> DeepCFRConfig:
|
||||||
if path is None:
|
if path is None:
|
||||||
@@ -45,6 +47,17 @@ def _with_overrides(config: DeepCFRConfig, overrides: dict[str, Any]) -> DeepCFR
|
|||||||
return DeepCFRConfig.model_validate(data)
|
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]:
|
def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]:
|
||||||
overrides: dict[str, Any] = {}
|
overrides: dict[str, Any] = {}
|
||||||
run_overrides = overrides.setdefault("run", {})
|
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
|
overrides.setdefault("evaluation", {})["games"] = args.eval_games
|
||||||
if args.no_save:
|
if args.no_save:
|
||||||
overrides.setdefault("checkpoint", {})["save_every_iteration"] = False
|
overrides.setdefault("checkpoint", {})["save_every_iteration"] = False
|
||||||
|
if args.exact_resume:
|
||||||
|
overrides.setdefault("checkpoint", {})["exact_resume"] = True
|
||||||
return overrides
|
return overrides
|
||||||
|
|
||||||
|
|
||||||
@@ -80,13 +95,14 @@ def train_command(args: argparse.Namespace) -> None:
|
|||||||
config = _load_config(args.config)
|
config = _load_config(args.config)
|
||||||
overrides = _train_overrides_from_args(args)
|
overrides = _train_overrides_from_args(args)
|
||||||
config = _with_overrides(config, overrides)
|
config = _with_overrides(config, overrides)
|
||||||
|
resume_path = _resolve_resume_path(config, args.resume)
|
||||||
trainer = DeepCFRTrainer(
|
trainer = DeepCFRTrainer(
|
||||||
config,
|
config,
|
||||||
config.rules.to_lost_cities_config(seed=config.run.seed),
|
config.rules.to_lost_cities_config(seed=config.run.seed),
|
||||||
device=args.device or config.run.device,
|
device=args.device or config.run.device,
|
||||||
)
|
)
|
||||||
if args.resume:
|
if resume_path:
|
||||||
trainer.load_checkpoint(args.resume)
|
trainer.load_checkpoint(resume_path)
|
||||||
trainer.train()
|
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("--traversals-per-iteration", type=int)
|
||||||
train.add_argument("--num-workers")
|
train.add_argument("--num-workers")
|
||||||
train.add_argument("--checkpoint-dir")
|
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("--device")
|
||||||
train.add_argument("--eval-every", type=int)
|
train.add_argument("--eval-every", type=int)
|
||||||
train.add_argument("--eval-games", type=int)
|
train.add_argument("--eval-games", type=int)
|
||||||
|
|||||||
@@ -218,6 +218,7 @@ class CheckpointConfig(StrictModel):
|
|||||||
save_iteration_interval: int = 0
|
save_iteration_interval: int = 0
|
||||||
save_latest_only: bool = False
|
save_latest_only: bool = False
|
||||||
progress_interval_seconds: float = 20.0
|
progress_interval_seconds: float = 20.0
|
||||||
|
exact_resume: bool = False
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def path(self) -> Path:
|
def path(self) -> Path:
|
||||||
|
|||||||
@@ -171,6 +171,7 @@ class DeepCFRTrainer:
|
|||||||
return {
|
return {
|
||||||
"config": self.config.to_dict(),
|
"config": self.config.to_dict(),
|
||||||
"game_config": self.game_config.to_snapshot(),
|
"game_config": self.game_config.to_snapshot(),
|
||||||
|
"resume_semantics": "networks_optimizers_iteration_only",
|
||||||
"iteration": self.iteration,
|
"iteration": self.iteration,
|
||||||
"input_dim": self.input_dim,
|
"input_dim": self.input_dim,
|
||||||
"action_size": self.action_size,
|
"action_size": self.action_size,
|
||||||
@@ -188,8 +189,17 @@ class DeepCFRTrainer:
|
|||||||
return save_checkpoint(path, self.checkpoint_payload(metrics))
|
return save_checkpoint(path, self.checkpoint_payload(metrics))
|
||||||
|
|
||||||
def load_checkpoint(self, path: str | Path) -> None:
|
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)
|
payload = load_checkpoint(path, device=self.device)
|
||||||
self.iteration = int(payload.get("iteration", 0))
|
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(
|
for network, state_dict in zip(
|
||||||
self.advantage_networks, payload["advantage_networks"], strict=True
|
self.advantage_networks, payload["advantage_networks"], strict=True
|
||||||
):
|
):
|
||||||
@@ -417,8 +427,7 @@ class DeepCFRTrainer:
|
|||||||
metrics.append(item)
|
metrics.append(item)
|
||||||
self._append_metrics(item, elapsed)
|
self._append_metrics(item, elapsed)
|
||||||
self._maybe_record_self_play_snapshot(iteration)
|
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):
|
if self._time_limit_reached(run_started):
|
||||||
break
|
break
|
||||||
iteration += 1
|
iteration += 1
|
||||||
@@ -445,7 +454,7 @@ class DeepCFRTrainer:
|
|||||||
|
|
||||||
def _save_iteration_checkpoints(self, iteration: int, item: IterationMetrics) -> None:
|
def _save_iteration_checkpoints(self, iteration: int, item: IterationMetrics) -> None:
|
||||||
checkpoint_dir = self.run_dir
|
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 / f"iteration_{iteration:05d}.pt", item)
|
||||||
self.save_checkpoint(checkpoint_dir / "latest.pt", item)
|
self.save_checkpoint(checkpoint_dir / "latest.pt", item)
|
||||||
|
|
||||||
|
|||||||
@@ -10,7 +10,10 @@ from coolrl_lost_cities.games.classic.deep_cfr.benchmark import (
|
|||||||
benchmark_traversal,
|
benchmark_traversal,
|
||||||
benchmark_traversal_modes,
|
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 (
|
from coolrl_lost_cities.games.classic.deep_cfr.cli import (
|
||||||
|
_RESUME_LATEST,
|
||||||
|
_resolve_resume_path,
|
||||||
_train_overrides_from_args,
|
_train_overrides_from_args,
|
||||||
_with_overrides,
|
_with_overrides,
|
||||||
)
|
)
|
||||||
@@ -77,6 +80,7 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None:
|
|||||||
"eval_every": None,
|
"eval_every": None,
|
||||||
"eval_games": None,
|
"eval_games": None,
|
||||||
"no_save": True,
|
"no_save": True,
|
||||||
|
"exact_resume": False,
|
||||||
},
|
},
|
||||||
)()
|
)()
|
||||||
config = load_config("configs/deep_cfr/deep_cfr_selfplay_full_depth_slot_playability.yaml")
|
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
|
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:
|
def test_deep_cfr_playability_encoding_extends_input_shape() -> None:
|
||||||
state = GameState.new_game(LostCitiesConfig(seed=61), seed=61)
|
state = GameState.new_game(LostCitiesConfig(seed=61), seed=61)
|
||||||
base_dim = input_dim(state)
|
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)
|
restored.load_checkpoint(latest)
|
||||||
|
|
||||||
assert latest.exists()
|
assert latest.exists()
|
||||||
|
assert load_checkpoint(latest)["resume_semantics"] == "networks_optimizers_iteration_only"
|
||||||
assert (checkpoint_dir / "config.json").exists()
|
assert (checkpoint_dir / "config.json").exists()
|
||||||
assert (checkpoint_dir / "metrics.jsonl").exists()
|
assert (checkpoint_dir / "metrics.jsonl").exists()
|
||||||
assert (checkpoint_dir / "runtime_progress.json").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")
|
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 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 "Iteration complete:" in train_log
|
||||||
|
assert "reservoir memories and RNG state are not restored" in train_log
|
||||||
assert restored.iteration == 1
|
assert restored.iteration == 1
|
||||||
assert "eval_random_games" in metrics[0].eval_metrics
|
assert "eval_random_games" in metrics[0].eval_metrics
|
||||||
assert "eval_random_play_action_rate" 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:
|
def test_deep_cfr_trainer_multiprocessing_smoke_run(tmp_path) -> None:
|
||||||
trainer = DeepCFRTrainer(
|
trainer = DeepCFRTrainer(
|
||||||
_deep_cfr_config(
|
_deep_cfr_config(
|
||||||
|
|||||||
Reference in New Issue
Block a user