はじめに
強化学習苦手の会 Advent Calendar 2020 19日目の @ashigirl96 です
去年の３月に大学を卒業してから仕事しつつも細々と強化学習の論文を読んできました。学部のときも今も独学で学んでいるので、最新の知識を手に入れるのも大変だなあと日々感じてました。そんなときに@sei_shinagawaさん主催の強化学習苦手の会が始まりました。もくもくと自習をする会…凄く良い…。ただ、私がコミュ障というのと、大して強化学習わからないのに参加していいのか躊躇して、もくもく会に参加できずに今に至ります。強化学習苦手の会 アドベントカレンダをきっかけに参加できればと思い筆を執りました。
今回話したいネタが、一年前に読んでもピンと来なかったMaximum a Posteriori Policy Optimization (MPO)とそのOn-Policy版のV-MPOです。私は数学苦手なので、踏み込んだ話よりこれらの論文を読んで「面白いな」と思ったポイントを共有していきたいです。（間違っていれば教えて下さい）
面白いと思った部分が
- Soft Actor CriticとMPOが兄弟？
- MPOの方策改善にEMアルゴリズムを使っている
- MPOとV-MPOの関係は？
です。
MPOとV-MPO
Soft Actor Critic と (V-)MPOは兄弟？
Soft Actor CriticとMPOはどちらもControl as Inferenceから発展したアルゴリズムです。[DL輪読会]Control as Inferenceと発展 がとても参考になるのでぜひ読んでみてください。
Control as Inferenceの前に基礎となるMDPから単純なOn-Policy方策勾配アルゴリズムの変遷について話したいと思います。
MDPと簡単な方策勾配
一般的なマルコフ決定過程(MDP)は、時刻$t$における状態$\mathbf{s}_{t}\in\mathcal{S}$ と行動$\mathbf{a}_{t} \in \mathcal{A}$ と報酬 ${r}_{t} = r(\mathbf{s}_{t} , \mathbf{a}_{t} )$を定義します。状態遷移確率を $p(\mathbf{s}_{t+1} | \mathbf{s}_{t}, \mathbf{a}_{t})$とします。エージェントはパラメトリック方策分布 $ \pi_{\theta}\left(\mathbf{a}_{t} \mid \mathbf{s}_{t}\right)$と定義し、軌跡$\tau=\left(s_{1}, a_{1}, \ldots\right)$に対する分布を
$$
p(\tau)=\rho\left(\mathbf{s}_{1}\right) \prod_{t=1}^{T} p\left(\mathbf{s}_{t+1} \mid \mathbf{s}_{t}, \mathbf{a}_{t}\right) \pi_{\theta}\left(\mathbf{a}_{t} \mid \mathbf{s}_{t}\right) \tag{1}
$$
とします。そして、(語弊がありますが)MDPでの目的はリターン（報酬の総和）を最大化する方策のパラメータベクトル$\theta$をみつけることです。
$$
\theta^{\star}=\arg \max_{\theta} E_{\tau \sim \pi_{\theta}}\left[ \sum_{t=1}^{T}r\left(\mathbf{s}_{t}, \mathbf{a}_{t}\right)\right] \tag{1}
$$
$J(\theta) = E_{\tau \sim \pi_{\theta}}\left[ \sum_{t=1}^{T}r\left(\mathbf{s}_{t}, \mathbf{a}_{t}\right)\right]$ を目的関数とします。この目的関数の最大化するためには、勾配を求めることが考えられます。それが
$$
\nabla_{\theta} J(\theta) = E_{\tau \sim \pi_{\theta}}\left[ \sum_{t=1}^{T}r\left(\mathbf{s}_{t}, \mathbf{a}_{t}\right) \nabla_{\theta} \ln \pi_{\theta}(\mathbf{a}_{t} \mid \mathbf{s}_{t}) \right] \tag{3}
$$
で、$\alpha$を学習率としたとき
$$
\theta_{k+1}=\theta_{k}+\alpha \nabla_{\theta} J\left(\pi_{\theta_{k}}\right) \tag{4}
$$
と、方策パラメータを更新していきます。
この$r\left(\mathbf{s}_{t}, \mathbf{a}_{t}\right)$をさまざまな関数に置き換えていき、一般化させたのがGAE(General Advantage Estimation)のモチベーションですよね。
最適性制御のための推論
一般的な強化学習が期待報酬を最大化する軌跡を見つけることに対して、Control as Inferenceでは、軌跡上の事前分布から、タスクが達成や方策が改善したなどの結果を条件付けし、この結果と一致する軌跡の事後分布を推定します。$(1)$に対して、最適性変数$O_t = { 0, 1}$を加えます。そのとき軌跡に対する尤度関数を
$$
p_{\pi}(O=1 \mid \tau) \propto \exp \left(\frac{\sum_{t} r_{t}}{\eta}\right) \tag{5}
$$
とします。このとき、 $\eta$ は温度です。
一般的な強化学習が$(2)$のようにリターンを最大化する方策を見つけるのに対して、Control as Inferenceでは、最適性変数の対数周辺尤度の下界を計算します。ここの計算の違いが、SACとMPOの生き別けれ(?)になります。（気持ちが伝わってほしいだけなので、式変形は色々省いていたりしてます; 簡単のため事前分布$\log p(\theta)$を省いたりしてます）
$$
\begin{aligned}
\log p_\theta(O=1) &= \log \int p(\tau) p(O=1 \mid \tau) d \tau \geq \int q(\tau)\left[\log p(O=1 \mid \tau)+\log p(\tau) - \log q(\tau)\right] d \tau \
&\overset{\text{SAC}}{=} \mathbb{E}_{\tau \sim q}\left[\sum_{t} r_{t} / \alpha\right] +\alpha \mathcal{H}(q (\tau)) = \mathbb{E}_{\tau \sim \pi_{\theta}}\left[\sum_{t} r_{t} / \alpha +\alpha \mathcal{H}(\pi_{\theta}(a_t | s_t)) \right] = \mathcal{J}(\pi_\theta)\
&\overset{\text{MPO}}{=} \mathbb{E}_{\tau \sim q}\left[\sum_{t} r_{t} / \alpha\right]-\mathrm{KL}\left(q(\tau) | p(\tau)\right) = \mathbb{E}_{\tau \sim q}\left[\sum_{t} r_{t} / \alpha -\mathrm{KL}\left(q(a | s_t) | \pi_{\theta}(a | s_t) \right) \right] = \mathcal{J}(q, \pi_\theta)
\end{aligned} \
$$
上記の式を語弊を恐れずに以下のように解釈することができます。
面白ポイント：もともと、一般的な強化学習に最適性変数を加えることによって、ベイズの枠組みで最適制御問題を確率的に推論できるようになりました。SACは、対数周辺尤度の下限に対して、$\log p(\tau)$ を対数一様分布と仮定することで無視し、変分分布を方策$\pi_{\theta}$として目的関数を求めています。それに対して、MPOは、確率分布$p$を方策$\pi_{\theta}$としています。そして、変分分布$q$をノンパラメトリックな分布とし、$\log p_\theta(O=1)$の下限がなるべくタイトになるような$q$を計算（推定 or 学習）してから、その変分分布によって再重み付けされたサンプルを用いて$\pi_\theta(a | s)$を学習します。
方策改善にEMアルゴリズムを使っている
上記した面白ポイントで言及している通り、
そして、変分分布$q$をノンパラメトリックな分布とし、$\log p_\theta(O=1)$の下限がなるべくタイトになるような$q$を計算（推定 or 学習）してから、その変分分布によって再重み付けされたサンプルを用いて$\pi_\theta(a | s)$を学習
する手段として
- E-step: $q$について $\mathcal{J}(q, \pi_\theta)$ を最適化する
- M-step: $\pi_\theta$について $\mathcal{J}(q, \pi_\theta)$ を最適化する
のステップそれぞれを交互に座標上昇して$\mathcal{J}$を最適化するEMアルゴリズムを使っています。私は、このControl as Inferenceの文脈から派生したからこそEM-アルゴリズムのようなベイズ推定がRLで使えるということがアツく感じます。
ここで、少し余談（ホントは書きたいがボリューム増えそうで心が折れた話）です。
E-stepでは $\mathrm{KL}\left(q(a | s_t) | \pi_{\theta}(a | s_t) \right)$を最小化するような$q$を見つけるわけなんですが、式変形をする途中で、式$(5)$の温度$\eta$に対して、以下のような制約付き最適化問題を考えるようになります。
$$
\begin{array}{l}
\max _{q} \int \mu_{q}(s) \int q(a \mid s) Q_{\theta_{i}}(s, a) d a d s \
\text {s.t.} \int \mu_{q}(s) \mathrm{KL}\left(q(a \mid s), \pi\left(a \mid s, \theta_{i}\right)\right) d a<\epsilon \
\iint \mu_{q}(s) q(a \mid s) d a d s=1
\end{array} \tag{6}
$$
上記の制約付き最適化問題から以下のラグランジュ方程式（緩和）を考え、$\eta$についての双対関数を導出します。これを最適化します。
$$
\begin{aligned}
L(q, \eta, \gamma)=& \int \mu_{q}(s) \int q(a \mid s) Q_{\theta_{i}}(s, a) \text { dads+ } \
& \eta\left(\epsilon-\int \mu_{q}(s) \int q(a \mid s) \log \frac{q(a \mid s)}{\pi\left(a \mid s, \theta_{i}\right)}\right)+\gamma\left(1-\iint \mu_{q}(s) q(a \mid s) d a d s\right)
\end{aligned}
$$
一年前この論文を読んだとき、ラグランジュ緩和がわからなかった（今もわからん）ので、現在わかってることを話してみます。（ここも怪しいので間違っていたら教えて下さい）
- $(6)$のような難しい制約問題がある
- すべての制約を最大化したい方程式にまとめる。→ラグランジュ緩和
- ラグランジュ緩和した方程式の最大値 $\geq$ 下の制約問題の最大値（となる）
→ 良いラグランジュ乗数を見つけたい
- ラグランジュ未定乗数法を使う
- MPOのように双対関数を導いて乗数について最大化を行う
このようにして、制約問題を最大化することができます。
私は今までこのような制約問題に対して、TRPOのようなペナルティ法を用いるアルゴリズムは見たことがありますが、双対問題から解くのは見たことがなく、面白ポイントだと思いました。
最適化数学的には、ラグランジュ関数のラグランジュ乗数を双対関数から簡単に見つけづらい→ペナルティ法を使うという論理だと思うので、逆に双対関数を使っているのが面白いんですかね。
V-MPOとMPOの関係は？
尤度関数が違う
MPOの尤度関数が$(5)$に対してV-MPOでは状態-行動に対する尤度関数を
$$
p_{\theta_{\text {old }}}(\mathcal{I}=1 \mid \tau) \propto \exp \left(\frac{\sum_{t} A^{\pi_{\theta \text { old }}(s_t, a_t)}}{\eta}\right) \tag{7}
$$
としてます。SACやMPOの$O_t$が行動に対して報酬が最大になるというバイナリイベントとして解釈できたのに対して、V-MPOの$\mathcal{I}_t$は前回の方策に対して現在の方策が相対的に改善したかというバイナリイベントとして解釈できます。方策が改善されれば$\mathcal{I}=1$で、されてなければ$\mathcal{I}=0$です。
方策評価が違う
MPOの方策評価に Retraceアルゴリズム(off-policy)
$$
\begin{align}
\mathcal{L}_{V}(\phi)&=\min _{\phi} \mathbb{E}_{\mu_{0}(s), b(a \mid s)}\left[\left(Q_{\theta_{i}}\left(s_{t}, a_{t}, \phi\right)-Q_{t}^{\text {ret }}\right)^{2}\right], \text { with } \
Q_{t}^{\text {ret }}&=Q_{\phi^{\prime}}\left(s_{t}, a_{t}\right)+\sum_{j=t}^{\infty} \gamma^{j-t}\left(\prod_{k=t+1}^{j} c_{k}\right)\left[r\left(s_{j}, a_{j}\right)+\mathbb{E}_{\pi\left(a \mid s_{j+1}\right)}\left[Q_{\phi^{\prime}}\left(s_{j+1}, a\right)\right]-Q_{\phi^{\prime}}\left(s_{j}, a_{j}\right)\right] \
c_{k}&=\min \left(1, \frac{\pi\left(a_{k} \mid s_{k}\right)}{b\left(a_{k} \mid s_{k}\right)}\right)
\end{align} \tag{8}
$$
を使っているのに対して、V-MPOの方策評価に $n$ステップ目的(on-policy)
$$
\mathcal{L}_{V}(\phi)=\frac{1}{2|\mathcal{D}|} \sum_{s_{t} \sim \mathcal{D}}\left(V_{\phi}^{\pi}\left(s_{t}\right)-G_{t}^{(n)}\right)^{2} \tag{9}
$$
を使っているのが、MPO(off-policy)とV-MPO(on-policy)の違いです。
V-MPOのAdvantage関数の計算
V-MPOでは軌跡のサンプルのうち、advantage関数の値が高い50%をE-Stepのバッチに用いることによって、大幅に学習が改善されたそうです。この考え方は共分散行列適応進化戦略(CMA-ES)のテクニックに似ているそうです。
数学的にこのようなことをして良いと言えるのかまだよくわかりません…
最後に
MPO、V-MPOは今までのような強化学習の方策改善とは違いEMスタイルを用いている点が面白いですよね。その土壌としてControl as Inferenceという変分ベイズが背景にあるので、ベイズの恩恵が受けられそうで、これからの強化学習の発展が楽しみです。
本当は、実装してみたり、式変形をもっと丁寧に書いたり、理解が歯抜けになっている部分をきちんと埋めてからこの記事を書きたかったですが、時間がなかったり、アカデミックじゃない個人でこのような論文を精読するにはなかなか骨が折れます…。せっかく強化学習苦手の会があるので、積極的に参加して話し合っていきたいです。
前回のアドベントカレンダQ-learningの収束性について書いたの2年前か… 
時間がすぎるの速いですね
もうすぐ一年が終わりますが、皆さん来年もよろしくお願いします。