Add current-vs-average eval diagnostic

scripts/eval_current_vs_average.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python
+"""Compare Deep CFR current regret-matching policy vs average strategy policy."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state
+from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
+
+from coolrl_lost_cities.games.classic.bots import build_bot
+from coolrl_lost_cities.games.classic.deep_cfr.config import config_from_dict
+from coolrl_lost_cities.games.classic.deep_cfr.evaluate import (
+    EvalRuntimeCounters,
+    PolicyEvalDiagnostics,
+    StrategyNetPolicy,
+    _accumulate_game_diagnostics,
+    _advance_opponent_turn,
+    _advance_policy_turn,
+    _EvalGame,
+    _finalize_if_done,
+)
+from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
+from coolrl_lost_cities.games.classic.policy import LostCitiesPolicy, PolicyInput
+
+
+class CurrentRegretPolicy(LostCitiesPolicy):
+    def __init__(
+        self,
+        advantage_networks: list[torch.nn.Module],
+        *,
+        device: torch.device | str,
+        encoding,
+        epsilon: float,
+        all_negative_fallback: str,
+    ) -> None:
+        self.advantage_networks = advantage_networks
+        self.device = torch.device(device)
+        self.encoding = encoding
+        self.epsilon = float(epsilon)
+        self.all_negative_fallback = all_negative_fallback
+        self.runtime = EvalRuntimeCounters()
+
+    def select_actions_batch(self, states: list[GameState]) -> list[tuple[int, float]]:
+        return [self.select_action(state) for state in states]
+
+    def select_action(self, state: GameState) -> tuple[int, float]:
+        started = time.perf_counter()
+        self.runtime.policy_turns += 1
+        legal_actions, distribution = self.action_distribution(state)
+        probs = distribution[legal_actions]
+        entropy = _entropy(probs)
+        unified = int(legal_actions[int(np.argmax(probs))])
+        self.runtime.policy_select_seconds += time.perf_counter() - started
+        return state.from_unified_action(unified), entropy
+
+    def action_distribution(self, state: GameState) -> tuple[np.ndarray, np.ndarray]:
+        started = time.perf_counter()
+        legal = np.asarray(state.unified_legal_mask(), dtype=bool)
+        legal_actions = np.flatnonzero(legal)
+        self.runtime.policy_legal_mask_seconds += time.perf_counter() - started
+        if len(legal_actions) == 0:
+            raise RuntimeError("no legal action available")
+
+        started = time.perf_counter()
+        info = encode_info_state(state, state.current_player, self.encoding)
+        self.runtime.policy_encoding_seconds += time.perf_counter() - started
+
+        started = time.perf_counter()
+        network = self.advantage_networks[int(state.current_player)]
+        with torch.inference_mode():
+            x = torch.as_tensor(info, dtype=torch.float32, device=self.device).unsqueeze(0)
+            advantages = network(x).squeeze(0).detach().cpu().numpy().astype(np.float32)
+        self.runtime.policy_network_seconds += time.perf_counter() - started
+
+        started = time.perf_counter()
+        distribution = self._regret_matching_distribution(advantages, legal, legal_actions)
+        self.runtime.policy_postprocess_seconds += time.perf_counter() - started
+        return legal_actions, distribution
+
+    def _regret_matching_distribution(
+        self,
+        advantages: np.ndarray,
+        legal: np.ndarray,
+        legal_actions: np.ndarray,
+    ) -> np.ndarray:
+        positive = np.where(legal, np.maximum(advantages, 0.0), 0.0).astype(np.float32)
+        total = float(positive.sum())
+        if total > self.epsilon:
+            return positive / total
+        distribution = np.zeros_like(advantages, dtype=np.float32)
+        if self.all_negative_fallback == "uniform":
+            distribution[legal_actions] = 1.0 / float(len(legal_actions))
+            return distribution
+        best = float(np.max(advantages[legal_actions]))
+        best_actions = legal_actions[advantages[legal_actions] == best]
+        distribution[int(best_actions[0])] = 1.0
+        return distribution
+
+    def act(self, obs_or_state: PolicyInput) -> int:
+        if not isinstance(obs_or_state, GameState):
+            legal = np.asarray(obs_or_state["legal_mask"], dtype=bool)
+            legal_actions = np.flatnonzero(legal)
+            if len(legal_actions) == 0:
+                raise RuntimeError("no legal action available")
+            return int(legal_actions[0])
+        action, _entropy_value = self.select_action(obs_or_state)
+        return action
+
+
+def evaluate_policy(
+    policy: CurrentRegretPolicy | StrategyNetPolicy,
+    config: LostCitiesConfig,
+    *,
+    games: int,
+    seed: int,
+    opponent: str,
+    max_steps: int,
+    batch_size: int,
+) -> dict[str, float | int]:
+    diagnostics = PolicyEvalDiagnostics()
+    started = time.perf_counter()
+    active_games: list[_EvalGame] = []
+    for index in range(games):
+        game_seed = seed + index
+        swap = index % 2 == 1
+        policy_player = 1 if swap else 0
+        opponent_policy = build_bot(opponent, seed=game_seed * 2 + (1 - policy_player))
+        policies = [opponent_policy, policy] if swap else [policy, opponent_policy]
+        active_games.append(
+            _EvalGame(
+                state=GameState.new_game(config, seed=game_seed),
+                policies=policies,
+                policy_player=policy_player,
+                diagnostics=PolicyEvalDiagnostics(games=1),
+                first_open_recoverable_by_color={},
+                game_index=index,
+                game_seed=game_seed,
+            )
+        )
+
+    while active_games:
+        pending_policy_games: list[_EvalGame] = []
+        next_active_games: list[_EvalGame] = []
+        for game in active_games:
+            if _finalize_if_done(game, max_steps=max_steps):
+                _accumulate_game_diagnostics(diagnostics, game.diagnostics)
+                continue
+            if game.state.current_player == game.policy_player:
+                pending_policy_games.append(game)
+            else:
+                _advance_opponent_turn(game)
+                if _finalize_if_done(game, max_steps=max_steps):
+                    _accumulate_game_diagnostics(diagnostics, game.diagnostics)
+                else:
+                    next_active_games.append(game)
+
+        chunk_size = max(1, int(batch_size))
+        for start in range(0, len(pending_policy_games), chunk_size):
+            chunk = pending_policy_games[start : start + chunk_size]
+            actions = policy.select_actions_batch([game.state for game in chunk])
+            for game, (action, entropy) in zip(chunk, actions, strict=True):
+                _advance_policy_turn(game, action, entropy)
+                if _finalize_if_done(game, max_steps=max_steps):
+                    _accumulate_game_diagnostics(diagnostics, game.diagnostics)
+                else:
+                    next_active_games.append(game)
+
+        active_games = next_active_games
+
+    diagnostics.runtime.accumulate(policy.runtime)
+    return diagnostics.to_dict(time.perf_counter() - started)
+
+
+def load_policies(
+    checkpoint: Path,
+    *,
+    device: torch.device | str,
+) -> tuple[StrategyNetPolicy, CurrentRegretPolicy, LostCitiesConfig, dict]:
+    payload = torch.load(checkpoint, map_location="cpu")
+    cfg = config_from_dict(payload["config"])
+    game_config = LostCitiesConfig(**payload["game_config"])
+    input_dim = int(payload["input_dim"])
+    action_size = int(payload["action_size"])
+
+    strategy_network = DeepCFRMLP.from_config(input_dim, action_size, cfg.network).to(device)
+    strategy_network.load_state_dict(payload["strategy_network"])
+    strategy_network.eval()
+
+    advantage_networks = [
+        DeepCFRMLP.from_config(input_dim, action_size, cfg.network).to(device) for _ in range(2)
+    ]
+    for network, state_dict in zip(advantage_networks, payload["advantage_networks"], strict=True):
+        network.load_state_dict(state_dict)
+        network.eval()
+
+    average_policy = StrategyNetPolicy(
+        strategy_network,
+        device=device,
+        seed=cfg.run.seed * 2,
+        encoding=cfg.encoding,
+    )
+    current_policy = CurrentRegretPolicy(
+        advantage_networks,
+        device=device,
+        encoding=cfg.encoding,
+        epsilon=cfg.traversal.regret_matching_epsilon,
+        all_negative_fallback=cfg.regret_matching.all_negative_fallback,
+    )
+    return average_policy, current_policy, game_config, payload
+
+
+def _entropy(probs: np.ndarray) -> float:
+    probs = probs[probs > 0.0]
+    if len(probs) == 0:
+        return 0.0
+    return float(-(probs * np.log(probs)).sum())
+
+
+def _checkpoint_iteration(path: Path, payload: dict) -> int:
+    if isinstance(payload.get("iteration"), int):
+        return int(payload["iteration"])
+    stem = path.stem
+    if stem.startswith("iteration_"):
+        return int(stem.removeprefix("iteration_"))
+    return -1
+
+
+def _select_metrics(result: dict[str, float | int]) -> dict[str, float | int]:
+    keys = (
+        "win_rate0",
+        "avg_score_diff0",
+        "avg_opened_colors",
+        "score_per_opened_color",
+        "bad_open_rate",
+        "good_open_rate",
+        "negative_expedition_rate",
+        "positive_expedition_rate",
+        "policy_entropy",
+        "play_action_rate",
+        "discard_action_rate",
+        "draw_deck_rate",
+        "draw_pile_rate",
+        "avg_game_length",
+        "max_step_timeouts",
+        "elapsed_seconds",
+    )
+    return {key: result[key] for key in keys if key in result}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("checkpoints", nargs="+", type=Path)
+    parser.add_argument("--opponent", action="append", default=None)
+    parser.add_argument("--games", type=int, default=100)
+    parser.add_argument("--seed", type=int, default=79_000)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--max-steps", type=int, default=10_000)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+
+    opponents = args.opponent or ["random", "safe_heuristic_strict"]
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+
+    rows: list[dict] = []
+    for checkpoint in args.checkpoints:
+        average_policy, current_policy, game_config, payload = load_policies(
+            checkpoint,
+            device=args.device,
+        )
+        iteration = _checkpoint_iteration(checkpoint, payload)
+        for policy_name, policy in (("average", average_policy), ("current", current_policy)):
+            for opponent in opponents:
+                result = evaluate_policy(
+                    policy,
+                    game_config,
+                    games=args.games,
+                    seed=args.seed + max(iteration, 0) * 1000,
+                    opponent=opponent,
+                    max_steps=args.max_steps,
+                    batch_size=args.batch_size,
+                )
+                row = {
+                    "checkpoint": str(checkpoint),
+                    "iteration": iteration,
+                    "policy": policy_name,
+                    "opponent": opponent,
+                    "games": args.games,
+                    **_select_metrics(result),
+                }
+                rows.append(row)
+                print(json.dumps(row, sort_keys=True))
+
+    args.output.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n")
+    print(f"wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()