MuZeroの論文を読むその６（探索）

付録B 探索

MuZeroで使用される探索アルゴリズムについて説明する。
私たちのアプローチは、信頼区間の上限(UCB; upper confidence bounds)のあるモンテカルロ木探索、単一エージェントドメインの最適な方策とゼロサムゲームのミニマックス価値関数に漸近的に収束するプランニングアプローチに基づいている*1。

探索木のすべてのノードは、内部状態𝑠に関連付けられている。
𝑠からの各行動𝑎には、一連の統計 $\{N(s,a), Q(s,a), P(s,a), R(s,a), S(s,a) \}$ を保存するエッジ $(s,a)$ がある。それぞれ訪問数N、平均価値Q、方策P、報酬R、および状態遷移Sを表す。

AlphaZeroと同様に、探索は3つの段階に分割され、多くのシミュレーションで繰り返される。

選択

各シミュレーションは内部ルート状態 $s^0$ から始まり、シミュレーションがリーフノード $s^l$ に到達すると終了する。
シミュレーションの各仮想時間ステップ $k = 1 \dots l$ に対して、信頼区間の上限を最大化することにより、内部状態 $s^{k-1}$ の保存された統計量に従って行動 $a^k$ が選択される*2 *3。

$\displaystyle \begin{equation} a^{k} = \\ \arg\max_{a}\bigg[ Q(s, a) + P(s, a) \cdot \frac{\sqrt{\sum_b N(s, b)}}{1 + N(s, a)} \bigg(c_1 + \log\Big(\frac{\sum_b N(s, b) + c_2 + 1}{c_2}\Big)\bigg) \bigg] \end{equation}$
$\displaystyle\tag{2}$

定数 $c_1$ および $c_2$ は、ノードがより頻繁にアクセスされるときに、価値 $Q(s, a)$ に対する事前分布 $P(s, a)$ の影響を制御するために使用される。
私たちの実験では、 $c_1 = 1.25$ および $c_2 = 19652$ である。

$k \lt l$ の場合、次の状態と報酬は、状態遷移と報酬テーブル $s^{k} = S(s^{k-1}, a^{k})$ , $r^{k} = R(s^{k-1}, a^{k})$ で検索される。

展開

シミュレーションの最後のタイムステップlで、報酬と状態がダイナミクス関数 $r^l, s^l = g_\theta(s^{l-1}, a^{l})$ によって計算され、対応するテーブル $R(s^{l-1}, a^{l}) = r^l, S(s^{l-1}, a^{l}) = s^l$ に格納される。
方策と価値は、予測関数 $\mathbf{p}^l, v^l = f_\theta(s^l)$ によって計算される。
状態 $s^l$ に対応する新しいノードが探索木に追加される。
新しく展開されたノードの各エッジ $(s^l, a)$ は、 $\{ N(s^l, a)=0, Q(s^l, a)=0, P(s^l, a)=\mathbf{p}^l \}$ に初期化される。
探索アルゴリズムは、シミュレーションごとにダイナミクス関数と予測関数をそれぞれ最大で1回呼び出すことに注意する。
計算コストはAlphaZeroと同じオーダーである。

バックアップ

シミュレーションの最後に、軌跡に沿った統計が更新される。
バックアップは、環境が(1とは異なる割引γを持っている)中間報酬を出すことができる場合に一般化されており、価値の見積もりには制限がない(ボードゲームでは、割引は1と想定されており、中間報酬はない)。
$k= l ... 0$ の場合、以下の(価値関数 $v^l$ からブートストラップされる)累積割引報酬の $l − k$ ステップ推定値を形成する。

$\displaystyle G^k = \sum_{\tau=0}^{l-1-k}{\gamma ^ \tau r_{k+ 1 +\tau}} + \gamma^{l - k} v^l \tag{3}$

$k = l ... 1$ の場合、シミュレーションパスの各エッジ $(s^{k-1}, a^{k})$ の統計を次のように更新する。

$\displaystyle \begin{equation} \begin{aligned} Q(s^{k-1},a^k) & := \frac{N(s^{k-1},a^k) \cdot Q(s^{k-1},a^k) + G^k}{N(s^{k-1},a^k) + 1} \\ N(s^{k-1},a^k) & := N(s^{k-1},a^k) + 1 \end{aligned} \end{equation} \tag{4}$

2人ゼロサムゲームでは、価値関数は[0,1]の間に制限されると想定される。
この選択により、pUCTルール（式2）を使用して、価値の推定値と確率を組み合わせることができる。
ただし、多くの環境では価値に制限がないため、pUCTルールを調整する必要がある。
簡単な解決策は、環境で観察できる最大スコアを使用して、価値を再スケーリングするか、pUCT定数を適切に設定することである*4。
ただし、どちらのソリューションもゲーム固有であり、MuZeroアルゴリズムに事前知識を追加する必要がある。
これを回避するために、MuZeroは、その時点までに探索木で観測された最小値と最大値を使用して、正規化されたQ値推定値 $\overline{Q} \in [0, 1$ ]を計算する。
選択の段階でノードに到達すると、アルゴリズムは、以下の式を使用してpUCTルールで使用されるエッジの正規化されたQ値を計算する。

$\displaystyle \begin{equation} \overline{Q}(s^{k-1}, a^k) = \frac{Q(s^{k-1}, a^k) - \min_{s, a \in Tree}Q(s, a)}{\max_{s,a \in Tree} Q(s, a) - \min_{s,a \in Tree} Q(s, a)} \end{equation} \tag{5}$

感想

学習した環境のモデルを使用してシミュレーションを行う際のアルゴリズムの詳細について述べられています。
AlphaZeroのpUCTを使用したMCTSシミュレーションとほとんど同じですが、バックアップする際の報酬に割引が適用されています。また、即時報酬にも対応しています。
これにより、ボードゲーム以外のAtariなどにも適用できるアルゴリズムになっています。
pUCTの式は、AlphaZeroの疑似コードにあった式と同じです。

次回は、ハイパーパラメータについてです。
(続く)

*1:http://ggp.stanford.edu/readings/uct.pdf

*2:http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.172.9450&rep=rep1&type=pdf

*3:Science | AAAS

*4:https://dke.maastrichtuniversity.nl/m.winands/documents/CGSameGame.pdf