ハムスターでもわかるTRPO ③tensorflow2での実装例

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

[TRPOシリーズ一覧]

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

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

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

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

f:id:horomary:20200826004751p:plain:w300

TRPOの更新式
実装
- コード全体
更新ステップの計算まで
ステップサイズの線形探索
Pendulum-v0 でのテスト結果
そしてPPOへ

TRPOの更新式

前回までの解説で、結局のところTRPOでは次式の通りに方策パラメータを更新すればよいことがわかりました。

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

$\displaystyle{ \text{where } }$

$\displaystyle{ g = \nabla_{\theta} L(\theta) | _{\theta_{old}} = {\frac{\nabla_{\theta}\pi_{\theta}(a | s ) | _{\theta_{old}}}{\pi_{\theta_{old}}(a | s )} A^{\pi}(s, a)} }$

$\displaystyle{ H ={\nabla_{\theta}{^2} {D}_{KL}(\pi_{\theta_{old}} || \pi_{\theta})|_{\theta{old}}} }$

$\displaystyle{ s = {H^{-1}g} }$

$\displaystyle{g}$ については、現在（更新前）のパラメータ周辺でのL(θ)の勾配ですので実装でとくに難しいことはありません。

問題は更新前/更新後方策関数のKLダイバージェンスのヘシアン $\displaystyle{H}$ とその逆行列とgの積 $\displaystyle{s=H^{-1}g}$ の算出です。
この計算処理がTRPOの実装が複雑になる原因となっています。

ヘシアンの耐えられない重さ

上述の更新式を愚直に実装しようとするとヘシアン（ヘッセ行列）の逆行列を計算する必要がありますが、ヘシアンの計算コストはパラメータ数の３乗オーダーで増加していきますので（参考：Hessian free）、深層学習ではとても実用的な計算速度にはなりません。

しかし、TRPO（に限らずヘシアンを利用する最適化手法）の更新式をよく見るとヘシアンそのものが計算に必要なのではなく、ヘシアンの逆行列とベクトルの積およびヘシアンとベクトルの積（ $\displaystyle{H^{-1}g}$ および $\displaystyle{Hs}$ ）さえわかれば事足りることに気づきます。

※方策関数のパラメータ数がNのときヘシアン $\displaystyle{H}$ は N×N行列で勾配 $\displaystyle{g}$ はN×１ベクトルなので $\displaystyle{s=H^{-1}g}$ は N×１ベクトル、したがって $\displaystyle{Hs}$ も N×１ベクトルであることに注意しましょう。

ヘシアンの逆行列とベクトルの積

ヘシアンの逆行列と勾配ベクトルの積を x と置きます。

$\displaystyle{ x = H^{-1}g }$

これを変形すると

$\displaystyle{ Hx = g }$

となり連立一次方程式 $\displaystyle{ Ax = b }$ の形になります。そして連立一次方程式の解ｘは共役勾配法によってよい近似解を得ることができます。

このトリックによりヘシアン逆行列の愚直な計算を回避することができます。

Hessian-vector product の計算

資料４：Efficiently Computing the Fisher Vector Product in TRPO より

$\displaystyle{ x = H^{-1}g }$ を近似するための共役勾配法アルゴリズム中ではヘシアンと任意のベクトルの積 $\displaystyle{ Hv }$ を計算する必要があります。

※共役勾配法についての詳細は過去記事を参照ください。
horomary.hatenablog.com

ヘシアンを真面目に計算してしまうとヘシアンの逆行列ほどではないにせよやはり計算量が多すぎるのでここでも計算トリックを使って、Hを陽に計算せずにHvを計算します。 具体的には”KLダイバージェンスの勾配と任意のベクトルｖの積の総和” についての勾配を計算することによって Hv が得られます。

f:id:horomary:20200805225547p:plain:w400 — 引用元： https://www.telesens.co/2018/06/09/efficiently-computing-the-fisher-vector-product-in-trpo/

これらのテクニックにより $\displaystyle{H^{-1}g}$ および $\displaystyle{Hs}$ を現実的な計算量で得ることができるようになります。

実装

openAI/baselinesの実装(tensorflow1.X) を参考に tensorflow2で実装しました。
baselines/trpo_mpi.py at master · openai/baselines · GitHub

実装全体はgithubへ： https://github.com/horoiwa/deep_reinforcement_learning_gallery

コード全体

1024ステップ分のトラジェクトリを取得済み、アドバンテージを計算済みの状態から方策関数を更新するコードのみ掲載します。

更新ステップの計算まで

fullstep = $\displaystyle{ \sqrt{\frac{2{\delta}}{{s^{\prime}}^{\mathrm{T}} {H }{s^{\prime}}}}{H^{-1}g} }$ の計算までは数式通りに実装するだけです。

tensorflow2.X だと GradientTapeのおかげでどこで勾配が流れるのかわかりやすいため、tensorflow1.X での実装に比べてだいぶわかりやすくなっています。

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

fullstep = $\displaystyle{ \sqrt{\frac{2{\delta}}{{s^{\prime}}^{\mathrm{T}} {H }{s^{\prime}}}}{H^{-1}g} }$ はあくまでテイラー展開による近似によって計算される値なので、このステップで更新した結果本当にL(θ)が改善するか、KL距離制約を満たしているかを確認します。

もしL(θ)が改善しない or KL距離制約を満たさないならばステップサイズを縮小します。既定の回数この処理を繰り返しても条件を満たさないならばこの回でのパラメータ更新は諦めてトラジェクトリを破棄します。

これがTRPOにおける Line search です。