TRPO(Trust Region Policy Optimization)まとめ

方策勾配法の学習の安定化にあたっては、TRPO(Trust Region Policy Optimization)やPPO(Proximal Policy Optimization)のようにステップ幅の調整が解決策になります。当記事ではTRPOとPPOについて取りまとめました。
「ゼロから作るDeep Learning④ー強化学習編」の$10.2.3$の「TRPO・PPO」やTRPO、PPOの論文などを参考に作成を行いました。

・用語/公式解説
https://www.hello-statisticians.com/explain-terms

・TRPO論文
・PPO論文

前提の確認

方策勾配法まとめ

方策勾配法の目的関数や勾配

軌道$\tau$の収益$G(\tau)$を最大にするような方策$\pi_{\theta}(A_t|S_t)$を得るにあたっては、下記の目的関数を$p$次元ベクトル$\theta$について最大化すればよい。
$$
\large
\begin{align}
J(\theta) &= \mathbb{E}_{\tau \sim \pi_{\theta}}[G(\tau)] \quad (1.1) \\
\theta & \in \mathbb{R}^{p}
\end{align}
$$

上記で定義した$J(\theta)$の勾配ベクトルは下記のように表される。
$$
\large
\begin{align}
\nabla_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum_{t=0}^{T} G(\tau) \nabla_{\theta} \log{\pi_{\theta}(A_t|S_t)} \right] \quad (1.2)
\end{align}
$$

詳しい導出は下記で取り扱った。

REINFORCEMENT、ベースライン、Actor-Critic

$(1.1)$式では収益$G(\tau)$を元に目的関数$J(\theta)$を定義し、勾配ベクトル$\nabla_{\theta} J(\theta)$を$(1.2)$式で表した。$(1.2)$式は収益$G(\tau)$に基づいて軌道の重み付けを行なったと大まかに解釈できるが、本来的には状態$S_t$後の収益のみを用いるのがより適切である。このように$(1.2)$式の$G(\tau)$については様々な表し方が考えられるので、$(1.2)$式を$G(\tau)$を一般化した$\Phi_{t}$を用いて下記のように再定義する。
$$
\large
\begin{align}
\nabla_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum_{t=0}^{T} \Phi_{t} \nabla_{\theta} \log{\pi_{\theta}(A_t|S_t)} \right] \quad (1.3)
\end{align}
$$

上記の一例であるベースライン付きREINFORCEは、$\Phi_{t}$を下記のように定義する。
$$
\large
\begin{align}
\Phi_{t} &= G_{t} – b(S_t) \quad (1.4) \\
G_t &= R_t + \gamma R_{t+1} + \gamma^{2} R_{t+2} + \cdots + \gamma^{T-t} R_{T} \quad (1.5)
\end{align}
$$

ここで上記の収益$G_t$にQ関数$Q(S_t,A_t)$、ベースライン$b(S_t)$に価値関数を用いると$\Phi_{t}$はアドバンテージ関数$A(S_t,A_t)$に一致する。
$$
\large
\begin{align}
\Phi_{t} &= Q(S_t,A_t) – V(S_t) \\
&= A(S_t,A_t) \quad (1.6)
\end{align}
$$

アドバンテージ関数は状態$S_t$における行動$A_t$の評価を単体で行うことができるように定義される関数である。また、オーソドックスなREINFORCEやActor-Criticなどベースライン付きREINFORCE以外の$\Phi_{t}$の定義は下記で詳しく取りまとめた。

共役勾配法

対称行列$A$と$A$の固有ベクトルを$\mathbf{v}_{1}, \, \mathbf{v}_{2}$と定義する。ここで$\mathbf{d}_{1}=\mathbf{v}_{1}$とおくとき、対称行列$A$について$\mathbf{d}_{1}$と共役であるベクトル$\mathbf{d}_{2}$は、下記の式に基づいて得られる。
$$
\large
\begin{align}
\mathbf{d}_{2} = \mathbf{v}_{2} – \frac{\mathbf{v}_{1}^{\mathrm{T}} A \mathbf{v}_{2}}{\mathbf{v}_{1}^{\mathrm{T}} A \mathbf{v}_{1}} \mathbf{v}_{1} \quad (1.7)
\end{align}
$$

TRPOでは最適化における各ステップでの修正の方向の計算にあたって、通常の勾配ベクトルではなく上記に基づいて計算を行なった$\mathbf{d}_{2}$を用いる。共役勾配法の式の詳しい解釈に関しては下記で詳しく取り扱った。

直線探索

$$
\large
\begin{align}
\mathbf{x}_{k+1} &= \mathbf{x}_{k} – \alpha_{k} \nabla f(\mathbf{x}_{k}) \quad (1.8) \\
\mathbf{x}_{k} & \in \mathbb{R}^{p} \\
\nabla &= \left(\begin{array}{c} \displaystyle \frac{\partial}{\partial x_1} \\ \vdots \\ \displaystyle \frac{\partial}{\partial x_p} \end{array} \right)
\end{align}
$$

上記の$(1.8)$式は勾配降下法の式であるが、このような式の$\alpha_{k}$を$f(\mathbf{x}_{k} – \alpha_{k} \nabla f(\mathbf{x}_{k}))$が最大値を取るように設定するのが直線探索(line search)である。

直線探索のプログラム作成にあたってはnp.argmaxなどを用いるとプログラムをシンプルに作成できる。詳しくは下記で取り扱った。

フィッシャー情報行列

$p$次元パラメータベクトル$\theta \in \mathbb{R}^{p}$を用いて尤度関数を$L(\theta|x)$のように定義するとき、フィッシャー情報行列(FIM; Fisher Information Matrix)の$I(\theta)$は下記のように定義される。
$$
\large
\begin{align}
I(\theta) &= E \left[ \frac{\partial}{\partial \theta} \log{L(\theta|x)} \frac{\partial}{\partial \theta^{\mathrm{T}}} \log{L(\theta|x)} \right] \\
&= E \left[ \left(\begin{array}{c} \displaystyle \frac{\partial \log{L(\theta|x)}}{\partial \theta_{1}} \\ \vdots \\ \displaystyle \frac{\partial \log{L(\theta|x)}}{\partial \theta_{p}} \end{array} \right) \left(\begin{array}{ccc} \displaystyle \frac{\partial \log{L(\theta|x)}}{\partial \theta_{1}} & \cdots & \displaystyle \frac{\partial \log{L(\theta|x)}}{\partial \theta_{p}} \end{array} \right) \right]
\end{align}
$$

詳しくは下記で取り扱った。

Hessian-Free Optimization

$$
\large
\begin{align}
\nabla = \frac{\partial}{\partial \theta} = \left(\begin{array}{c} \displaystyle \frac{\partial}{\partial \theta_1} \\ \vdots \\ \displaystyle \frac{\partial}{\partial \theta_p} \end{array} \right)
\end{align}
$$

上記のように方向微分の演算子$\nabla$を定義するとき、ヘッセ行列$H$とベクトル$\mathbf{u}$の積$H \mathbf{u}$について下記のような近似が成立する。
$$
\large
\begin{align}
H \mathbf{u} & \simeq \frac{\nabla f(\theta + \varepsilon \mathbf{u}) – \nabla f(\theta)}{\varepsilon} \quad (1.10)
\end{align}
$$

上記の式の導出については下記で詳しく取り扱った。

TRPO

TRPOの目的関数と制約

TRPO(Trust Region Policy Optimization)では下記のように目的関数(objective function)と制約を定義する。
$$
\large
\begin{align}
\mathrm{maximize}_{\theta} &: \, \mathbb{E}_{t} \left[ \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t|s_t)} A_{t} \right] \quad (2.1) \\
\mathrm{subject \, to} &: \, \mathbb{E}_{t} \left[ KL[\pi_{\theta_{\mathrm{old}}}(\cdot|s_t),\pi_{\theta}(\cdot|s_t)] \right] \leq \delta \quad (2.2) \\
\theta & \in \mathbb{R}^{p}, \, \theta_{\mathrm{old}} \in \mathbb{R}^{p}
\end{align}
$$

$(2.1)$式は$(1.1)$式の$G(\tau)$をアドバンテージ関数で置き換え、さらに重点サンプリング(importance sampling)を行うにあたって式変形を行なったものである。$(2.1)$式のように定義する目的関数をTRPOやPPOの論文ではsurrogate objectiveということも抑えておくと良い。

上記の最適化にあたっては、重点サンプリングを元に目的関数の近似値を計算し、近似値の計算結果を元に共役勾配法直線探索(line search)に基づく繰り返し演算を行う。詳細は次項以降で取り扱う。

最適化問題の定式化

前項の「TRPOの目的関数と制約」の$(2.1), \, (2.2)$式を下記のように式の簡略化を行う。
$$
\large
\begin{align}
\mathrm{Maximize} & : \, L(\theta) \quad (2.1)’ \\
\mathrm{subject \, to} & : \, \overline{D}_{KL}(\theta_{\mathrm{old}}, \theta) \leq \delta \quad (2.2)’ \\
\theta & \in \mathbb{R}^{p}, \, \theta_{\mathrm{old}} \in \mathbb{R}^{p}
\end{align}
$$

TRPOでは上記の最適化問題を下記のように制約式ではなくペナルティを用いて再定義する。
$$
\large
\begin{align}
\mathrm{Maximize} & : \, L(\theta) – \lambda \overline{D}_{KL}(\theta_{\mathrm{old}}, \theta) \quad (2.3)
\end{align}
$$

制約やペナルティに用いられるKLダイバージェンスは下記のように二次近似(quadratic approximation)を行う。
$$
\large
\begin{align}
\overline{D}_{KL}(\theta_{\mathrm{old}}, \theta) & \simeq \frac{1}{2} (\theta_{\mathrm{old}}, \theta)^{\mathrm{T}} A (\theta_{\mathrm{old}} – \theta) \quad (2.4) \\
A_{ij} &= \frac{\partial^{2}}{\partial \theta_{i} \partial \theta_{j}} \overline{D}_{KL}(\theta_{\mathrm{old}}, \theta) \quad (2.5) \\
\theta & \in \mathbb{R}^{p}, \, \theta_{\mathrm{old}} \in \mathbb{R}^{p}
\end{align}
$$

フィッシャー情報行列は前節の$(1.9)$式のように定義されることが一般的であるが、TRPO論文では$(2.5)$式によって$A$の要素の推定が行われる。同様に目的関数$L(\theta)$について一次近似を行い方向微分を計算すると、$Ax=g$形式の方程式が得られるので、共役勾配法を用いて解けばよい。

共役勾配法と線形探索

「共役勾配法」「Hessian-Free Optimization」で確認した計算に基づいてパラメータ$\theta \in \mathbb{R}^{p}$をUpdateする方向の$\mathbf{d}_{t+1}$が得られた際に、$\theta_{t+1} = \theta_{t} + \beta \mathbf{d}_{t+1}$によって$\theta_{t+1}$の計算を行う。この際に制約式における$\delta$は下記に対応する。
$$
\large
\begin{align}
\delta = \overline{D}_{KL}(\theta_{\mathrm{old}}, \theta) \simeq \frac{1}{2}(\beta \mathbf{d}_{t+1})^{\mathrm{T}} A (\beta \mathbf{d}_{t+1}) = \frac{1}{2} \beta^{2} \mathbf{d}_{t+1}^{\mathrm{T}} A \mathbf{d}_{t+1} \quad (2.6)
\end{align}
$$

$(2.6)$式を$\beta$について解くと下記が得られる。
$$
\large
\begin{align}
\delta & \simeq \frac{1}{2} \beta^{2} \mathbf{d}_{t+1}^{\mathrm{T}} A \mathbf{d}_{t+1} \quad (2.6) \\
\beta^{2} & \simeq \frac{2 \delta}{\mathbf{d}_{t+1}^{\mathrm{T}} A \mathbf{d}_{t+1}} \\
\beta & \simeq \sqrt{\frac{2 \delta}{\mathbf{d}_{t+1}^{\mathrm{T}} A \mathbf{d}_{t+1}}} \quad (2.7)
\end{align}
$$

したがって、線形探索にあたっての$\beta$による探索幅の最大値を$\displaystyle \sqrt{\frac{2 \delta}{\mathbf{d}_{t+1}^{\mathrm{T}} A \mathbf{d}_{t+1}}}$に設定することで、$(2.2)$式の制約の範囲で線形探索を行うことができる。

「TRPO(Trust Region Policy Optimization)まとめ」への1件の返信

コメントは受け付けていません。