跳棋项目小记

点击即玩：https://cc0.ylxdzsw.com

项目历史

项目开始于 2021 年 4 月，彼时 AlphaZero 还有一定热度，我也跟风学了学强化学习，打算做点东西练手，就搞这个没什么人搞过的跳棋。当时也花了不少精力，调来调去，也试了很多技术，然而最后训练出来的 AI 效果很差，根本就是乱走，我就放弃了，这个项目就搁置了。

一晃两年就过去了，疫情来了又走了，ChatGPT 犹如一股泥石流席卷了机器学习各个领域。我突然想起这个被遗忘的项目，决定把它补完。又调了一个星期，效果依然很差，不过总算是能跑起来了。这里记录一下这个项目中踩过的坑和学到的东西，给它盖上棺材板。

项目分成四部分，第一块是用 Rust 实现的游戏逻辑和搜索算法，分别编译成 .so 和 wasm 供另外两个部分使用。采用 Rust 是因为打算实现并发的纯搜索算法对战用来生成大量数据供模型学习。这块一定程度上成功了，在我的电脑上不到一天跑出了 10 亿回合，超过 20 多 G 的对战数据。然而，最后训练出来的 AI 依然是个智障。第二块是前端界面，作为我对这个项目的基本要求，整个游戏包括神经网络模型推理都要在不借助后端服务器的情况下完成，最终的结果放在 github page 上点击即玩。第三块是 Python 实现的模型训练，导出到前端使用。选择 Python 的理由是没得选。第四块是用来生成跳棋棋盘结构相关信息的一次性脚本，包括棋盘邻接矩阵，SVG 图等，用 Julia 实现的，支持多种大小的棋盘，这样可以先用小棋盘测试 AI，靠谱再训练大棋盘。

前端本来还有支持使用 peer.js 进行联机对战的宏伟计划，不过搁置后再重启的时候没这热情了，就删掉了。前端的推理这块，我先尝试了 tch-rs 期望能够静态连接 Emscripten 一块编译成 wasm。果然，折腾了大半天实在搞不定。然后我又尝试了 tensorflow.js，然而它却并跑不了 Keras 的 Transformer 模型。我本来都打算放弃了，最后挣扎着试了一下邪道的先 PyTorch 导出成 onnx 再用 onnxruntime-web。没想到出乎意料的简单，只遇到了 PyTorch 1.8 导出的模型没有问题，1.10 导出的会报错，以及直接 new 出来的模型能用，但是 torch.jit.script 过的模型不行，以及用 onnx opset 11 导出的能用，14 的不行三个小问题而已。而且这几个问题都是两年前遇到的，最新版本下已经都没有了。

跳棋实现

常见的围棋象棋都是方形棋盘，用一个二维数组就能轻易表示，而且也天然的适合 CNN 网络。然而跳棋是六角形格子，而且还不是一个完整的六边形，而是一个奇葩的星形结构。一个方式是采用斜的坐标系，但是这样坐标并不能很好的反映两个点之间的距离。我选择的方法是把棋盘上每个格子赋了一个独立的 id，然后用一个邻接矩阵存放这些格子之间的相邻关系。

一个标准大小的跳棋“局面”用 21 个字节就可以表示，第一个字节表示现在谁走，接着 10 个字节是先手玩家 (p1) 的棋子的格子 id，再后面是后手玩家的。得益于它小而且长度固定，我直接采用一个定长数组实现，用的时候也可以随便复制，大幅降低了代码复杂度。而且由于定长数组放在栈上，没有堆内存申请释放计数等开销，性能相当好。

跳棋特点分析

Complete Information: 跳棋双方可以看到完整的棋盘，因此我们上面的 21 个字节的信息对双方都是完全可知的。完全信息也意味着接下来的下法只取决于这 21 个字节的信息，而与游戏之前如何到达这个局面的历史无关。

Deterministic: 在同一个局面下，走同样的一步棋，总会造成同样的后续局面，不存在随机性。

Zero-sum: 游戏只有两方，且只有一方能赢，因此让自己获得优势等价于给对手创造劣势。我们对每个状态只需要记录先手玩家的胜率，每一步先手玩家会尝试最大化这个数，而后手玩家会选择最小化它。

Model-based: 严格来说不是跳棋本身的特点，而是我们采用的方法的类型。Model-based 指对于给定的一个局面，我们可以预知采取每一个动作会造成的后续局面（或者概率，如果不是 deterministic 的游戏）。换句话说，AI 在下棋的时候可以调用跳棋模拟器。

基本思路

综合上面的特点，可以发现跳棋其实比一般的强化学习问题要稍微简单一点，在于我们可以直接把动作 $a$ 的概念给去掉。AI 每一步不再是选择“怎么走”，而是直接选择“想要到达的后续局面”。强化学习里的几个关键变量，客观的包括状态 $s$ ，每个状态下可以进行的动作 $a$ ，每个状态对应的回报 $r_{s}$ ，以及在状态 $s$ 下采取动作 $a$ 可能到达的后续状态的转移概率表 $P (s, a)$ 。跟 AI 采取的策略相关的则包括 AI 在某个状态下选择某个动作的概率 $π (s, a)$ ，在这个策略下每个状态的价值 $V_{π} (s)$ ，以及每个状态下每个动作的价值 $Q_{π} (s, a)$ 。在我们的完全信息且确定性的游戏里，我们可以直接去掉 $P$ 和 $Q$ ，专注优化 $V$ 即可。 $V$ 可以用如下公式计算：

$V (s) = s^{'} max {r_{s^{'}} + γV (s^{'})}$

对于极小规模的游戏，动态规划（假设没有环）或者列表迭代算这个式子是可行的，然而对于有一定规模的游戏，比如标准棋盘的跳棋，按上面21个字节的方法计算，可能的状态有 $2 \times (10 121) \times (10 111) = 13042418121078104783313040752$ 种，显然迭代不完。解决这个问题基本的思路就是用神经网络参数化的函数 $f_{θ} (s)$ 来拟合 $V (s)$ 。原理是：1、不同的状态之间有一些有一定的相似性，例如两个状态 $s_{1}$ 和 $s_{2}$ 只相差一个无关紧要的棋子。神经网络可能可以发现这些相似性，这样在更新 $θ$ 来调整 $f_{θ} (s_{1})$ 的同时， $f_{θ} (s_{2})$ 也会进行相应的调整。对于迭代中没有遇到过的状态，也能根据与它相似的状态给出一个不错的估计。2、虽然状态空间很大，但是很多状态在合理的对局中根本不会出现，例如把子走到棋盘左右两边角落里。因此在迭代计算的时候，我们不是直接循环遍历或者等概率抽取所有状态来更新，而是通过模拟对局，只更新对局中出现过的状态，这样可以在有限的时间内优先拟合常见状态。

模型结构

我首先尝试了 Transformer，目前大有大一统之势的网络结构。我把棋盘格子 id 作为类似自然语言中的 token id，其中后手玩家的 id 在棋盘格子 id 的基础上再加棋盘大小，这样就能把先手玩家和后手玩家的棋子区分开。游戏下一步谁走这个信息也额外再增加两个 id 来表示。这样给模型的输入就是 21 个 token，用 embedding 映射到 256 维的向量，之后就直接调 PyTorch 的 TransformerEncoderLayer 层，最后分出两个头，分别输出策略和价值。策略是一个 121 * 121 的大矩阵，表示把棋盘格子 i 里的棋子走到 j 的概率权重。不合法的走法都在 softmax 前 mask 掉了。

后面发现效果不好，只好返璞归真，搬出了最基础的 MLP。输入是一个 1 + 2 * 121 的 0-1 向量，其中第一个元素表示下一步谁走，接下来 121 个元素是先手玩家的棋子，哪个格子上有哪个元素就是 1，其它的就是 0，这样这里面总是只有 10 个 1。后面的 121 个元素则是后手玩家的子。这个简单的网络最后效果比 Transformer 好些，虽然依然很菜。

第一次尝试

我先是严格照搬 AlphaZero，虽然上面分析了我们没有引入 $Q$ 的必要，但是为了防止在简化的过程中出错，我也还是照着搬了过来。然而，跑了一会儿我就意识到了问题：直接以随机状态开始对弈，也就是 AlphaZero 说的 tabula rasa，下几百万回合都走不到一个终局。原因是跳棋跟象棋、围棋不一样，围棋每一步往棋盘上上放子，越下越多，象棋越吃子越少，即使随机乱下也有很大概率可以下完。跳棋只是移动棋子，光靠随机游走把所有棋子移入对手的基地违反热力学定律。然后我想着加规则，到了 80 回合还没结束就强制结束，其中在目标区域棋子更多的玩家获胜。训练了很久，AI 依然是随机乱下。我认为是由于我 MCTS 的轮数太少，无法得出有效的结果来学习。然而每一轮 MCTS 都要进行几十次神经网络计算，实在算力不及谷歌，只好放弃。

训练时还遇到一个百思不得其解的事情，就是 AI 总是倾向于往左边角落里走。调了很久才意识到是我 MCTS expand 的时候没有打乱子节点的顺序，id 比较小的总是排在前面，在 PUCT 分数相同的情况下总是被优先选择，就造成了这个情况。

第二次尝试

搁置了两年后，我决定不再从头自我对弈，而是先用基本的 AI 收集 trace 去监督训练。我先设计了一个评估函数，计算每个子到对手的基地的角落格子的六角距离之和，然后用对手的距离减自己的距离就可以得到局面得分。借助这个得分，我实现了 greedy 和 alphabeta 两个搜索算法，寻找得分最高的后续局面。greedy 顾名思义，直接遍历当前局面的后续局面，挨个计算得分，然后走得分最高的那个。然而直接这样操作变数太少，每次都会走一样的路线，收集不到有意义的 trace。所以我又加了一点随机性：不是简单的选择最高得分，而是把后续局面的得分算一个 softmax，然后依照概率抽样。

alphabeta 则复杂一点，它的基本想法是越靠近终局，得分估计越准，因此在搜索的时候不光看下一步的得分，而是会再多走几步，优化最后一步的得分。因为要搜索的次数跟探索的深度 d 呈指数关系，所以其实也搜不了多长。对跳棋来说，10 个子的标准棋盘 d=3 就已经有点卡了，d=4 则需要很有耐心。alphabeta 会记录目前遇到的最优解来进行剪枝，在我看来和 ILP 的 branch-and-bound 极为相似，可以说就是对于零和游戏 min-max 结构的搜索树的特化版。

这两个 baseline 意外的强，我自己标准棋盘下不过 alphabeta。于是我就用它们进行了上千万次对局，然后用这些 trace 去训练 MLP。拟合效果还算可以，然而实际跑起来依然不尽人意。MCTS+Model 勉强能和 Alphabeta+heuristic 打个平手，而 Greedy+Model 和 Alphabeta+Model 让人看了血压暴涨，经常在对面基地门口逡巡不进。

其它一些尝试

我也试了一下增加一个规则，每一步棋必须让总距离和减小，也就是只能前进，不能平移或者后退。这样总算是可以让 AI 下到终局了，然而效果仍然一般。

我还试了一下纯 MCTS，设计了一个超大的搜索树，用一个 RWLock 来锁整棵树，然后每个节点上用 Atomic 来存访问计数和 p1 获胜计数。这样在 MCTS 搜索过程中，只有增加节点的时候需要拿 write-lock，只是更新计数的时候直接拿 read-lock 然后通过原子操作来更新，在 128 线程并发搜索下效果依然不错。然而即使搜到最后内存撑爆，还是搜不出什么靠谱的对局，基本都是 xjb 乱走。

发现 Transformer 拟合欠佳之后，我也试了一下 xgboost，效果意外不错，但是实在太慢，数据量一上去就跑不动了。此外它没办法在前端使用，虽然有一个基于 Ecmascripten 的移植项目，但是已经几年没维护了。

其它一些发现

Alphabeta 算法在无法取胜的情况下，会用一个棋子呆在自己基地里来阻止对面取胜，非常赖皮。这个问题估计得修改游戏规则，不过我没想好该怎么改。

处于劣势的时候基于模型的 AI 会倾向于走奇葩走位，比如走到两边的角落里。我认为原因是中间的常见走法在训练的时候已经发现处于劣势，而奇葩走位在训练的时候没见过，预测的胜率可能会是比较中庸的 50%，相比之下比常见走法更高。要解决这个问题可能需要改动训练的损失函数或者 regularization 方法，让模型倾向于对没见过的场面给出负面预测而不是 50% 胜率的预测。

纯 MCTS 在接近终盘的情况下，效果其实相当不错。

其它一些坑

我的搜索算法在 Rust 里实现，搜索的过程中需要调用神经网络来预测给定局面的得分。然而用的神经网络库 onnxruntime-web 只提供了 async 的 API，在 wasm 里没法简单的调用。另一方面，Rust 暂时没有类似 Python 或者 Zig 那样随时 yield 返回控制权的方法，因此只能改写搜索算法，让它接受一个得分表。每次需要评估局面得分的时候就去表里查，查不到直接返回缺哪些，前端运行神经网络之后填进表里重新再调用一次。这样虽然性能差了一点，需要重复许多操作，不过至少不用手写 continuation。

其它一些想法

直接预测给定局面的胜率可能有点困难，一个另外的想法是用模型预测 heuristic 给出的判断对不对，类似 boosting 那样去拟合 heuristic 的 redisual。这样对于没有见过的局面，模型应该会倾向于不推翻 heuristic 的预测，因此即使在训练不到位的情况下也有一个不错的默认值。