🌿

Batch Monte Carlo Tree Search

に公開
AI
強化学習
AlphaZero
mcts
tech

はじめに

AlphaZeroを実装する過程でMCTSとGPUとの相性の悪さを痛感し、改善手法を検索してみた。
その結果、辿り着いた論文がとても素晴らしいので解読メモを残そうと思う。

MCTSのアルゴリズムで一番の問題点は、探索・状態評価・逆伝播の過程で1つの盤面ずつ評価を行っていくということである。たとえば、これがCPU上の簡易な計算で高速に評価できる盤面であるならば、それは問題にはならない。
しかし、これがGPU上でニューラルネットワークを用いた計算を行う必要があるならば話は変わってくる。GPU上での計算は、並列計算をそこそこのスピードで大量に行うことを長所としている。しかし、MCTSの盤面評価は愚直に実装すると、盤面評価を一つずつ行ってしまい、長所を活かしていないどころか、GPUの計算時間よりもテンソルを送受信する時間のほうがかかってしまう。

そこでよく取られる手法が、推論プロセスと探索プロセスを並列に実行し、推論のみをバッチ化して並列に行う手法である。このバッチ化 / 並列化というのがGPU推論を行う上でのキーワードである。実際、Batch-MCTSの前にも、Root並列化・Leaf並列化・Tree並列化など様々な手法が提案されていた。

更にこの手法では、本来の探索（1盤面ずつ評価するアルゴリズム）を守りながらも、異なる枝・異なる深さにある盤面評価をバッチ化してGPUに移譲し、高速に評価を行うのが特徴である。

リンク集

論文

記事

前提知識

  • AlphaZeroなどに用いられるPUCTの探索アルゴリズムについての深い理解
  • 強化学習の概念理解

提案手法

論文の提案手法について、要点や単盤面評価版のPUCTとの相違点を紹介する。
この論文の提案は主に以下の5つである。

  • Batch MCTS : batch化する盤面の探索方法
  • Virtual Mean : 未知盤面探索用に仮に配置する状態価値
  • \muFPU : 未探索盤面におけるQ値の代替スタート値
  • Second Move Heuristic : 探索途中で覆せない評価値が出た場合に他の盤面を探索する
  • Last Iteration : 指定の探索回数をこなしたあとに余分に複数回探索を行う

特徴

  • 盤面評価をバッチ化する
  • 2種類の探索木を用いて評価を行う
    • 未知盤面探索木 : 未知盤面を探索するための木、未知盤面発見時の特別な状態更新を行うために用意する
    • 盤面評価木 : 盤面の評価に用いる木

データ構造

論文に記載されている式は、筆者の実装とは表記が異なったため、筆者の実装に表記を合わせて解説する。
※筆者の実装は、メンダコさんの記事を参考に行っているため、メンダコさんの実装に合わせる目的もある。

この手法では、MCツリーに加えてTranspositionテーブルを実装する。

探索木

AlphaZeroで一般的に用いられる探索木である。

  • V(s) : 状態価値
  • \pi(s) : 事前確率(一つ前の時点で学習した方策)
  • N(s, a) : sにおいてaを選んだ回数
    • N(s, a) = \sum_a N(s, a) : sを通過した回数、最初に訪れた時(expand)は含まない
  • S(s, a) : sにおいてaを選んだ時に逆伝搬された価値の合計

※通常評価用も探索用も同じ構造

Transpositionテーブル

推論評価で得た価値と方策を蓄積する。

  • V(s)
  • \pi(s)

アルゴリズム

アルゴリズムは以下の2つのフェーズの繰り返しで進行する。

通常探索

盤面評価木を用いて行う。
ただし、盤面の展開時の評価はTranspositionテーブルの値を用いて行う。
もし、テーブルに無い未知盤面を発見した場合は、未知盤面探索に移行する。

未知盤面探索

未知盤面探索において重要な点は、未知盤面の逆伝播用の価値の計算方法と訪問回数の更新である。なお、ここからの更新では、未知盤面探索用の木を用いて行う。
手順は以下の通りである。

  • 通常と同様の探索を行う
  • Transpositionテーブルが既知の盤面であれば通常通りその価値を逆伝播する
  • 未知盤面であれば評価待ちバッチに追加し、終端盤面からUnknownを返却する
  • Unknownが伝播されてきた時は、Virtual Meanによる更新を行う
  • 上の探索を繰り返し、未知盤面が設定した閾値まで溜まったらモデル推論による評価を行い、Transpositionテーブルに結果を記録する
  • 通常探索に戻る

Virtual Meanによる更新

  • \mu = \frac{S(s, a)}{N(s, a)}
  • S(s, a) = S(s, a) + \mu \times vl
  • N(s, a) = N(s, a) + vl

この更新式では、パスの探索回数のみ増加し、訪問した盤面の評価値は変化しない。
通過する盤面の評価値を歪めずに探索回数のみを増加させる目的でS(s,a)の更新も行う。

例えば、探索前の評価値は以下のように算出できる。
Q(s, a) = \frac{S(s, a)}{N(s, a)}

Virtual Meanを用いて更新を行った後の盤面の評価値も等しくなる。
Q'(s, a) = \frac{S(s, a) + \mu \times vl}{N(s, a) + vl} = \frac{N(s,a)S(s,a) + S(s,a)vl}{N(s,a)(N(s,a) + vl)} = \frac{S(s,a)(N(s,a) + vl)}{N(s,a)(N(s,a) + vl)} = Q(s, a)

\muFPUの使用

FPUとは、sにおいて探索がまだ行われていない行動aの評価値平均Q(s,a)の代わりに用いる値である。

Q(s,a) = \frac{S(s,a)}{N(s,a)}
※N(s,a) = 0のときにこの式は使えないため、代わりの値を設定する。

FPUはいろんな種類がある

  • constant FPU : 固定値を設定する
  • best mean FPU : sにおける他のQ値の最大値を置く
  • \muFPU : sにおけるaの選択割合をかけたQ値の合計を用いる

\muFPUの算出式は以下の通りである。
\mu FPU = \sum_a \frac{N(s, a)}{N(s)} Q(s, a) = \frac{\sum_a S(s,a)}{N(s)}

これを用いることで、探索が遅れている盤面と探索済みの盤面の評価の乖離を防ぐことが目的である。

Second Move Heuristic

探索が偏ってしまうと、特定の盤面が多く探索されてしまい他の盤面についての探索が進まないことが多くなる。
この手法は、もう残りの探索回数では覆せないほどの探索回数を行った盤面が発生したら、そこを除外して残りの探索を行うというものである。

  • rb : 残り探索回数
  • N(s,a_{best}) > N(s,a_{2nd}) + rbの場合に、a_{best}を探索から除外

Last Iteration

この手法は、指定回数の探索が終わった後、最後に通常評価用木を未知盤面探索木に再度複製する。その後、未知盤面探索を行いその探索木から取った方策を使用する。
※ただし、発見した未知盤面のモデル推論は行わない。

これを行うことで、Transpositionテーブルには登録されたが使用されていない状態をなるべく多く探して方策に反映することが目的である。
※全て使えるわけではない。

手順を以下に示す。

  • 指定回数の探索を終える
  • 通常評価木を未知盤面探索木に複製
  • 未知盤面探索木で未知盤面が指定の個数見つかるまで探索を行う
    • この時、Transpositionテーブルにある評価値をより多く反映できる
  • 未知盤面探索木で方策評価を行う (\pi(a|s) = \frac{N(s,a)}{N(s)})
  • 未知盤面探索木を破棄する

この手順を行う際のvlの値をvllとすると、方策が歪まないようにvll=1に設定するのが良い。

結果

Inference Speed

Batch-MCSTの難しいところは、Batchを大きくすれば良いという訳ではないということである。
なるべく、通常評価と未知盤面探索のフェーズイテレーションが頻繁である方法で、かつなるべく多くの盤面を探索できる手法で size=32 良いのではないかという記述があった。

Size Batches per second Inferences per second
1 38.20 38
2 36.60 73
4 36.44 146
8 33.31 267
16 32.92 527
32 31.10 995
64 26.00 1664
128 18.32 2345

G. Inference Speed より

Virtual Mean

size=32のバッチを用いた場合のPUCTにおいて、普通のPUCTとの対戦結果を比較している。
Virtual LossとVirtual Meanの比較を行った結果、Virtual Meanの方が勝率が高いことが分かった。

※ Bはバッチ数、Batchはバッチサイズ

Virtual Loss

Penalty vl B Batch Nodes Inference Winrate
Virtual Loss 1 8 32 24.47 23.17 0.1300
Virtual Loss 2 8 32 24.37 24.46 0.1525
Virtual Loss 3 8 32 24.11 25.16 0.2075
Virtual Loss 4 8 32 23.87 25.53 0.2025
Virtual Loss 5 8 32 23.91 25.72 0.1600
Penalty vl B Batch Nodes Inference Winrate
Virtual Loss 1 32 32 166.09 28.08 0.7725
Virtual Loss 2 32 32 157.69 28.25 0.7900
Virtual Loss 3 32 32 151.02 28.30 0.7800
Virtual Loss 4 32 32 144.45 28.19 0.7550

Virtual Mean

Penalty vl B Batch Nodes Inference Winrate
Virtual Mean 1 8 32 46.45 20.22 0.2625
Virtual Mean 2 8 32 43.75 21.64 0.3025
Virtual Mean 3 8 32 41.63 22.10 0.3100
Virtual Mean 4 8 32 40.41 22.53 0.2400
Penalty vl B Batch Nodes Inference Winrate
Virtual Mean 1 32 32 612.02 26.63 0.9700
Virtual Mean 2 32 32 619.07 27.83 0.9675
Virtual Mean 3 32 32 593.91 28.20 0.9500

C. The Virtual Mean より

The Second Move Heuristic

全ての探索予算レベルで、一貫してSecond Move Heuristicが有効であった。

Budget Winrate
32 0.5925
64 0.6350
128 0.6425
256 0.5925
512 0.6250
1024 0.5600

E. The Second Move Heuristic

Last Iteration

U=0の場合よりもU=40の方が勝率が高いことが読み取れ、Last Iterationの効果があることが分かる。
また、U=10のように未知盤面の発見回数が少ない場合は vll=1 よりも vll=3 の方が勝率が高いが、U=40のように未知盤面の発見回数が多い場合は vll=1 の方が勝率が高いことが分かる。

U vl vll B Batch P Nodes Inference Winrate
10 3 3 8 32 64 109.02 21.90 0.4975
10 3 1 8 32 64 75.09 22.04 0.4450
40 3 3 8 32 64 232.09 22.17 0.5100
40 3 1 8 32 64 129.90 22.02 0.5275
0 1 1 32 32 512 729.41 28.88 0.6275
40 1 3 32 32 512 962.84 28.81 0.6650
40 1 1 32 32 512 835.07 28.86 0.6800

D. Last Iteration

おわりに

論文自体はBatch MCTSというネーミングだが、状態のbatch化以外にも別々に利用できる便利な手法が多く提案されている論文だった。
ゲーム探索のみならず、様々な場面で活用が見受けられるため、この手法による改善は非常に有用ではないだろうか。

Discussion