ChainerでTrust Region Policy Optimization (TRPO)

# TL;DR

Chainer v3でHessian-vector productが実装できるようになったのでTRPOを実装しました。

![out.gif](https://qiita-image-store.s3.amazonaws.com/0/41196/93c8bade-2e10-2438-ad5f-8ce582996fbe.gif)

# Trust Region Policy Optimization

[Trust Region Policy Optimization (TRPO)](https://arxiv.org/abs/1502.05477)はICML2015で発表された深層強化学習アルゴリズムです。2015年はDeepMindのDQNがNatureに掲載された年であり、Nature版DQNと同じくらい古いアルゴリズムといえますが、行動空間が連続値の場合の深層強化学習では未だにSOTAの1つです。

TRPOは次の制約付き最適化問題を解いて確率的方策$\pi_\theta$のパラメータを更新します。

```math
\text{maximize}_\theta ~ L_{\theta_\text{old}}(\theta)\\
\text{subject to} ~ \bar{D}_\text{KL}(\theta_\text{old},\theta) \le \delta,\\
\text{where} ~ L_{\theta_\text{old}}(\theta) = \mathbb{E}_{s \sim \rho_{\theta_\text{old}}, a \sim \pi_{\theta_\text{old}}} \Big[ \frac{\pi_\theta(a|s)}{\pi_{\theta_{old}}(a|s)} A_{\theta_\text{old}}(s,a) \Big],\\
\bar{D}_\text{KL}(\theta_\text{old},\theta) = \mathbb{E}_{s \sim \rho_{\theta_\text{old}}} [ D_\text{KL}(\pi_{\theta_{old}}(\cdot|s)||\pi_\theta(\cdot|s))].
```

ここで最大化している$L_{\theta_\text{old}}(\theta)$は、真に最大化したい値（$\pi_\theta$に従った場合の期待収益）の$\theta \simeq \theta_\text{old}$における近似になっており、KLダイバージェンスの制約はその近似の信頼領域を表しています。詳しくは元論文や、自分の[ICML2015読み会発表資料](https://www.slideshare.net/mooopan/trust-region-policy-optimization)などを参照してください。

上の制約付き最適化問題を具体的にどうやって解くかというと、まず$\theta \simeq \theta_\text{old}$において$L$、$\bar{D}_\text{KL}$をそれぞれ1次近似、2次近似します：

```math
L_{\theta_\text{old}}(\theta) \approx g \cdot (\theta - \theta_\text{old}) + L_{\theta_\text{old}}(\theta_\text{old}),\\
\bar{D}_\text{KL}(\theta_\text{old},\theta) \approx \frac{1}{2} (\theta - \theta_\text{old})^T 
 A(\theta - \theta_\text{old}),\\
\text{where} ~ g = \nabla_\theta L_{\theta_\text{old}}(\theta)\Bigr|_{\theta=\theta_\text{old}},\\
 A_{ij} = \frac{\partial}{\partial \theta_i} \frac{\partial}{\partial \theta_j} \bar{D}_\text{KL}(\theta_\text{old},\theta)\Bigr|_{\theta=\theta_\text{old}}.
```

近似後の制約付き最適化問題を解き、解となる新しいパラメータ$\theta_\text{new}$を求めます：

```math
\theta_\text{new} = \theta_\text{old} + \beta s,\\
\text{where} ~ s = A^{-1} g,\\
\beta = \sqrt{2 \delta / s^T A s}.
```

ここで求めた更新量$\beta s$はあくまで元の問題の近似解なので、以下を満たすかどうか実際に計算して確認し、どちらも満たすまで$\beta$を減らしていきます：

```math
L_{\theta_\text{old}}(\theta_\text{new}) > L_{\theta_\text{old}}(\theta_\text{old}),\\
\bar{D}_\text{KL}(\theta_\text{old},\theta_\text{new}) < \delta.
```

こうして最終的に得られた更新量がTRPOの1イテレーションの更新量になります。

# Hessian-vector productの実装

TRPOの更新量の計算のうち、$s = A^{-1} g$を求めるには、Conjugate Gradient法を使用します。そのためには、Hessian-vector product（ヘッセ行列とベクトルの積）を計算する必要があります。ヘッセ行列を求めてから直接その逆行列を求めてもいいのですが、パラメータ数が多くなると現実的ではなくなります。

Chainerではv2まで勾配の勾配を計算することができなかったので、Hessian-vector productの計算には有限差分法を用いるしかなかったのですが、v3からは勾配の勾配の計算がサポートされたのでそれを使って効率的に求めることができます。$y$の$\theta$についてのヘッセ行列とベクトル$u$の積は、$\nabla_\theta (u^T \nabla_\theta y)$として求めることができるので、これをそのままChainerで実装すると以下のようなコードになります。

```python
import chainer
import chainer.functions as F
import numpy as np


def flatten_and_concat_variables(vs):
    # Variableをつなげて1つのベクトルにする関数
    return F.concat([F.flatten(v) for v in vs], axis=0)

# パラメータa, bからなるモデルを考える
a = chainer.Parameter(initializer=0, shape=2)
b = chainer.Parameter(initializer=0, shape=2)
params = [a, b]
size = a.size + b.size

# この式の、パラメータについてのヘッセ行列を考える
y = F.sum((2 * a - b) ** 2)

# yの勾配をつなげて1つのベクトルにする（後でさらに勾配を計算するのでenable_double_backprop=Trueが必要）
gy = chainer.grad([y], params, enable_double_backprop=True)
flat_gy = _flatten_and_concat_variables(gy)
print('gy:', flat_gy.data)

# ランダムなベクトルに対して、ヘッセ行列とベクトルの積を求めたい
vec = np.random.rand(size).astype(np.float32)
print('random vector', vec)

# ヘッセ行列を計算せずに、ヘッセ行列とベクトルの積だけ求める！
hvp = chainer.grad([F.sum(flat_gy * vec)], params)
flat_hvp = _flatten_and_concat_variables(hvp)
print('Hessian-vector product (smart)', flat_hvp.data)

def compute_hessian(flat_gy):
    # 愚直にヘッセ行列を計算する関数（計算の確認のみに使用）
    ggy = []
    for i in range(flat_gy.size):
        ggyi = chainer.grad([flat_gy[i]], params)
        flat_ggyi = _flatten_and_concat_variables(ggyi)
        ggy.append(flat_ggyi.data)
    return np.stack(ggy)

h = compute_hessian(flat_gy)
print('Hessian:', h)
print('Hessian-vector product (dumb)', h.dot(vec))


# 結果は一致する
np.testing.assert_allclose(h.dot(vec), flat_hvp.data)
```

出力は以下のようになり、Hessian-vector productが計算できているのが確認できます。

```
gy: [ 0.  0. -0. -0.]
random vector [ 0.94696099  0.29959747  0.1657497   0.05266208]
Hessian-vector product (smart) [ 6.91268921  2.18613148 -3.4563446  -1.09306574]
Hessian: [[ 8.  0. -4. -0.]
 [ 0.  8. -0. -4.]
 [-4.  0.  2. -0.]
 [ 0. -4. -0.  2.]]
Hessian-vector product (dumb) [ 6.91268921  2.18613148 -3.4563446  -1.09306574]
```

# TRPOの実装

TRPOは[ChainerRL](https://github.com/chainer/chainerrl)のエージェントとして実装しました（まだmasterにはマージされてません）。

https://github.com/muupan/chainerrl/blob/trpo/chainerrl/agents/trpo.py
https://github.com/muupan/chainerrl/blob/trpo/examples/gym/train_trpo_gym.py

試しにOpenAI GymのHumanoid-v1を学習させてみました。アドバンテージの推定のために方策だけでなく状態価値関数も別にニューラルネットで学習しており、設定は隠れ層のユニット数を400にしている以外は上の`train_trpo_gym.py`と同一です。

![400units_100M.png](https://qiita-image-store.s3.amazonaws.com/0/41196/f6222573-1fec-5b55-e5f2-9cab55dd52b4.png)
学習曲線です。縦軸は100回の評価エピソードの平均収益です。

![out.gif](https://qiita-image-store.s3.amazonaws.com/0/41196/93c8bade-2e10-2438-ad5f-8ce582996fbe.gif)
元気よく走り回っています。