この記事は，以下の論文の解説です．
Fully Parameterized Quantile Function for Distributional Reinforcement Learning (NIPS 2019)
記事内容では，強化学習の基礎的な知識を前提としています．
また，記事中の図は全て論文からの引用です．
不備がございましたら，ご指摘頂けると幸いです．
また，この論文で提案する手法(FQF)の再現実装をGithubで公開しています．アルゴリズムがわかりやすいように工夫して，コメントも詳しく書いていますので，理解の助けになればと思います．ぜひ参照してみてください．
概要
この論文は，近年の分布型強化学習の進展に基づき，連続な価値分布の分位点関数を効率的・柔軟に近似可能な分布型強化学習の手法を提案した論文になります．
主に先行研究の結果がベースになっていますので，ご存じない方は先にお読み頂くことをオススメします(特にIQN)．
C51: A Distributional Perspective on Reinforcement Learning (論文，解説記事)
QR-DQN: Distributional Reinforcement Learning with Quantile Regression (論文，解説記事)
IQN: Implicit Quantile Networks for Distributional Reinforcement Learning (論文，解説記事)
近年の分布型強化学習では，価値分布をよりよく近似するためにどのようにモデル化するかが重要な課題となっていました．価値分布 $Z(x,a)$ は将来に得られる価値の確率分布のことを指し，価値分布の期待値は価値関数 $Q(x,a)$ に一致します．
C51では，価値の上限・下限を事前に定め，一様に離散化した価値における確率を出力するニューラルネットワークを用いました．すなわち，$N$個のビンで離散化した価値分布をモデル化しました．
QR-DQNでは，価値分布の累積確率分布の逆関数(=分位点関数)を考え，一様に離散化した累積確率における価値(分位点)を出力するニューラルネットワークを用いました．すなわち，$N$個のビンで離散化した分位点関数をモデル化しており，価値の上限・下限がハイパーパラメータであるというC51の課題を解決しました．しかし，C51もQR-DQNも固定した点での確率または価値を用いて近似を行っており，近似した分布の表現力が制限されてしまうという課題がありました．
IQNでは，QR-DQNのように固定した累積確率ではなく，一様分布から累積確率をサンプリングし，累積確率から価値への写像をニューラルネットワークでモデル化しました．すなわち，十分な数のサンプリングとそれに伴う写像を行うことで，分位点関数を暗にモデル化しました．IQNでは，十分なネットワークの表現力と無限のサンプリングを行うことで，分位点関数全体を正確に近似できると指摘されています．しかし，実際に無限のサンプリングを行うことは出来ないという課題が生じます．
この論文では，IQNの「一様分布から累積確率をサンプリングする」部分を修正し，より良い近似を行うための累積確率を出力する確率提案ネットワークを用いることを提案しています．IQNと同様の役割を持つ分位点ネットワークと共に用いることで，より良い価値分布の近似が可能になると期待されます．
準備
以下では，通常の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^\pi(x,a) =_D R(x,a) + \gamma Z^\pi(X', A')
この分布ベルマンオペレータは，p-Wasserstein距離の上界 $\bar d_p$ に関して縮小であることが示されています(C51の論文付録参照)．すなわち，分布ベルマンオペレータを適用していけば，方策 $\pi$ に関する価値分布に指数関数的に収束していくことが示されています．また，価値分布 $Z_1,Z_2$ の間の $\bar d_p$ は以下の式で表されます．ただし，$W_p$ はp-Wasserstein距離を表します．
\bar d_p(Z_1,Z_2) = \sup_{x,a} W_p(Z_1(x,a),Z_2(x,a))
QR-DQN
QR-DQNでは，価値分布そのものではなく分位点関数をモデル化することを考え，$N$ 個のビンで一様に離散化した累積確率 $\tau_i = i ;/; N$ での価値 $\theta_i(x,a)$ をニューラルネットワークで表現しました．これは価値分布を，一様な $N$ 個のデルタの混合分布で近似していることと等価です．
Z_\theta(x,a) = \frac{1}{N}\sum_{i=1}^N \delta_{\theta_i(x,a)}
$\theta_i(x,a)$ が累積確率 $\hat \tau_i =(\tau_{i-1} + \tau_i) ; / ; 2$ のときの真の価値(分位点)と等しくなるとき，真の価値分布と近似した価値分布の1-Wasserstein距離が最小となることが示されています(上図，QR-DQNの論文参照)．
この $\theta_i$ の推定には，下記の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}) &= |\tau - \mathbb I_{\{\delta_{ij}<0\}}|\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}
この分位点回帰を繰り返すことで，近似した価値分布の期待値が真の価値関数 $Q(x,a)$ に収束していくことが示されています(QR-DQNの論文参照)．
IQN
IQNでは，固定した累積確率 $\tau_i$ ではなく，一様分布からサンプリングした累積確率 $\tau \sim \mathcal U([0,1])$ を用いて分位点関数を近似することを提案しました．それに伴い，サンプリングされた累積確率(に対応する潜在表現)から価値(分位点)への写像をニューラルネットワークでモデル化しました(このネットワーク自体もIQNと呼びます)．
十分なネットワークの表現力と無限のサンプリングにより，IQNは分位点関数全体を正確に近似できると指摘されています．しかし，実際に無限のサンプリングを行うことは出来ないという課題があります．
Fully parameterized Quantile Function (FQF)
この論文では，FQFと呼ばれる分布型強化学習の手法を提案します．
FQFでは，確率提案ネットワーク(Fraction Proposal Network) : $P_{\omega_1}$ と分位点ネットワーク(Quantile Value Network) : $F_{Z,\omega_2}^{-1}$ の2つをニューラルネットワークでモデル化することで，分位点関数全体を近似します．
分位点関数の近似
FQFでは，$N$ 個の累積確率と，それらの累積確率での価値(分位点)をそれぞれモデル化します．すなわち，分位点関数の $x$ 軸における点(累積確率)と $y$ 軸における点(価値)の両方を可変とすることで，分位点関数の近似をよりよく行うことを目指します．
近似した価値分布は，$N$ 個のデルタ関数の重み付き平均で表されます．ただし，累積確率 $\tau_i (i\in[0,N])$ は $\tau_i < \tau_{i+1}$ を満たし，$\tau_0 = 0, \tau_N = 1$ とします．$\theta_i$ に関しては後述します．
Z_{\theta,\tau} (x,a) = \sum_{i=0}^{N-1} (\tau_{i+1}-\tau_i)\delta_{\theta_i(x,a)}
また，ある確率分布変数 $Z$ の分位点関数を $F_Z^{-1}$ とします．分位点関数 $F_Z^{-1}$ から，前述のように $\theta_i$ と $\tau_i$ を用いて近似した関数 $F_Z^{-1,\theta,\tau}$ への写像 $\Pi_{\theta,\tau}$ は，以下の式で表されます．ただし，$H$ はステップ関数で，$H_\tau(\omega) = H(\omega-\tau)$ を表します．
F_Z^{-1,\theta.\tau}(\omega) = \Pi^{\theta,\tau} F_Z^{-1}(\omega) = \theta_0 + \sum_{i=0}^{N-1} (\theta_{i+1} - \theta_i) H_{\tau_{i*1}}(\omega)
上図は分位点関数(青色)とその近似(灰色)の例を示しており， $\theta_i$ と $\tau_i$ を用いて近似することで $(\tau_i,\theta_i)$ の点で段差が生じていることが分かります．FQFでは，まず確率提案ネットワークを用いて累積確率の集合 $\tau$ を生成し，分位点ネットワークを用いてその累積確率における分位点 $\theta$ を求めます．
あとで説明するように，分位点ネットワークの学習は1-Wasserstein距離 $W_1$ の最小化により行います．$\theta_i$ と $\tau_i$ を用いて近似した分位点関数 $F_Z^{-1,\theta,\tau}$ と真の分位点関数 $F_Z^{-1}$ の $W_1$ は，以下の式で表されます．
W_1(Z,\theta,\tau) = \sum_{i=0}^{N-1} \int_{\tau_i}^{\tau_{i+1}}|F_Z^{-1}(\omega)-\theta_i|d\omega
(注: $W_1$ は確率分布間の距離を表しますが，分位点関数と確率分布は等価なので，分位点関数に対応する確率分布間の $W_1$ のことを単に分位点関数間の $W_1$ と呼んでいます．)
上図では，いい感じの累積確率を利用した場合(左)とIQNのようにランダムにサンプリングした累積確率を利用した場合(右)の $W_1$ の例を示しています．いい感じの累積確率を利用することで，よりよく分位点関数を近似できることが分かります．FQFでは，よりよく分位点関数を近似できる累積確率を生成することを目標に，確率提案ネットワークを学習します．
確率提案ネットワークの学習
まず $\tau$ を固定したときの最適な $\theta$ を見つけることを考えます．
補題1
$\tau_{i-1}<\tau_i$ を満たす任意の $\tau_1,\cdots,\tau_{N-1} \in [0,1]$ について，$W_1$ を最小化する $\theta$ は以下の式で表される．ただし，$\tau_0 = 0, \tau_N = 1$ とする．
\theta_i = F_Z^{-1}(\frac{\tau_i + \tau_{i+1}}{2})
この補題1では， $\theta_i$ が累積確率 $\hat \tau_i = (\tau_i + \tau_{i+1}) ; / ; 2$ における真の分位点 $F_Z^{-1}(\hat \tau_i)$ に等しいとき，真の価値分布と近似した価値分布の $W_1$ は最小になることを示しています．これを $W_1$ の式に代入すると，以下の命題1を得ることができます(証明は論文付録参照)．
命題1
任意の単調増加する分位点関数 $F_Z^{-1}$ について，$F_z^{-1}$ と $F_Z^{-1,\tau}$ の間の1-Wasserstein lossを以下のように定義する．
W_1(Z,\tau) = \sum_{i=0}^{N-1}\int_{\tau_i}^{\tau_{i+1}} |F_Z^{-1}(\omega)-F_Z^{-1}(\hat \tau_i)|d\omega
このとき，勾配 $\frac{\partial W_1}{\partial \tau_i}$ は以下の式で与えられる．( $\forall i \in (0,1)$ )
\frac{\partial W_1}{\partial \tau_i} = 2F_Z^{-1}(\tau_i)-F_Z^{-1}(\hat\tau_i)-F_Z^{-1}(\hat\tau_{i-1})
さらに，$\tau_{i-1}<\tau_{i+1}$ となる任意の $\tau_{i-1},\tau_{i+1}\in[0,1]$ について，$\frac{\partial W_1}{\partial \tau_i} = 0$ となる $\tau_i \in (\tau_{i-1},\tau_{i+1})$ が必ず存在する．
この命題1では，$W_1$ を最小化するために $\tau_i$ をどの方向に動かしたらいいかを表す勾配は正確に計算できることを示しています．一般的に $W_1$ をバイアスなしで計算することは難しいのですが，この命題1の勾配を用いることで，それを直接計算せずに $W_1$ を最小化するような $\tau_i$ を最適化することができます．
パラメータ $\omega_1$ を持つ確率提案ネットワーク $P_{\omega_1}$ を，この勾配をもとに確率的勾配降下法を用いて $W_1$ の最小化を行うことで，$\omega_1$ が収束することが保証されています．実際には真の分位点関数 $F_Z^{-1}$ を知ることはできないので，代わりにパラメータ $\omega_2$ を持つ分位点ネットワーク $F_{Z,\omega_2}^{-1}$ を用いて勾配を計算します．
FQFの価値分布の期待値，すなわち価値関数は以下の式で表されます．ただし，$\tau_0=0,\tau_N=1$ です．
Q(x,a) = \sum_{i=0}^{N-1}(\tau_{i+1}-\tau_i) F_{Z,\omega_2}^{-1} (\hat \tau_i)
分位点ネットワークの学習
以下では，適切に選ばれた累積確率 $\tau_i$ のもとで，分位点関数を分位点ネットワーク $F_{Z,\omega_1}^{-1}$ を用いて近似することを考えます．ただし，分位点ネットワークの学習は，累積確率の選び方以外はIQNと(大体)同じなので，概要のみを説明していきます．
まず $\tau_i$ が選ばれたとき，近似した分布と真の分布の $W_1$ を最小化する $\theta_i$ は，$\theta_i = F_Z^{-1}(\hat \tau_i)$ で与えられました(補題1)．しかし，真の分位点関数 $F_Z^{-1}$ は知ることができないので，最適分布ベルマンオペレータを適用後の近似した分位点関数 $\mathcal T F_{Z,\omega_1}^{-1}$ を推定のターゲットとします．分位点回帰を用いて $W_1( F_{Z,\omega_1}^{-1},\mathcal T F_{Z,\omega_1}^{-1})$ を最小化するように $\omega_1$ を更新していくことで，近似した価値分布の期待値が( $\tau_i$ が固定された際の)価値関数 $Q(x,a)$ に収束していきます(QR-DQNの論文参照)．
また，分位点回帰は以下の流れで行います．$Z$ を $(x_t,a_t)$ に関する価値，$Z'$ を $(x_{t+1},a_{t+1})$ に関する価値とすると，それぞれ累積確率 $\hat \tau_i,\hat \tau_j$ を用いたときのTD誤差は以下の式で表されます．
\delta_{ij}^t = r_t + \gamma F_{Z',\omega_1}^{-1} (\hat \tau_i) - F_{Z,\omega_1}^{-1}(\hat \tau_j)
分位点回帰では，以下のquantile Huber loss : $\rho_\tau^\kappa$ を用いて学習を行います．
\begin{align}
\rho_\tau^\kappa(\delta_{ij}) &= |\tau - \mathbb I_{\{\delta_{ij}<0\}}|\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}
この分位点回帰を，現在の分位点関数とターゲットの分位点関数の両方に関して十分な数の $\tau$ を用いて行うことで，分布全体が更新されていきます．従って，分位点ネットワークの損失関数は，以下の式で表すことができます．
\mathcal L(x_t,a_t,r_t,x_{t+1}) = \frac{1}{N}\sum_{i=0}^{N-1}\sum_{j=0}^{N-1}\rho_{\hat \tau_j}^\kappa (\delta_{ij}^t)
(注: IQNでは $\tau_i, \tau_j$ において分位点の値を計算するのに対して，FQFでは，補題1に沿って $\hat \tau_i, \hat \tau_j$ において分位点の値を計算しているという違いがあります．)
まとめると，FQFは以下のようなアルゴリズムになります．
実装
FQFでは，DQNと同様の構造のCNN $\psi : \mathcal X\to \mathbb R^d$ で状態を潜在ベクトルへと変換します．
確率提案ネットワーク $P_{\omega_2}: \mathbb R^d\to (0,1)^N$ は1層の全結合層で構成され，状態の潜在ベクトル $\psi(x)$ を入力として受け取り，累積確率の提案 $\tau_i$ を出力します．ただし，累積確率 $\tau_i$ は，$\tau_{i-1} < \tau_i$ かつ $\tau_0 = 0,\tau_N=1$ を満たす必要がありました．確率提案ネットワークの $N$ 次元出力の最後にSoftmax関数を適用すると，その出力 $q\in \mathbb R^N$ は，$q_i \in (0,1),; i\in [0,N-1]$ と $\sum_{i=0}^{N-1} q_i = 1$ を満たします．よって，$\tau_i = \sum_{j=0}^{i-1}q_j,; i \in [0,N]$ とすることで，この条件を満たすことができます．
また，IQNと同様に，累積確率 $\tau$ は潜在表現 $\phi(\tau)$ に変換して用います．ただし，$\omega_{ij}$ と $b_j$ はニューラルネットワークのパラメータです．
\phi_j(\tau) = {\rm ReLU}\left( \sum_{i=0}^{N-1}\cos(i\pi\tau)\omega_{ij} + b_j\right)
その後，状態の潜在ベクトル $\psi(x)$ との要素積を取り，分位点ネットワーク $F_{Z,\omega_2}^{-1}$ に入力することで，価値(分位点)を計算します．
F_Z^{-1}(\tau) \simeq F_{Z,\omega_2}^{-1} (\psi(x)\odot \phi(\tau))
よって，FQFで近似した価値分布の期待値(=価値関数 $Q$ )は，以下の式で計算されます．
Q = \frac{1}{N}\sum_{i=0}^{N-1}(\tau_{i+1}-\tau_i)F_Z^{-1}(\frac{\tau_i+\tau_{i+1}}{2})
検証
Arcade Learning Environment(ALE)のAtariの環境で検証．比較手法としては，IQNやRainbowなどを用いました．
上図は，55ゲームの人間のスコアを基準とした性能を示しています．FQFは，Prioritized Experience Replay等のテクニックを用いていないのにも関わらず，様々なテクニックを併用した手法(Rainbow)も含め，他の全手法を上回る性能を示していることがわかります．
一方，FQFの学習速度はIQNに比べて20%ほど遅く，IQNでは $\tau$ のサンプル数を増やしてもそこまで速度が落ちないのに対し，FQFは $\tau$ のサンプル数の上昇に伴い速度が大きく落ちることが指摘されています．
まとめ
この論文では，近年の分布型強化学習の進展に基づき，**連続な価値分布をより効率的・柔軟に近似する手法(FQF)**を提案しました．FQFでは，価値分布の分位点関数の両軸(累積確率，価値)をパラメトライズすることで，より良い近似を行えることを検証実験により示しました．この検証実験では，様々なテクニックを併用した手法(Rainbow)の性能を上回る性能を示しました．
FQFの課題としては，(個人的には)以下のようなものがあると考えました．
- 命題1では分位点関数の単調増加性を仮定しているものの，分位点ネットワークでは単調増加性は保証されず，確率提案ネットワークの学習が不安定になってしまう点
- IQNに比べて計算コストが大きい点