ハムスターでもわかるTRPO ②制約付き最適化問題をどう解くか

強化学習初学者の鬼門であるTRPO(Trust region policy optimization)を丁寧に解説し、tensorflow2で実装します（その②）。

[TRPOシリーズ一覧]

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

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

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

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

f:id:horomary:20200821225540j:plain:w400 — https://petponder.com/how-to-care-for-your-pet-djungarian-hamster

前回のまとめ
目的関数L(θ)の1次近似
KLダイバージェンスの2次近似
近似後の制約付き最大化問題
そして実装へ
補足：自然勾配法との関係

前回のまとめ

TRPOでは方策パラメータθの更新を、更新前/更新後方策のKL距離制約つきの目的関数L(θ)最大化問題と捉えます。前回は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] \tag{1}}$

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

制約付き最大化問題

制約付き最大化問題なのでラグランジュの未定乗数法で解けないか、と考えます。

ラグランジュの未定乗数法（例：2変数の場合）

制約付き最大化問題
$\displaystyle{ {\text{maximize}} \ f(\theta_1, \theta_2) \\ {\text{subject to } \ g(\theta_1, \theta_2) = 0 } }$
の解θは
$\displaystyle{ L(\theta_1, \theta_2, \lambda) = f(\theta_1, \theta_2) - \lambda{g(\theta_1, \theta_2)} }$ について
$\displaystyle{ \frac{\partial{L}}{\partial{\theta_1}} = \frac{\partial{L}}{\partial{\theta_2}} = \frac{\partial{L}}{\partial{\lambda}} =0 }$

※ラグランジュの未定乗数法についての解説はリンク先を参照：
ラグランジュの未定乗数法と例題 | 高校数学の美しい物語

これをそのまま適用すれば良さそうに思ってしまいますが、しかし残念ながら話はそう簡単ではなく、 $\displaystyle{{\pi_{\theta_{new}}(a | s )}}$ 内でのパラメータ同士の絡み合いがきつすぎて実用的にはとても解けそうにありません。

そこで目的関数L(θ)および制約式についてテイラー展開による近似を行い単純化していきます。

具体的には目的関数L(θ)は一次までのテイラー展開、KLダイバージェンスについては二次までのテイラー展開を行うことで近似式とします。

テイラー展開（2次まで）

たとえば $\displaystyle{ f(\theta) }$ について $\displaystyle{ \theta_{old} }$ 周辺で2次までのテイラー展開を行うと下式のようになります。

$\displaystyle{ f(\theta) | _{\theta_{old}} \simeq f(\theta_{old}) + f^{\prime}(\theta_{old})(\theta - \theta_{old}) + \frac{1}{2}f^{\prime\prime}(\theta_{old})(\theta - \theta_{old})^{2} }$

目的関数L(θ)の1次近似

※注意：見やすさのためにこれ以降は $\displaystyle{\theta_{new}}$ を単に $\displaystyle{\theta}$ と表記します。 $\displaystyle{\theta_{old}}$ はそのまま $\displaystyle{\theta_{old}}$ の表記です。

L(θ)は $\displaystyle{\theta_{old}}$ 周辺での1次までのテイラー展開により近似します。

$\displaystyle{ L(\theta) \simeq L(\theta_{old}) + g^{\mathrm{T}}(\theta - \theta_{old}) }$
$\displaystyle{ \text{where } g = \nabla_{\theta} L(\theta) | _{\theta_{old}} }$

$\displaystyle{ \nabla_{\theta} L(\theta) | _{\theta_{old}} }$ は $\displaystyle{ \theta = \theta_{old} }$ におけるL(θ)の勾配です。

また、 $\displaystyle{L(\theta)}$ の右辺第一項 $\displaystyle{L(\theta_{old})}$ は定数なので最大化問題では無視してよいことに注意してください。

一次近似とはずいぶん大胆な近似なので、そんなに近似して大丈夫か？と思います。
実際大丈夫じゃないこともあるのでTRPOではパラメータ更新を確定する前に、更新後のパラメータが期待報酬を改善するか、KLダイバージェンス制約を満たしているかを確認します。この処理についての詳細はステップサイズの線形探索で説明します。

KLダイバージェンスの2次近似

資料３： Taylor expansion of KL

L(θ)と同様にKL距離の制約式も $\displaystyle{\theta_{old}}$ 周辺でテイラー展開していきます。KL距離は2次まで展開します。

$\displaystyle{ {D_{KL}(\pi_{\theta_{old}} || \pi_{\theta})} \simeq {D_{KL}(\pi_{\theta_{old}} || \pi_{\theta_{old}}) } + {\nabla_\theta{D}_{KL}(\pi_{\theta_{old}} || \pi_{\theta})|_{\theta{old}}} ^{\mathrm{T}} (\theta - \theta_{old}) + \frac{1}{2}(\theta - \theta_{old})^{\mathrm{T}} {H }(\theta - \theta_{old}) }$
$\displaystyle{ \text{where } H ={\nabla_{\theta}{^2} {D}_{KL}(\pi_{\theta_{old}} || \pi_{\theta})|_{\theta{old}}} }$

まず右辺第一項 $\displaystyle{{D_{KL}(\pi_{\theta_{old}} || \pi_{\theta_{old}}) }}$ は0なので消えます。KLダイバージェンスの定義から当然ですね。

さらに右辺第二項 $\displaystyle{{\nabla_\theta{D}_{KL}(\pi_{\theta_{old}} || \pi_{\theta})|_{\theta{old}}} ^{\mathrm{T}} (\theta - \theta_{old})}$ も0になって消えるのですがあまり直感的ではないですね。
そういうものだと放置したくない方は資料３： Taylor expansion of KLの証明もしくは f-divergenceと汎関数微分 - れおなちずむを参照ください。

さて、右辺第一項、第二項が消えた結果、結局のところ制約式は次式になります。

$\displaystyle{ {D_{KL}(\pi_{\theta_{old}} || \pi_{\theta})} \simeq \frac{1}{2}(\theta - \theta_{old})^{\mathrm{T}} {H }(\theta - \theta_{old}) }$

HはKLダイバージェンスの二階微分でありすなわちヘッセ行列（Hessian）です。

近似後の制約付き最大化問題

Efficiently Computing the Fisher Vector Product in TRPO – Telesens より

テイラー展開による近似によって目的関数と制約関数がかなり単純化されました。

$\displaystyle{ \underset{\theta} {\text{argmax }} g^{\mathrm{T}}(\theta - \theta_{old}) }$
$\displaystyle{ {\text{subject to }} \frac{1}{2}(\theta - \theta_{old})^{\mathrm{T}} {H }(\theta - \theta_{old}) \leq \delta }$

この制約付き最大化問題についてラグランジュ関数Gを作ります。
※ラグランジュ関数はLで表記されることが多いですがすでにLは使っているのでGで表記します

$\displaystyle{ G = g^{\mathrm{T}}(\theta - \theta_{old}) - \lambda({\frac{1}{2}(\theta - \theta_{old})^{\mathrm{T}} {H }(\theta - \theta_{old}) - \delta}) }$

ここで $\displaystyle{(\theta - \theta_{old})=s}$ とおくと

$\displaystyle{ G = g^{\mathrm{T}}s - \lambda({\frac{1}{2}s^{\mathrm{T}} {H }s - \delta}) }$

ここで、 $\displaystyle{s=(\theta - \theta_{old})}$ はパラメータの更新方向となることに注意してください。

パラメータの更新方向と更新サイズ

まずはGをsについて微分します。

$\displaystyle{ \frac{\partial{G}}{\partial{s}} = {g} - {\lambda}{Hs} = 0 }$

よって、

$\displaystyle{ s = \theta - \theta_{old} = \frac{1}{\lambda}{H^{-1}g} }$

$\displaystyle{\lambda}$ が解けていないので適切な更新サイズはわかりませんが、パラメータを $\displaystyle{ {H^{-1}g} }$ 方向に更新すればよいことはわかりました。

そこで、パラメータの更新方向 $\displaystyle{ {H^{-1}g} }$ を $\displaystyle{ s^{\prime}}$ 、更新サイズを $\displaystyle{ \beta}$ と置くと、適切な更新サイズsは、 $\displaystyle{ s = \beta{s^{\prime}} }$ と表せます。

$\displaystyle{\beta}$ はKL距離制約を満たすように決めればよいので、

$\displaystyle{ \frac{1}{2}{s}^{\mathrm{T}} {H }{s} = \frac{1}{2}\beta{s^{\prime}}^{\mathrm{T}} {H }\beta{s^{\prime}} = \delta }$

より更新サイズβが

$\displaystyle{ \beta = \sqrt{\frac{2{\delta}}{{s^{\prime}}^{\mathrm{T}} {H }{s^{\prime}}}} }$

と求まります。

θnewはどうなるか

ここまでで、方策パラメータθは更新方向を $\displaystyle{s^{\prime} = H^{-1}g }$ にとり、
更新サイズを $\displaystyle{ \beta = \sqrt{\frac{2{\delta}}{{s^{\prime}}^{\mathrm{T}} {H }{s^{\prime}}}} }$ に設定することで、KL制約を満たしつつL(θ)を最大化できることがわかりました。

つまり更新後のパラメータ $\displaystyle{\theta_{new}}$ は、

$\displaystyle{ \theta_{new} = \theta_{old} + \beta{s}^{\prime} = \theta_{old} + \sqrt{\frac{2{\delta}}{{s^{\prime}}^{\mathrm{T}} {H }{s^{\prime}}}}{H^{-1}g} \tag{3}}$

とすればよいということになります。

ステップサイズの線形探索

理想的には（３）式で更新すればよいのですが、テイラー展開による近似の影響で(3)式通りに更新しても [ KL制約を満たさない or L(θ) が改善されない]、ということがあります。そこでTRPOではまず最大のステップサイズ、つまり（３）式通りの更新で [KL制約を満たすか、L(θ)が改善されるか] を確認します。

もしそうでないのなら更新サイズβを縮小して同じことを繰り返します。たとえばbaselinesのTRPO実装では0.5倍ずつβを小さくしていって条件を満たすステップサイズを探索しています。

そして実装へ

TRPOの理論の説明はここまでで終わりですが、実際はパラメータが非常に多い深層学習で $\displaystyle{ {H^{-1}} }$ を計算するのは非常に計算コストが高い（パラメータ数の3乗オーダー）という実用上の問題が残っています。しかし $\displaystyle{ {H^{-1}} }$ そのものではなく $\displaystyle{ {H^{-1}}g }$ であれば共役勾配法によって良い近似解を得ることができます。