从零开始：用Python实现AlphaZero五子棋AI（附完整代码）

从零开始用Python实现AlphaZero五子棋AI附完整代码最近几年深度强化学习在游戏AI领域取得了令人瞩目的成就从AlphaGo到AlphaZero这些里程碑式的项目不仅展示了算法的强大也极大地降低了我们普通人复现顶尖技术的门槛。很多开发者朋友都跃跃欲试想亲手搭建一个能和自己对弈的智能体但面对复杂的理论公式和庞大的代码库往往不知从何下手。这篇文章我就想和你聊聊如何抛开那些让人望而生畏的数学推导直接用Python把AlphaZero的核心思想“捏”出来最终得到一个能运行、能学习、能和你下五子棋的AI。我们的目标非常明确不追求在理论层面做到极致还原而是聚焦于工程实现。我会带你一步步构建蒙特卡洛树搜索MCTS、一个轻量化的残差网络ResNet以及自博弈Self-Play的训练循环。整个过程就像搭积木我们会先理解每一块积木的作用然后用代码把它们拼装起来。你不需要是强化学习专家只要对Python和PyTorch有基本了解就能跟上。最终你会得到一套完整的、可运行的代码并且能清晰地看到AI是如何从一张“白纸”通过自我对弈逐渐成长为一名“五子棋高手”的。这不仅是学习一个算法更是一次完整的AI项目实战体验。1. 环境准备与项目框架搭建在动手写核心算法之前我们需要先把“舞台”搭好。一个清晰的项目结构能让你在后续的编码和调试中事半功倍。这里我们不使用任何复杂的框架就用最基础的Python库来构建一切。首先确保你的开发环境已经安装了必要的依赖。我强烈建议使用Anaconda创建一个独立的虚拟环境避免包版本冲突。# 创建并激活一个名为alphazero_gomoku的conda环境 conda create -n alphazero_gomoku python3.8 conda activate alphazero_gomoku # 安装核心依赖 pip install torch torchvision torchaudio pip install numpy pip install matplotlib # 用于可视化训练过程接下来我们来规划项目的目录结构。在你的工作区新建一个文件夹比如alphazero_gomoku并在里面创建如下几个文件alphazero_gomoku/ ├── game.py # 定义五子棋游戏规则和逻辑 ├── mcts.py # 实现蒙特卡洛树搜索 ├── network.py # 定义策略价值网络ResNet ├── self_play.py # 实现自博弈数据收集 ├── train.py # 主训练脚本 └── evaluate.py # 模型评估与对弈脚本这个结构非常直观每个文件职责单一。game.py是基础它定义了AI所要交互的世界——五子棋的棋盘、落子规则、胜负判定。我们先从这里开始。一个简单的五子棋游戏类需要包含以下核心方法get_init_board(): 返回初始的空棋盘状态。get_valid_moves(board, player): 获取当前玩家所有合法的落子位置在空位落子。get_next_state(board, player, action): 执行落子动作返回新的棋盘状态和下一个玩家。get_game_ended(board, player): 判断游戏是否结束返回胜利1、失败-1、平局0或未结束None。get_canonical_form(board, player): 将棋盘转换为当前玩家的视角对于零和博弈很重要。在实现时棋盘可以用一个n x n的二维数组比如n15表示用1代表黑子-1代表白子0代表空位。胜负判断就是检查横、竖、斜四个方向是否有连续五个同色棋子。注意在实现get_canonical_form时一个常见的技巧是无论当前是黑子还是白子都统一从“当前玩家”的视角来看棋盘。即如果当前是白子值为-1走棋就把整个棋盘乘以-1这样白子就变成了1黑子变成了-1网络总是学习“我”值为1如何对抗“对手”值为-1。2. 构建策略价值网络轻量级ResNetAlphaZero的强大很大程度上源于其使用的深度残差网络ResNet。它能够从原始的棋盘状态中同时提取出“局面好坏”价值和“走哪里好”策略两种信息。对于五子棋这种中等复杂度的游戏我们不需要像原论文那样用几十甚至上百层的网络一个几层的小型ResNet就足够了。我们的网络输入是棋盘的“图像”。对于15x15的五子棋我们可以构造一个[2, 15, 15]的张量作为输入通道。第一个通道是“我的棋子”位置值为1第二个通道是“对手的棋子”位置值为1。这种表示比单纯一个通道用1和-1表示更有利于卷积网络学习。网络结构大致分为三部分特征提取主干由若干个残差块Residual Block堆叠而成。每个残差块通常包含两个卷积层、批归一化BatchNorm和ReLU激活函数并通过一个捷径连接Shortcut Connection相加。策略头Policy Head接在主干之后通过卷积和全连接层输出一个长度为15*151的向量对应所有棋盘位置加上一个“放弃”动作的概率分布。这里使用Softmax进行归一化。价值头Value Head同样接在主干之后通过卷积、全连接层和Tanh激活函数输出一个标量范围在[-1, 1]之间表示当前玩家视角下的胜率估计1为必胜-1为必败。在network.py中我们可以这样定义网络import torch import torch.nn as nn import torch.nn.functional as F class ResidualBlock(nn.Module): def __init__(self, num_channels): super().__init__() self.conv1 nn.Conv2d(num_channels, num_channels, kernel_size3, padding1) self.bn1 nn.BatchNorm2d(num_channels) self.conv2 nn.Conv2d(num_channels, num_channels, kernel_size3, padding1) self.bn2 nn.BatchNorm2d(num_channels) def forward(self, x): residual x out F.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) out residual out F.relu(out) return out class AlphaZeroNet(nn.Module): def __init__(self, board_size15, num_res_blocks5, num_channels128): super().__init__() self.board_size board_size # 初始卷积层 self.conv_input nn.Conv2d(2, num_channels, kernel_size3, padding1) self.bn_input nn.BatchNorm2d(num_channels) # 残差塔 self.res_blocks nn.ModuleList([ResidualBlock(num_channels) for _ in range(num_res_blocks)]) # 策略头 self.conv_policy nn.Conv2d(num_channels, 2, kernel_size1) self.bn_policy nn.BatchNorm2d(2) self.fc_policy nn.Linear(2 * board_size * board_size, board_size * board_size) # 价值头 self.conv_value nn.Conv2d(num_channels, 1, kernel_size1) self.bn_value nn.BatchNorm2d(1) self.fc_value1 nn.Linear(board_size * board_size, 64) self.fc_value2 nn.Linear(64, 1) def forward(self, x): # 特征提取 s F.relu(self.bn_input(self.conv_input(x))) for block in self.res_blocks: s block(s) # 策略输出 p F.relu(self.bn_policy(self.conv_policy(s))) p p.view(-1, 2 * self.board_size * self.board_size) p self.fc_policy(p) p F.log_softmax(p, dim1) # 输出log概率便于计算交叉熵损失 # 价值输出 v F.relu(self.bn_value(self.conv_value(s))) v v.view(-1, self.board_size * self.board_size) v F.relu(self.fc_value1(v)) v torch.tanh(self.fc_value2(v)) # 输出范围[-1, 1] return p, v这个网络有几个关键设计点批归一化BatchNorm加速训练并提升稳定性在残差网络中几乎是标配。策略头输出Log-Softmax这与后续使用的负对数似然损失NLLLoss相匹配是分类问题的标准做法。价值头输出Tanh将价值约束在[-1, 1]之间符合棋类游戏胜负的数值表示。网络准备好后我们就可以进入AlphaZero的“大脑”——蒙特卡洛树搜索了。3. 实现蒙特卡洛树搜索MCTSMCTS是AlphaZero进行决策的核心引擎。它不像传统搜索算法那样暴力计算所有可能性而是通过模拟Simulation的方式有选择地探索更有希望的行棋路径并不断更新对每个位置的评估。这个过程可以类比为一个人在下棋时的“心算”先在脑子里尝试几种走法看看每种走法后续可能的发展然后选择看起来最好的那一个。MCTS的一次模拟包含四个步骤选择Selection、扩展Expansion、模拟Simulation在AlphaZero中由网络完成、回溯Backpropagation。我们会为树中的每个节点存储以下信息属性名数据类型描述statenumpy array该节点对应的棋盘状态规范形式。parentMCTSNode指向父节点的引用。childrendict字典键为动作值为对应的子节点对象。visit_countint该节点被访问的总次数N。total_valuefloat所有模拟回溯到该节点的价值总和W。prior_probfloat先验概率P由策略网络给出。在mcts.py中我们首先定义节点类然后实现搜索过程。选择Selection从根节点开始递归地选择子节点直到到达一个尚未完全展开的节点叶子节点或未探索的节点。选择的标准是PUCTPolynomial Upper Confidence Trees算法它平衡了利用Exploitation和探索Exploration。def select_child(self, c_puct1.0): 根据PUCT公式选择最优的子节点。 best_score -float(inf) best_action None best_child None for action, child in self.children.items(): # PUCT公式: Q c_puct * P * sqrt(N_parent) / (1 N_child) # Q child.total_value / (child.visit_count 1e-8) q_value child.total_value / (child.visit_count 1e-8) u_value c_puct * child.prior_prob * math.sqrt(self.visit_count) / (child.visit_count 1) score q_value u_value if score best_score: best_score score best_action action best_child child return best_action, best_child公式中的c_puct是一个超参数控制探索的强度。q_value代表了该动作的历史平均胜率利用u_value则鼓励访问次数少的动作探索。扩展Expansion与模拟Simulation当选择过程到达一个叶子节点游戏未结束且该节点有未尝试的合法动作时我们进行扩展。从策略网络获取该状态所有合法动作的先验概率P为每个合法动作创建一个新的子节点并将其prior_prob初始化为网络给出的概率。然后我们使用价值网络V对这个新状态进行一次评估得到这个叶子节点的价值v。注意在AlphaZero中我们不再进行传统的随机模拟直到终局而是用神经网络的一次前向传播来替代这极大地提高了效率。回溯Backpropagation拿到叶子节点的价值v后我们沿着选择路径反向更新所有祖先节点的信息。这里有一个关键点在零和博弈中从父节点到子节点玩家发生了切换。因此回溯的价值需要取反以始终表示当前节点玩家的视角价值。def backpropagate(self, value): 从当前节点开始向上回溯更新祖先节点的访问次数和总价值。 node self while node is not None: node.visit_count 1 node.total_value value # 这里的value已经是当前节点玩家的视角价值 value -value # 对于父节点价值视角相反 node node.parent完成一次完整的MCTS搜索意味着我们从根节点开始执行了固定次数比如800次的“选择-扩展-模拟-回溯”循环。最终根节点下每个动作的访问次数N就构成了一个改进后的策略分布π。我们通常根据这个分布经过温度参数τ调整来采样最终的动作或者在训练时直接将其作为策略网络的训练目标。4. 自博弈循环与模型训练有了游戏环境、神经网络和MCTS我们就可以让AI自己和自己下棋Self-Play并利用产生的数据来训练网络了。这是整个项目最核心的循环也是AI从“无知”到“强大”的学习过程。自博弈和数据收集的流程如下初始化重置游戏到初始状态将当前状态s设为MCTS的根节点。单步决策以当前状态s为根运行固定次数的MCTS模拟例如800次。MCTS过程中需要策略价值网络为新的叶子节点提供先验概率P和价值v。模拟结束后根据根节点各子节点的访问次数计算策略分布ππ(a) N(a)^(1/τ) / sum(N(b)^(1/τ))τ为温度参数训练早期τ1增加探索后期τ-0趋于贪婪。根据分布π采样一个动作a执行早期可以加入狄利克雷噪声以增加探索。记录三元组(s, π, z)其中z在游戏结束时才被确定胜1负-1平0。状态推进执行动作a得到新状态s并将MCTS的根节点移动到对应于s的节点如果该节点已存在于树中可以复用其子树大幅提升效率。循环与存储重复步骤2-3直到游戏结束。将本局游戏中记录的所有(s, π, z)存储到经验回放缓冲区Replay Buffer中。在self_play.py中实现这个循环。经验回放缓冲区通常有容量限制当满的时候会丢弃最旧的数据这保证了训练数据的新鲜度。接下来是训练部分。我们从经验回放缓冲区中随机采样一个批次Batch的数据(s, π, z)然后计算损失函数来更新网络参数。AlphaZero的损失函数由三部分组成def compute_loss(policy_pred, value_pred, target_pi, target_z): # 策略损失预测策略分布与MCTS搜索结果的交叉熵 policy_loss -torch.sum(target_pi * policy_pred) / target_pi.size(0) # 价值损失预测价值与最终胜负结果的均方误差 value_loss F.mse_loss(value_pred, target_z) # 总损失 total_loss policy_loss value_loss return total_loss, policy_loss, value_loss这里policy_pred是网络输出的log概率所以直接用负对数似然计算交叉熵。target_pi是MCTS产生的策略分布访问次数归一化。target_z是游戏结果。训练循环train.py的大致结构如下# 初始化网络、优化器、经验缓冲区等 model AlphaZeroNet() optimizer torch.optim.Adam(model.parameters(), lr0.001, weight_decay1e-4) replay_buffer ReplayBuffer(capacity100000) for iteration in range(num_iterations): # 阶段1自博弈收集数据 print(fIteration {iteration}: Self-Playing...) new_data self_play(model, num_games100) # 自博弈100局 replay_buffer.add(new_data) # 阶段2训练网络 print(fIteration {iteration}: Training...) for epoch in range(num_training_epochs): # 从缓冲区采样批次数据 batch_states, batch_target_pis, batch_target_zs replay_buffer.sample(batch_size512) # 前向传播 policy_pred, value_pred model(batch_states) # 计算损失 total_loss, policy_loss, value_loss compute_loss(policy_pred, value_pred, batch_target_pis, batch_target_zs) # 反向传播与优化 optimizer.zero_grad() total_loss.backward() optimizer.step() # 可以记录损失值用于监控 # 可选阶段3评估新模型 if iteration % eval_interval 0: win_rate evaluate(model, previous_model) # 与上一代模型对弈 if win_rate threshold: previous_model model # 更新为新的最佳模型这个“自博弈-训练”的循环会不断进行。随着网络能力的提升MCTS的模拟质量会更高产生的数据(π, z)也会更准确从而进一步训练出更好的网络。这是一个典型的“数据驱动模型模型优化数据”的正反馈过程。5. 实战调试与性能优化技巧理论实现和实际跑通之间往往隔着无数个Bug和性能瓶颈。在这一部分我想分享几个在实现过程中容易遇到的问题和优化策略这些都是我踩过坑之后总结出来的经验。常见问题与调试方法训练不收敛价值预测始终接近0检查点首先确认游戏逻辑game.py的胜负判断是否正确。可以写一些单元测试模拟一些必胜、必败的棋局看get_game_ended的返回值是否符合预期。检查点确认MCTS的回溯逻辑是否正确特别是价值取反value -value那一步。一个简单的测试是让AI在必胜局面下运行MCTS看根节点下最佳动作的Q值是否接近1。检查点检查经验缓冲区中存储的target_z游戏结果是否正确。确保在存储时z是从状态s对应的当前玩家视角记录的胜利1或失败-1。策略网络输出过于均匀没有重点可能原因MCTS的探索参数c_puct设置过大导致搜索过于分散访问次数N分布均匀。可以尝试适当调小c_puct。可能原因在自博弈的早期策略先验P加入了过多的狄利克雷噪声。可以随着训练进行逐渐减少噪声的强度。调试方法可视化MCTS搜索后的策略分布π。在某个典型的中盘局面打印出访问次数最多的前5个动作及其概率看是否集中在人类直觉上的“好点”附近。训练速度慢性能瓶颈往往在MCTS。一次搜索要进行数百次网络前向传播每次扩展和模拟都需要。确保在mcts.py中对同一个状态S网络前向传播只执行一次并将结果(P, v)缓存起来避免重复计算。使用GPU确保你的PyTorch安装了CUDA版本并将模型.to(device)。MCTS中的批量预测batch_size 1可以显著提升效率。例如在一次搜索中将所有需要评估的叶子状态收集成一个批次一次性送入网络而不是逐个评估。复用搜索树在自博弈的每一步新的根节点往往是上一步某个子节点。保留这棵子树而不是从头开始建树可以节省大量模拟次数。一些提升AI强度的实用技巧数据增强五子棋棋盘具有旋转和翻转的对称性。在将数据(s, π, z)存入缓冲区前可以对状态s和对应的策略π需要同步变换动作索引进行随机旋转和翻转能有效增加数据多样性提升泛化能力。异步训练这是一个高级优化。可以设计一个主进程负责训练网络多个工作进程并行进行自博弈。工作进程定期从主进程拉取最新的网络参数主进程则持续收集工作进程产生的数据并更新网络。这能极大加快数据收集速度。温度调度Temperature Scheduling在自博弈中前期使用较高的温度如τ1采样动作鼓励探索随着训练进行逐渐降低温度如τ0.1使AI的选择更确定、更贪婪。在最终对弈时通常设置τ0直接选择访问次数最多的动作。定期评估与保存像上面的训练循环所示定期让最新训练的模型与之前保存的最佳模型进行一定数量的对弈比如100局。如果胜率达到一定比例如55%就用新模型替换旧的最佳模型。这保证了我们始终使用一个稳定的强模型进行自博弈数据收集避免因模型突然变差而产生低质量数据。实现一个完整的AlphaZero项目是一次充满挑战但收获巨大的旅程。它迫使你去深入理解MCTS的每个步骤、神经网络的输入输出、以及自博弈这个巧妙的训练范式。当你第一次看到自己训练的AI开始走出一些有模有样的棋形甚至能防住你的“三三”进攻时那种成就感是无与伦比的。代码的细节很多关键在于动手去写去调试。我提供的代码框架和思路是一个起点你可以在此基础上尝试不同的网络结构、调整超参数、甚至修改游戏规则比如尝试六子棋观察AI行为的变化这本身就是最有趣的学习过程。