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

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

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:

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, prior))
-        pending.append((info.astype(np.float32), legal_mask, pi, player))
         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

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 AlphaZeroLogitsView, AlphaZeroNet
-from .network import 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(
+            result = evaluate_strategy_network(
-                self.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(
+def _entropy(distribution: np.ndarray) -> float:
-    network: AlphaZeroNet,
+    probs = np.asarray(distribution, dtype=np.float64)
-    config: LostCitiesConfig,
+    probs = probs[probs > 0.0]
-    *,
+    if len(probs) == 0:
-    opponent: str,
+        return 0.0
-    games: int,
+    return float(-(probs * np.log(probs)).sum())
-    seed: int,
+
-    device: torch.device | str,
+
-    encoding=None,
+def _kl_divergence(target: np.ndarray, prior: np.ndarray | None) -> float:
-    max_steps: int = 10_000,
+    if prior is None:
-    mcts_config=None,
+        return 0.0
-) -> dict[str, float | int]:
+    pi = np.asarray(target, dtype=np.float64)
-    rng = random.Random(seed)
+    p = np.asarray(prior, dtype=np.float64)
-    score_diffs: list[int] = []
+    mask = pi > 0.0
-    wins = losses = draws = 0
+    if not np.any(mask):
-    policy_actions = play_actions = 0
+        return 0.0
-    for game_index in range(games):
+    return float((pi[mask] * (np.log(pi[mask]) - np.log(np.clip(p[mask], 1.0e-12, 1.0)))).sum())
-        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),
-    }

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])