From b5550d3840241626ec6cb74a361940bb3cfcb3f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Fri, 8 May 2026 17:14:21 +0900 Subject: [PATCH] Add current-vs-average eval diagnostic --- scripts/eval_current_vs_average.py | 305 +++++++++++++++++++++++++++++ 1 file changed, 305 insertions(+) create mode 100644 scripts/eval_current_vs_average.py diff --git a/scripts/eval_current_vs_average.py b/scripts/eval_current_vs_average.py new file mode 100644 index 0000000..32fcc07 --- /dev/null +++ b/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()