在传统的强化学习模型中，价值函数（value function）输出的是每个动作（action）在给定状态（state）下的价值 / 期望回报（expected return)。Distributional RL 则假设这个价值是一个随机变量，传统价值函数的输出是这个随机变量的均值。也就是说，传统价值函数的目标是近似地估计价值的期望，而 distributional RL 的目标是近似地估计价值的分布（概率密度函数）。
什么是 Distributional RL
Bellman’s Equation 描述了当前状态 x 下采用一个动作 a 的期望回报（即价值）等于期望奖励加未来的期望回报
$$Q(x, a) = \mathbb{E} R(x, a) + \gamma \mathbb{E} Q(X, A)$$
所以，学习价值的概率分布显然要比只学习它的期望更加靠谱。例如，假设 $(x, a)$ 的真实的概率分布是一个关于 $a$ 的取值对称双峰分布，那么单单预测一个固定的期望回报显然不能很好的评估当前状态的价值，在训练过程中也容易陷入 local minimum （一个峰）。而一个设置良好的概率分布受到的影响会小得多，比如我就假定价值分布是一个有两个中心的高斯混合模型。
Distributional RL 这个概念和 Bellman’s equation 一样“古老”，去年（2017年）DeepMind 在 ICML 上发表的一篇文章 A Distributional Perspective on Reinforcement Learning 将这个概念带到了一个新的高度。理论推导和实验都表明，论文中 Distributional RL 模型 C51 相比传统 DQN 在收敛速度和结果上都取得了显著的提升。因此，DeepMind 在这之后的 DQN 文章（如 Rainbow）大都将 distributional RL 当作标准配置。下面将详细讲一下论文的核心概念。
A Distributional Perspective on Reinforcement Learning（论文）
Distributional Bellman’s Equation
论文定义了一个随机变量 $Z$，它的均值是传统价值函数的输出，它的概率分布称为价值分布（value distribution），并由此得到 distributional Bellman’s equation
$$Z(x, a) = R(x, a) + \gamma Z(X’, A’)$$
传统的 Bellman’s equation 描述的是当前价值和未来的价值之间的关系。Distributional Bellman’s equation 描述的则是当前价值分布和未来价值分布之间的关系。公式中等号在这里的意思是，等号两边的随机变量是同分布的。
Parametric Distribution
要学习这个随机变量 $Z$，首先要确定它的概率分布是什么形式。之前有研究使用高斯分布，本文中使用的是参数分布（parametric distribution）的形式，即分布的形状由一组可学习的参数控制。具体的，价值分布是由 N 个参数 ${\theta_i}$ 控制的一个离散分布，它的支撑（support，概率不为0的点）由 $\{z_i\}$ 构成，也就是均匀分布在 $V_{\min}$ 和 $V_{\max}$ 中的N个点，概率分布定义如下
$$Z_\theta (x, a) = z_i \text{ w.p. } p_i (x, a) = \frac{e^{\theta_i(x, a)}}{\sum_je^{\theta_j(x, a)}} $$
模型的目标就是学习每个支撑的参数函数 $\{\theta_i\}$。与 value-based 模型类似，价值分布训练的过程就是给定一个 $\text{transition} (x, a, r, x’)$，要让 $Z(x, a)$ 和 $\mathcal{T}Z = r + \gamma Z(x’, \pi(x’))$ 的分布尽可能相同，其中 $\pi$ 是 greedy 或者 $\epsilon$ greedy policy。
Categorical Projection
上述的分布形式有一个问题，就是在训练的过程中 distributional Bellman’s equation 右边计算得到的值往往不在这些支撑上（由于 reward 和 $\gamma$）， 导致我们无法更新支撑的概率参数。本文就提出了一个很巧妙的方法解决这个问题，文中称为 categorical projection。因为计算得到的 sample Bellman update $\hat{\mathcal{T}}Z_{\theta}$ 应该和 $Z_{\theta}$ 同分布，给定 $\text{transition} (x, a, r, x’)$，$\hat{\mathcal{T}}z_j = \mathcal{T}Z_{\theta_j}(x, a)$ 的概率和 $z_j = Z_{\theta_j}(x’, \pi(x’))$ 的概率 $p_{j} (x’, \pi(x’))$ 相同。由于 $\hat{\mathcal{T}}z_j$ 有可能不是支撑，那么我们就把 $p_j (x’, \pi(x’))$，也就是它的概率按比例分配到距离它最近的支撑上，这也是论文中公式（7）的含义。
$$(\Phi\hat{\mathcal{T}}Z_\theta(x, a))_i = \sum^{N-1}_{j=0} \left[1-\frac{|[\hat{\mathcal{T}}z_j]^{V_{\text{MAX}}}_{V_{\text{MIN}}} - z_i|}{\Delta z}\right]^1_0 p_j(x’, \pi(x’))$$
最后，根据 distributional Bellman’s equation，分配过后得到的分布 $\Phi \hat{\mathcal{T}}Z_{\theta}$ 和 $Z_{\theta}$ 应该是同分布的，因此以 KL divergence 为损失函数，我们可以通过梯度下降的方法来训练模型。上述过程在论文中被称为 Categorical Algorithm，用此方法训练的 DQN 模型被称为 Categorical DQN，模型C51的含义就是有51个支撑的 Categorical DQN。
网络结构及 Tensorflow 实现
模型的网络结构与一般的DQN类似，区别在于：传统 DQN 最后一层输出的是一个 $N$ 维的向量，表示 $N$ 个动作对应的价值，Categorical DQN 输出的是 $N \times M$ 维的向量，表示的是 $N$ 个动作在 $M$ 个价值分布的支撑上的概率。Google 新发布的强化学习研究框架 Dopamine 中有一个很简洁的实现，还附带了一个例子说明每一步计算是做什么的，可以加深对算法的理解。
其他形式的 Distributional RL
除了 Categorical DQN 这种类型，研究人员还探索了不同形式的 Distributional RL：
- Distributed Distributional Deterministic Policy Gradients (D4PG，论文)。D4PG 将 categorical 这种形式用在 Actor-Critic 类强化学习算法的 Critic 网络上。D4PG 的想法是很自然且合理的，因为 Critic 网络本身就是一个价值判断网络，与 DQN 中的 value function 起到相同的作用。这篇文章说明 distributional RL 的思想也适用于 Continuous Control 领域，效果还不错。
- Implicit Quantile Networks for Distributional Reinforcement Learning (IQN，论文，代码)。Categorical 的方法虽然将预测期望回报拓展到了预测价值分布，但是模型在做决策时仍然使用的是期望回报（见上小节 Algorithm 1 里的前3行）。并且，Categorical 方法要确定一组先验且离散的支撑。通过估计分布的百分点函数 （quantile function）， IQN 可以输出连续的期望回报，这种灵活性使得它的性能相比 C51 有很大的提升。IQN 是之前的 QR-DQN 模型（论文）的加强版。