#!/usr/bin/env python3 """ eval_elo.py — Deep CFR 模型对抗评估脚本 让两个不同轮次的 checkpoint 模型 (Model A / Model B) 进行 100,000 局对战， 计算 Model A 对 Model B 的百手赢率 (bb/100)。 === 核心设计 === 1. 多进程并发: ProcessPoolExecutor(mp_context='spawn'), 22 Workers 2. 绝对公平: 50,000 局 Model A 做 P0(小盲), 50,000 局 Model A 做 P1(大盲) 3. 严格推理: 使用 avg_strategy + legal_mask 归一化 + numpy.random.choice 采样 4. bb/100 指标: (总筹码收益 / BB) / (总局数 / 100) 用法: python eval_elo.py --model_a ckpt_10000.pt --model_b ckpt_5000.pt python eval_elo.py --model_a checkpoints/ckpt_iter_10000.pt --model_b checkpoints/ckpt_iter_5000.pt --num_games 10000 """ # ── 必须在 import torch 之前，锁死 C++ 线性代数库多线程，防止 spawn 模式下 OpenMP 死锁 ── import os os.environ["OMP_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" os.environ["OPENBLAS_NUM_THREADS"] = "1" import argparse import random import multiprocessing as mp from typing import Dict, List, Tuple import numpy as np import torch import pyspiel from concurrent.futures import ProcessPoolExecutor, as_completed torch.set_num_threads(1) # ── 将 poker/ 根目录加入 sys.path ── import sys _POKER_DIR = os.path.abspath(os.path.dirname(__file__)) if _POKER_DIR not in sys.path: sys.path.insert(0, _POKER_DIR) from env_adapter import ( HUNL_SB_BB_STRING, BetTranslator, extract_env_state, CFR_ACTIONS, NUM_CFR_ACTIONS, STACK_NORMALIZE, ) from card_model.config import PAD_TOKEN, BOARD_SIZE from card_model.data_generator import extract_cards_from_state from card_model.model import CardModel from cfr_net import CFRNetwork, CARD_DIM, ENV_DIM, NUM_ACTIONS # ───────────────────── 常量 ───────────────────── STREET_NORMALIZE = 3.0 BIG_BLIND = 100 # BB = 100 筹码 NUM_WORKERS = 22 # 并行 Worker 数 CARD_MODEL_CHECKPOINT = os.path.join(_POKER_DIR, "card_model", "data", "best_card_model.pt") # ───────────────────── 全局 Worker 状态 ───────────────────── _WORKER_STATE: Dict = {} # ───────────────────── 特征构建 ───────────────────── def build_env_features(env_info: dict) -> torch.Tensor: """将 extract_env_state 返回的字典归一化为 5 维 env_features Tensor。""" features = [ env_info["pot"] / STACK_NORMALIZE, env_info["p0_stack"] / STACK_NORMALIZE, env_info["p1_stack"] / STACK_NORMALIZE, env_info["street"] / STREET_NORMALIZE, float(env_info["position"]), ] return torch.tensor(features, dtype=torch.float32) def build_card_features(card_model: CardModel, state) -> torch.Tensor: """ 使用 CardModel 从当前 state 提取 50 维胜率直方图 (card_features)。 必须使用 extract_cards_from_state(state, player_id=current_player) 提取当前行动玩家的底牌，绝不能看错牌！ """ current_player = state.current_player() hole_cards, board_cards = extract_cards_from_state(state, player_id=current_player) model_device = next(card_model.parameters()).device x_hole = torch.tensor([hole_cards], dtype=torch.int64, device=model_device) # [1, 2] padded_board = board_cards + [PAD_TOKEN] * (BOARD_SIZE - len(board_cards)) x_board = torch.tensor([padded_board], dtype=torch.int64, device=model_device) # [1, 5] with torch.no_grad(): _, pred_histogram = card_model(x_hole, x_board) # [1, 50] return pred_histogram.squeeze(0).cpu() # ───────────────────── Worker 初始化 ───────────────────── def _init_worker( model_a_state_dict: dict, model_b_state_dict: dict, card_state_dict: dict, ) -> None: """ ProcessPoolExecutor 的 Worker 初始化函数。 在每个 Worker 进程启动时调用一次，创建该进程专属的： - OpenSpiel 游戏实例 (标准 P0=SB, P1=BB) - CPU 上的两个 CFRNetwork（Model A / Model B，eval 模式） - CPU 上的 CardModel（eval 模式） - BetTranslator 所有模型权重通过 state_dict 参数传入（CPU Tensor）， 避免 CUDA 跨进程问题。 """ global _WORKER_STATE # 每个 Worker 自行创建 OpenSpiel 游戏实例（不可跨进程传递） _WORKER_STATE["game"] = pyspiel.load_game(HUNL_SB_BB_STRING) # P0=SB(50), P1=BB(100) # 在 CPU 上创建 Model A 的 CFRNetwork 并加载权重 cfr_net_a = CFRNetwork( card_dim=CARD_DIM, env_dim=ENV_DIM, num_actions=NUM_ACTIONS, ) cfr_net_a.load_state_dict(model_a_state_dict) cfr_net_a.eval() _WORKER_STATE["cfr_net_a"] = cfr_net_a # 在 CPU 上创建 Model B 的 CFRNetwork 并加载权重 cfr_net_b = CFRNetwork( card_dim=CARD_DIM, env_dim=ENV_DIM, num_actions=NUM_ACTIONS, ) cfr_net_b.load_state_dict(model_b_state_dict) cfr_net_b.eval() _WORKER_STATE["cfr_net_b"] = cfr_net_b # 在 CPU 上创建 CardModel 并加载权重 # CardModel 是共享的——两个模型用相同的牌面编码器 card_model = CardModel() card_model.load_state_dict(card_state_dict) card_model.eval() _WORKER_STATE["card_model"] = card_model # BetTranslator 是无状态对象，可以直接创建 _WORKER_STATE["translator"] = BetTranslator() # ───────────────────── 单步推理 ───────────────────── def _choose_action( cfr_net: CFRNetwork, card_model: CardModel, translator: BetTranslator, state, ) -> int: """ 使用 avg_strategy 进行动作采样，返回引擎原生 action ID。 严格规则： 1. 使用 extract_cards_from_state(state, player_id=current_player) 提取底牌 2. 调用 get_strategy() 后丢弃第一个返回值，使用第二个返回值 avg_strategy 3. 绝对禁止 argmax！用 avg_strategy * legal_mask 归一化后 numpy.random.choice 采样 4. 用 BetTranslator.cfr_to_engine_action() 转换为引擎动作 """ current_player = state.current_player() # 提取环境特征 env_info = extract_env_state(state, translator) env_features = build_env_features(env_info) # [5] card_features = build_card_features(card_model, state) # [50] legal_mask = env_info["legal_mask"] # list[int], 长度5 # 构造网络输入（加 batch 维度） card_input = card_features.unsqueeze(0) # [1, 50] env_input = env_features.unsqueeze(0) # [1, 5] legal_mask_tensor = torch.tensor([legal_mask], dtype=torch.float32) # [1, 5] # ── 核心：调用 get_strategy，丢弃 current_strategy，只用 avg_strategy ── with torch.no_grad(): _, avg_strategy = cfr_net.get_strategy(card_input, env_input, legal_mask_tensor) # 取出 [5] 向量，转为 numpy avg_strat = avg_strategy.squeeze(0).cpu().numpy() # [5], numpy array # ── avg_strategy 与 legal_mask 相乘，重新归一化 ── legal_mask_np = np.array(legal_mask, dtype=np.float64) masked_probs = avg_strat.astype(np.float64) * legal_mask_np prob_sum = masked_probs.sum() if prob_sum > 1e-9: masked_probs = masked_probs / prob_sum else: # 极端兜底：在合法动作上均匀分布 num_legal = int(legal_mask_np.sum()) if num_legal > 0: masked_probs = legal_mask_np / num_legal else: masked_probs[1] = 1.0 # 最后兜底：CALL # ── 使用 numpy.random.choice 按概率分布采样 ── cfr_action_idx = int(np.random.choice(NUM_CFR_ACTIONS, p=masked_probs)) # ── 通过 BetTranslator 将 CFR 离散动作映射为引擎原生 action ID ── engine_action = translator.cfr_to_engine_action(state, cfr_action_idx) return engine_action # ───────────────────── Worker 对战函数 ───────────────────── def worker_play_games( num_games: int, model_a_is_p0: bool, ) -> float: """ Worker 进程中的对战函数。 执行 num_games 局对战，返回 Model A 的总筹码收益。 始终使用标准游戏 (P0=SB, P1=BB)，通过 nets 列表的顺序实现位置轮换： - model_a_is_p0=True: Model A = P0(SB), Model B = P1(BB) - model_a_is_p0=False: Model A = P1(BB), Model B = P0(SB) 在 OpenSpiel 中，state.returns()[0] 始终是 P0 的收益， state.returns()[1] 始终是 P1 的收益。 因此需要根据 model_a_is_p0 将 P0/P1 收益映射到 Model A 的收益。 """ global _WORKER_STATE cfr_net_a = _WORKER_STATE["cfr_net_a"] cfr_net_b = _WORKER_STATE["cfr_net_b"] card_model = _WORKER_STATE["card_model"] translator = _WORKER_STATE["translator"] game = _WORKER_STATE["game"] if model_a_is_p0: # Model A 坐在 P0(SB) 位置 nets = [cfr_net_a, cfr_net_b] else: # Model A 坐在 P1(BB) 位置 nets = [cfr_net_b, cfr_net_a] total_chips_won_by_a = 0.0 for _ in range(num_games): state = game.new_initial_state() while not state.is_terminal(): # ── Chance Node: 随机发牌 ── if state.is_chance_node(): outcomes = state.chance_outcomes() action_list, prob_list = zip(*outcomes) chance_action = random.choices(action_list, weights=prob_list, k=1)[0] state.apply_action(chance_action) continue # ── Player Node: 由对应模型推理选动作 ── current_player = state.current_player() current_net = nets[current_player] engine_action = _choose_action(current_net, card_model, translator, state) state.apply_action(engine_action) # ── 终局：收集 Model A 的筹码收益 ── returns = state.returns() if model_a_is_p0: # Model A 是 P0 total_chips_won_by_a += returns[0] else: # Model A 是 P1 total_chips_won_by_a += returns[1] return total_chips_won_by_a # ───────────────────── 加载模型权重 ───────────────────── def load_cfr_state_dict(checkpoint_path: str) -> dict: """ 加载 CFRNetwork 的 state_dict。 支持两种 checkpoint 格式： - 训练存档格式: {"model_state_dict": ..., "optimizer_state_dict": ..., "iteration": ...} - 纯模型格式: 直接的 state_dict（如 export_model.py 导出的） """ print(f"[加载] 正在读取: {checkpoint_path}") ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) if isinstance(ckpt, dict) and "model_state_dict" in ckpt: state_dict = ckpt["model_state_dict"] iter_info = ckpt.get("iteration", "未知") print(f"[加载] 训练存档格式, iteration={iter_info}") else: state_dict = ckpt print(f"[加载] 纯模型权重格式") # 确保所有 Tensor 在 CPU 上 return {k: v.cpu() for k, v in state_dict.items()} def load_card_state_dict(checkpoint_path: str) -> dict: """加载 CardModel 的 state_dict。""" print(f"[加载] CardModel: {checkpoint_path}") ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) if isinstance(ckpt, dict) and "model_state_dict" in ckpt: state_dict = ckpt["model_state_dict"] else: state_dict = ckpt return {k: v.cpu() for k, v in state_dict.items()} # ───────────────────── 主函数 ───────────────────── def main(): parser = argparse.ArgumentParser( description="Deep CFR 模型对抗评估 — 计算 Model A 对 Model B 的 bb/100" ) parser.add_argument( "--model_a", type=str, required=True, help="Model A 的 checkpoint 路径" ) parser.add_argument( "--model_b", type=str, required=True, help="Model B 的 checkpoint 路径" ) parser.add_argument( "--num_games", type=int, default=100_000, help="总对战局数 (默认 100,000)" ) parser.add_argument( "--num_workers", type=int, default=NUM_WORKERS, help=f"并行 Worker 数 (默认 {NUM_WORKERS})" ) parser.add_argument( "--card_model", type=str, default=CARD_MODEL_CHECKPOINT, help=f"CardModel 权重路径 (默认 {CARD_MODEL_CHECKPOINT})" ) args = parser.parse_args() total_games = args.num_games num_workers = args.num_workers # ── 1. 加载模型权重 ── model_a_state_dict = load_cfr_state_dict(args.model_a) model_b_state_dict = load_cfr_state_dict(args.model_b) card_state_dict = load_card_state_dict(args.card_model) # ── 2. 分配任务：半数局 Model A 做 P0，半数做 P1 ── games_per_side = total_games // 2 # 每一侧均匀分配给 num_workers 个 Worker games_per_worker_per_side = games_per_side // num_workers remainder_per_side = games_per_side % num_workers # 构建 Worker 任务列表: (num_games, model_a_is_p0) tasks: List[Tuple[int, bool]] = [] # Model A 做 P0 (小盲) 的任务 for i in range(num_workers): games = games_per_worker_per_side + (1 if i < remainder_per_side else 0) if games > 0: tasks.append((games, True)) # Model A 做 P1 (大盲) 的任务 for i in range(num_workers): games = games_per_worker_per_side + (1 if i < remainder_per_side else 0) if games > 0: tasks.append((games, False)) actual_total = sum(t[0] for t in tasks) actual_a_as_p0 = sum(t[0] for t in tasks if t[1]) actual_a_as_p1 = sum(t[0] for t in tasks if not t[1]) print(f"\n{'='*70}") print(f" Deep CFR 模型对抗评估") print(f" Model A: {args.model_a}") print(f" Model B: {args.model_b}") print(f" 总局数: {actual_total} (A做P0: {actual_a_as_p0}, A做P1: {actual_a_as_p1})") print(f" Workers: {num_workers} | BB = {BIG_BLIND} 筹码") print(f"{'='*70}\n") # ── 3. 启动多进程对战 ── spawn_ctx = mp.get_context('spawn') with ProcessPoolExecutor( max_workers=num_workers, initializer=_init_worker, initargs=(model_a_state_dict, model_b_state_dict, card_state_dict), mp_context=spawn_ctx, ) as executor: # 提交所有任务 futures = [] for task_idx, (num_games_task, model_a_is_p0) in enumerate(tasks): future = executor.submit(worker_play_games, num_games_task, model_a_is_p0) futures.append((future, task_idx, num_games_task, model_a_is_p0)) # 收集结果 total_chips_won_by_a = 0.0 completed = 0 for future, task_idx, num_games_task, model_a_is_p0 in futures: try: chips = future.result() total_chips_won_by_a += chips completed += 1 position_str = "A=P0(SB)" if model_a_is_p0 else "A=P1(BB)" if completed % 5 == 0 or completed == len(futures): print(f" [进度] {completed}/{len(futures)} 任务完成 " f"({position_str}, {num_games_task}局, " f"Model A 筹码: {chips:+.0f})") except Exception as e: print(f" [警告] 任务 {task_idx} 执行失败: {e}") # 失败的任务不影响整体，仅跳过 # ── 4. 计算 bb/100 ── # bb/100 = (总筹码收益 / BB) / (总局数 / 100) # = (total_chips_won / BIG_BLIND) / (total_games / 100) # = total_chips_won * 100 / (BIG_BLIND * total_games) bb_per_100 = (total_chips_won_by_a / BIG_BLIND) / (actual_total / 100) # ── 5. 打印结果 ── sign = "+" if bb_per_100 >= 0 else "" print(f"\n{'='*70}") print(f" 评估完成!") print(f" 总局数: {actual_total}") print(f" Model A 总筹码收益: {total_chips_won_by_a:+.0f}") print(f" 经过 {actual_total:,} 局对抗，Model A 对 Model B 的百手赢率为: " f"{sign}{bb_per_100:.1f} bb/100") print(f"{'='*70}") if __name__ == "__main__": main()