3. Distributional Reinforcement Learning with Quantile Regression

C51算法理论上用Wasserstein度量衡量两个累积分布函数间的距离证明了价值分布的可行性，但在实际算法中用KL散度对离散支持的概率进行拟合，不能作用于累积分布函数，不能保证Bellman更新收敛；且C51算法使用价值分布的若干个固定离散支持，通过调整它们的概率来构建价值分布。

而分位数回归(quantile regression)的distributional RL对此进行了改进。首先，使用了C51的“转置”，即固定若干个离散支持的均匀概率，调整离散支持的位置；引入分位数回归的思想，近似地实现了Wasserstein距离作为损失函数。

Quantile Distribution

假设$\mathcal{Z}_Q$是分位数分布空间，可以将它的累积概率函数均匀分为$N$等分，即$\tau_0,\tau_1...,\tau_N(\tau_i=\frac{i}{N},i=0,1,..,N)$。使用模型$\theta:\mathcal{S}\times \mathcal{A}\to \mathbb{R}^N$来预测分位数分布$Z_\theta \in \mathcal{Z}_Q$，即模型$\{\theta_i (s,a)\}$将状态-动作对$(s,a)$映射到均匀概率分布上。$Z_\theta (s,a)$的定义如下

\[Z_\theta (s,a):=\frac{1}{N}\sum_{i=1}^N \delta_{\theta_i(s,a)} \tag{1}
\]

其中，$\delta_z$表示在$z\in\mathbb{R}$处的Dirac函数

与C51算法相比，这种做法的好处：

不再受预设定的支持限制，当回报的变化范围很大时，预测更精确
取消了C51的投影步骤，避免了一些先验知识
使用分位数回归，可以近似最小化Wassertein损失，梯度下降不再有偏

Quantile Approximation

Quantile Projection

使用1-Wassertein距离对随机价值分布$Z\in \mathcal{Z}$到$\mathcal{Z}_Q$的投影进行量化：

\[\mathcal{\Pi}_{W_1}Z:={\arg\min}_{{Z_\theta}\in\mathcal{Z}_Q}W_1(Z,Z_\theta)
\]

假设$Z_\theta$的支持集为$\{\theta_1,...,\theta_N \}$，那么

\[W_1(Z,Z_\theta)=\sum_{i=1}^N \int_{\tau_{i-1}}^{\tau_i} |F_Z^{-1}(w)-\theta_i|dw
\]

其中，$\tau_i,\tau_{i-1}\in[0,1]$论文指出，当$F_Z^{-1}$是逆累积分布函数时，$F_Z^{-1}((\tau_{i-1}+\tau_i)/2)$最小。因此，量化中点为$\mathcal{\hat\tau_i}=\frac{\tau_{i-1}+\tau_i}{2}(1\le i\le N)$，且最小化$W_1$的支持$\theta_i=F_Z^{-1}(\mathcal{\hat\tau_i})$。如下图

【注】C51是将回报空间(横轴)均分为若干个支持，然后求Bellman算子更新后回报落在每个支持上的概率，而分位数投影是将累积概率(纵轴)分为若干个支持(图中是4个支持)，然后求出对应每个支持的回报值；图中阴影部分的面积和就是1-Wasserstein误差。

Quantile Regression

建立分位数投影后，需要去近似分布的分位数函数，需要引入分位数回归损失。对于分布$Z$和一个给定的分位数$\tau$，分位数函数$F_Z^{-1}(\tau)$的值可以通过最小化分位数回归损失得到

\[\mathcal{L}_{\text{QR}}^\tau(\theta):=\mathbb{E}_{\hat Z\sim Z}[\rho_\tau (\hat Z -\theta)],\quad \text{where} \quad \rho_\tau (u)=u(\tau-\delta_\{u<0\}),\forall u\in\mathbb{R}
\]

最终，整体的损失函数为

\[\sum_{i=1}^N \mathbb{E}_{\hat Z\sim Z}[\rho_{\hat{\tau}_i} (\hat Z -\theta)]
\]

但是，分位数回归损失在0处不平滑。论文进一步提出了quantile Huber loss：

\[\mathcal{L}_{\mathcal{K}}(u)=
\begin{cases}
& \frac{1}{2}u^2,\quad\quad\quad\quad \text{if} |u|\le \mathcal{K} \\
& \mathcal{K}(|u|-\frac{1}{2}\mathcal{K}),\,\, \text{otherwise}
\end{cases}
\]

\[\rho_{\tau}^{\mathcal{K}}(u)=|\tau-\delta_{\{u<0\}}|\mathcal{L}_{\mathcal{K}}(u)
\]

Implement

QR TD-Learning

QRTD算法(quantile regression temporal difference learning algorithm)的更新

\[\theta_i(s)\leftarrow \theta_i(s)+\alpha (\hat{\mathcal{\tau}}_i-\delta_{\{r+\gamma z^\prime < \theta_i (s) \}})
\]

$a\sim\pi (\cdot|s),r\sim R(s,a),s^\prime\sim P(\cdot|s,a),z^\prime\sim Z_\theta(s^\prime)$

其中，$Z_\theta$是由公式(1)给出的分位数分布，$\theta_i (s)$是状态$s$下$F_{Z^\pi (s)}^{-1}(\mathcal{\hat \tau}_i)$的估计值。

QR-DQN

QR-DQN算法伪代码

Append

1. Dirac Delta Function

\[\delta_a (x)=\delta (x-a)=0,(x\neq 0) \quad且\quad \int_{-\infty}^\infty \delta_a (x)d_x=1
\]

References

Will Dabney, Mark Rowland, Marc G. Bellemare, Rémi Munos. Distributional Reinforcement Learning with Quantile Regression. 2017.

Distributional RL