MuZeroの論文を読むその２（MuZeroアルゴリズム）

続きです。

MuZeroアルゴリズム

MuZeroアルゴリズムについて詳しく説明する。
予測は、各タイムステップtで、 $k=1 \ldots K$ ステップのそれぞれについて、過去の観測 $O_{1}, \ldots, O_{t}$ および将来の行動 $a_{t+1}, \dots, a_{t+k}$ を条件とするパラメーター $\theta$ を使用したモデル $\mu \theta$ によって行われる。
モデルは、3つの将来の量を予測する：

方策 $\mathbf{p}_{t}^{k} \approx \pi\left(a_{t+k+1} | o_{1}, \ldots, o_{t}, a_{t+1}, \ldots, a_{t+k}\right)$
価値関数 $v_{t}^{k} \approx \mathbb{E}[u_{t+k+1}+\gamma u_{t+k+2}+\ldots | o_{1}, \ldots, o_{t}, a_{t+1}, \ldots, a_{t+k}]$
即時報酬 $r_{t}^{k} \approx u_{t+k}$

ここで、 $u$ は真の観測報酬、 $\pi$ は実際の行動の選択に使用される方策、 $\gamma$ は環境の割引関数である。

モデルの目的

内部的には、各タイムステップtで、モデルは表現関数、ダイナミクス関数、および予測関数の組み合わせによって表される。
ダイナミクス関数 $r^{k}, s^{k}=g_{\theta}\left(s^{k-1}, a^{k}\right)$ は、各仮想ステップkで即時報酬 $r_k$ および内部状態 $s_k$ を計算する再帰プロセスである。
これは、特定の状態と行動に対して予想される報酬と状態遷移を計算するMDPモデルの構造を反映している*1。
ただし、モデルベースRLへの従来のアプローチとは異なり、この内部状態 $s_k$ には環境状態のセマンティクスが関連付けられていない。
これは単にモデル全体の隠れ状態であり、その唯一の目的は関連する将来の量（方策、価値、および報酬）を正確に予測することである。
この論文では、ダイナミクス関数は決定論的に表される。確率的遷移への拡張は将来の作業のために残されている。
方策と価値の関数は、AlphaZeroの方策と価値の結合ネットワークと同様に、予測関数 $\mathbf{p}^{k}, v^{k}=f_{\theta}\left(s^{k}\right)$ によって内部状態 $s_k$ から計算される。
「ルート」状態 $s^0$ は、過去の観測値 $s^{0}=h_{\theta}\left(o_{1}, \dots, o_{t}\right)$ をエンコードする表現関数を使用して初期化される。
繰り返すが、これには将来の予測の支援を超える特別なセマンティクスはない。

モデルを使用したプランニング

このようなモデルが与えられると、過去の観測 $o_{1}, \dots, o_{t}$ が与えられた場合に、仮想の将来の軌跡 $a^{1}, \ldots, a^{k}$ を探索できる。
たとえば、単純な探索では、価値関数を最大化するkステップ行動シーケンスを単純に選択できる。
より一般的には、ダイナミクス関数によって引き起こされる内部報酬と状態空間に、MDPプランニングアルゴリズムを適用できる。
具体的には、私たちはAlphaZeroの探索に似たMCTS アルゴリズムを使用する。これは、単一エージェントドメインと中間報酬を許可するように一般化されている（Methodsを参照）。
各内部ノードで、現在のモデルパラメーター $\theta$ によって生成された方策、価値、および報酬の推定値を利用する。
MCTS アルゴリズムは、推奨方策 $\pi_t$ および推定価値 $v_t$ を出力する。
それから、行動 $a_{t+1} \sim \pi_{t}$ が選択される。

モデルの訓練

モデルのすべてのパラメータは、各仮想ステップkの方策、価値、および報酬を、kの実際のタイムステップが経過した後に観測された対応するターゲット値に正確に一致するように一緒に訓練される。
AlphaZeroと同様に、改善された方策ターゲットはMCTS探索によって生成される。
最初の目標は、予測方策 $\mathbf{p}_{t}^{k}$ と探索方策 $\pi_{t+k}$ の間の誤差を最小化することである。
また、AlphaZeroと同様に、ゲームまたはMDPをプレイすることにより、改善された価値目標が生成される。
ただし、AlphaZeroとは異なり、探索価値 $z_{t}=u_{t+1}+\gamma u_{t+2}+\ldots+\gamma^{n-1} u_{t+n}+\gamma^{n} \nu_{t+n}$ から未来へのnステップをブートストラップすることにより、割引と中間報酬を伴う長いエピソードを許可する。
ボードゲームの最終結果{負け，引き分け，勝ち}は、エピソードの最終ステップで発生する $u_{t} \in\{-1,0,+1\}$ の報酬として扱われる。
具体的には、2番目の目標は、予測価値 $v_{t}^{k}$ と目標価値 $z_{t+k}$ の間の誤差を最小化することである。
報酬目標は、単に観測された報酬である。
したがって、3番目の目標は、予測された報酬 $r_{t}^{k}$ と観測された報酬 $u_{t+k}$ の間の誤差を最小化することである。
最後に、L2正則化項も追加され、全体的な損失になる。

$\displaystyle l_{t}(\theta)=\sum_{k=0}^{K} l^{r}\left(u_{t+k}, r_{t}^{k}\right)+l^{v}\left(z_{t+k}, v_{t}^{k}\right)+l^{p}\left(\pi_{t+k}, \mathbf{p}_{t}^{k}\right)+c\|\theta\|^{2} \tag{1}$

ここで、 $l^r$ 、 $l^v$ 、および $l^p$ は、それぞれ報酬、価値、および方策の損失関数である。
補足図S2は、MuZeroアルゴリズムがどのように計画、行動、学習するかを制御する式をまとめたものである。

補足図S2

$\displaystyle {\text { Model }} \\ \left.\begin{array}{c} {s^{0}=h_{\theta}\left(o_{1}, \ldots, o_{t}\right)} \\ {r^{k}, s^{k}=g_{\theta}\left(s^{k-1}, a^{k}\right)} \\ {\mathbf{p}^{k}, v^{k}=f_{\theta}\left(s^{k}\right)}\end{array}\right\} \mathbf{p}^{k}, v^{k}, r^{k}=\mu_{\theta}\left(o_{1}, \ldots, o_{t}, a^{1}, \ldots, a^{k}\right)$

$\displaystyle {\text { Search }} \\ \begin{array}{c}{\nu_{t}, \pi_{t}=M C T S\left(s_{t}^{0}, \mu_{\theta}\right)} \\ {a_{t} \sim \pi_{t}}\end{array}$

$\displaystyle {\text { Learning Rule }} \\ \begin{array}{c}{\mathbf{p}_{t}^{k}, v_{t}^{k}, r_{t}^{k}=\mu_{\theta}\left(o_{1}, \ldots, o_{t}, a_{t+1}, \ldots, a_{t+k}\right)}\end{array} \\ z_{t}=\left\{\begin{array}{ll}{u_{T}} & {\text { for games }} \\ {u_{t+1}+\gamma u_{t+2}+\ldots+\gamma^{n-1} u_{t+n}+\gamma^{n} \nu_{t+n}} & {\text { for general MDPs }}\end{array}\right. \\ l_{t}(\theta)=\sum_{k=0}^{K} l^{r}\left(u_{t+k}, r_{t}^{k}\right)+l^{v}\left(z_{t+k}, v_{t}^{k}\right)+l^{p}\left(\pi_{t+k}, p_{t}^{k}\right)+c\|\theta\|^{2}$

$\displaystyle {\text { Losses }} \\ l^{r}(u, r)=\left\{\begin{array}{ll}{0} & {\text { for games }} \\ {\phi(u)^{T} \log \mathbf{r}} & {\text { for general MDPs }}\end{array}\right. \\ l^{v}(z, q)=\left\{\begin{array}{ll}{(z-q)^{2}} & {\text { for games }} \\ {\phi(z)^{T} \log \mathbf{q}} & {\text { for general MDPs }}\end{array}\right. \\ l^{p}(\pi, p)=\pi^{T} \log \mathbf{p}$

図1：学習モデルを使用したプランニング、行動、および訓練

f:id:TadaoYamaoka:20191121214604p:plain

（A）MuZeroがモデルを使用してプランニングする方法

モデルは、表現、ダイナミクス、および予測のための3つの接続されたコンポーネントで構成されている。
前の隠れ状態 $s^{k-1}$ と候補行動 $a^k$ が与えられると、ダイナミクス関数 $g$ は即時報酬 $r^k$ と新しい隠れ状態 $s^k$ を生成する。
方策 $p^k$ および価値関数 $v^k$ は、予測関数 $f$ によって隠れ状態 $s^k$ から計算される。
初期隠れ状態 $s^0$ は、過去の観測（囲碁の盤面やAtariの画面など）を表現関数 $h$ に渡すことで取得される。

（B）MuZeroが環境でどのように動作するか

Aで説明したように、モンテカルロ木探索は各タイムステップtで実行される。
行動 $a_{t+1}$ は、ルートノードからの各アクションの訪問回数に比例する探索方策 $\pi_t$ からサンプリングされる。
環境は行動を受け取り、新しい観測値 $o_{t+1}$ と報酬 $u_{t+1}$ を生成する。
エピソードの終わりに、軌跡データはリプレイバッファに保存される。

（C）MuZeroがモデルを訓練する方法

軌跡は、リプレイバッファからサンプリングされる。
最初のステップでは、表現関数 $h$ は、選択された軌跡から過去の観測 $o_{1}, \dots, o_{t}$ を入力として受け取る。
その後、Kステップのモデルが再帰的に展開される。
各ステップkで、ダイナミクス関数 $g$ は入力として前のステップからの隠れ状態 $s^{k-1}$ と実際の行動[tex:a_{t+k}を受け取る。
表現、ダイナミクス、および予測関数のパラメーターは、3つの量を予測するために、backpropagation-through-timeによりエンドツーエンドで一緒に訓練される。
3つの量：方策 $\mathbf{p}^{k} \approx \pi_{t+k}$ 、価値関数 $v^{k} \approx z_{t+k}$ 、および報酬 $r_{t+k} \approx u_{t+k}$
ここで、 $z_{t+k}$ はサンプルの収益（最終報酬（ボードゲーム）またはnステップ収益（Atari））である。

感想

MuZeroのアルゴリズムは図1を見るとわかりやすいです。

図1(A)を見ると、表現関数 $h$ 、ダイナミクス関数 $g$ 、予測関数 $f$ を使用してMCTSでシミュレーションを行っています。
UCBの式は付録で解説があるので、別途読んでいきます。
AlphaZeroでは手を指した後の局面は、実際にゲームの盤面に手を適用して進めていますが、MuZeroでは、ダイナミクス関数 $g$ を使用して次の局面の隠れ状態を取得する点が異なっています。

図1(B)の行動の選択は、AlphaZeroと同じく訪問回数に応じて選択されています。

図1(C)の学習は、リプレイバッファから取得した軌跡について、開始状態から予測関数 $f$ とダイナミクス関数 $g$ に従いKステップ分展開した系列と、実際にMCTSに従い選択した軌跡が一致するように学習されています。再帰的な関数になるため、リカレントニューラルネットワークで使用されるbackpropagation-through-time(BPTT)を使用してエンドツーエンドで学習されています。実装する場合は、この部分が難しそうです。

(続く)

*1:https://dl.acm.org/citation.cfm?id=528623