1187 lines
45 KiB
Python
1187 lines
45 KiB
Python
"""
|
||
mccfr_trainer.py — MCCFR 自我博弈训练流水线
|
||
|
||
基于外部采样 (External Sampling MCCFR) 的 Deep CFR 训练器。
|
||
核心流程:
|
||
1. traverse(state, traversing_player) 递归遍历博弈树，
|
||
计算每个信息集的反事实遗憾值 (Counterfactual Regret)
|
||
2. 将 [info_state, legal_mask, regrets, strategy] 存入 CFRBuffer
|
||
3. 从 Buffer 采样 mini-batch，训练 CFRNetwork 的 Regret Head 和 Policy Head
|
||
|
||
=== 多进程并行数据生成 ===
|
||
|
||
阶段 A (Data Generation) 已改造为多进程版本：
|
||
- 使用 ProcessPoolExecutor(mp_context=spawn) 并发执行
|
||
- spawn 启动方式避免 fork 继承 CUDA Context 导致死锁
|
||
- 进程池跨 iteration 复用，通过函数参数直接传递权重（内存安全）
|
||
- 每个 Worker 批量处理多局，大幅减少 IPC 消息数量
|
||
- Worker 返回 (utility_list, experience_list)，经验以 Python 基本类型传输
|
||
|
||
=== 外部采样 MCCFR 核心逻辑 ===
|
||
|
||
MCCFR 是一种近似求解大规模扩展式博弈纳什均衡的算法。
|
||
"外部采样" 指的是：在对方回合，根据当前策略采样 1 个动作继续；
|
||
在遍历者回合，遍历所有合法动作以精确计算 regret。
|
||
|
||
设 v(I, a) 为在信息集 I 采取动作 a 后的期望收益，
|
||
v(I) = sum_a(σ(I,a) * v(I,a)) 为加权期望收益，
|
||
则反事实遗憾值:
|
||
regret(I, a) = v(I, a) - v(I)
|
||
|
||
当遍历者反复经历相同信息集时，累计 regret 越大的动作
|
||
越应该被优先选择——这就是 Regret Matching 的核心直觉。
|
||
"""
|
||
|
||
# ── 必须在 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 random
|
||
import sys
|
||
import os
|
||
import glob
|
||
import tempfile
|
||
import uuid
|
||
import multiprocessing as mp
|
||
from typing import List, Tuple, Dict, Any
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import pyspiel
|
||
|
||
torch.set_num_threads(1)
|
||
|
||
|
||
# ── 将 poker/ 根目录加入 sys.path，方便导入同目录模块 ──
|
||
_POKER_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
if _POKER_DIR not in sys.path:
|
||
sys.path.insert(0, _POKER_DIR)
|
||
|
||
# ── 将项目根目录加入 sys.path，方便导入 card_model 包 ──
|
||
_PROJECT_DIR = os.path.dirname(_POKER_DIR)
|
||
if _PROJECT_DIR not in sys.path:
|
||
sys.path.insert(0, _PROJECT_DIR)
|
||
|
||
from env_adapter import (
|
||
HUNL_FULLGAME_STRING,
|
||
BetTranslator,
|
||
extract_env_state,
|
||
CFR_ACTIONS,
|
||
NUM_CFR_ACTIONS,
|
||
STACK_NORMALIZE,
|
||
)
|
||
from card_model.config import NUM_BINS, 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
|
||
from cfr_buffer import CFRBuffer
|
||
|
||
|
||
|
||
# ───────────────────── 终极版超参数 (128G RAM + 24G VRAM, 24核) ─────────────────────
|
||
|
||
# ── 大规模并行与迭代 ──
|
||
NUM_ITERATIONS = 30000 # 总训练迭代次数（Deep CFR: 高频迭代，少局数）
|
||
GAMES_PER_ITER = 1000 # 每次迭代搜索的局数（快速进入更新阶段）
|
||
# NUM_WORKERS: 留 2 核给主进程（GPU 训练 + 日志 + OS），避免上下文切换惩罚
|
||
# 24 核 - 2 = 22 个 worker，每个 worker 纯 CPU 密集型遍历
|
||
NUM_WORKERS = 22
|
||
|
||
# ── 巨型经验池 (吞噬你的 128G 内存) ──
|
||
# 1条样本大约 1KB 多。500 万条样本约占 6~8 GB 内存。
|
||
# 设为 1000 万，充分利用你的系统内存，避免灾难性遗忘 (Catastrophic Forgetting)
|
||
BUFFER_MAX_SIZE = 10_000_000
|
||
MIN_BUFFER_SIZE_FOR_TRAIN = 200_000 # 初期攒够 2 万条就开始训练
|
||
|
||
# ── 显存利用 (24G VRAM) ──
|
||
# CFR 网络极其小 (3层 MLP, ~100K 参数)。BS=32768 时:
|
||
# Forward 中间激活: ~33 MB, Backward: ~33 MB, Optimizer: ~0.8 MB
|
||
# 总计 < 100 MB，在 24G VRAM 上毫无压力
|
||
# 超大 Batch Size 能让策略收敛平滑，减少策略震荡
|
||
TRAIN_BATCH_SIZE = 16384
|
||
TRAIN_STEPS_PER_ITER = 64 # 每次迭代后，网络更新 64 步
|
||
|
||
# ── 优化器细节 ──
|
||
LEARNING_RATE = 5e-4 # BS 变大，LR 稍微降一点，求稳
|
||
WEIGHT_DECAY = 1e-4 # 防止过拟合
|
||
|
||
|
||
# 训练
|
||
#NUM_ITERATIONS = 100 # 总训练迭代次数
|
||
#GAMES_PER_ITER = 50 # 每个 iteration 跑多少局自对弈
|
||
#TRAIN_BATCH_SIZE = 256 # 训练 mini-batch 大小
|
||
#TRAIN_STEPS_PER_ITER = 20 # 每 iteration 训练多少步
|
||
#MIN_BUFFER_SIZE_FOR_TRAIN = 1000 # Buffer 至少有多少条数据才开始训练
|
||
#LEARNING_RATE = 1e-3 # AdamW 学习率
|
||
#WEIGHT_DECAY = 1e-4 # AdamW 权重衰减
|
||
|
||
# 归一化常量（与 env_features 的 5 维对应）
|
||
# STACK_NORMALIZE = 20000.0 已从 env_adapter 导入（与 Botzone 初始筹码对齐）
|
||
STREET_NORMALIZE = 3.0 # 街道归一化因子 (0/3, 1/3, 2/3, 3/3)
|
||
|
||
CLIP_GRAD_NORM = 1.0 # 梯度裁剪阈值，防止训练不稳定
|
||
|
||
# Card Model 权重路径（设为 None 则使用随机初始化）
|
||
CARD_MODEL_CHECKPOINT = "card_model/data/best_card_model.pt"
|
||
|
||
# 设备
|
||
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
# ── 检查点配置 ──
|
||
CHECKPOINT_DIR = os.path.join(_POKER_DIR, "checkpoints")
|
||
CHECKPOINT_INTERVAL = 50 # 每 50 个 iteration 保存一次
|
||
KEEP_LAST_N_CHECKPOINTS = 3 # 本地最多保留最近的 3 个检查点文件，防止磁盘撑爆
|
||
|
||
|
||
# ───────────────────── 辅助函数 ─────────────────────
|
||
|
||
def build_env_features(env_info: dict) -> torch.Tensor:
|
||
"""
|
||
将 extract_env_state 返回的字典归一化为 5 维 env_features Tensor。
|
||
|
||
5 维内容:
|
||
[pot/20000, p0_stack/20000, p1_stack/20000, street/3.0, position]
|
||
|
||
Args:
|
||
env_info: extract_env_state() 返回的字典
|
||
|
||
Returns:
|
||
[5] 的 FloatTensor
|
||
"""
|
||
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)
|
||
|
||
|
||
# ── CardModel 结果缓存 ──
|
||
# 在扑克博弈树中，只要没有发新牌，牌面概率绝对不变。
|
||
# 以 (hole_cards_tuple, board_cards_tuple) 为 key 缓存 CardModel 输出，
|
||
# 避免在同一个发牌阶段（如 Flop）反复推理。
|
||
CARD_CACHE: Dict[str, torch.Tensor] = {}
|
||
|
||
|
||
def build_card_features(
|
||
card_model: CardModel,
|
||
state,
|
||
) -> torch.Tensor:
|
||
"""
|
||
使用 CardModel 从当前 state 提取 50 维胜率直方图 (card_features)。
|
||
|
||
带字典缓存：只要 hole_cards 和 board_cards 不变，直接返回缓存结果。
|
||
在每次 iteration 开始前由主循环调用 CARD_CACHE.clear() 防止内存泄漏。
|
||
|
||
流程:
|
||
1. 调用 extract_cards_from_state(state) 获取 hole_cards, board_cards
|
||
2. 构造 cache_key，命中则直接返回
|
||
3. 未命中则构造 Tensor 进行推理（在模型所在设备上进行）
|
||
4. 返回 [50] 的 FloatTensor（CPU）
|
||
|
||
Args:
|
||
card_model: 已加载权重的 CardModel（eval 模式）
|
||
state: OpenSpiel State
|
||
|
||
Returns:
|
||
[50] 的 FloatTensor（脱离计算图，在 CPU 上）
|
||
"""
|
||
global CARD_CACHE
|
||
|
||
hole_cards, board_cards = extract_cards_from_state(state)
|
||
|
||
# 将卡牌列表转为可哈希的 key
|
||
cache_key = f"{tuple(hole_cards)}_{tuple(board_cards)}"
|
||
|
||
if cache_key in CARD_CACHE:
|
||
return CARD_CACHE[cache_key]
|
||
|
||
# 构造输入 Tensor（使用 card_model 当前所在设备）
|
||
model_device = next(card_model.parameters()).device
|
||
x_hole = torch.tensor([hole_cards], dtype=torch.int64, device=model_device) # [1, 2]
|
||
# board_cards 不足 5 张时用 PAD_TOKEN 填充
|
||
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]
|
||
|
||
# 脱离计算图，转到 CPU，去掉 batch 维度
|
||
result = pred_histogram.squeeze(0).cpu()
|
||
CARD_CACHE[cache_key] = result
|
||
return result
|
||
|
||
|
||
# ───────────────────── 核心：外部采样遍历（单进程版本，保留兼容） ─────────────────────
|
||
|
||
def traverse(
|
||
state,
|
||
traversing_player: int,
|
||
card_model: CardModel,
|
||
cfr_net: CFRNetwork,
|
||
buffer: CFRBuffer,
|
||
translator: BetTranslator,
|
||
depth: int = 0,
|
||
) -> float:
|
||
"""
|
||
外部采样 MCCFR 的博弈树遍历函数（单进程版本）。
|
||
|
||
递归地遍历博弈树，在遍历者回合计算每个动作的反事实遗憾值 (Regret)，
|
||
并将经验 [info_state, legal_mask, regrets, strategy] 存入 Buffer。
|
||
|
||
=== 三种节点类型的处理 ===
|
||
|
||
1. Terminal Node (终局):
|
||
直接返回 traversing_player 的收益值 state.returns()[traversing_player]
|
||
|
||
2. Chance Node (发牌):
|
||
按概率随机采样一个发牌动作，递归继续
|
||
|
||
3. Player Node (玩家行动):
|
||
提取信息集特征，通过 CFRNetwork 获取当前策略 σ(I,·)
|
||
然后分两种情况:
|
||
|
||
a) 对方回合 (current_player != traversing_player):
|
||
根据 σ(I,·) 的概率分布采样 1 个动作，执行后递归
|
||
→ 这是"外部采样"的含义：只采样 1 条路径，大幅减少计算量
|
||
|
||
b) 遍历者回合 (current_player == traversing_player):
|
||
遍历所有合法动作 a:
|
||
- clone state，执行动作 a，递归得到 v(I, a)
|
||
计算加权期望收益: v(I) = Σ_a σ(I,a) * v(I,a)
|
||
计算遗憾值: regret(I, a) = v(I,a) - v(I)
|
||
将经验存入 Buffer
|
||
|
||
Args:
|
||
state: OpenSpiel 的 State 对象
|
||
traversing_player: 当前遍历者 (0 或 1)
|
||
card_model: Card Model（eval 模式），用于提取牌面特征
|
||
cfr_net: CFR 策略网络（eval 模式），用于获取当前策略
|
||
buffer: CFR 经验回放池
|
||
translator: BetTranslator，CFR 动作 ↔ 引擎动作转换
|
||
|
||
Returns:
|
||
float: 当前节点对 traversing_player 的期望收益
|
||
"""
|
||
|
||
# ── 深度熔断保护：超过极限深度，强制截断博弈树 ──
|
||
if depth >= 40:
|
||
return 0.0
|
||
|
||
# ── 1. Terminal Node: 终局，直接返回收益 ──
|
||
if state.is_terminal():
|
||
return float(state.returns()[traversing_player])
|
||
|
||
# ── 2. 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]
|
||
child_state = state.clone()
|
||
child_state.apply_action(chance_action)
|
||
utility = traverse(child_state, traversing_player, card_model, cfr_net,
|
||
buffer, translator, depth + 1)
|
||
del child_state
|
||
return utility
|
||
|
||
# ── 3. Player Node: 玩家行动 ──
|
||
current_player = state.current_player()
|
||
|
||
env_info = extract_env_state(state)
|
||
env_features = build_env_features(env_info)
|
||
card_features = build_card_features(card_model, state)
|
||
legal_mask = env_info["legal_mask"]
|
||
|
||
info_state = torch.cat([card_features, env_features], dim=0)
|
||
|
||
# 使用 cfr_net 当前所在设备（数据生成阶段为 CPU）
|
||
net_device = next(cfr_net.parameters()).device
|
||
legal_mask_tensor = torch.tensor(
|
||
[legal_mask], dtype=torch.float32, device=net_device
|
||
)
|
||
card_input = card_features.unsqueeze(0).to(net_device)
|
||
env_input = env_features.unsqueeze(0).to(net_device)
|
||
|
||
with torch.no_grad():
|
||
current_strategy, _ = cfr_net.get_strategy(
|
||
card_input, env_input, legal_mask_tensor
|
||
)
|
||
|
||
strategy_list = current_strategy.squeeze(0).cpu().tolist()
|
||
|
||
# ── 3a. 对方回合 ──
|
||
if current_player != traversing_player:
|
||
legal_indices = [i for i, m in enumerate(legal_mask) if m == 1]
|
||
if not legal_indices:
|
||
legal_indices = [1]
|
||
|
||
probs = [strategy_list[i] for i in legal_indices]
|
||
prob_sum = sum(probs)
|
||
if prob_sum > 0:
|
||
probs = [p / prob_sum for p in probs]
|
||
else:
|
||
probs = [1.0 / len(legal_indices)] * len(legal_indices)
|
||
|
||
sampled_idx = random.choices(legal_indices, weights=probs, k=1)[0]
|
||
engine_action = translator.cfr_to_engine_action(state, sampled_idx)
|
||
child_state = state.clone()
|
||
child_state.apply_action(engine_action)
|
||
|
||
utility = traverse(child_state, traversing_player, card_model, cfr_net,
|
||
buffer, translator, depth + 1)
|
||
del child_state
|
||
return utility
|
||
|
||
# ── 3b. 遍历者回合 ──
|
||
legal_indices = [i for i, m in enumerate(legal_mask) if m == 1]
|
||
if not legal_indices:
|
||
return 0.0
|
||
|
||
action_utilities = [0.0] * NUM_CFR_ACTIONS
|
||
|
||
for a in legal_indices:
|
||
child_state = state.clone()
|
||
engine_action = translator.cfr_to_engine_action(child_state, a)
|
||
child_state.apply_action(engine_action)
|
||
action_utilities[a] = traverse(
|
||
child_state, traversing_player, card_model, cfr_net,
|
||
buffer, translator, depth + 1
|
||
)
|
||
del child_state
|
||
|
||
node_utility = 0.0
|
||
for a in legal_indices:
|
||
node_utility += strategy_list[a] * action_utilities[a]
|
||
|
||
regrets = [0.0] * NUM_CFR_ACTIONS
|
||
for a in legal_indices:
|
||
regrets[a] = (action_utilities[a] - node_utility) / STACK_NORMALIZE
|
||
|
||
buffer.add(
|
||
info_state=info_state,
|
||
legal_mask=legal_mask,
|
||
regrets=regrets,
|
||
strategy=strategy_list,
|
||
)
|
||
|
||
return node_utility
|
||
|
||
|
||
# ───────────────────── 多进程 Worker 基础设施 ─────────────────────
|
||
#
|
||
# Worker 进程通过 _init_worker 初始化，在其中自行创建：
|
||
# - pyspiel 游戏实例（不可跨进程序列化）
|
||
# - CPU 上的 CFRNetwork（从主进程传入的 state_dict 加载权重）
|
||
# - CPU 上的 CardModel（从主进程传入的 state_dict 加载权重）
|
||
# - BetTranslator（轻量无状态对象）
|
||
#
|
||
# 权重热更新机制：
|
||
# 模型参数极小（< 500KB），直接通过 worker_traverse_batch 函数参数传递，
|
||
# 无需临时文件中转，彻底消除并发读写冲突和 BrokenProcessPool 风险。
|
||
#
|
||
# IPC 安全机制：
|
||
# - 使用 spawn 启动方式，避免 fork 继承 CUDA Context 导致死锁
|
||
# - Worker 返回经验全部使用 Python 基本类型（list[float]/list[int]）
|
||
# - 每个 Worker 批量处理多局游戏，减少 IPC 消息数量，防止管道阻塞
|
||
|
||
# Worker 进程的全局状态
|
||
_WORKER_STATE: Dict[str, Any] = {}
|
||
|
||
|
||
def _init_worker(
|
||
cfr_state_dict: dict,
|
||
card_state_dict: dict,
|
||
) -> None:
|
||
"""
|
||
ProcessPoolExecutor 的 worker 初始化函数。
|
||
|
||
在每个 worker 进程启动时调用一次，创建该进程专属的：
|
||
- OpenSpiel 游戏实例
|
||
- CPU 上的 CFRNetwork（eval 模式）
|
||
- CPU 上的 CardModel（eval 模式）
|
||
- BetTranslator
|
||
|
||
传入的 state_dict 中所有 Tensor 已通过 .cpu() 确保在 CPU 上，
|
||
避免 CUDA 跨进程问题。
|
||
|
||
权重热更新：
|
||
每次 iteration 主进程通过 worker_traverse_batch 的 cfr_state_dict 参数
|
||
直接传递最新权重，Worker 在 batch 开始时 load_state_dict 加载，
|
||
彻底消除临时文件并发读写冲突。
|
||
|
||
Args:
|
||
cfr_state_dict: 主进程 CFRNetwork 的 state_dict（CPU Tensor）
|
||
card_state_dict: 主进程 CardModel 的 state_dict（CPU Tensor）
|
||
"""
|
||
global _WORKER_STATE
|
||
|
||
# 每个 worker 自行创建 OpenSpiel 游戏实例（不可跨进程传递）
|
||
_WORKER_STATE["game"] = pyspiel.load_game(HUNL_FULLGAME_STRING)
|
||
|
||
# 在 CPU 上创建 CFRNetwork 并加载主进程的最新权重
|
||
cfr_net = CFRNetwork(
|
||
card_dim=CARD_DIM,
|
||
env_dim=ENV_DIM,
|
||
num_actions=NUM_ACTIONS,
|
||
)
|
||
cfr_net.load_state_dict(cfr_state_dict)
|
||
cfr_net.eval()
|
||
_WORKER_STATE["cfr_net"] = cfr_net
|
||
|
||
# 在 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 worker_traverse_batch(
|
||
game_indices: List[int],
|
||
cfr_state_dict: dict,
|
||
) -> Tuple[List[Tuple[int, float]], str]:
|
||
"""
|
||
多进程版本的批量 MCCFR 遍历函数，在 Worker 进程中执行。
|
||
|
||
一个 Worker 一次处理多局游戏（由 game_indices 指定），
|
||
大幅减少 IPC 消息数量，防止管道被海量小消息撑爆。
|
||
|
||
权重通过函数参数直接传入（模型 < 500KB），无需临时文件中转，
|
||
彻底消除并发读写冲突和 BrokenProcessPool 风险。
|
||
|
||
经验数据通过临时文件回传，避免 multiprocessing.Pipe 序列化爆炸：
|
||
Worker 将经验列表 torch.save 到临时 .pt 文件，仅返回文件路径。
|
||
主进程 torch.load 读取后立即删除临时文件。
|
||
|
||
Args:
|
||
game_indices: 要处理的游戏索引列表（长度约 GAMES_PER_ITER/NUM_WORKERS）
|
||
cfr_state_dict: 主进程当前 iteration 的 CFRNetwork 权重（CPU Tensor）
|
||
|
||
Returns:
|
||
game_results: [(game_idx, utility), ...] 每局的索引和收益
|
||
exp_file_path: 经验临时文件路径，主进程需 torch.load 读取后删除
|
||
"""
|
||
global _WORKER_STATE
|
||
global CARD_CACHE
|
||
|
||
# ── 清除本 Worker 上个批次的卡牌缓存，彻底封杀内存泄露！ ──
|
||
CARD_CACHE.clear()
|
||
|
||
# ── 直接通过参数加载最新权重，无需文件中转 ──
|
||
_WORKER_STATE["cfr_net"].load_state_dict(cfr_state_dict)
|
||
|
||
game = _WORKER_STATE["game"]
|
||
cfr_net = _WORKER_STATE["cfr_net"]
|
||
card_model = _WORKER_STATE["card_model"]
|
||
translator = _WORKER_STATE["translator"]
|
||
|
||
game_results: List[Tuple[int, float]] = []
|
||
all_experiences: List[Tuple] = []
|
||
|
||
for game_idx in game_indices:
|
||
traversing_player = game_idx % 2
|
||
state = game.new_initial_state()
|
||
|
||
experiences: List[Tuple] = []
|
||
utility = _traverse_worker(
|
||
state, traversing_player, card_model, cfr_net,
|
||
experiences, translator,
|
||
)
|
||
|
||
game_results.append((game_idx, utility))
|
||
all_experiences.extend(experiences)
|
||
|
||
#print(f"Worker 跑完了一局! (Game {game_idx})", flush=True)
|
||
|
||
|
||
# ── 将经验落盘，避免 IPC 管道序列化爆炸 ──
|
||
# 直接 return all_experiences 会导致 multiprocessing.Pipe 被撑爆
|
||
# （单次可达数百 MB 的 Python list 序列化数据），引发 OS OOM 和 BrokenProcessPool。
|
||
# 改用临时文件：torch.save 落盘 → 仅传输文件路径字符串 → 主进程读取后删除。
|
||
exp_file_path = os.path.join(
|
||
tempfile.gettempdir(),
|
||
f"cfr_exp_{uuid.uuid4().hex}.pt",
|
||
)
|
||
torch.save(all_experiences, exp_file_path)
|
||
del all_experiences # 立即释放 Worker 内存，避免文件缓存双倍占用
|
||
|
||
return game_results, exp_file_path
|
||
|
||
|
||
def _traverse_worker(
|
||
state,
|
||
traversing_player: int,
|
||
card_model: CardModel,
|
||
cfr_net: CFRNetwork,
|
||
experiences: List[Tuple],
|
||
translator: BetTranslator,
|
||
depth: int = 0,
|
||
) -> float:
|
||
"""
|
||
Worker 进程内部的递归遍历函数。
|
||
|
||
与 traverse() 逻辑完全一致，唯一区别是将经验追加到
|
||
本地 experiences 列表（而非全局 CFRBuffer），且存入的
|
||
info_state 使用 list[float] 而非 Tensor，以便跨进程传输。
|
||
|
||
内存安全保证：
|
||
- 所有 Tensor 操作在 torch.no_grad() 下进行
|
||
- info_state_tensor.tolist() 后，原始 Tensor 即可被 GC 回收
|
||
- 递归过程中 state.clone() 产生的对象在递归返回后也被 GC
|
||
- 最终返回的 experiences 只包含 Python 基本类型，
|
||
主进程 torch.tensor() 新建时无 grad_fn，buffer.add() 安全
|
||
|
||
设备说明：
|
||
Worker 中所有模型始终在 CPU 上。Tensor 设备由模型参数所在
|
||
设备自动推导（next(cfr_net.parameters()).device），无需显式传递。
|
||
|
||
Args:
|
||
state: OpenSpiel State
|
||
traversing_player: 当前遍历者
|
||
card_model: CardModel（eval，CPU）
|
||
cfr_net: CFRNetwork（eval，CPU）
|
||
experiences: 本地经验暂存列表
|
||
translator: BetTranslator
|
||
|
||
Returns:
|
||
float: 当前节点对 traversing_player 的期望收益
|
||
"""
|
||
# ── 深度熔断保护：超过极限深度，强制截断博弈树 ──
|
||
if depth >= 40:
|
||
return 0.0
|
||
|
||
# ── 1. Terminal Node ──
|
||
if state.is_terminal():
|
||
return float(state.returns()[traversing_player])
|
||
|
||
# ── 2. 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]
|
||
child_state = state.clone()
|
||
child_state.apply_action(chance_action)
|
||
utility = _traverse_worker(child_state, traversing_player, card_model, cfr_net,
|
||
experiences, translator, depth + 1)
|
||
del child_state
|
||
return utility
|
||
|
||
# ── 3. Player Node ──
|
||
current_player = state.current_player()
|
||
|
||
env_info = extract_env_state(state)
|
||
env_features = build_env_features(env_info)
|
||
card_features = build_card_features(card_model, state)
|
||
legal_mask = env_info["legal_mask"]
|
||
|
||
info_state_tensor = torch.cat([card_features, env_features], dim=0)
|
||
|
||
# 使用 cfr_net 当前所在设备（Worker 中始终为 CPU）
|
||
net_device = next(cfr_net.parameters()).device
|
||
legal_mask_tensor = torch.tensor(
|
||
[legal_mask], dtype=torch.float32, device=net_device
|
||
)
|
||
card_input = card_features.unsqueeze(0).to(net_device)
|
||
env_input = env_features.unsqueeze(0).to(net_device)
|
||
|
||
with torch.no_grad():
|
||
current_strategy, _ = cfr_net.get_strategy(
|
||
card_input, env_input, legal_mask_tensor
|
||
)
|
||
|
||
strategy_list = current_strategy.squeeze(0).cpu().tolist()
|
||
|
||
# ── 3a. 对方回合 ──
|
||
if current_player != traversing_player:
|
||
legal_indices = [i for i, m in enumerate(legal_mask) if m == 1]
|
||
if not legal_indices:
|
||
legal_indices = [1]
|
||
|
||
probs = [strategy_list[i] for i in legal_indices]
|
||
prob_sum = sum(probs)
|
||
if prob_sum > 0:
|
||
probs = [p / prob_sum for p in probs]
|
||
else:
|
||
probs = [1.0 / len(legal_indices)] * len(legal_indices)
|
||
|
||
sampled_idx = random.choices(legal_indices, weights=probs, k=1)[0]
|
||
engine_action = translator.cfr_to_engine_action(state, sampled_idx)
|
||
child_state = state.clone()
|
||
child_state.apply_action(engine_action)
|
||
|
||
utility = _traverse_worker(child_state, traversing_player, card_model, cfr_net,
|
||
experiences, translator, depth + 1)
|
||
del child_state
|
||
return utility
|
||
|
||
# ── 3b. 遍历者回合 ──
|
||
legal_indices = [i for i, m in enumerate(legal_mask) if m == 1]
|
||
if not legal_indices:
|
||
return 0.0
|
||
|
||
action_utilities = [0.0] * NUM_CFR_ACTIONS
|
||
|
||
for a in legal_indices:
|
||
child_state = state.clone()
|
||
engine_action = translator.cfr_to_engine_action(child_state, a)
|
||
child_state.apply_action(engine_action)
|
||
action_utilities[a] = _traverse_worker(
|
||
child_state, traversing_player, card_model, cfr_net,
|
||
experiences, translator, depth + 1,
|
||
)
|
||
del child_state
|
||
|
||
node_utility = 0.0
|
||
for a in legal_indices:
|
||
node_utility += strategy_list[a] * action_utilities[a]
|
||
|
||
regrets = [0.0] * NUM_CFR_ACTIONS
|
||
for a in legal_indices:
|
||
regrets[a] = (action_utilities[a] - node_utility) / STACK_NORMALIZE
|
||
|
||
# 存入本地经验列表（info_state 转为 list[float] 以便跨进程传输）
|
||
# .tolist() 后 info_state_tensor 可被 GC 回收，不会造成内存累积
|
||
info_state_list = info_state_tensor.tolist()
|
||
experiences.append((info_state_list, legal_mask, regrets, strategy_list))
|
||
|
||
return node_utility
|
||
|
||
|
||
# ───────────────────── 检查点保存与恢复 ─────────────────────
|
||
|
||
def save_checkpoint(
|
||
iteration: int,
|
||
cfr_net: CFRNetwork,
|
||
buffer: CFRBuffer,
|
||
optimizer,
|
||
checkpoint_dir: str,
|
||
keep_last_n: int,
|
||
) -> None:
|
||
"""
|
||
原子保存检查点，并清理旧文件。
|
||
|
||
原子操作流程：
|
||
1. 先写入临时文件 ckpt_tmp.pt
|
||
2. 写入成功后 os.replace 重命名为 ckpt_iter_{iteration}.pt
|
||
3. 复制一份 latest_ckpt.pt 作为默认恢复入口
|
||
4. 清理超过 keep_last_n 的旧检查点
|
||
|
||
Args:
|
||
iteration: 当前迭代次数
|
||
cfr_net: CFR 策略网络
|
||
buffer: CFR 经验回放池
|
||
optimizer: AdamW 优化器（可能为 None）
|
||
checkpoint_dir: 检查点保存目录
|
||
keep_last_n: 保留最近 N 个检查点
|
||
"""
|
||
os.makedirs(checkpoint_dir, exist_ok=True)
|
||
|
||
# 构造保存字典
|
||
save_dict = {
|
||
"iteration": iteration,
|
||
"model_state_dict": {k: v.cpu() for k, v in cfr_net.state_dict().items()},
|
||
"buffer_state_dict": buffer.state_dict(),
|
||
}
|
||
if optimizer is not None:
|
||
# 优化器状态中的 Tensor 可能在 GPU 上，需转 CPU
|
||
opt_state = optimizer.state_dict()
|
||
opt_state_cpu = {}
|
||
for k, v in opt_state.items():
|
||
if isinstance(v, torch.Tensor):
|
||
opt_state_cpu[k] = v.cpu()
|
||
elif isinstance(v, dict):
|
||
opt_state_cpu[k] = _recursive_tensor_to_cpu(v)
|
||
else:
|
||
opt_state_cpu[k] = v
|
||
save_dict["optimizer_state_dict"] = opt_state_cpu
|
||
|
||
# 原子保存：先写临时文件，再 rename
|
||
ckpt_path = os.path.join(checkpoint_dir, f"ckpt_iter_{iteration}.pt")
|
||
tmp_path = os.path.join(checkpoint_dir, "ckpt_tmp.pt")
|
||
|
||
torch.save(save_dict, tmp_path)
|
||
os.replace(tmp_path, ckpt_path)
|
||
|
||
# 更新 latest_ckpt.pt（用 copy 而非 symlink，更跨平台）
|
||
latest_path = os.path.join(checkpoint_dir, "latest_ckpt.pt")
|
||
# os.replace 是原子的，先用临时文件再替换
|
||
tmp_latest = os.path.join(checkpoint_dir, "latest_ckpt_tmp.pt")
|
||
torch.save(save_dict, tmp_latest)
|
||
os.replace(tmp_latest, latest_path)
|
||
|
||
print(f"[Checkpoint] 已保存: iter={iteration}, buffer={len(buffer)}")
|
||
|
||
# 清理旧检查点（保留最近 keep_last_n 个）
|
||
_cleanup_old_checkpoints(checkpoint_dir, keep_last_n)
|
||
|
||
|
||
def _recursive_tensor_to_cpu(d: dict) -> dict:
|
||
"""递归地将字典中所有 Tensor 转到 CPU。"""
|
||
out = {}
|
||
for k, v in d.items():
|
||
if isinstance(v, torch.Tensor):
|
||
out[k] = v.cpu()
|
||
elif isinstance(v, dict):
|
||
out[k] = _recursive_tensor_to_cpu(v)
|
||
elif isinstance(v, list):
|
||
out[k] = [
|
||
item.cpu() if isinstance(item, torch.Tensor) else item
|
||
for item in v
|
||
]
|
||
else:
|
||
out[k] = v
|
||
return out
|
||
|
||
|
||
def _cleanup_old_checkpoints(checkpoint_dir: str, keep_last_n: int) -> None:
|
||
"""删除旧的检查点文件，只保留最近 keep_last_n 个。"""
|
||
pattern = os.path.join(checkpoint_dir, "ckpt_iter_*.pt")
|
||
files = glob.glob(pattern)
|
||
if len(files) <= keep_last_n:
|
||
return
|
||
|
||
# 按文件名中的 iteration 数字排序
|
||
def _extract_iter(path: str) -> int:
|
||
basename = os.path.basename(path)
|
||
# ckpt_iter_123.pt -> 123
|
||
num = basename.replace("ckpt_iter_", "").replace(".pt", "")
|
||
try:
|
||
return int(num)
|
||
except ValueError:
|
||
return -1
|
||
|
||
files_with_iter = [(f, _extract_iter(f)) for f in files if _extract_iter(f) >= 0]
|
||
files_with_iter.sort(key=lambda x: x[1])
|
||
|
||
# 删除多余的旧文件
|
||
to_delete = files_with_iter[:-keep_last_n]
|
||
for path, it in to_delete:
|
||
os.remove(path)
|
||
print(f"[Checkpoint] 已删除旧检查点: iter={it}")
|
||
|
||
|
||
def _optimizer_state_to_device(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
|
||
"""将优化器状态中所有 Tensor 移到指定设备。
|
||
|
||
PyTorch 的 optimizer.load_state_dict() 不会自动将 Tensor 移到
|
||
当前参数所在的设备，需要手动处理。这对于从 CPU 检查点恢复到
|
||
GPU 上的场景至关重要。
|
||
|
||
Args:
|
||
optimizer: PyTorch 优化器
|
||
device: 目标设备
|
||
"""
|
||
for state in optimizer.state.values():
|
||
for k, v in state.items():
|
||
if isinstance(v, torch.Tensor):
|
||
state[k] = v.to(device)
|
||
|
||
|
||
def load_checkpoint(
|
||
checkpoint_path: str,
|
||
cfr_net: CFRNetwork,
|
||
buffer: CFRBuffer,
|
||
) -> dict:
|
||
"""
|
||
从检查点文件恢复训练状态。
|
||
|
||
Args:
|
||
checkpoint_path: 检查点文件路径
|
||
cfr_net: CFR 策略网络（将被加载权重）
|
||
buffer: CFR 经验回放池（将被加载数据）
|
||
|
||
Returns:
|
||
dict: 包含 iteration 和 optimizer_state_dict（如果有）的字典
|
||
"""
|
||
print(f"[Checkpoint] 正在加载: {checkpoint_path}")
|
||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||
|
||
# 恢复模型权重
|
||
cfr_net.load_state_dict(ckpt["model_state_dict"])
|
||
|
||
# 恢复 Buffer
|
||
buffer.load_state_dict(ckpt["buffer_state_dict"])
|
||
|
||
result = {
|
||
"iteration": ckpt.get("iteration", 0),
|
||
}
|
||
if "optimizer_state_dict" in ckpt:
|
||
result["optimizer_state_dict"] = ckpt["optimizer_state_dict"]
|
||
|
||
print(f"[Checkpoint] 恢复成功: iter={result['iteration']}, buffer={len(buffer)}")
|
||
return result
|
||
|
||
|
||
# ───────────────────── 训练一步 ─────────────────────
|
||
|
||
def train_step(
|
||
cfr_net: CFRNetwork,
|
||
buffer: CFRBuffer,
|
||
optimizer: torch.optim.Optimizer,
|
||
device: torch.device,
|
||
) -> Tuple[float, float, float]:
|
||
"""
|
||
从 Buffer 采样一个 mini-batch，训练 CFRNetwork 一步。
|
||
|
||
Loss 组成:
|
||
1. Regret Loss: MSE(预测 regret, 目标 regret)，只在 legal_mask=1 的位置计算
|
||
2. Policy Loss: MSE(预测策略, 目标策略)，拟合 Buffer 中存储的策略分布
|
||
|
||
总 Loss = Regret Loss + Policy Loss
|
||
|
||
Args:
|
||
cfr_net: CFR 策略网络
|
||
buffer: CFR 经验回放池
|
||
optimizer: AdamW 优化器
|
||
device: 运行设备
|
||
|
||
Returns:
|
||
(total_loss, regret_loss, policy_loss): 三个 float 标量
|
||
"""
|
||
cfr_net.train()
|
||
|
||
info_states, legal_masks, target_regrets, target_strategies = buffer.sample(
|
||
TRAIN_BATCH_SIZE
|
||
)
|
||
|
||
info_states = info_states.to(device)
|
||
legal_masks = legal_masks.to(device)
|
||
target_regrets = target_regrets.to(device)
|
||
target_strategies = target_strategies.to(device)
|
||
|
||
card_features = info_states[:, :CARD_DIM]
|
||
env_features = info_states[:, CARD_DIM:]
|
||
|
||
pred_regrets, pred_policy_logits = cfr_net(card_features, env_features)
|
||
|
||
# ── Loss 1: Regret Loss ──
|
||
regret_diff = (pred_regrets - target_regrets) ** 2
|
||
regret_loss = (regret_diff * legal_masks).sum() / legal_masks.sum().clamp(min=1.0)
|
||
|
||
# ── Loss 2: Policy Loss ──
|
||
masked_logits = pred_policy_logits.masked_fill(legal_masks == 0, float("-inf"))
|
||
pred_strategy = F.softmax(masked_logits, dim=-1)
|
||
|
||
all_illegal = (legal_masks.sum(dim=-1, keepdim=True) == 0)
|
||
num_legal = legal_masks.sum(dim=-1, keepdim=True).clamp(min=1.0)
|
||
uniform = legal_masks / num_legal
|
||
pred_strategy = torch.where(all_illegal, uniform, pred_strategy)
|
||
|
||
policy_diff = (pred_strategy - target_strategies) ** 2
|
||
policy_loss = (policy_diff * legal_masks).sum() / legal_masks.sum().clamp(min=1.0)
|
||
|
||
total_loss = regret_loss + policy_loss
|
||
|
||
optimizer.zero_grad()
|
||
total_loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(cfr_net.parameters(), max_norm=CLIP_GRAD_NORM)
|
||
optimizer.step()
|
||
|
||
return (
|
||
total_loss.item(),
|
||
regret_loss.item(),
|
||
policy_loss.item(),
|
||
)
|
||
|
||
|
||
# ───────────────────── 主训练循环 ─────────────────────
|
||
|
||
def main():
|
||
"""
|
||
MCCFR 训练主循环（多进程版本）。
|
||
|
||
每个 iteration 分两个阶段:
|
||
阶段 A (Data Generation) — 多进程并行:
|
||
- 使用 ProcessPoolExecutor(mp_context=spawn) 并发执行
|
||
- 进程池跨 iteration 复用，通过共享内存 dict 热更新权重
|
||
- 每个 Worker 批量处理多局，减少 IPC 消息数量
|
||
- 主进程汇总所有 Worker 的经验到全局 CFRBuffer
|
||
|
||
阶段 B (Network Training) — 单进程 GPU:
|
||
- 如果 Buffer 中数据量足够，进行 mini-batch 训练
|
||
- 更新 Regret Head 和 Policy Head
|
||
|
||
日志输出:
|
||
- 每 iteration 打印 Buffer 大小、Loss 变化
|
||
"""
|
||
# ── 强制使用 spawn 启动方式，防止 fork 继承 CUDA Context 死锁 ──
|
||
# Linux 默认 fork，fork 后子进程继承父进程的 CUDA Context，
|
||
# 但 CUDA 运行时内部锁状态不一致，任何后续 CUDA 操作都会永久死锁。
|
||
# spawn 从零启动 Python 解释器，完全隔离 CUDA 状态。
|
||
# 必须在 if __name__ == "__main__" 保护内、任何 CUDA 操作之前调用。
|
||
mp.set_start_method('spawn', force=True)
|
||
|
||
print("=" * 70)
|
||
print(" MCCFR Self-Play Training Pipeline (Multi-Process)")
|
||
print(" External Sampling + Deep CFR")
|
||
print(f" NUM_WORKERS = {NUM_WORKERS}")
|
||
print(f" Start method = spawn")
|
||
print("=" * 70)
|
||
|
||
# ── 1. 初始化 OpenSpiel 环境 ──
|
||
game = pyspiel.load_game(HUNL_FULLGAME_STRING)
|
||
translator = BetTranslator()
|
||
print(f"\n[环境] 游戏加载成功: {HUNL_FULLGAME_STRING}")
|
||
|
||
# ── 2. 初始化 Card Model（先放 CPU，数据生成阶段用 CPU 推理） ──
|
||
card_model = CardModel() # 初始在 CPU
|
||
if CARD_MODEL_CHECKPOINT and os.path.isfile(CARD_MODEL_CHECKPOINT):
|
||
ckpt = torch.load(CARD_MODEL_CHECKPOINT, map_location="cpu", weights_only=False)
|
||
card_model.load_state_dict(ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt)
|
||
print(f"[CardModel] 已加载权重: {CARD_MODEL_CHECKPOINT}")
|
||
else:
|
||
print(f"[CardModel] 使用随机初始化（未提供 checkpoint）")
|
||
card_model.eval()
|
||
|
||
# ── 3. 初始化 CFR Network（先放 CPU，数据生成阶段用 CPU 推理） ──
|
||
cfr_net = CFRNetwork(
|
||
card_dim=CARD_DIM,
|
||
env_dim=ENV_DIM,
|
||
num_actions=NUM_ACTIONS,
|
||
) # 初始在 CPU
|
||
total_params = sum(p.numel() for p in cfr_net.parameters())
|
||
print(f"[CFRNetwork] 参数量: {total_params:,}，初始设备: CPU")
|
||
|
||
# ── 4. 初始化 Buffer ──
|
||
buffer = CFRBuffer(max_size=BUFFER_MAX_SIZE)
|
||
print(f"[Buffer] 容量: {BUFFER_MAX_SIZE:,}")
|
||
|
||
# Optimizer 延迟到首次阶段 B 时创建（此时 cfr_net 已移至 GPU）
|
||
# 但如果从检查点恢复且存档中有 optimizer_state_dict，则提前初始化
|
||
optimizer = None
|
||
|
||
# ── 检查点恢复 ──
|
||
# 检查 latest_ckpt.pt 是否存在，存在则自动恢复
|
||
start_iteration = 1
|
||
latest_ckpt_path = os.path.join(CHECKPOINT_DIR, "latest_ckpt.pt")
|
||
if os.path.isfile(latest_ckpt_path):
|
||
ckpt_info = load_checkpoint(latest_ckpt_path, cfr_net, buffer)
|
||
start_iteration = ckpt_info["iteration"] + 1
|
||
|
||
# 如果存档中有 optimizer 状态，必须现在就初始化并恢复
|
||
# 原因：optimizer 的动量缓冲区绑定在参数对象上，
|
||
# 必须在 cfr_net 移到 GPU 后、任何训练之前加载
|
||
if "optimizer_state_dict" in ckpt_info:
|
||
# 先将网络移至 GPU（optimizer 的参数必须在正确设备上）
|
||
cfr_net.to(DEVICE)
|
||
optimizer = torch.optim.AdamW(
|
||
cfr_net.parameters(),
|
||
lr=LEARNING_RATE,
|
||
weight_decay=WEIGHT_DECAY,
|
||
)
|
||
# 加载优化器状态（可能包含 GPU 上的 Tensor，需映射到正确设备）
|
||
optimizer.load_state_dict(ckpt_info["optimizer_state_dict"])
|
||
# 将优化器状态中的所有 Tensor 移到与参数相同的设备
|
||
_optimizer_state_to_device(optimizer, DEVICE)
|
||
# 网络移回 CPU，等阶段 A 再用
|
||
cfr_net.to("cpu")
|
||
print(f"[Optimizer] 已从检查点恢复优化器状态, device={DEVICE}")
|
||
|
||
print(f"[Checkpoint] 将从 Iteration {start_iteration} 继续训练")
|
||
else:
|
||
print("[Checkpoint] 未找到检查点，从头开始训练")
|
||
|
||
# ── 5. 初始化 CPU state_dict（给 Worker 首次初始化用） ──
|
||
# 必须在创建进程池之前转为 CPU，避免 CUDA Tensor 跨进程传输
|
||
cfr_state_dict_cpu = {k: v.cpu() for k, v in cfr_net.state_dict().items()}
|
||
card_state_dict_cpu = {k: v.cpu() for k, v in card_model.state_dict().items()}
|
||
|
||
# ── 6. 创建进程池（跨 iteration 复用，避免重复初始化开销） ──
|
||
# mp_context=spawn 确保子进程不继承 CUDA Context
|
||
# 使用 spawn context 创建 ProcessPoolExecutor，而非全局 set_start_method
|
||
spawn_ctx = mp.get_context('spawn')
|
||
|
||
print(f"\n{'='*70}")
|
||
print(f" 开始训练: {NUM_ITERATIONS} iterations, "
|
||
f"每 iter {GAMES_PER_ITER} 局自对弈, "
|
||
f"{NUM_WORKERS} workers 并行")
|
||
print(f"{'='*70}")
|
||
|
||
# 将 GAMES_PER_ITER 局均匀分配给 NUM_WORKERS 个 Worker
|
||
# 每个 Worker 一次处理一批游戏，减少 IPC 消息数量
|
||
game_indices_all = list(range(GAMES_PER_ITER))
|
||
# 计算每个 Worker 分配的游戏索引
|
||
chunks: List[List[int]] = [[] for _ in range(NUM_WORKERS)]
|
||
for i, gi in enumerate(game_indices_all):
|
||
chunks[i % NUM_WORKERS].append(gi)
|
||
|
||
with ProcessPoolExecutor(
|
||
max_workers=NUM_WORKERS,
|
||
initializer=_init_worker,
|
||
initargs=(cfr_state_dict_cpu, card_state_dict_cpu),
|
||
mp_context=spawn_ctx,
|
||
) as executor:
|
||
|
||
for iteration in range(start_iteration, NUM_ITERATIONS + 1):
|
||
|
||
# ──────────── 阶段 A: 数据生成 — 多进程 ────────────
|
||
|
||
# 动态设备切换：数据生成阶段，模型放 CPU，避免 Batch=1 的 GPU PCIe 延迟
|
||
card_model.to("cpu")
|
||
cfr_net.to("cpu")
|
||
cfr_net.eval()
|
||
|
||
# 清除 CardModel 缓存，防止内存泄漏
|
||
CARD_CACHE.clear()
|
||
|
||
# 获取当前网络权重，直接通过函数参数传递给 Worker（无需临时文件）
|
||
cfr_state_dict_cpu = {k: v.cpu() for k, v in cfr_net.state_dict().items()}
|
||
|
||
# 提交批量任务（每个 Worker 处理一批游戏，权重直接参数传递）
|
||
futures = {}
|
||
for worker_id, chunk in enumerate(chunks):
|
||
if not chunk:
|
||
continue
|
||
future = executor.submit(worker_traverse_batch, chunk, cfr_state_dict_cpu)
|
||
futures[future] = worker_id
|
||
|
||
# 收集结果并汇总经验
|
||
game_results = []
|
||
total_experiences = 0
|
||
|
||
for future in as_completed(futures):
|
||
worker_id = futures[future]
|
||
try:
|
||
worker_game_results, exp_file_path = future.result()
|
||
except Exception as e:
|
||
print(f"[警告] Worker {worker_id} 执行失败: {e}")
|
||
continue
|
||
|
||
# 记录对局收益
|
||
for game_idx, utility in worker_game_results:
|
||
traversing_player = game_idx % 2
|
||
game_results.append((traversing_player, utility))
|
||
|
||
# 从临时文件加载经验（避免 IPC 管道爆炸）
|
||
try:
|
||
experiences = torch.load(exp_file_path, map_location="cpu", weights_only=False)
|
||
except Exception as e:
|
||
print(f"[警告] Worker {worker_id} 经验文件加载失败: {e}")
|
||
continue
|
||
finally:
|
||
# 无论加载成功与否，立即删除临时文件，防止 /tmp 堆积
|
||
if os.path.isfile(exp_file_path):
|
||
os.remove(exp_file_path)
|
||
|
||
# 将 worker 返回的经验写入全局 Buffer
|
||
for info_state_list, legal_mask, regrets, strategy in experiences:
|
||
buffer.add(
|
||
info_state=info_state_list,
|
||
legal_mask=legal_mask,
|
||
regrets=regrets,
|
||
strategy=strategy,
|
||
)
|
||
|
||
total_experiences += len(experiences)
|
||
del experiences # 立即释放，避免主进程内存峰值
|
||
|
||
# 统计本 iteration 数据生成的概况
|
||
p0_utils = [u for p, u in game_results if p == 0]
|
||
p1_utils = [u for p, u in game_results if p == 1]
|
||
avg_p0 = sum(p0_utils) / len(p0_utils) if p0_utils else 0.0
|
||
avg_p1 = sum(p1_utils) / len(p1_utils) if p1_utils else 0.0
|
||
|
||
# ──────────── 阶段 B: 网络训练 — GPU ────────────
|
||
|
||
# 动态设备切换：训练阶段，cfr_net 移至 GPU 利用大 Batch 高效训练
|
||
cfr_net.to(DEVICE)
|
||
|
||
# 首次进入训练阶段时创建 optimizer（此时参数已在 GPU 上）
|
||
if optimizer is None:
|
||
optimizer = torch.optim.AdamW(
|
||
cfr_net.parameters(),
|
||
lr=LEARNING_RATE,
|
||
weight_decay=WEIGHT_DECAY,
|
||
)
|
||
print(f"[Optimizer] AdamW lr={LEARNING_RATE}, weight_decay={WEIGHT_DECAY}, device={DEVICE}")
|
||
|
||
avg_total_loss = 0.0
|
||
avg_regret_loss = 0.0
|
||
avg_policy_loss = 0.0
|
||
trained_steps = 0
|
||
|
||
if len(buffer) >= MIN_BUFFER_SIZE_FOR_TRAIN:
|
||
for step in range(TRAIN_STEPS_PER_ITER):
|
||
total_loss, regret_loss, policy_loss = train_step(
|
||
cfr_net, buffer, optimizer, DEVICE
|
||
)
|
||
avg_total_loss += total_loss
|
||
avg_regret_loss += regret_loss
|
||
avg_policy_loss += policy_loss
|
||
trained_steps += 1
|
||
|
||
avg_total_loss /= trained_steps
|
||
avg_regret_loss /= trained_steps
|
||
avg_policy_loss /= trained_steps
|
||
|
||
# ──────────── 打印日志 ────────────
|
||
|
||
buffer_size = len(buffer)
|
||
loss_str = (
|
||
f"Total={avg_total_loss:.6f} "
|
||
f"Regret={avg_regret_loss:.6f} "
|
||
f"Policy={avg_policy_loss:.6f}"
|
||
if trained_steps > 0
|
||
else "训练未开始（Buffer 不足）"
|
||
)
|
||
|
||
print(
|
||
f"[Iter {iteration:3d}/{NUM_ITERATIONS}] "
|
||
f"Buffer={buffer_size:6d} "
|
||
f"Exp={total_experiences:6d} "
|
||
f"P0_util={avg_p0:+.1f} P1_util={avg_p1:+.1f} "
|
||
f"Loss: {loss_str}"
|
||
)
|
||
|
||
# ──────────── 定时保存检查点 ────────────
|
||
|
||
if iteration % CHECKPOINT_INTERVAL == 0:
|
||
# 保存时网络在 GPU 上，save_checkpoint 会将权重转到 CPU
|
||
save_checkpoint(
|
||
iteration=iteration,
|
||
cfr_net=cfr_net,
|
||
buffer=buffer,
|
||
optimizer=optimizer,
|
||
checkpoint_dir=CHECKPOINT_DIR,
|
||
keep_last_n=KEEP_LAST_N_CHECKPOINTS,
|
||
)
|
||
|
||
# ── 8. 训练结束，保存最终检查点 ──
|
||
save_checkpoint(
|
||
iteration=NUM_ITERATIONS,
|
||
cfr_net=cfr_net,
|
||
buffer=buffer,
|
||
optimizer=optimizer,
|
||
checkpoint_dir=CHECKPOINT_DIR,
|
||
keep_last_n=KEEP_LAST_N_CHECKPOINTS,
|
||
)
|
||
# 同时保存一份到根目录（向后兼容）
|
||
save_path = os.path.join(_POKER_DIR, "cfr_net_checkpoint.pt")
|
||
torch.save({
|
||
"model_state_dict": {k: v.cpu() for k, v in cfr_net.state_dict().items()},
|
||
"optimizer_state_dict": optimizer.state_dict() if optimizer else {},
|
||
"iteration": NUM_ITERATIONS,
|
||
"buffer_size": len(buffer),
|
||
}, save_path)
|
||
print(f"[保存] 模型已保存到: {save_path}")
|
||
|
||
print(f"\n{'='*70}")
|
||
print(f" 训练完成！共 {NUM_ITERATIONS} iterations")
|
||
print(f" 最终 Buffer 大小: {len(buffer)}")
|
||
print(f"{'='*70}")
|
||
|
||
|
||
# ───────────────────── 入口 ─────────────────────
|
||
|
||
if __name__ == "__main__":
|
||
main()
|