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

強化学習初学者の鬼門であるTrust Region Policy Optimization (TRPO、信頼領域ポリシー最適化)を丁寧に解説し、tensorflow2で実装します。

[TRPOシリーズ一覧]

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

② ハムスターでもわかるTRPO ②制約付き最適化問題をどう解くか - どこから見てもメンダコ

③ ハムスターでもわかるTRPO ③tensorflow2での実装例 - どこから見てもメンダコ

関連： TRPOにおけるHessian-vector-productと共役勾配法 - どこから見てもメンダコ

はじめに

TRPO(Trust Region Policy Optimization) は自然方策勾配法の派生であり、DQN, A3Cと並んで近年の深層強化学習の最重要手法のひとつです。

PPO, ACKTRなど派生手法も多く、強化学習を学ぶ上では避けては通れない手法なのですが、私のように数学力がハムスターレベルの人間にはなかなかに理解が困難で大変苦労したので同じような人のために解説と実装を残します。実用的にはTRPOではなく、実装が簡単で安定性も良好な後継手法PPOを使えばよいかもしれませんが、理解できるとはぐれメタルのような経験値が入るので逃げずに頑張る価値はあると思います。

※ 方策勾配定理はわかるけどTRPOはよくわからない人向けの解説です。方策勾配定理についての詳細説明はありません。

f:id:horomary:20200714004440p:plain:w300

はじめに
方策勾配法とその問題点
Trust Region Policy Optimization
- KL制約項について
  - 方策パラメータθの小さな変化≠方策関数の出力の小さな変化
  - KL距離は方策関数の変化の大きさの指標
- 目的関数：L(θ)
② 近似と実装編へつづく
- 補足： conservative policy iteration

※この解説は下記の資料を基に作成しています

[1502.05477] Trust Region Policy Optimization : TRPO論文。

資料１：TRPO論文著者の講義スライド

資料２：UCバークレーの講義資料。

資料３：カーネギーメロン大学の講義資料。

資料４：個人ブログ。実装方法が詳細に解説されている

動画１：資料１の講義動画

方策勾配法とその問題点

資料１：Two Limitations of “Vanilla” Policy Gradient Methods
資料２：P4, Policy Gradients Review

割引報酬和の期待値を表す関数である $\displaystyle{ J(\theta) }$ の勾配 $\displaystyle{ g }$ は下式のようになるというのが方策勾配定理でした。

$\displaystyle{ g = \nabla J(\theta) = E_{{\tau}\sim\pi_\theta}\left[ {\sum_{t=0}^{\infty}{\gamma^{t}\nabla_\theta\log\pi_{\theta}(a_t | s_t )A^{\pi_\theta}(s_t , a_t )}}\right] }$

※ $\displaystyle{{\tau}}$ ：トラジェクトリ（状態とアクションのシーケンス）
※ $\displaystyle{{\tau}\sim\pi_\theta}$ ：トラジェクトリは方策 $\displaystyle{\pi_\theta}$ に従って集められた、の意味

方策勾配法の問題点は、 $\displaystyle{ J(\theta) }$ を改善するための方策パラメータθの更新方向の正しさしか保証してくれず、適切なステップ幅がわからないことです。

よって方策勾配定理に従って方策関数パラメータθを更新しても、パラメータを大きく更新しすぎると方策が劣化することがあります。さらに悪いことに、劣化した方策関数は報酬の少ない悪いサンプルを集めるようになるので次の更新ではさらに方策が悪化するという負のスパイラルに陥ります。この問題はステップ幅（学習率）を十分に小さく設定することで回避できますが、小さくしすぎると今度は学習が進みません。

この問題へのアプローチがTRPO（Trust region policy optimization, 信頼領域ポリシーの最適化）です。方策勾配定理がパラメータの更新方向しか保証してくれないのに対して、TRPOではその名前の通り適切な更新方向とともに方策が改善すると信頼できる適切な更新サイズも保証してくれます。

Trust Region Policy Optimization

TRPOでは方策パラメータθの更新を、更新前/更新後方策間のKL距離制約つき目的関数L(θ)最大化問題と捉えます。

$\displaystyle{ \underset{\theta_{new}} {\text{argmax }} L(\theta_{new}) = E_{{s}\sim{d^{\pi{θold}}}, \ {a}\sim\pi_{θold} }\left[ \frac{\pi_{\theta_{new}}(a | s )}{\pi_{\theta_{old}}(a | s )} A^{\pi}(s, a) \right] }$

$\displaystyle{ {\text{subject to }} E\left[ {D_{KL}(\pi_{\theta_{old}} || \pi_{\theta_{new}}) } \right] \leq \delta }$

更新前の方策関数 $\displaystyle{ \pi_{\theta_{old}} }$ と更新後の方策パラメータ $\displaystyle{ \pi_{\theta_{new}} }$ のKL距離が $\displaystyle{ \delta }$ 以下という制約のもとで、目的関数 $\displaystyle{ L(\theta) }$ を最大化する $\displaystyle{ \theta_{new}}$ を求める、というのがTRPOにおける方策の更新であり、勾配降下法で方策を更新する方策勾配法とはまったく異なったスキームとなっています。方策勾配法では勾配方向を選択してから固定ステップ幅*1で更新するのに対し、TRPOではKL距離で定義した信頼領域半径の中で目的関数を最大化できるように勾配方向とステップ幅を決定します。もっとわかりやすく言うとTRPOでは方策の更新にOptimizer(AdamとかRMSPropとか)を使わず毎回の更新ごとに制約付き最大化問題を解きます。

目的関数L(θ)がどこから湧いたのかなど混乱しますが、目的関数についてはひとまず置いといてまずはKL距離制約の意味から考えてみましょう。

KL制約項について

方策パラメータθの小さな変化≠方策関数の出力の小さな変化

前述の通り方策パラメータθを大きく更新しすぎると方策関数の出力が急激に変化し、結果として劣化する恐れがあります。これは逆に、方策関数の出力が大きく変化しないならばパラメータθを大きく更新しても問題ない、ということでもあります。パラメータθを安全に大きく更新できるならば少ないサンプルでも学習が効率的に進むので非常に有用です。

f:id:horomary:20200819223257p:plain:w700 — 方策パラメータθの小さな変化≠方策関数の出力の小さな変化
(資料２, The Problem is More Than Step Size より)

KL距離は方策関数の変化の大きさの指標

方策関数の更新前/更新後の出力のKL距離は、方策関数の変化の大きさの良い指標となります。
なぜなら方策関数はある状態sにおいてアクションaを選択する確率分布を出力します。そしてKL距離は２つの確率分布間の距離尺度であるからです。

KL距離の制約下で方策パラメータθを更新することで、一回の更新での方策関数の大きすぎる変質を防ぎつつ可能な最大の更新サイズで方策パラメータθを更新することができるようになります。直感的には、KL距離を制約に使うことで更新前パラメータ $\displaystyle{\theta_{old}}$ 周辺での方策関数の変化が急斜面なら慎重に、緩斜面なら大胆にパラメータ更新するようになるので安全にサンプル効率を向上させる効果がある、と理解するとよいのではないでしょうか。

目的関数：L(θ)

前述の通り、TRPOのパラメータ更新は目的関数L(θ)の制約付き最大化問題です。

$\displaystyle{ \underset{\theta_{new}} {\text{argmax }} L(\theta_{new}) = E_{{s}\sim{d^{\pi{θold}}}, \ {a}\sim\pi_{θold} }\left[ \frac{\pi_{\theta_{new}}(a | s )}{\pi_{\theta_{old}}(a | s )} A^{\pi}(s, a) \right] }$

このL(θ)は後継手法である PPO にも登場しますが、この式は一体どこから湧いて出たのか、と混乱します。

θ_old からθ_newへの更新による報酬期待値の変化を考える

資料２, P12, Proof of Relative Policy Performance Identity
資料３ Relating objectives of two policies

あるタイミングで方策パラメータθをθ_old からθ_newに更新したときの期待報酬和Jの改善量を考えると(1)式が導かれます。（証明は上記資料を参照）

$\displaystyle{ J(\pi_{\theta_{new}}) - J(\pi_{\theta_{old}}) = E_{{s}\sim{d^{\pi{θnew}}}, \ {a}\sim\pi_{θnew}}\left[ A^{\pi}(s, a) \right] \tag{1}}$

(1)式から、方策パラメータθの更新による方策改善の最大化するためには、上式の右辺を最大化するようなθnewを求めればよい、ということがわかります。しかし、(1)式の右辺はまだ計算不可能です。なぜならば学習に利用できるサンプルは更新前ポリシー $\displaystyle{\pi_{\theta_{old}}}$ によって収集されたものなので

$\displaystyle{ {s}\sim {d^{\pi_{\theta_{old}}}}, \ {a}\sim\pi_{\theta_{old}} }$
※ $\displaystyle{ d^{\pi_{\theta_{old}}} }$ は $\displaystyle{ {\pi_{\theta_{old}}} }$ に従ってサンプルを集めた時の各状態sの出現確率

であるのに対して、(1)式右辺で必要なサンプルは更新後のポリシー $\displaystyle{\pi_{\theta_{new}}}$ で収集されたものであること、つまり

$\displaystyle{ {s}\sim {d^{\pi_{\theta_{new}}}}, \ {a}\sim\pi_{\theta_{new}} }$
※ $\displaystyle{ d^{\pi_{\theta_{new}}} }$ は $\displaystyle{ {\pi_{\theta_{new}}} }$ に従ってサンプルを集めた時の各状態sの出現確率

であるためです。この問題は重点サンプリング法（Importance Sampling）というトリックを使うことで解決します。

重点サンプリング法

重点サンプリング（Importance Sampling）は、関心のある分布とは異なる分布から生成されたサンプルから、特定の分布の特性を推定する一般的な手法です。 Importance sampling - Wikipedia

ある確率分布Pから生成されたサンプルxから計算されるf(x)の平均値（など統計量）を計算したいが、手元には確率分布Qから生成されたサンプルしかない、というときに使用します。on-policy手法を疑似的にoff-policyにするトリックとも表現できます。

$\displaystyle{ E_{x\sim{p}}\left[f(x)\right] = E_{x\sim{q}}\left[\frac{p(x)}{q(x)} f(x)\right] \tag{2}}$

余談ですが重点サンプリングは分子動力学シミュレーションなんかでもよく使う手法で、例えば知りたいのはあるタンパク質の0℃での平均構造だが、低温だとシミュレーション効率が悪いので50℃でシミュレーションした結果を重点サンプリングで重みづけして0℃での平均構造を推定する、というような利用をします。

目的関数L(θ)の導出

(1)式に対して、(2)式と同様に重点サンプリング法を適用しactionの確率分布を $\displaystyle{ {\pi_{\theta_{old}}} }$ に挿げ替えます。

$\displaystyle{ \ E_{{s}\sim{d^{\pi{θnew}}}, \ {a}\sim\pi_{θnew}}\left[ A^{\pi}(s, a) \right] }$

$\displaystyle{ = E_{{s}\sim{d^{\pi{θnew}}}, \ {a}\sim\pi_{θold} }\left[ \frac{\pi_{\theta_{new}}(a | s )}{\pi_{\theta_{old}}(a | s )} A^{\pi}(s, a) \right] }$

$\displaystyle{ d^{\pi_{\theta_{new}}} }$ については、ポリシーが十分に近いならば各状態ｓの出現確率は更新前後でほとんど変わらないだろう、という仮定*2を置いて、 $\displaystyle{ d^{\pi_{\theta_{new}}} }$ ≒ $\displaystyle{ d^{\pi_{\theta_{new}}} }$ と近似してしまいます。すると、

$\displaystyle{ = E_{{s}\sim{d^{\pi{θold}}}, \ {a}\sim\pi_{θold} }\left[ \frac{\pi_{\theta_{new}}(a | s )}{\pi_{\theta_{old}}(a | s )} A^{\pi}(s, a) \right] }$

$\displaystyle{ = L(\theta_{new}) }$

となり最大化する目的関数L(θ)が導出できました。

② 近似と実装編へつづく

ここままでで、TRPOの方策更新とは目的関数L(θ)の制約付き最大化問題を解くことである、ということとL(θ)の意味を確認しました。

しかし、この最大化問題を解くためには目的関数および制約式にいくつかの近似を行う必要がありますので、②ではこの解説を行います。

horomary.hatenablog.com

補足： conservative policy iteration

TRPO論文のキモは Approximately Optimal Approximate Reinforcement Learningが提唱したconservative policy iteration （保守的なポリシー更新による単調改善）が混合ポリシー更新でなくても成立することを証明したことであり、この貢献こそが今後もTRPOが重要な手法でありつづける根拠です。しかし、この話から始めるととんでもなく分かりにくくなるので（そしてTRPO論文がわかりにくい直接の原因でもあるので）、本稿では触れていません。参考： Summary: Conservative Policy Iteration | by Zac Wellmer | Arxiv Bytes | Medium

*1:Adamとか使うなら実際は固定ステップではないがわかりやすさのため

*2:この近似ゆえにあまりにも離れたpolicy間で重点サンプリングを行うのは危険