MuZero

# 概要

　強化学習のアルゴリズムには、モデルフリーとモデルベースがあります。モデルフリーの有名なアルゴリズムは、深層Q学習(DQN)、アクター・クリティック(Actor Critic) などがあります。モデルベースの有名なアルゴリズムは、AlphaGo、AlphaZeroなどがあります。

　モデルフリーなアルゴリズムは、どのような行動すれば、高い報酬が得られるか探索する必要があり、サンプリングを多く行わなければならず、学習に時間がかかることも知られています。反対に、モデルベースなアルゴリズムは、サンプル効率が良いことが知られています。

　今回の記事では、atariのゲームもできるモデルベースなアルゴリズ（おそらくモデルフリーとの組み合わせ？）、MuZeroについて紹介したいと思います。ただし、私はAlphaZeroの論文を読んだことがなく、MuZeroの最大のポイントはルールの知識がなくても学習ができることですが、その辺をちゃんと理解していません。今回の記事は、浅くMuZeroについて紹介したいと思います。

論文は以下を引用しています。

https://arxiv.org/abs/1911.08265

以下の論文の著者が

https://arxiv.org/abs/2102.12924

tensorflowのコードを公開しています。

https://github.com/kaesve/muzero

# ３つのニューラルネットワーク
　MuZeroは、表現関数・dynamics関数・予測関数と呼ばれる３つのニューラルネットワークで構成される。

```math
\begin{align}
\mbox{representation} \ \ :& \ \ s^0=h_{\theta}(o_1,o_2,\ldots,o_t) \\
\mbox{dynamics} \ \ :&  \ \ r^k,s^k= g_{\theta}(s^{k-1},a^k)          \\
\mbox{prediction} \ \ :&  \ \ p^k,v^k= f_{\theta}(s^k)              \\
\end{align}
```
$a,s,r,p,v$は、順に行動・隠れ状態・報酬・ポリシー・価値関数である。上付き添え字$k$は、モンテカルロ木の探索ステップを表す。$o_1,o_2,\ldots,o_t$は、過去の観測状態を表す。$\theta$はニューラルネットワークの学習パラメータである。

　表現関数$ h_{\theta}$は隠れ状態を、dynamics関数 $g_{\theta}$は報酬と次の隠れ状態を、予測関数 $f_{\theta}$は、ポリシーと価値を返す関数である。この関数を使い、計画・行動・学習について説明する。

# MuZero のアルゴリズム

![スクリーンショット 2021-07-16 183223.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1521272/ee223b16-7858-4fa8-d356-4081121e95a2.png)

この図は、MuZeroの論文から引用している。

　MuZeroのアルゴリズムは、時間ステップを$t$としてモンテカルロ木探索で計算する際の時間ステップを $k=1,\ldots,K$ とする。過去の状態$o_1,\ldots,o_t$ と未来の行動 $a_{t+1},\ldots,a_{t+k}$ から、ポリシー、価値関数、即時報酬

```math
\begin{align}
\mbox{policy} \ \ :& \ \ p_t^k \approx \pi\left(a_{t+k+1} |o_1,\ldots,o_t,a_{t+1},\ldots,a_{t+k}  \right) \\
\mbox{value function} \ \ :&  \ \ v_t^k \approx \mathrm{E}\left[u_{t+k+1}+\gamma u_{t+k+2}+\ldots  | o_1,\ldots,o_t,a_{t+1},\ldots,a_{t+k} \right]         \\
\mbox{immediate reward} \ \ :&  \ \ r_t^k \approx u_{t+k}             \\
\end{align}
```
を予測する（ように学習を行う）。$u_t$は観測された報酬である。

　まずは、計画（図A）について説明する。過去の状態 $o_1,\ldots,o_t$ を表現関数 $h_{\theta}$ に入力し、初期隠れ状態$s^0$を出力させる。前の隠れ状態 $s^{k-1}$ と次の行動 $a^k$ をdynamics関数 $g_{\theta}$ に代入し、（即時）報酬 $r^k$ と隠れ状態 $s^k$ を出力させる。そして、隠れ状態 $s^k$ を予測関数 $f_{\theta}$ に代入し、ポリシー $p^k$ と価値$v^k$ を得る。行動 $a^k$ の選択については、後に説明する Upper confidence bound を使い計算される。

　また隠れ状態は、正規化を行い

```math
\begin{align}
s \leftarrow \frac{s-\min(s)}{\max(s)-\min(s)}
\end{align}
```
とする。


　次に、行動（図B）について説明する。モンテカルロ木探索でポリシーを計算し、ポリシーから行動 $a_t$ を選択する。行動 $a_t$ の選択については、図Aにおける行動 $a^k$とは異なり、モンテカルロ木のノードに滞在した数によって選択される。

観測状態 $o_{t+1}$ と報酬 $u_{t+1}$ が得られるので、その結果をreplay Buffer に保存する。

　最後に、学習（図C）について説明する。replay Bufferから学習に使うデータのサンプリングを行う。Hypothetical Steps(unroll step ?) を $k=1,\ldots,K$ とし、隠れ状態 $s^{k-1}$ と実際に選択した行動 $a^k$ をdynamics関数 $g_{\theta}$ に代入し、報酬 $r^k$ と隠れ状態 $s^k$ を出力させ、これを繰り返す。
　$z_{t+k}$は最終的な報酬であり、$z_t$はブーストラップを行い

```math
\begin{align}
z_t = u_{t+1}+\gamma u_{t+2}+\cdots+\gamma^{n-1} u_{t+n} + \gamma^n \nu_{t+n}
\end{align}
```
と計算される。$\nu_t$は、モンテカルロ木探索で推定された報酬（価値関数）である。

　ブーストラップについて簡単に説明すると、通常のQ学習で使われる目標値を、次の状態 $s^{\prime}$とすれば

```math
\begin{align}
r + \gamma \underset{a}{\max} Q(s^{\prime},a)
\end{align}
```
と書ける。ブーストラップされた目標値は、$n$ステップ後の状態 $s_n$として

```math
\begin{align}
\sum_{t=0}^{n-1} \gamma^t r_t + \gamma^n \underset{a}{\max} Q(s_n,a)
\end{align}
```
とする。$n$ステップ後のQ値を使い計算される。

　Hypothetical Step $k$ のポリシー・価値関数・報酬は、モンテカルロ木探索で得られたsearch policy$\pi_{t+k}$と、ブーストラップされた探索価値関数$z_{t+k}$、観測された報酬$u_{t+k}$とすれば、

```math
\begin{align}
p^k & \approx \pi_{t+k} \\
v^k & \approx z_{t+k}        \\
r_t^k & \approx u_{t+k}             \\
\end{align}
```
と予測されるように、誤差逆伝播法（再帰ニューラルネットで使われるBackpropagation Through Time）を使いパラメータ$\theta$を更新する。つまり、サンプルされた$t$ステップのデータにおける誤差関数は、Hypothetical Step $k$ のポリシー・価値関数・報酬を使い

```math
\begin{align}
l_t(\theta) =\sum_{k=0}^{K} \left\{l^r(u_{t+k},r_t^k) + l^v(z_{t+k},v_t^k)  + l^p(\pi_{t+k},p_t^k) \right\} + c\|\theta \|^2
\end{align}
```
と計算される。$l^r(u,r)$は、観測された報酬と予測された報酬との誤差であり、
$l^v(z,v)$は、ブーストラップされた探索価値関数と予測された価値関数の誤差であり、
$l^p(\pi,p)$は、search policyと予測ポリシーの誤差である。これらは、自乗誤差かクロスエントロピーが使われる。


# モンテカルロ木探索

　モンテカルロ木探索について説明する。ある状態$s$における行動$a$とする。それぞれの$(a,s)$において、滞在数・平均価値関数・ポリシー・報酬・次の状態のテーブル

```math
\begin{align}
\{N(s,a),Q(s,a),P(s,a),R(s,a),S(s,a)\}
\end{align}
```
をつくる。

　行動$a^k$は、Upper confidence bound から計算される。

```math
\DeclareMathOperator*{\argmax}{arg\,max} %argmax
\begin{align}
a^k =\argmax_{a} \left[\bar{Q}(s,a) +P(s,a) \frac{\sqrt{\sum_b N(s,b)}}{1+N(s,a)} \left(c_1 +\log\left\{\frac{\sum_b N(s,b)+c_2+1 }{c_2} \right\} \right)\right] 
\end{align}
```
$c_1=1.25,c_2=19652$ である。また、$\bar{Q}(s,a)$は、規格化した値であり

```math
\DeclareMathOperator*{\min}{min\ } %min 
\DeclareMathOperator*{\max}{max\ } %max 
\begin{align}
\bar{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{align}
```
である。

図Aにおいて、dynamics関数から報酬・隠れ状態$r^l,s^l = g_{\theta}(s^{l-1},a^l)$ が得られ、予測関数からポリシー・価値関数 $p^l,v^l=f_{\theta}(s^{l})$ が得られる。これらをテーブルに保存する。

```math
\begin{align}
R(s^{l-1},a^l)=r^l , S(s^{l-1},a^l)=s^l ,P(s^l,a^l)=p^l
\end{align}
```
新しいノード（新しい状態$s^l$）の場合は、

```math
\begin{align}
N(s^{l},a^l)=0 , Q(s^{l},a^l)=0 ,P(s^l,a^l)=p^l
\end{align}
```
と初期化する。

　$k=l,l-1,\ldots,0$として、累積割引報酬をブーストラップを行い

```math
\begin{align}
G^k=\sum_{\tau=0}^{l-1-k} \gamma^{\tau} r_{k+1+\tau} +\gamma^{l-k}v^l
\end{align}
```
と計算する。$k=0,1,\ldots,l$として、$Q,N$のテーブルを

```math
\begin{align}
Q(s^{k+1},a^k) &= \frac{N(s^{k+1},a^k){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{align}
```
と更新する。

# Visit count distribution
 滞在数の確率分布は、温度$T$、そして$N(\alpha)=N(s^0,\alpha)$として

```math
\begin{align}
P_{\alpha}=\frac{N(\alpha)^{1/T} }{\sum_{b}N(b)^{1/T}}
\end{align}
```
で与えられる。例えば、$T\to \infty $では一様分布になる。つまり、温度を下げていくことで、滞在数の影響が大きくなる。論文では、温度を最初の500k(1k=1000 つまり 500000)ステップは$T=1$とし、次の250kステップは$T=0.5$、残りの250kステップは$T=0.25$と温度を下げていく。
 
　滞在数の確率分布を使い、図Bにおける行動$a_{t+1}\sim P_{\alpha}$を選択する。確率$P_{\alpha}$は、search policy $\pi_{t+k}$である。


# prioritized replay

　モンテカルロ木探索で推定された価値関数$\nu$とブーストラップされた探索価値関数$z$の値の開きが大きいデータを重点的に学習に使いたいため、priority $p(i)$を

```math
\begin{align}
p(i)=\frac{p_i^\alpha}{\sum_k p_k^\alpha} \ \ , \ \ p_i = |\nu_i - z_i |
\end{align}
```
としてデータをサンプリングする。学習が不安定にならないように、Importance sampling weightを

```math
\begin{align}
w_i = \left(\frac{1}{N} \frac{1}{P(i)} \right)^{\beta}
\end{align}
```
と定義する。

# scale gradient

　図Cにおいて、Hypothetical Steps(unroll step) の計算の際、各ステップにおいて、勾配の大きさが同程度になるように、誤差に$1/K$を掛けスケーリングする。また、dynamics関数に隠れ状態$s$を入力する前に、隠れ状態$s$に$1/2$を掛けスケーリングする。

# 価値関数・報酬のスケーリング

　価値関数・報酬は以下の関数を使いスケーリングする。

```math
\begin{align}
h(z) &= \rm{sgn}(z)\left(\sqrt{|z|+1}-1 \right) +\epsilon z \\
h^{-1}(z)&= \rm{sgn}(z) \left\{\left(\frac{\sqrt{1+4\epsilon(|z|+1+\epsilon)}-1 }{2\epsilon}\right)^2  -1 \right\} \\
\epsilon=10^{-3}
\end{align}
```
　dynamics関数と予測関数から出力される価値関数と報酬は、ソフトマックス関数から出力され（出力サイズは2*support_size+1）、この出力から期待値を計算し、$h^{-1}$を使いスケーリングされスカラー値となった価値関数と報酬が得られる。

　学習を行う際は、スカラー値の価値関数と報酬を、$h$を使いスケーリングを行い、出力サイズを2*support_size+1に戻し誤差関数の計算がされる。


# 簡単な学習（CartPole-v1）
　
　強化学習において、比較的簡単な CartPole-v1 だと学習できた。

　![resulte_.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1521272/58c600f3-1cf0-b40c-d31f-57014760d6b2.gif)


# 最後に

　今後は、atari などのゲームを試してみようかな。