林子豪的 PKM
← 返回 Blog

AlphaZero 的搜索：MCTS

How Monte Carlo Tree Search connects game states, legal moves, policy priors, and value estimates.

8 min read系列 · game-ai
alphazeromctssearchdraughts

上一篇里的 bitboard 解决的是局面表示和走法生成：哪些格子被占住，哪些棋子能走，哪些棋子能吃，king 的 ray 要怎么截断。它让程序可以很快回答一个局面里的规则问题。

MCTS 接上的问题不一样。它不关心一个棋子具体是通过 << 7 还是 ray table 算出来的；它只关心：给定当前局面，哪些动作合法，走了某个动作以后局面变成什么，以及这个新局面大概对谁更好。

所以从 MCTS 开始，bitboard 会退到实现层。搜索层看到的是一个游戏接口：

legal_moves(state)       -> moves
next_state(state, move)  -> state'
terminal_value(state)    -> nil or value
encode(state)            -> neural-network input
move_index(move)         -> policy output index

这里的 state 可以内部用 bitboard 存；move 也可以是 source、landing、capture path 这些结构。但 MCTS 本身不应该写 draughts 的规则。它只负责在这些合法动作之间分配搜索预算。

Search tree

MCTS 建的是一棵局面树。root 是当前真实局面；从 root 出发，每条边是一种合法动作；动作后面的 child node 是走完这步以后的新局面。

在普通代码里，可以把统计量存在 node 上，也可以存在 edge 上。AlphaZero 的写法更自然地把统计量理解成 state, action 这一条边上的数据：

N(s,a) 是从局面 s 选择动作 a 的访问次数。W(s,a) 是这条边累计回传的价值总和。Q(s,a) 是平均价值：

Q(s,a) = \frac{W(s,a)}{N(s,a)}

P(s,a) 是神经网络一开始给这个动作的先验概率。

这些量会随着每次 simulation 更新。搜索刚开始时，P 来自神经网络；搜索跑得越多，N 和 Q 越能反映树搜索后的判断。

下面这个图不是一个完整的 draughts 规则树，而是一个 MCTS 统计树。P,Q,N,U,\operatorname{score} 都是 edge 统计量；它们表示的是从 parent 选择这条 move 之后得到的统计结果。图里的 edge 粗细按同一个 parent 下的 \operatorname{score} 相对大小归一化，不能跨 parent 全局比较。

{:title "32 次 simulation 之后的 root"
 :description "每个 child 显示的是 parent 选择这条 move 时的 edge 统计量。"
 :edge-metric :score
 :edge-note "边的粗细只按同一个 parent 下的 score 相对大小归一化；它表示局部选择偏好，不能跨 parent 全局比较。"
 :highlight-path ["a3-b4" "b6-c5"]
 :tree ["root"
        {:node {:label "s0"
                :role "root"
                :n 32
                :note "Current real position"}
         :moves {"a3-b4"
                 {:node {:label "s1"
                         :role "expanded"}
                  :edge {:n 18
                         :q 0.42
                         :p 0.31
                         :u 0.08
                         :score 0.50}
                  :moves {"b6-a5"
                          {:node {:label "s3"
                                  :role "leaf"}
                           :edge {:n 12
                                  :q -0.18
                                  :p 0.52
                                  :u 0.06
                                  :score -0.12}
                           :moves {}}
                          "b6-c5"
                          {:node {:label "s4"
                                  :role "leaf"}
                           :edge {:n 6
                                  :q 0.06
                                  :p 0.21
                                  :u 0.14
                                  :score 0.20}
                           :moves {}}}}
                 "c3-d4"
                 {:node {:label "s2"
                         :role "expanded"}
                  :edge {:n 9
                         :q 0.17
                         :p 0.44
                         :u 0.15
                         :score 0.32}
                  :moves {"f6-e5"
                          {:node {:label "s5"
                                  :role "leaf"}
                           :edge {:n 9
                                  :q -0.03
                                  :p 0.37
                                  :u 0.19
                                  :score 0.16}
                           :moves {}}}}
                 "g3-h4"
                 {:node {:label "s6"
                         :role "candidate"}
                  :edge {:n 5
                         :q 0.05
                         :p 0.08
                         :u 0.25
                         :score 0.30}
                  :moves {}}}}]
 :layout-opts {:node-width 188
               :node-height 126
               :column-gap 76
               :row-gap 24
               :padding 36}}

图中高亮的 a3-b4 -> b6-c5 可以理解成某一次 simulation 选中的 path。等 leaf value 算出来以后，backup 阶段更新的也是这条 path 上的 edge。

这里有一个容易混淆的点：MCTS 最后真正用来选动作的，通常不是原始神经网络的 P，也不是单纯的 Q，而是搜索以后 root 下面每个动作的访问次数 N。神经网络先给一个方向，MCTS 再用树搜索把这个方向修正成一个更可靠的动作分布。

One simulation

一次 MCTS simulation 可以理解成从 root 往下走一条路径，然后把 leaf 的评估结果传回 root。

它不是每次都从头枚举整棵树。每次只走一条路径。重复很多次以后，访问次数自然会集中到更有希望的分支上。

伪代码大概是：

def run_simulation(root):
    node = root
    state = root.state
    path = []

    while node.expanded and terminal_value(state) is None:
        move, edge = select_child(node)
        path.append(edge)

        state = next_state(state, move)
        node = edge.child

    value = terminal_value(state)

    if value is None:
        policy, value = neural_net(encode(state))
        moves = legal_moves(state)
        expand(node, moves, policy)

    backup(path, value)

这段代码里，select_child 是 MCTS 的核心：已经展开过的节点，下一步应该走哪个 child？expand 是把一个没见过的局面展开成若干合法动作。backup 是把最后得到的 value 传回路径上的每条边。

经典 MCTS 在 leaf 之后会随机 rollout 到终局，再用胜负当作 value。AlphaZero 不做随机 rollout。它直接用神经网络评估 leaf：

policy, value = neural_net(encode(state))

policy 负责告诉搜索“哪些动作看起来值得试”。value 负责告诉搜索“这个局面对当前 player-to-move 大概有多好”。

Selection

Selection 要在两个目标之间折中。

一个目标是 exploitation：已经搜过很多次、平均结果很好的动作，应该继续看。这个部分由 Q 表示。

另一个目标是 exploration：一个动作可能还没搜过几次，但神经网络觉得它很有希望，也应该给机会。这个部分由 P 和访问次数控制。

AlphaZero 用的选择公式通常写成 PUCT：

\begin{aligned}
\operatorname{score}(s,a) &= Q(s,a) + U(s,a) \\
U(s,a) &= c_{\mathrm{puct}}\,P(s,a)\,
\frac{\sqrt{\sum_{b} N(s,b)}}{1 + N(s,a)}
\end{aligned}

每次从当前节点往下走时，就选择 \operatorname{score}(s,a) 最大的动作。

Q(s,a) 越大，说明这个动作过去搜索出来的平均结果越好。P(s,a) 越大，说明神经网络一开始更看好这个动作。N(s,a) 越小，分母越小，探索项越大；同一个动作被访问得越多，探索奖励就会下降。

\sum_{b} N(s,b) 是当前节点下所有动作的总访问次数。父节点被访问得越多，搜索越有理由继续扩展一些还没试够的动作。

c_{\mathrm{puct}} 控制探索强度。它越大，搜索越愿意相信 policy prior，给低访问动作更多机会；它越小，搜索越偏向已经有较高 Q 的动作。

Expansion and policy mask

Selection 会停在一个还没展开过的局面，或者一个终局。

如果是终局，就不需要问神经网络。游戏规则可以直接给出结果：

terminal_value(state) -> -1, 0, or 1

这里的 value 是从当前 player-to-move 的视角看。如果当前轮到的一方已经输了，value 就是 -1。如果是平局，就是 0。

如果不是终局，就调用神经网络：

policy, value = neural_net(encode(state))

然后用游戏规则列出合法动作：

moves = legal_moves(state)

神经网络的 policy head 通常是固定大小的向量，但每个局面的合法动作数量不同。所以展开节点时，需要把非法动作 mask 掉，只保留当前局面的合法动作：

priors = {}

for move in legal_moves(state):
    index = move_index(move)
    priors[move] = policy[index]

normalize(priors)

这就是 move_index 出现的地方。MCTS 可以处理任意结构的 move，但神经网络输出必须是固定维度。两者之间需要一个稳定映射。

在 draughts 里，这个设计会遇到一个具体问题：连续吃子算一个动作，还是每一跳算一个动作？

更简单的做法是让 legal_moves(state) 返回完整的一回合动作。普通走子是 source 到 target；吃子则是一条完整 capture path。这样 MCTS 的每条边都对应“当前玩家做完一次决策后，轮到对手”。回传 value 时，每走一层就变号。

如果把连续吃子的每一跳都暴露给 MCTS，也可以做，但 state 里必须额外表示“当前仍然是同一个棋子继续吃，不能切换到其他棋子，也还没有轮到对手”。这会让搜索树里的 turn semantics 变复杂。除非有很强的理由，先把完整回合作为 action 会更干净。

Backup

神经网络给出的 value 是 leaf state 对当前 player-to-move 的价值。可是路径往上走一层，player-to-move 会换成对手。对手觉得好的局面，对上一层玩家就是坏的。

所以在双人零和、双方轮流行动的游戏里，backup 时每回到上一层都要变号：

def backup(path, leaf_value):
    value = leaf_value

    for edge in reversed(path):
        value = -value

        edge.N += 1
        edge.W += value
        edge.Q = edge.W / edge.N

假设 leaf 的 value 是 +0.8，意思是 leaf 局面对 leaf 当前要走的一方很好。那么它对上一手刚走完的人就是 -0.8，再往上一层又变回 +0.8。

这个符号约定非常重要。Q(s,a) 必须表示“在 state s 里，当前 player-to-move 选择 action a 的平均价值”。只要一直保持这个约定，selection 里的 Q + U 才是在为当前要走的人选好动作。

Choosing the real move

一次 simulation 只更新一条路径。实际下棋前，会对同一个 root 跑很多次 simulation，例如几百次或几千次。

跑完以后，root 下每个合法动作都有一个访问次数：

N(s_0,a)

这些访问次数就是 MCTS 产出的改进策略。训练时，AlphaZero 不只是选择访问次数最多的动作；它会把访问次数转成一个 policy target：

\pi(a \mid s_0) =
\frac{N(s_0,a)^{1/\tau}}
{\sum_{b} N(s_0,b)^{1/\tau}}

\tau 是 temperature。\tau 大时，访问次数分布会更平，采样更有探索性。\tau 接近 0 时，分布会集中到访问次数最多的动作，接近 \arg\max。

自对弈早期通常需要探索，所以会按 \pi 采样。真正对弈或自对弈后期，可以直接选访问次数最高的动作：

a^* = \arg\max_a N(s_0,a)

这也是 AlphaZero 里很关键的一步：神经网络原始输出的是 P，搜索以后得到的是 \pi。训练时让神经网络去拟合 \pi，相当于让网络学习“经过搜索修正后的自己”。

The loop

AlphaZero 的 MCTS 可以看成一个 policy improvement operator。

神经网络先给出粗略判断：

f_\theta(s) = (p, v)

MCTS 用游戏规则和 PUCT 搜索，把这个粗略 policy 改进成 root visit distribution：

\operatorname{MCTS}(s, f_\theta) \rightarrow \pi

自对弈执行动作，直到终局，得到最终结果 z。训练样本就是：

(\operatorname{encode}(s), \pi, z)

下一轮训练让神经网络同时学习两件事：

\begin{aligned}
\text{policy head} &\rightarrow \pi \\
\text{value head} &\rightarrow z
\end{aligned}

所以 AlphaZero 不是单纯地“神经网络下棋”。更准确地说，是神经网络和搜索互相增强：神经网络给 MCTS 直觉，MCTS 用规则和模拟预算修正直觉，然后把修正后的结果变成新的训练目标。

Bitboard 在这个循环里的位置也就清楚了。它不是 AlphaZero 的核心算法，但它决定了 legal_moves 和 next_state 能不能足够快。MCTS 会反复调用这两个函数；如果走法生成慢，搜索次数就上不去；搜索次数上不去，\pi 的质量就会差。

下一步再看神经网络时，重点就不是“网络怎么直接选一步棋”，而是两个输出头分别服务什么：policy head 给 MCTS prior，value head 替代随机 rollout。搜索树和神经网络之间的接口，就是 AlphaZero 算法真正闭合的地方。

One simulation as steps

{:title "一轮 simulation 的分步视图"
 :description "这个示例运行实际 MCTS：Tic-Tac-Toe game protocol + uniform evaluator。"
 :game :tic-tac-toe
 :evaluator :uniform
 :state {:board [:x :x nil
                 :o :o nil
                 nil nil nil]
         :turn :x}
 :initial-simulations 1
 :cpuct 1.5
 :temperature 1.0
 :edge-metric :score
 :edge-note "Edge width 按同一父节点下的 PUCT score 归一化；按“下一轮”可以继续增加 simulation。"
 :layout-opts {:node-width 188
               :node-height 126
               :column-gap 76
               :row-gap 24
               :padding 36}}