強化学習 as Inference： Maximum a Posteriori Policy Optimizationの実装

方策が最適である確率の下界をＥＭアルゴリズムっぽく最大化する強化学習手法 Maximum a Posteriori policy Optimization (ICLR2018) をBipedalWalker-v3向けにtensorflow2で実装します。

openreview.net

はじめに
最適制御確率の下界を導出する
MPOとは：最適制御確率の下界JをEMアルゴリズムっぽく最大化
- E-step：変分分布ｑの最適化、あるいはノンパラ版TRPO
- M-step：方策パラメータθの更新
BipedalWalker-2dでの学習結果
- 雑記

※コード全文：
GitHub - horoiwa/deep_reinforcement_learning_gallery: Deep reinforcement learning with tensorflow2

※MPOの実装はdeepmind/acmeを参考にしています
GitHub - deepmind/acme: A library of reinforcement learning components and agents

はじめに

方策勾配法：劣悪なサンプル効率と不安定な更新

深層強化学習において、連続値行動環境をコントロールする方策（Policy）を訓練するためには方策勾配法が広く使われています。しかし方策勾配法は基本的にオンポリシーなのでサンプルを使い捨てるゆえの劣悪なサンプル効率と、勾配の分散の大きさ*1ゆえの不安定なネットワーク更新という２つの困難が産業応用の壁となっています。

Trust Region Policy Optimization (2015)では信頼領域法の導入により方策勾配法の安定性を大きく高めることに成功しましたが、オンポリシーゆえのサンプル効率の悪さは課題として残ります。源泉かけ流しのごとくサンプルを消費するオンポリシーアルゴリズムはシミュレータならよいですが実機を使うロボティクスなんかでは時間/金銭コストが高すぎます。

horomary.hatenablog.com

Maximum a Posteriori Policy Optimization ：確立推論の枠組みで方策を最適化

Maximum a Posteriori Policy Optimization (MPO) は方策勾配法ではなく Control as Inference、すなわち確率推論のフレームワークで制御ポリシーを訓練する手法です。 MPOはオフポリシーアルゴリズムゆえにサンプル効率が良好であり、かつTRPOのような信頼領域法を用いることでロバストな更新を実現するという実用的な手法です。論文の発表はICLR2018ですが、DeepMindの2022年のNature論文 "深層強化学習によるトカマク型核融合炉の制御 " でMPOが採用されていることからも実用性の高さがうかがえます。

また、MPOにおけるポリシー更新は本質的に教師あり学習のためオフラインセッティングと相性がよいことも重要なポイントです。

www.nature.com

Control as inference では行動の最適性を確率分布で表現する

Control as inferenceの枠組みでもっとも重要なコンセプトは最適性確率変数Oです。例えばあるトラジェクトリτが与えられた時、そのトラジェクトリが最適トラジェクトリである確率はP(O=1 | τ)、そうでない確率はP(O=0 | τ)と表現されます。同様に状態s_tにおいてアクションa_t が最適行動である確率は P(O_t=1 | s_t, a_t)となります。最適性変数Oの導入の最大のメリットはMDPにおける"最適な制御"をグラフィカルに表現できるようになることです。

最適性変数Oにより行動の最適性を明示的に確率分布で表現できるために環境の不確実性を自然に扱えるようになります。また、確率推論のさまざまなツールを利用可能になるのも大きなメリットであり、実際にMPOではEMアルゴリズムやELBOなどを活用しています。

↓しっかり学びたい方はSergey Levineの講義動画へ↓

www.youtube.com

PDF： http://rail.eecs.berkeley.edu/deeprlcourse-fa17/f17docs/lecture_11_control_and_inference.pdf

Control as InferenceとしてのSoft Actor-Critic

連続値コントロール環境ではTRPO/PPOと並んで大人気の Soft Actor-Critic も Control as Inference フレームワークの手法と解釈することができます*2。しかしSAC論文内ではControl as Inference観点からの説明が乏しいので詳細を知りたい方はやはりSergey Levineの講義動画（上記）を見ましょう。

雑に説明するならSACもMPOも方策関数の形状をQ関数に似せることを目的としているので似たような手法という感覚です。なお、実装の観点ではSACがDDPGの派生みたいな感じになっているのに対してMPOは既存の有力手法とははっきり異なる実装になっています。

horomary.hatenablog.com

最適制御確率の下界を導出する

MPOの目的は方策πで行動決定を実行したときにそれが最適制御である確率 Pπ(O=1) を最大化することです。

$\displaystyle { \log p_\pi (O=1 ) }$

このままではナンセンスなのでより具体化すると、

$\displaystyle { \log p_\pi (O=1 ) = \log \int p_\pi( \tau )p( O=1 | \tau ) d\tau }$

すなわち、最適制御である確率＝トラジェクトリτが方策πに従って生成されるときにトラジェクトリτが最適である確率の期待値 です。

つづいてVAEなどでもお馴染みのイエンゼンの不等式より任意の確率分布q(τ)について、

$\displaystyle { \log \int p_\pi( \tau )p( O=1 | \tau ) d\tau \geqq \int q(\tau) ( \log p( O=1 | \tau ) + \log \frac{p_\pi( \tau )}{q(\tau)} ) d \tau }$

$\displaystyle { = \int q(\tau) \log p( O=1 | \tau) d\tau + \int q(\tau) \log \frac{p_\pi( \tau )}{q(\tau)} ) d \tau }$

第一項はトラジェクトリτがqに従って生成されるときに、トラジェクトリが最適である確率の期待値です。ここで、トラジェクトリが最適である確率 p(O=1 | τ) はトラジェクトリτの報酬和と指数比例する と想定すると、

$\displaystyle { \int q(\tau) \log p( O=1 | \tau) d\tau + \int q(\tau) \log \frac{p_\pi( \tau )}{q(\tau)} ) d \tau }$

$\displaystyle { = E_{\tau \sim q} \left[ \log \exp \sum_{t} \frac{r_t}{\alpha} \right] + \int q(\tau) \log \frac{p_\pi( \tau )}{q(\tau)} ) d \tau }$

$\displaystyle { = E_{\tau \sim q} \left[ \sum_{t} \frac{r_t}{\alpha} \right] + \int q(\tau) \log \frac{p_\pi( \tau )}{q(\tau)} ) d \tau }$

報酬和が高いほど最適トラジェクトリである確率が指数的に高まるというのは、まあそりゃそうだけどもという感じ。ここまで方策のパラメータ（≒ニューラルネットの重み）を明示せずにPπ(τ)と表記してきたのでPπ(τ) = P(τ | θ)P(θ) と書き直すと、

$\displaystyle { = E_{\tau \sim q} \left[ \sum_{t} \frac{r_t}{\alpha} \right] + \int q(\tau) \log \frac{p( \tau | \theta )p(\theta)}{q(\tau)} ) d \tau }$

$\displaystyle { = E_{\tau \sim q} \left[ \sum_{t} \frac{r_t}{\alpha} \right] + \int q(\tau) \log \frac{p( \tau | \theta )}{q(\tau)} ) d \tau + \log p(\theta) }$

第二項はKLダイバージェンス $\displaystyle { KL( q(\tau) || p_\pi( \tau | \theta) ) }$ そのものなのですが、トラジェクトリτは扱いが困難なので、やや強引ながらトラジェクトリのKLダイバージェンスは各状態行動ステップのKLダイバージェンスに分解できると想定すると、

$\displaystyle { = E_{\tau \sim q} \left[ \sum_{t} \frac{r_t}{\alpha} - KL(q( \cdot | s_t) || \pi( \cdot | s_t, \theta)) \right] + \log p(\theta) }$

$\displaystyle { = J(q, \theta )}$

ここまでで、方策πで行動決定を実行したときにそれが最適制御である確率 Pπ(O=1)の下界Jを、任意の確率分布ｑと方策パラメータθで表現することができました。

MPOとは：最適制御確率の下界JをEMアルゴリズムっぽく最大化

※煩雑になるので割引率γ=1として省略

MPOとは上で導出された最適制御である確率Pπ(O=1)の下界J(q, θ)をEMアルゴリズムっぽい方法で最大化する手法です。すなわち、Estep: θを定数と見なしてｑについてＪを最大化 → Mstep: ｑを固定してθについてＪを最大化を繰り返します。

$\displaystyle { J(q, \theta ) = E_{\tau \sim q} \left[ \sum_{t}^{\infty} \frac{r_t}{\alpha} - KL(q( \cdot | s_t) || \pi( \cdot | s_t, \theta)) \right] + \log p(\theta) }$

E-step：変分分布ｑの最適化、あるいはノンパラ版TRPO

Eステップでは方策パラメータθ を定数とみなし、下界Ｊを最大化するような方策分布qを算出します。

$\displaystyle { \text{arg} \max_{q} J(q, \theta ) }$

$\displaystyle { = \text{arg} \max_{q} E_{\tau \sim q} \left[ \sum_{t}^{\infty} \frac{r_t}{\alpha} - KL(q(\cdot | s_t) || \pi(\cdot | s_t, \theta)) \right] + \log p(\theta) }$

logp(θ)はE-stepでは定数扱いなのでJの最大化には寄与しないので無視できます。さらに温度パラメータαは非負定数なのでJに掛けてもargmax_qの結果は変わらないため、

$\displaystyle { = \text{arg}\max_{q} E_{\tau \sim q} \left[ \sum_{t}^{\infty} r_t - \alpha KL(q(\cdot | s_t) || \pi(\cdot | s_t, \theta)) \right] }$

さらにトラジェクトリ期待値Ｅτ~q を状態行動ステップ期待値Eμ(s)Ea~qに書き直すと、

$\displaystyle { = \text{arg}\max_{q} E_{\mu (s)} \left[ E_{a \sim q} \left[ \sum_{t}^{\infty} r_t \right] - \alpha KL(q(\cdot | s_t) || \pi(\cdot | s_t, \theta)) \right] }$

ここで初登場するμ(s)というのは定常分布から状態sがサンプリングされるという意味ですが、実装的にはリプレイバッファからとってくるの意味なのであんま気にしなくてOK。ところで時刻t...∞までの報酬和の期待値というのはつまり 状態行動価値Q(s, a) であるので、

$\displaystyle { = \text{arg}\max_{q} E_{\mu (s)} \left[ E_{a \sim q} \left[ Q(s_t, a_t) \right] - \alpha KL(q(\cdot | s_t) || \pi(\cdot | s_t, \theta)) \right] }$

つまりJ(q, θ)を最大化する方策分布qとは、[第二項] 現在の方策πから離れすぎないというＫＬ制約のもとで、[第一項] 状態行動価値Q(s, a)が大きな行動を出力する方策分布ということになります。このソフト制約付き最適化問題をハード制約つき最適化問題の形式に書き直したものが以下です。

この制約付き最適化問題において変分分布qを更新後方策、πを更新前方策と見ればほぼTRPOと同じ問題になります。*3。ゆえにTRPOは実質的にEステップのみのMPOと見ることができるます。

ハムスターでもわかるTRPO ①基本編 - どこから見てもメンダコ

しかし、TRPOではｑをパラメタライズ（≒ニューラルネットで表現）してハード制約付き最適化問題を解いたのに対して、MPOではノンパラでハード制約付き最適化問題を解きます。

論文Appendix D.2 E-Step より、期待値オペレータを積分形式に書き直す＆確率分布なので総和１の制約を追加した制約付き最適化問題

をラグランジュの未定乗数法で解いた結果が次式（論文：式8）となります。

つまり最適方策ｑとは現在の方策πを状態行動価値Qのボルツマン分布で重みづけしたものというのがこの制約付き最適化問題の答えとなります。

ここで、温度パラメータη*は次式を最小化するηを求めることによって得られます。

M-step：方策パラメータθの更新

M-stepではE-stepで求めたqを固定して方策パラメータθを更新します。

$\displaystyle { \text{arg} \max_{\theta} J(q, \theta ) }$

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim q} \left[ \sum_{t}^{\infty} \frac{r_t}{ \alpha} \right] - KL(q(\cdot | s_t) || \pi(\cdot | s_t, \theta)) \right] + \log p(\theta) }$

argmaxθに影響しない項を消去すると、

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ q(\cdot | s_t) \log \pi(\cdot | s, \theta) \right] + \log p(\theta) }$

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim q} \left[ \log \pi( a | s, \theta) \right] \right] + \log p(\theta) }$

第一項はqにしたがってアクションがサンプリングされたときのlogπの期待値であるので、重点サンプリング法でπからアクションがサンプリングされたときの期待値に挿げ替えると、

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim \pi} \left[ \frac{q(a | s, \theta)}{\pi(a | s, \theta)} \log \pi( a | s, \theta) \right] \right] + \log p(\theta) }$

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim \pi} \left[ \frac{ \pi(a | s, \theta) \exp \frac{Q(s,a)}{\eta} }{\pi(s | a, \theta)} \log \pi( a | s, \theta) \right] \right] + \log p(\theta) }$

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim \pi} \left[ \exp \left( \frac{Q(s,a)}{\eta} \right) \log \pi( a | s, \theta) \right] \right] + \log p(\theta) }$

第一項はexp(Q(s, a)/η)と現在方策πのクロスエントロピーを最大化するということなので、直感的にはボルツマン分布で重みづけされたQ(s,a)に方策πができるだけ似るように更新しよう、と解釈できます。

第二項は方策パラメータの事前分布であるので任意の分布を設定できます。更新前パラメータθiを平均、フィッシャー情報行列F/λを共分散とする多変量正規分布

を事前分布に設定し、さらにargmaxθに影響しない定数項を消去する *4 と、

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim \pi} \left[ \exp \left( \frac{Q(s,a)}{\eta} \right) \log \pi( a | s, \theta) \right] \right] - \lambda (\theta_i - \theta)^{T} F^{-1}_{\theta i} (\theta_i - \theta) }$

ここで、第二項は更新前方策πθiと更新後方策πθのKLダイバージェンスの二次までのテイラー展開であると解釈できます*5。

よって最終的にM-stepの目的関数は、

$\displaystyle { = \text{arg}\max_{\theta} E_{\mu (s)} \left[ E_{a \sim \pi} \left[ \exp \left( \frac{Q(s,a)}{\eta} \right) \log \pi( a | s, \theta) - \lambda KL(\pi_{\theta i}(a | s)) || \pi_{\theta }(a | s)) \right] \right] }$

結局M-stepもソフトKL制約つき最適化問題に帰着しました。E-stepではハードKL制約問題に書き直してまじめに解きましたが、M-stepではこのままソフト制約問題としてそのまま最大化（実装的には-1をかけて最小化）してしまいます。

ただし、λは制約違反の大きさに応じて適応的に更新していくことに注意してください。すなわち、KL制約がハイパラとして設定する許容値εを上回るようならλを大きくしてKLペナルティを増大させ、そうでないならλを縮小します。Soft Actor-Criticの温度パラメータ自動調整とやっていることはほとんど同じです。