AlphaZero實戰：從零學下五子棋（附代碼）

2個多月前，AlphaGo Zero橫空出世，完全從零開始，僅通過自我對弈就能天下無敵，瞬間刷爆朋友圈，各路大神分分出來解讀，驚嘆於其思想的簡單、效果的神奇。很快就有大神放出了開源版的AlphaGo Zero，但是只有代碼，沒有訓練出來的模型，因為據大神推算，在普通消費級的電腦上想訓練出AlphaGo Zero的模型需要1700年！然而DeepMind在AlphaGo Zero的論文里只強調運行的時候需要4個TPU，而完全沒有提及訓練過程的最大計算需求在於生成self-play數據，還引起了一點小爭議。還好，過了不到兩個月，在12月初，DeepMind就在Arxiv上低調放出了更加通用的AlphaZero的論文。AlphaZero幾個小時就征服圍棋、國際象棋和日本將棋的壯舉再次驚嘆世人，但同時DeepMind大方公開的self-play階段使用的5000個TPU也讓大家紛紛感嘆，原來是「貧窮限制了我們的想像力」！

扯得有點遠了，讓我們回到這篇文章的正題：AlphaZero實戰，通過自己動手從零訓練一個AI，去體會AlphaZero自我對弈學習成功背後的關鍵思想和一些重要技術細節。這邊選擇了五子棋作為實踐對象，因為五子棋相對比較簡單，大家也都比較熟悉，這樣我們能更專註於AlphaZero的訓練過程，同時也能通過親自對陣，來感受自己訓練出來的AI慢慢變強的過程。經過實踐發現，對於在6*6的棋盤上下4子棋這種情況，大約通過500~1000局的self-play訓練（2小時），就能訓練出比較靠譜的AI；對於在8*8的棋盤上下5子棋這種情況，通過大約2000~3000局自我對弈訓練（2天），也能得到比較靠譜的AI。所以雖然貧窮，但我們還是可以去親身感受最前沿成果的魅力！完整代碼以及4個訓練好的模型已經上傳到了github：https://github.com/junxiaosong/AlphaZero_Gomoku

我們先來看兩局訓練好的AI模型（3000局self-play訓練得到）對弈的情況，簡單感受一下：

從上面的對局樣例可以看到，AI已經學會了怎麼下五子棋，知道什麼時候要去堵，怎麼樣才能贏，按我自己對陣AI的感受來說，要贏AI已經不容易了，經常會打平，有時候稍不留神就會輸掉。這裡有一點需要說明，上面展示的兩局AI對弈中，AI執行每一步棋的時候分別只執行了400次和800次MCTS模擬，進一步增大模擬次數能夠顯著增強AI的實力，參見AlphaZero論文中的Figure 2（註：AlphaZero在訓練的時候每一步只執行800次MCTS simulations，但在評估性能的時候每一步棋都會執行幾十萬甚至上百萬次MCTS模擬）。

下面，我結合AlphaZero演算法本身，以及github上的具體實現，從自我對局和策略價值網路訓練兩個方面來展開介紹一下整個訓練過程，以及自己實驗過程中的一些觀察和體會。

自我對局（self-play）

完全基於self-play來學習進化是AlphaZero的最大賣點，也是整個訓練過程中最關鍵也是最耗時的環節。這裡有幾個關鍵點需要說明：

1. 使用哪個模型來生成self-play數據？

在AlphaGo Zero版本中，我們需要同時保存當前最新的模型和通過評估得到的歷史最優的模型，self-play數據始終由最優模型生成，用於不斷訓練更新當前最新的模型，然後每隔一段時間評估當前最新模型和最優模型的優劣，決定是否更新歷史最優模型。而到了AlphaZero版本中，這一過程得到簡化，我們只保存當前最新模型，self-play數據直接由當前最新模型生成，並用於訓練更新自身。直觀上我們可能會感覺使用當前最優模型生成的self-play數據可能質量更高，收斂更好，但是在嘗試過兩種方案之後，我們發現，在6*6棋盤上下4子棋這種情況下，直接使用最新模型生成self-play數據訓練的話大約500局之後就能得到比較好的模型了，而不斷維護最優模型並由最優模型生成self-play數據的話大約需要1500局之後才能達到類似的效果，這和AlphaZero論文中訓練34小時的AlphaZero勝過訓練72小時的AlphaGo Zero的結果也是吻合的。個人猜測，不斷使用最新模型來生成self-play數據可能也是一個比較有效的exploration手段，首先當前最新模型相比於歷史最優模型一般不會差很多，所以對局數據的質量其實也是比較有保證的，同時模型的不斷變化使得我們能覆蓋到更多典型的數據，從而加快收斂。

2. 如何保證self-play生成的數據具有多樣性？

一個有效的策略價值模型，需要在各種局面下都能比較準確的評估當前局面的優劣以及當前局面下各個action的相對優劣，要訓練出這樣的策略價值模型，就需要在self-play的過程中儘可能的覆蓋到各種各樣的局面。前面提到，不斷使用最新的模型來生成self-play數據可能在一定程度上有助於覆蓋到更多的局面，但僅靠這麼一點模型的差異是不夠的，所以在強化學習演算法中，一般都會有特意設計的exploration的手段，這是至關重要的。在AlphaGo Zero論文中，每一個self-play對局的前30步，action是根據正比於MCTS根節點處每個分支的訪問次數的概率採樣得到的（也就是上面Self-play示意圖中的 ，有點類似於隨機策略梯度方法中的探索方式），而之後的exploration則是通過直接加上Dirichlet noise的方式實現的（ ,有點類似於確定性策略梯度方法中的探索方式）。在我們的實現中，self-play的每一步都同時使用了這兩種exploration方式，Dirichlet noise的參數取的是0.3，即 , 同時 .

3. 始終從當前player的視角去保存self-play數據

在self-play過程中，我們會收集一系列的 數據， 表示局面， 是根據MCTS根節點處每個分支的訪問次數計算的概率， 是self-play對局的結果，其中 和 需要特別注意從每一步的當前player的視角去表示。比如 中用兩個二值矩陣分別表示兩個player的棋子的位置，那麼可以是第一個矩陣表示當前player的棋子位置，第二個矩陣表示另一個player的棋子位置，也就是說第一個矩陣會交替表示先手和後手player的棋子位置，就看 局面下誰是當前player。 也類似，不過需要在一個完整的對局結束後才能確定這一局中每一個 中的 ，如果最後的勝者是 局面下的當前player，則 ,如果最後的敗者是 局面下的當前player，則 ,如果最後打平，則 .

4. self-play數據的擴充

圍棋具有旋轉和鏡像翻轉等價的性質，其實五子棋也具有同樣的性質。在AlphaGo Zero中，這一性質被充分的利用來擴充self-play數據，以及在MCTS評估葉子節點的時候提高局面評估的可靠性。但是在AlphaZero中，因為要同時考慮國際象棋和將棋這兩種不滿足旋轉等價性質的棋類，所以對於圍棋也沒有利用這個性質。而在我們的實現中，因為生成self-play數據本身就是計算的瓶頸，為了能夠在算力非常弱的情況下儘快的收集數據訓練模型，每一局self-play結束後，我們會把這一局的數據進行旋轉和鏡像翻轉，將8種等價情況的數據全部存入self-play的data buffer中。這種旋轉和翻轉的數據擴充在一定程度上也能提高self-play數據的多樣性和均衡性。

策略價值網路訓練

所謂的策略價值網路，就是在給定當前局面 的情況下，返回當前局面下每一個可行action的概率以及當前局面評分的模型。前面self-play收集到的數據就是用來訓練策略價值網路的，而訓練更新的策略價值網路也會馬上被應用到MCTS中進行後面的self-play，以生成更優質的self-play數據。兩者相互嵌套，相互促進，就構成了整個訓練的循環。下面還是從幾個點分別進行說明：

1. 局面描述方式

在AlphaGo Zero中，一共使用了17個 的二值特徵平面來描述當前局面，其中前16個平面描述了最近8步對應的雙方player的棋子位置，最後一個平面描述當前player對應的棋子顏色，其實也就是先後手。在我們的實現中，對局面的描述進行了極大的簡化，以 的棋盤為例，我們只使用了4個 的二值特徵平面，其中前兩個平面分別表示當前player的棋子位置和對手player的棋子位置，有棋子的位置是1，沒棋子的位置是0. 然後第三個平面表示對手player最近一步的落子位置，也就是整個平面只有一個位置是1，其餘全部是0. 第四個平面，也就是最後一個平面表示的是當前player是不是先手player，如果是先手player則整個平面全部為1，否則全部為0. 其實在最開始嘗試的時候，我只用了前兩個平面，也就是雙方的棋子的位置，因為直觀感覺這兩個平面已經足夠表達整個完整的局面了。但是後來在增加了後兩個特徵平面之後，訓練的效果有了比較明顯的改善。個人猜想，因為在五子棋中，我方下一步的落子位置往往會在對手前一步落子位置的附近，所以加入的第三個平面對於策略網路確定哪些位置應該具有更高的落子概率具有比較大的指示意義，可能有助有訓練。同時，因為先手在對弈中其實是很佔優勢的，所以在局面上棋子位置相似的情況下，當前局面的優劣和當前player到底是先手還是後手十分相關，所以第四個指示先後手的平面可能對於價值網路具有比較大的意義。

2. 網路結構

在AlphaGo Zero中，輸入局面首先通過了20或40個基於卷積的殘差網路模塊，然後再分別接上2層或3層網路得到策略和價值輸出，整個網路的層數有40多或80多層，訓練和預測的時候都十分緩慢。所以在我們的實現中，對這個網路結構進行了極大的簡化，最開始是公共的3層全卷積網路，分別使用32、64和128個 的filter，使用ReLu激活函數。然後再分成policy和value兩個輸出，在policy這一端，先使用4個 的filter進行降維，再接一個全連接層，使用softmax非線性函數直接輸出棋盤上每個位置的落子概率；在value這一端，先使用2個 的filter進行降維，再接一個64個神經元的全連接層，最後再接一個全連接層，使用tanh非線性函數直接輸出 之間的局面評分。整個策略價值網路的深度只有5~6層，訓練和預測都相對比較快。

3. 訓練目標

前面提到，策略價值網路的輸入是當前的局面描述 ，輸出是當前局面下每一個可行action的概率 以及當前局面的評分 ，而用來訓練策略價值網路的是我們在self-play過程中收集的一系列的 數據。根據上面的策略價值網路訓練示意圖，我們訓練的目標是讓策略價值網路輸出的action概率 更加接近MCTS輸出的概率 ,讓策略價值網路輸出的局面評分 能更準確的預測真實的對局結果 。從優化的角度來說，我們是在self-play數據集上不斷的最小化損失函數： ,其中第三項是用於防止過擬合的正則項。既然是在最小化損失函數，那麼在訓練的過程中，如果正常的話我們就會觀察到損失函數在慢慢減小。下圖展示的是一次在 棋盤上進行五子棋訓練的過程中損失函數隨著self-play局數變化的情況，這次實驗一共進行了3050局對局，損失函數從最開始的4點幾慢慢減小到了2.2左右。

在訓練過程中，除了觀察到損失函數在慢慢減小，我們一般還會關注策略價值網路輸出的策略（輸出的落子概率分布）的entropy的變化情況。正常來講，最開始的時候，我們的策略網路基本上是均勻的隨機輸出落子的概率，所以entropy會比較大。隨著訓練過程的慢慢推進，策略網路會慢慢學會在不同的局面下哪些位置應該有更大的落子概率，也就是說落子概率的分布不再均勻，會有比較強的偏向，這樣entropy就會變小。也正是由於策略網路輸出概率的偏向，才能幫助MCTS在搜索過程中能夠在更有潛力的位置進行更多的模擬，從而在比較少的模擬次數下達到比較好的性能。下圖展示的是同一次訓練過程中觀察到的策略網路輸出策略的entropy的變化情況。

另外，在漫長的訓練過程中，我們最希望看到的當然是我們訓練的AI正在慢慢變強。所以雖然在AlphaZero的演算法流程中已經不再需要通過定期評估來更新最優策略，在我們的實現中還是每隔50次self-play對局就對當前的AI模型進行一次評估，評估的方式是使用當前最新的AI模型和純的MCTS AI（基於隨機rollout）對戰10局。pure MCTS AI最開始每一步使用1000次模擬，當被我們訓練的AI模型10:0打敗時，pure MCTS AI就升級到每一步使用2000次模擬，以此類推，不斷增強，而我們訓練的AlphaZero AI模型每一步始終只使用400次模擬。在上面那次3050局自我對局的訓練實驗中，我們觀察到：

• 經過550局，AlphaZero VS pure_MCTS 1000 首次達到10:0
• 經過1300局，AlphaZero VS pure_MCTS 2000 首次達到10:0
• 經過1750局，AlphaZero VS pure_MCTS 3000 首次達到10:0
• 經過2450局，AlphaZero VS pure_MCTS 4000 取得8勝1平1負
• 經過2850局，AlphaZero VS pure_MCTS 4000 取得9勝1負。

OK，到這裡整個AlphaZero實戰過程就基本介紹完了，感興趣的小夥伴可以下載我github上的代碼進行嘗試。為了方便大家直接和已經訓練好的模型進行對戰體驗，我專門實現了一個純numpy版本的策略價值前向網路，所以只要裝了python和numpy就可以直接進行人機對戰啦，祝大家玩的愉快！^_^

參考文獻：

1. AlphaZero: Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm
2. AlphaGo Zero: Mastering the game of Go without human knowledge

推薦閱讀：