この記事は，以下の論文の解説です．
Implicit Quantile Networks for Distributional Reinforcement Learning (2018)
記事内容では，強化学習の基礎的な知識を前提としています．
また，記事中の図は全て論文からの引用です．
不備がございましたら，ご指摘頂けると幸いです．
また，この論文で提案する手法(IQN)の再現実装をGithubで公開しています．アルゴリズムがわかりやすいように工夫して，コメントも詳しく書いていますので，理解の助けになればと思います．ぜひ参照してみてください．
概要
この論文は，「価値分布の分位点関数を暗にモデル化し，CVaRなどのリスク指標に基づいた方策を学習できる分布型強化学習の手法を提案した」論文になります．
主に先行研究の結果がベースになっていますので，ご存じない方は先にお読み頂くことをオススメします．
C51: A Distributional Perspective on Reinforcement Learning (論文，解説記事)
QR-DQN: Distributional Reinforcement Learning with Quantile Regression (論文，解説記事)
分布型強化学習の手法は，価値分布のモデル化の仕方と学習の仕方の2つで特徴付けることができます．C51では，固定した価値での確率を出力する関数(=離散化した価値分布)をモデル化し，KLダイバージェンスの最小化により学習を行いました．課題として，
- 固定した点のみでモデル化しているため，モデルの表現力が制限されること
- 価値の上限・下限を事前に定める必要があること
- 縮小性が示されているWasserstein距離ではなく，KLダイバージェンスを最小化していること
などがありました．
QR-DQNでは，固定した(累積)確率での価値を出力する関数(=離散化した分位点関数)をモデル化し，分位点回帰を用いてWasserstein距離の最小化により学習を行いました．C51の課題のうち，後者2つは解決したものの，離散化によりモデルの表現力が制限される点は改善されていません．
この論文では，QR-DQNをベースに，離散化した分位点関数をモデル化するのではなく累積確率から価値への連続な写像(=分位点関数全体)を暗にモデル化することで，この課題を解決していきます．
準備
以下では，通常のMDP $<\mathcal X, \mathcal A, R, P, \gamma>$ を考えます．$R$ は報酬関数を表し，確率変数です．また，方策 $\pi$ に従ったときの収益を $Z^\pi = \sum_{t=0}^\infty \gamma^t R_t$ と表します．確率変数 $Z^\pi$ の従う確率分布を価値分布と呼びます．また，価値分布のことも単に $Z^{\pi}$ と表すことにします．
分布ベルマンオペレータ $\mathcal T^\pi$ は，以下の式で表されます．ただし，$=_D$ は両辺の確率変数が従う確率分布が等しいことを表します．
\mathcal T^\pi Z(x,a) =_D R(x,a) + \gamma Z(X', A')
同様に，最適分布ベルマンオペレータ $\mathcal T$ は，以下の式で表されます．これは，状態 $x$ で行動 $a$ を選択したあとは，$Z$ に基づく貪欲方策に従って行動選択し続けたときの将来的な価値を表していると考えることができます．
\mathcal T Z(x,a) =_D R(x,a) + \gamma Z(X', arg\max_{a'} \mathbb E[Z(X', a')])
分布ベルマンオペレータ $\mathcal T^\pi$ はp-Wasserstein距離の上界に関して縮小であるものの，最適分布ベルマンオペレータ $\mathcal T$ はいかなる計量に関しても縮小にならないことがC51の論文で示されています．ただし，最適分布ベルマンオペレータを繰り返し適用することで，価値分布の期待値が最適価値関数に収束することは示すことができます．
p-Wasserstein距離
p-Wasserstein距離： $W_p (p\in [1,\infty])$ は，累積分布関数の逆関数(=分位点関数)の間の$L_p$ノルムで表されます．よって，確率分布 $U, Y$ の間の$W_p$は，以下の式で表されます．ただし，$F_Y^{-1}(\omega)=\inf {y \in \mathbb R : \omega \le F_Y(y)}$ は，確率変数 $Y$ の分位点関数を表します．
W_p(U,V) = \left(\int_0^1 \mid F_Y^{-1}(\mathcal \omega) - F_U^{-1}(\omega)\mid^p d\omega\right)^{1/p}
以下では，p-Wasserstein距離を単に $W_p$ と呼びます．また，$W_p$ の上界を $\bar d_p$ と表します．
QR-DQN
QR-DQNの論文では，価値分布の適切な累積確率( $\hat \tau_i$ )における分位点関数の値(=分位点)を計算することで，分布ベルマンオペレータ $\mathcal T^\pi$ とある射影 $\Pi_{W_1}$ を組み合わせた操作が $\bar d_\infty$ に関して縮小となることを示しました．また，任意の累積確率における分位点を推定することができる分位点回帰を用いることで，価値分布 $Z$ に対する射影 $\Pi_{W_1} Z$ を計算できることを示しました．
従って，分位点回帰を用いて $Z_\theta = \Pi_{W_1}\mathcal T^\pi Z_\theta$ となるようにパラメータの更新を繰り返すことで，$\bar d_\infty(\Pi_{W_1}\mathcal T^\pi Z_\theta, Z_\theta)$ が縮小していき，$Z_\theta$ が真の価値分布 $Z$ に収束します．また，$\mathcal T^\pi$ の代わりに $\mathcal T$ を用いることで，分位点関数の期待値 $\mathbb E[Z_\theta]$ が最適価値関数 $Q^{*}$ に収束していくことが期待されます．
QR-DQNでは，価値分布 $Z$ を一様なディラック関数の混合分布でモデル化しました．
Z_\theta(x,a) = \frac{1}{N}\sum_{i=1}^N\delta_{\theta_i(x,a)}
ただし，$\theta_i$ は累積確率 $\hat \tau_i = \frac{\tau_{i-1} + \tau_i}{2} (i=1,\cdots,N)$ での分位点の推定値で，$\tau_i = i / N$ は一様な累積確率です．
この $\theta_i$ の推定には，下記の閾値 $\kappa$ を持つquantile Huber lossを用います(分位点回帰)．ただし，推定値のターゲットとして $\mathcal T \theta'_j = r + \gamma \theta'_j(x', \pi(x'))$ を用い，TD誤差を $\delta_{ij} = \mathcal T \theta'_j - \theta_i $ と表します．
\begin{align}
\rho_\tau^\kappa(\delta_{ij}) &= \mid\tau - \mathcal I_{\{\delta_{ij}<0\}}\mid \frac{\mathcal L_\kappa(\delta_{ij})}{\kappa} \\
\mathcal L_\kappa(\delta_{ij})&= \begin{cases}\frac{1}{2} \delta_{ij}^2 \;\;\;\;\;\;\;\;\;\;\:\;\;\;\;\; |\delta_{ij}| \le \kappa \\\kappa(|\delta_{ij}|-\frac{1}{2}\kappa) \;\;\;\;\;{\rm otherwise}\end{cases}
\end{align}
実際には，任意の $i,j \in [1,N]$ についてこの損失関数を最小化します．これは $Z_\theta = \Pi_{W_1}\mathcal T Z_{\theta'}$ となるようにパラメータ $\theta$ を更新していることになり，分位点関数の期待値 $\mathbb E[Z_\theta(x,a)]$ が最適価値関数 $Q^{*}(x,a)$ に収束していくことが期待されます．
歪みリスク指標
これまでの分布型強化学習では，価値分布を推定しているのにも関わらず，価値分布の期待値(=価値関数 $Q$ )の最大化に基づいた方策を学習していました．すなわち，以下のような方策を獲得することを目標をしていました．
\pi(x) = arg\max_{a} \mathbb E[Z(x,a)]
しかし，価値分布にはリスクなどの有益な情報が含まれており，期待値を最大化するだけでなく，リスクを考慮した方策を考えることもできそうな気がします．ここでは，単調増加関数 $h: [0,1]\to [0,1]$ を歪みリスク指標と呼び，以下のような方策を考えます．
\pi(x) = arg\max_a\int_{-\infty}^\infty a \frac{\partial}{\partial z}(h \circ F_{Z(x,a)})(z)dz
この方策は，累積分布関数 $F_Z$ を $h$ で重み付けした(歪ませた)上で期待値をとり，それを最大化する方策と考えることができます．$h$ に凸関数を用いたときは，価値が低くなる行動の重みを低くすることになるので risk-seeking といい，凹関数を用いたときは risk-averse と呼びます．また $h$ に恒等写像を用いたとき，すなわち従来の強化学習のときは risk-neutral と呼びます．
このように，歪みリスク指標を用いることでリスク考慮した方策を考えることができます．
Implicit Quantile Networks (IQN)
この論文で提案する手法，**Implicit Quantile Networks (IQN)**について説明していきます．
以下では，$F_Z^{-1}(\tau)$ を $\tau\in[0,1]$ における確率変数 $Z$ の分位点関数の値とし，単に $Z_\tau$ と表します．また，価値分布の分位点関数 $F_Z^{-1}$ のことも単に $Z$ と表します．このとき，一様分布からのサンプル $\tau \sim \mathcal U([0,1])$ を用いると，$Z_\tau(x,a) \sim Z(x,a)$ となることがわかります(注: この右辺は分位点関数ではなく確率分布を表します)．
この論文では，離散化した分位点関数をモデル化するのではなく，分位点関数全体を暗にモデル化します．つまり，一様分布からのサンプル $\tau \sim \mathcal U([0,1])$ から $Z_\tau(x,a)$ への写像をモデル化します．これにより，十分な数のサンプル $\tau$ について $Z_\tau(x,a)$ を計算することで，分位点関数 $F_Z^{-1}(x,a)$ を暗に表現することができます．
また，$\beta: [0,1] \to [0,1]$ を歪みリスク指標とします．恒等写像の場合はrisk-neutralとなります．このとき，リスク考慮した $Z(x,a)$ の期待値は，以下の式で表されます．
Q_\beta(x,a) := \mathbb E_{\tau\sim\mathcal U([0,1])}[Z_{\beta(\tau)}(x,a)]
この期待値は，歪みリスク指標 $\beta$ で累積分布関数を重み付けした期待値 $Q = \int_0^1 F_Z^{-1}(\tau)d\beta(t)$ と等価になります．すなわち，リスク考慮した方策に対応する価値関数を考えていることになります．従って，リスク考慮した方策を以下のように表すことができます．
\pi_\beta (x) = arg\max_{a\in \mathcal A} Q_\beta(s,a)
学習
2つの独立なサンプル $\tau, \tau' \sim \mathcal U([0,1])$ と方策 $\pi_\beta$ について，時間 $t$ でのTD誤差を以下のように表します．
\delta_t^{\tau,\tau'} = r_t + \gamma Z_{\tau'} (x_{t+1},\pi_\beta(x_{t+1})) - Z_\tau(x_t,a_t)
上式は，価値分布からのサンプルを用いてTD誤差を計算していることを意味します．十分な数のサンプルでこのTD誤差を計算することで，価値分布全体に関するTD誤差を計算していることになります．
従って，損失関数は以下の式で表せます．ただし，$\rho_{\tau_i}^\kappa$ はquantile Huber lossで，$N,N'$ は損失関数を推定するための独立なサンプル数です．
\mathcal L (x_t,a_t,r_t,x_{t+1}) = \frac{1}{N'}\sum_{i=1}^{N}\sum_{j=1}^{N'}\rho_{\tau_i}^\kappa (\delta_t^{\tau,\tau'})
この損失関数は，十分な数($N,N'$)のサンプルで分位点回帰を行うことを意味しています．よって，分位点関数全体が最適化されていき，価値分布の期待値が最適価値関数 $Q^{*}(x,a)$ に収束していくことが期待できます．
分位点回帰を用いて学習した $Z_{\beta}(x,a)$ と $K$ 個の独立なサンプル $\tilde \tau_k \sim \mathcal U([0,1])$ を用いて，リスク考慮した方策を以下のように近似することができます．
\hat \pi_\beta(x) = arg\max_{a\in \mathcal A} \frac{1}{K}\sum_{k=1}^K Z_{\beta(\tilde \tau_k)}(x,a)
実装
DQNで用いられた2つのネットワーク $\psi: \mathcal X \to \mathbb R^d$ (CNN) と $f: \mathbb R^d \to \mathbb R^{|A|}$ (MLP) を考えます．DQNでは，価値関数を $Q(x,a) \simeq f(\psi(x))_a$ のように推定していました．IQNでは，さらにもう1つのネットワーク $\phi: [0,1] \to \mathbb R^d$ を考えます．この $\phi$ はサンプリングした点 $\tau$ の特徴量を計算する関数を表します．これらを用いて，価値分布を以下のように近似します．ただし，$\odot$ はアダマール積です．
Z_\tau(x,a) \simeq f(\psi(x) \odot \phi(\tau))_a
IQNでは，これら3つのネットワークを用いて価値分布 $Z(x,a)$ をモデル化しました．ただし，サンプリングした点の特徴を計算する $\phi$ には，以下のような写像を用いました．( $n=64$ )
\phi_j(\tau) = {\rm ReLU}(\sum_{i=0}^{n-1}\cos(\pi i \tau)w_{ij} + b_{ij})
ハイパーパラメータ$N$ は分位点回帰を行う点の数を表しており，IQNのサンプル効率に関係すると考えられます．また $N'$ は分位点回帰の際のターゲットの点の数を表しており，勾配の推定の分散に関係すると考えられます．
下図は，$N,N'$ を変化させてAtariの6環境を学習させたときの1000万ステップの平均スコア(左)と最終スコア(右)です．これより，$N$ は大きい方が性能が上がるが，$N'$ はある程度大きければあまり影響しないことがわかります．よって，以下の実験では全て $N=64, N'=8$ を用いました．また，$K=32$ を用いました．
リスク考慮型強化学習
歪みリスク指標 $\beta$ を恒等写像から別の写像に変換することで，リスク考慮した方策を学習することができます．この記事では説明しませんが，IQNではいくつかのリスク指標(e.g. CVaR)に基づいた学習を行い，検証を行っています．気になる方は，論文を参照してみてください．
検証
Atari-57ベンチマークで検証します．比較対象はQR-DQNやRainbowなどです．ここで注意すべきは，RainbowはPrioritized Replay等のテクニックを多用しているのに対し，IQNではそれらを用いていない点です．また， IQNではrisk-neutralな方策を用いて学習しています．
上表・図から，IQNはQR-DQNを大きく上回り，Rainbowにも匹敵する性能を示すことがわかりました．
まとめ
この論文では，QR-DQNを拡張し，離散化した分位点関数ではなく分位点関数全体を学習する手法 IQNを提案しました．また検証実験では，IQNはQR-DQNを大きく上回る性能を示しました．
今後解決すべき問題としては
- サンプルベースの分位点回帰を用いた強化学習では，収束を保証することができるのか
- QR-DQNで示した(適切に)離散化した分位点関数の縮小性が，IQNのような連続な関数でも成り立つのか
- リスク考慮した方策に関して，通常の強化学習のように収束を保証することができるのか
などが挙げられます．