From 25a3fba53fb0152aae40f94013eff92d46dcf2cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Sun, 10 May 2026 23:04:05 +0900 Subject: [PATCH] Use Deep CFR diagnostics for IS-MCTS eval Wrap AlphaZeroNet with a logits-only view so IS-MCTS training evaluation can call evaluate_strategy_network and emit the same full diagnostic metric set as Deep CFR. Adds root prior capture and per-iteration MCTS entropy, value error, and policy-vs-search KL metrics. Tests: uv run python -m pytest tests/games/classic/ismcts/ -x; uv run python -m pytest tests/games/classic/test_deep_cfr_trainer.py -x; uv run lost-cities-ismcts train --config configs/ismcts/mini.yaml --set run.experiment_name=ismcts-metrics-smoke --set run.max_iterations=2 --set training.games_per_iter=2 --- .../games/classic/ismcts/network.py | 12 ++ .../games/classic/ismcts/replay_buffer.py | 1 + .../games/classic/ismcts/self_play.py | 22 ++- .../games/classic/ismcts/trainer.py | 140 +++++++++--------- tests/games/classic/ismcts/test_ismcts.py | 81 +++++++++- 5 files changed, 184 insertions(+), 72 deletions(-) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/network.py b/src/coolrl_lost_cities/games/classic/ismcts/network.py index 00ff0cc..007fe88 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/network.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/network.py @@ -70,3 +70,15 @@ class AlphaZeroNet(nn.Module): probs = torch.softmax(logits, dim=-1).masked_fill(~legal_mask.bool(), 0.0) normalizer = probs.sum(dim=-1, keepdim=True).clamp_min(1.0e-12) return probs / normalizer + + +class AlphaZeroLogitsView(nn.Module): + """Expose AlphaZeroNet's policy logits as a one-argument module.""" + + def __init__(self, net: AlphaZeroNet) -> None: + super().__init__() + self.net = net + + def forward(self, x: torch.Tensor) -> torch.Tensor: + logits, _value = self.net(x, legal_mask=None) + return logits diff --git a/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py b/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py index dbec70b..a55aa33 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py @@ -14,6 +14,7 @@ class ReplaySample: pi_target: np.ndarray v_target: float player: int + prior: np.ndarray | None = None class ReplayBuffer: diff --git a/src/coolrl_lost_cities/games/classic/ismcts/self_play.py b/src/coolrl_lost_cities/games/classic/ismcts/self_play.py index c092ded..b390bc8 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/self_play.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/self_play.py @@ -60,10 +60,23 @@ def play_self_play_game( max_steps: int = 10_000, ) -> list[ReplaySample]: state = GameState.new_game(game_config, seed=rng.randrange(2**31)) - pending: list[tuple[np.ndarray, np.ndarray, np.ndarray, int]] = [] + pending: list[tuple[np.ndarray, np.ndarray, np.ndarray, int, np.ndarray]] = [] steps = 0 while not state.terminal and steps < max_steps: player = int(state.current_player) + legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool) + info = encode_info_state(state, player, encoding) + with torch.inference_mode(): + x = torch.as_tensor(info[None, :], dtype=torch.float32, device=device) + mask = torch.as_tensor(legal_mask[None, :], dtype=torch.bool, device=device) + prior = ( + network.policy_distribution(x, mask) + .squeeze(0) + .detach() + .cpu() + .numpy() + .astype(np.float32) + ) searcher = IsMctsSearcher( network, mcts_config, @@ -72,20 +85,18 @@ def play_self_play_game( rng=random.Random(rng.randrange(2**31)), ) visits = searcher.search(state, player) - legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool) pi = visit_distribution(visits, state.action_size, temperature=temperature) if pi.sum() <= 0: legal_actions = np.flatnonzero(legal_mask) pi[legal_actions] = 1.0 / len(legal_actions) - info = encode_info_state(state, player, encoding) - pending.append((info.astype(np.float32), legal_mask, pi, player)) + pending.append((info.astype(np.float32), legal_mask, pi, player, prior)) action = select_from_distribution(pi, rng) state.apply_unified_action(action) steps += 1 final_diff0 = float(state.score_diff(0)) samples: list[ReplaySample] = [] - for info, legal_mask, pi, player in pending: + for info, legal_mask, pi, player, prior in pending: value = final_diff0 if player == 0 else -final_diff0 samples.append( ReplaySample( @@ -94,6 +105,7 @@ def play_self_play_game( pi_target=pi.astype(np.float32), v_target=value, player=player, + prior=prior, ) ) return samples diff --git a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py index f57de03..bbad61d 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py @@ -10,13 +10,12 @@ import numpy as np import torch from torch import nn -from coolrl_lost_cities.games.classic.bots import build_bot from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim +from coolrl_lost_cities.games.classic.deep_cfr.evaluate import evaluate_strategy_network from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig -from .mcts import IsMctsSearcher -from .network import AlphaZeroNet +from .network import AlphaZeroLogitsView, AlphaZeroNet from .replay_buffer import ReplayBuffer, ReplaySample from .self_play import play_self_play_game @@ -32,6 +31,7 @@ class IterationMetrics: self_play_seconds: float train_seconds: float eval_metrics: dict[str, float | int] + mcts_metrics: dict[str, float] def to_dict(self) -> dict[str, float | int]: data: dict[str, float | int] = { @@ -44,6 +44,7 @@ class IterationMetrics: "time/self_play_seconds": self.self_play_seconds, "time/train_seconds": self.train_seconds, } + data.update(self.mcts_metrics) data.update(self.eval_metrics) return data @@ -107,6 +108,7 @@ class IsMctsTrainer: self.network.eval() sp_started = time.perf_counter() added = 0 + iteration_samples: list[ReplaySample] = [] for _ in range(self.config.training.games_per_iter): samples = play_self_play_game( self.network, @@ -119,8 +121,10 @@ class IsMctsTrainer: max_steps=self.config.evaluation.max_steps, ) self.buffer.add(samples) + iteration_samples.extend(samples) added += len(samples) self_play_seconds = time.perf_counter() - sp_started + mcts_metrics = self._compute_mcts_metrics(iteration_samples) train_started = time.perf_counter() losses = [] @@ -140,6 +144,7 @@ class IsMctsTrainer: self_play_seconds=self_play_seconds, train_seconds=train_seconds, eval_metrics=eval_metrics, + mcts_metrics=mcts_metrics, ) def _train_batch(self, batch: list[ReplaySample]) -> tuple[float, float, float]: @@ -181,27 +186,72 @@ class IsMctsTrainer: def _evaluate(self, iteration: int) -> dict[str, float | int]: opponents = self.config.evaluation.opponents_for_iteration(iteration) + if ( + not opponents + and self.config.evaluation.eval_every > 0 + and self.config.run.max_iterations is not None + and iteration >= self.config.run.max_iterations + ): + opponents = self.config.evaluation.opponents if not opponents: return {} self.network.eval() results: dict[str, float | int] = {} + logits_view = AlphaZeroLogitsView(self.network) for opponent in opponents: - result = evaluate_policy( - self.network, + result = evaluate_strategy_network( + logits_view, self.game_config, - opponent=opponent, games=self.config.evaluation.games, seed=self.config.run.seed + iteration * 1000, + opponent=opponent, device=self.device, encoding=self.config.encoding, max_steps=self.config.evaluation.max_steps, - mcts_config=self.config.mcts, + batch_size=self.config.evaluation.batch_size, ) key = opponent.replace("-", "_") for metric_key, value in result.items(): results[f"eval/{key}/{metric_key}"] = value return results + def _compute_mcts_metrics(self, samples: list[ReplaySample]) -> dict[str, float]: + if not samples: + return { + "mcts/avg_visit_entropy": 0.0, + "mcts/value_prediction_error": 0.0, + "mcts/policy_mcts_kl": 0.0, + } + entropies = [_entropy(sample.pi_target) for sample in samples] + policy_kls = [ + _kl_divergence(sample.pi_target, sample.prior) + for sample in samples + if sample.prior is not None + ] + info = torch.as_tensor( + np.stack([sample.info_state for sample in samples]), + dtype=torch.float32, + device=self.device, + ) + legal = torch.as_tensor( + np.stack([sample.legal_mask for sample in samples]), + dtype=torch.bool, + device=self.device, + ) + target = torch.as_tensor( + [sample.v_target for sample in samples], + dtype=torch.float32, + device=self.device, + ) + with torch.inference_mode(): + _logits, value_pred = self.network(info, legal) + value_error = nn.functional.mse_loss(value_pred, target) + return { + "mcts/avg_visit_entropy": float(np.mean(entropies)) if entropies else 0.0, + "mcts/value_prediction_error": float(value_error.item()), + "mcts/policy_mcts_kl": float(np.mean(policy_kls)) if policy_kls else 0.0, + } + def _append_metrics(self, metrics: IterationMetrics) -> None: with self.metrics_path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(metrics.to_dict(), sort_keys=True) + "\n") @@ -229,62 +279,20 @@ class IsMctsTrainer: return (time.perf_counter() - started) / 60.0 >= self.config.run.max_minutes -def evaluate_policy( - network: AlphaZeroNet, - config: LostCitiesConfig, - *, - opponent: str, - games: int, - seed: int, - device: torch.device | str, - encoding=None, - max_steps: int = 10_000, - mcts_config=None, -) -> dict[str, float | int]: - rng = random.Random(seed) - score_diffs: list[int] = [] - wins = losses = draws = 0 - policy_actions = play_actions = 0 - for game_index in range(games): - policy_player = game_index % 2 - policies = [ - build_bot(opponent, seed=seed + game_index), - build_bot(opponent, seed=seed + game_index), - ] - state = GameState.new_game(config, seed=seed + game_index) - for _ in range(max_steps): - if state.terminal: - break - current = int(state.current_player) - if current == policy_player: - searcher = IsMctsSearcher( - network, - mcts_config or IsMctsConfig().mcts, - device=device, - encoding=encoding, - rng=random.Random(rng.randrange(2**31)), - ) - visits = searcher.search(state, current) - unified = max(visits, key=visits.get) - action = state.from_unified_action(unified) - else: - action = policies[current].act(state) - if current == policy_player and state.phase == "card": - policy_actions += 1 - if action % 2 == 0: - play_actions += 1 - state.apply_action(action) - diff = state.score_diff(policy_player) - score_diffs.append(diff) - wins += int(diff > 0) - losses += int(diff < 0) - draws += int(diff == 0) - return { - "games": games, - "wins0": wins, - "wins1": losses, - "draws": draws, - "win_rate0": wins / max(1, games), - "avg_score_diff0": float(np.mean(score_diffs)) if score_diffs else 0.0, - "play_action_rate": play_actions / max(1, policy_actions), - } +def _entropy(distribution: np.ndarray) -> float: + probs = np.asarray(distribution, dtype=np.float64) + probs = probs[probs > 0.0] + if len(probs) == 0: + return 0.0 + return float(-(probs * np.log(probs)).sum()) + + +def _kl_divergence(target: np.ndarray, prior: np.ndarray | None) -> float: + if prior is None: + return 0.0 + pi = np.asarray(target, dtype=np.float64) + p = np.asarray(prior, dtype=np.float64) + mask = pi > 0.0 + if not np.any(mask): + return 0.0 + return float((pi[mask] * (np.log(pi[mask]) - np.log(np.clip(p[mask], 1.0e-12, 1.0)))).sum()) diff --git a/tests/games/classic/ismcts/test_ismcts.py b/tests/games/classic/ismcts/test_ismcts.py index 65193be..644ccea 100644 --- a/tests/games/classic/ismcts/test_ismcts.py +++ b/tests/games/classic/ismcts/test_ismcts.py @@ -11,7 +11,7 @@ from coolrl_lost_cities.games.classic.ismcts.config import IsMctsConfig, MctsCon from coolrl_lost_cities.games.classic.ismcts.determinization import sample_determinization from coolrl_lost_cities.games.classic.ismcts.info_set import canonical_info_set_key from coolrl_lost_cities.games.classic.ismcts.mcts import IsMctsSearcher -from coolrl_lost_cities.games.classic.ismcts.network import AlphaZeroNet +from coolrl_lost_cities.games.classic.ismcts.network import AlphaZeroLogitsView, AlphaZeroNet from coolrl_lost_cities.games.classic.ismcts.replay_buffer import ReplayBuffer, ReplaySample from coolrl_lost_cities.games.classic.ismcts.self_play import play_self_play_game from coolrl_lost_cities.games.classic.ismcts.trainer import IsMctsTrainer @@ -69,6 +69,16 @@ def test_network_shapes_and_mask() -> None: assert torch.all(probs[~mask] == 0) +def test_logits_view_adapter() -> None: + state = GameState.new_game(mini_config(), seed=6) + dim = input_dim(state) + net = AlphaZeroNet(dim, state.action_size, hidden_size=16, num_layers=1) + logits_view = AlphaZeroLogitsView(net) + x = torch.as_tensor(encode_info_state(state, 0)[None, :], dtype=torch.float32) + logits = logits_view(x) + assert logits.shape == (1, state.action_size) + + def test_mcts_prior_drives_visits() -> None: state = GameState.new_game(mini_config(), seed=7) dim = input_dim(state) @@ -114,6 +124,7 @@ def test_self_play_game_returns_signed_targets() -> None: assert samples assert {sample.player for sample in samples} <= {0, 1} assert all(sample.pi_target.sum() > 0 for sample in samples) + assert all(sample.prior is not None for sample in samples) def test_trainer_one_iteration_smoke(tmp_path) -> None: @@ -142,3 +153,71 @@ def test_trainer_one_iteration_smoke(tmp_path) -> None: metrics = trainer.train() assert len(metrics) == 1 assert (tmp_path / "metrics.jsonl").exists() + + +def test_trainer_emits_full_eval_metrics(tmp_path) -> None: + config = IsMctsConfig.model_validate( + { + "run": {"max_iterations": 1, "seed": 12, "device": "cpu"}, + "rules": { + "n_colors": 3, + "n_ranks": 5, + "n_handshakes": 1, + "hand_size": 4, + "bonus_threshold": 4, + }, + "network": {"hidden_size": 16, "num_layers": 1}, + "mcts": {"n_simulations": 2}, + "training": {"games_per_iter": 1, "gradient_steps_per_iter": 1, "batch_size": 8}, + "checkpoint": {"save_every": 0}, + "evaluation": { + "eval_every": 1, + "games": 2, + "opponents": ["random"], + "num_workers": 1, + "max_steps": 80, + }, + } + ) + trainer = IsMctsTrainer( + config, + config.rules.to_lost_cities_config(seed=config.run.seed), + run_dir=tmp_path, + ) + metrics = trainer.train()[0].to_dict() + assert "eval/random/avg_opened_colors" in metrics + assert "eval/random/bad_open_rate" in metrics + assert "eval/random/per_game_negative_expeditions" in metrics + + +def test_trainer_emits_mcts_metrics(tmp_path) -> None: + config = IsMctsConfig.model_validate( + { + "run": {"max_iterations": 1, "seed": 13, "device": "cpu"}, + "rules": { + "n_colors": 3, + "n_ranks": 5, + "n_handshakes": 1, + "hand_size": 4, + "bonus_threshold": 4, + }, + "network": {"hidden_size": 16, "num_layers": 1}, + "mcts": {"n_simulations": 2}, + "training": {"games_per_iter": 1, "gradient_steps_per_iter": 1, "batch_size": 8}, + "checkpoint": {"save_every": 0}, + "evaluation": {"eval_every": 0, "num_workers": 1, "max_steps": 80}, + } + ) + trainer = IsMctsTrainer( + config, + config.rules.to_lost_cities_config(seed=config.run.seed), + run_dir=tmp_path, + ) + metrics = trainer.train()[0].to_dict() + for key in ( + "mcts/avg_visit_entropy", + "mcts/value_prediction_error", + "mcts/policy_mcts_kl", + ): + assert key in metrics + assert np.isfinite(metrics[key])