TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

MuZeroの論文を読む その2(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^rl^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)を使用してエンドツーエンドで学習されています。実装する場合は、この部分が難しそうです。

(続く)