[論文解説] IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures

この記事は，以下の論文の解説です．

[IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://arxiv.org/abs/1802.01561) (ICML 2018)

ただし，アルゴリズムを解説することを目的とし，検証実験に関しては触れていません．(めんどくさかったため)

記事内容では，強化学習の基礎的な知識を前提としています．
不備がございましたら，ご指摘頂けると幸いです．


## 概要

この論文では， **学習の安定性や効率を損なうことなく，数千のマシン上で学習を行うことのできる分散強化学習の手法 = IMPALA**，およびIMPALAでも利用されている **off-policyでの学習を実現するためのアルゴリズム = V-trace** を提案しています．

## IMPALA

![aaaa.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/248552/5eccc7c8-9f4b-f19f-0377-46699687bcfc.png)

図 : https://deepmind.com/blog/article/impala-scalable-distributed-deeprl-dmlab-30 から引用．

IMPALAでは，Actor-Criticの設定を用いて，方策 $\pi$ と価値関数 $V^\pi$ を学習します．IMPALAのアーキテクチャでは，データ収集のみを行う複数のプロセス(Actor)と，*off-policy* で学習を行う1個以上のプロセス(Learner)で構成されます(上図)．

Actorは $n$ ステップごとにデータ収集を行います．まず自身の方策 $\mu$ をLearnerの方策 $\pi$ に更新し，$n$ ステップ間データ収集を行います．その後，収集した経験データ(状態，行動，報酬)，方策の分布 $\mu(a_t|x_t)$ ，LSTMの初期状態をLearnerに送ります．一方，Learnerは複数のActorから送られてきたデータを用いて，繰り返し学習を行います．

このデータ収集と学習が分離されたシンプルな構成により，以下のメリットがあります．

- Learnerは各Actorが収集したデータを利用できるので，GPUを活用してデータ並列化できる．またデータ収集を学習から分離することで，LearnerのGPU利用効率が向上する．
- Actorを多数のマシンに分散させることが容易となる．またActorはそれぞれ非同期なので，1ステップにかかる時間にばらつきがある環境でも，同期のための待ち時間が生じない．
- 勾配情報に比べて経験データはサイズが小さいので，通信性能が向上する．

しかし，この構成では **データ収集時の方策 $\mu$ と学習している方策 $\pi$ が必ずしも一致しない** という問題が生じます．そこで，この方策のズレを補正する *V-trace* と呼ばれるアルゴリズムを提案しています．*V-trace* を用いることで，サンプル効率を損なうことなく，IMPALAアーキテクチャが持つ高いスループット性能を獲得することができます．さらに，IMPALAでは **複数Learnerを用いることで，複数GPUを効率的に活用することができる** というメリットがあります(上図右)．



## V-trace

データ収集と学習が分離した構成では，データを収集する方策と学習する方策の間にズレが生じるため，*off-policy* での学習がキモとなります．そこで，この**方策のズレを補正する**ために *V-trace* と呼ばれるアルゴリズムを提案していきます．

以下では，無限時間のMDPにおいて，期待割引報酬和(価値関数) $V^\pi(x):=\mathbb E_\pi[\sum_{t\ge0} \gamma^t r_t]$ を最大化する方策 $\pi$ を見つけることを考えます．ここで，$\gamma \in [0,1)$ は割引率，方策は確率的方策 $a_t \sim \pi(\cdot|x_t)$ とします．以下では，行動方策 $\mu$ で収集されたデータを用いて，学習方策 $\pi$ の価値関数 $V^\pi$ を学習することを考えます．



### V-traceオペレータ

$n$ ステップの *V-trace* オペレータ $\mathcal R^n$ を以下のように定義します．

```math
\mathcal R^n V(x_s) := V(x_s) + \mathbb E_\mu[
\sum_{t=s}^{s+n-1} \gamma^{t-s}(c_s\cdots c_{t-1})\delta_tV
]
```

ここで，$\delta_tV := \rho_{t}(r_t+ \gamma V(x_{t+1})-V(x_t))$ は，*Importance Sampling* (IS)により重み付けられたTD誤差を表します．また，$\rho_i = \min(\bar \rho,\frac{\pi(a_i|x_i)}{\mu(a_i|x_i)})$ ，$c_i = \min(\bar c, \frac{\pi(a_i|x_i)}{\mu(a_i|x_i)})$ は，クリップされたISの重み係数を表し，クリップの閾値は $\bar \rho \ge \bar c$ を満たすこととします．

ここで，*on-policy* ( $\mu = \pi$ ) の場合には，

```math
\begin{align}
\mathcal R^n V(x_s) &= V(x_s) + \mathbb E_\mu[
\sum_{t=s}^{s+n-1} \gamma^{t-s} (r_t+ \gamma V(x_{t+1})-V(x_t))] \\
&= \mathbb E_\mu[
\sum_{t=s}^{s+n-1}\gamma^{t-s} r_t + \gamma^n V(x_{s+n})]
\end{align}
```

と変形でき，***on-policy* の場合の *V-trace* は，オンライン学習時の $n$ ステップベルマンオペレータに一致する**ことがわかります．

*V-trace* において，2つの異なるIS重み係数の閾値 $\bar \rho$ と $\bar c$ は異なる役割を果たします．

まず **重み係数 $\rho_i$ の閾値 $\bar \rho$ は，*V-trace* オペレータの唯一の不動点を定義している** と考えることができます．関数近似誤差を生じない表形式の場合には，*V-trace* オペレータは，以下の式で表される方策 $\pi^{\bar \rho}$ の価値関数 $V^{\pi_{\bar \rho}}$ を唯一の不動点として持ちます(実際に計算すると，確認できます)．

```math
\pi_{\bar \rho}(a|x) := \frac{\min(\bar \rho \mu(a|x), \pi(a|x))}{\sum_b\min(\bar \rho \mu(b|x), \pi(b|x))}
```

よって，$\bar \rho$ が無限のとき，*V-trace* オペレータは方策 $\pi$ の価値関数 $V^\pi$ を唯一の不動点として持ち．$\bar \rho < \infty$ の場合には，$\pi$ と $\mu$ の間の方策の価値関数を不動点として持つことになります．**ISを行うと，行動方策 $\mu$ の確率密度が低いデータにおいて重み係数が非常に大きくなり，推定分散が大きくなってしまう**ので，$\bar \rho$ でクリップすることで分散を抑制します．よって，$\bar \rho$ が大きいときほど *off-policy* の学習における分散が大きく(一方でバイアスが小さく)，$\bar \rho$ が小さくなるにつれ分散が小さく(バイアスが大きく)なります．また，$\rho_i$ は $c_i$ と異なり時系列で掛け合わせる操作を行わないので，時系列によって発散するような挙動は生じません．

次に **重み係数 $c_i$ の閾値 $\bar c$ は，*V-trace* オペレータの収束の速さを制御している** と考えることができます．重み係数の掛け合わせ $(c_s\cdots c_{t-1})$ は，時刻 $t$ でのTD誤差 $\delta_t V =\rho_t (r_t + \gamma V(x_{t+1}) - V(x_t))$ が時刻 $s$ の価値関数 $V(x_s)$ の更新にどの程度寄与するかを評価しています．$c_i$ は時系列で掛け合わせる操作を伴うため発散しやすく，分散抑制のために重み係数 $c_i$ をクリップすることが重要です．ここで，**$\bar c$ の大きさは *V-trace* オペレータの不動点(学習が収束する点)には影響を与えない** ので，より分散を抑制するために $\bar \rho$ よりも小さな値に設定することが望ましいです．

実際には，*V-trace* は以下のように再帰的に計算することができます．

```math
\mathcal R^n V(x_t) = V(x_t) + + \mathbb E_{\mu}[
\delta_t V+ \gamma c_t (\mathcal R^n V(x_{t+1}) - V(x_{t+1}))
]
```


### V-trace Actor-Critic

方策 $\pi_\omega$ をパラメータ $\omega$ で，価値関数 $V_\theta$ をパラメータ $\theta$ で関数近似します．経験データは行動方策 $\mu$ で収集されたものとします．

価値関数の学習には，*V-trace* オペレータにおけるTD誤差を損失関数として用います．

```math
L_\theta = (\mathcal R^n V(x_s) - V_\theta(x_s))^2
```

勾配は容易に計算可能で，以下の式で表せます．

```math
\nabla_\theta L_\theta = (\mathcal R^n V(x_s) - V_\theta(x_s)) \nabla_\theta V_\theta(x_s)
```

また，方策勾配定理とISにより，方策 $\pi_{\bar\rho}$ の勾配は以下の式で表せます．

```math
E_{a_s\sim \mu}[\frac{\pi_{\bar \rho}(a_s|x_s)}{\mu(a_s|x_s)} \nabla_\omega \log \pi_{\bar \rho}(a_s|x_s) (q_s - b(x_s)) | x_s]
```

ここで，$q_s = r_s + \gamma \mathcal R^n V(x_{s+1})$ は *V-trace* オペレータの下での行動価値関数 $Q^{\pi_\omega}(x_s,a_s)$ の推定値を表し，$b(x_s)$ は分散を抑制するための状態依存のベースライン関数を表します．クリッピングよるバイアスが極めて小さい( $\bar \rho$ が十分大きい)場合，上述の勾配は $\pi_{\omega}$ の方策勾配の良い推定値であると考えられます．よって，ベースライン関数に$V_\theta(x_s)$ を用いることで，以下の方策勾配を得ます．

```math
\nabla_\omega L_\omega = \rho_s \nabla_\omega \log \pi_\omega(a_s|x_s) (
r_s + \gamma \mathcal R^n V(x_{s+1}) - V_\theta(x_s))
```

また，方策 $\pi_\omega$ が局所解に収束してしまうのを防ぐため，エントロピー損失を加えることも考えられます．

```math
L_{\rm ent} = - \sum_a \pi_\omega(a|x_s)\log \pi_\omega(a|x_s)
```

IMPALAでは，これら3つの損失関数を用いて学習を行います．

(検証実験の解説の需要がある方は，コメントお願いします．)