C51

经典强化学习与值分布强化学习的区别

由于状态转移的随机性，状态表示的混叠效应（编码状态时带来的信息丢失）以及函数逼近的引入（使用函数表示状态价值带来的信息丢失），状态价值通常具有一定随机性，并可以用某种分布来表示。经典强化学习（DQN算法等）使用单个值的形式预测状态价值可能会忽略掉许多有效信息。因此值分布强化学习旨在拟合出状态价值的分布，从而提高预测准确性。值得一提的是，PG类算法在应用于连续动作时通常也会以某种分布作为输出。但这种情况下输出的分布为动作的概率分布，主体为动作；而值分布强化学习输出的分布为给定状态-动作对情况下其动作价值的分布，主体为价值。

Bellman算子

算子，即operator，接受函数输入，并输出函数。若$Q(s,a)$是一个函数，如果$s$是一个$n$维的变量，$a$是一个$m$维的变量，那么动作价值函数$Q$是一个$\mathbb{R}^n \times \mathbb{R}^m \rightarrow \mathbb{R}$的函数。我们可以引入Bellman算子$\mathcal{T}$来表示Bellman方程。对于某一个具体策略$\pi$以及状态转移矩阵$P$，我们有：

$\mathcal { T } ^ { \pi } Q ( x , a ) : = \mathbb { E } R ( x , a ) + \gamma \underset { P , \pi } { \mathbb { E } Q } \left( x ^ { \prime } , a ^ { \prime } \right)$

其中$R(s,a)$是环境的一部分，其输入输出与$Q$相同，为$\mathbb{R}^n \times \mathbb{R}^m \rightarrow \mathbb{R}$。可以看出，Bellman算子接收输入函数$Q$，并输出函数$Q'$。

从Bellman算子的角度来描述策略评估，其实就是在找一个函数$Q$，满足$\mathcal{T}^\pi Q = Q$，也就是在找算子的不动点（经过算子运算后函数不发生改变）。具体可以参考这一篇文章。

以相同的手法我们可以定义最优Bellman算子，其用来表示Bellman最优方程：

$\mathcal { T } Q ( x , a ) : = \mathbb { E } R ( x , a ) + \gamma \mathbb { E } _ { P } \max _ { a ^ { \prime } \in \mathcal { A } } Q \left( x ^ { \prime } , a ^ { \prime } \right)$

同样的，使用最优Bellman算子做策略评估，就是在找算子$\mathcal{T}$的不动点。

同样由这篇文章，可以证明算子$\mathcal{T}, \mathcal{T}^\pi$都满足$\gamma -contraction$的条件，即有：

$\left\| \mathcal{T}^\pi Q_1 - \mathcal{T}^\pi Q_2 \right\| _ { \infty } \leq \gamma \left\| Q_1 - Q_2 \right\| _ { \infty } \\$

$\left\| \mathcal{T} Q_1 - \mathcal{T} Q_2 \right\| _ { \infty } \leq \gamma \left\| Q_1 - Q_2 \right\| _ { \infty }$

于是根据Contraction Mapping Theorem，上述两个算子都具有唯一的不动点。即有:

$\begin{array} { l } \mathcal { T } ^ { \pi \infty } Q = Q ^ { \pi } \\ \mathcal { T } ^ { \infty } Q = Q ^ { * } \end{array}$

KL散度

对于分布$p,q$，它们的KL散度定义为:

$\mathrm { KL } ( p \| q ) = \int p ( x ) \log \frac { p ( x ) } { q ( x ) } d x$

KL散度能衡量两个分布之间的差异，但是不具有对称性。其可以用于离散形式的计算，有：

$\mathrm { KL } ( p \| q ) = \sum _ { i = 1 } ^ { N } p \left( x _ { i } \right) \log \frac { p \left( x _ { i } \right) } { q \left( x _ { i } \right) } = \sum _ { i = 1 } ^ { N } p \left( x _ { i } \right) \left[ \log p \left( x _ { i } \right) - \log q \left( x _ { i } \right) \right]$

Distributional DQN

Distributional DQN将传统DQN中的value function换成了value distribution。原来的DQN中的值函数为$Q(s,a)$，它接收一个状态动作对$s,a$，并输出实数格式的动作价值$Q$。

与其相对，Distributional DQN接收一个状态动作对$x,a$，并输出一个值分布$Z(x,a)=P(v | x,a)$，这个分布描绘了状态动作对$x,a$所有可能的价值$v$的范围以及可能性。易得，$Q(x,a) = \mathbb{E}[Z(x,a)]$。

针对值分布，可以定义状态转移算子：

$\begin{aligned} P ^ { \pi } Z ( x , a ) & : = Z \left( X ^ { \prime } , A ^ { \prime } \right) \\ X ^ { \prime } & \sim P ( \cdot \mid x , a ) , A ^ { \prime } \sim \pi \left( \cdot \mid X ^ { \prime } \right) \end{aligned}$

这个算子是$\mathcal { Z } \rightarrow \mathcal { Z }$，也就是接收一个分布，并输出另一个分布。可以看出，算子接收当前状态动作对$x,a$生成的值分布$Z(x,a)$，根据状态转移矩阵$P$以及当前策略$\pi$生成下一时刻状态动作对$x',a'$生成的值分布$Z(x',a')$。

在值分布情况下的Bellman算子就可以被定义为：

$\mathcal { T } ^ { \pi } Z ( x , a ) : \stackrel { D } { = } R ( x , a ) + \gamma P ^ { \pi } Z ( x , a )$

与状态转移算子$P^\pi$相同，值分布下的Bellman算子同样也是$\mathcal{Z} \rightarrow \mathcal{Z}$。可视化后的情况如下图：

在DQN中，我们计算出一个\(r + \gamma \underset{a^*} \max Q(x', a^*)\)，并使当前$Q(x,a)$尽量靠近这个值。相对应的，我们可以在Distributional DQN中定义最优Bellman算子$\mathcal{T} Z$，其满足：

$\mathcal { T } Z = \mathcal { T } ^ { \pi } Z \text { for } \pi \in \mathcal { G } _ { Z } \\ \mathcal { G } _ { Z } : = \left\{ \pi : \sum _ { a } \pi ( a \mid x ) \mathbb { E } Z ( x , a ) = \max _ { a ^ { \prime } \in \mathcal { A } } \mathbb { E } Z \left( x , a ^ { \prime } \right) \right\}$

$\mathcal { G } _ { Z }$看起来很复杂，但实际上指的是令策略$\pi$取最优策略。因此，在Distributional DQN中，我们需要让当前值分布$Z(x,a)$尽可能的靠近$\mathcal{T} Z(x,a)$这个分布。

C51定义

由于$Z(x,a)$分布并不一定符合某种特殊的预设模式，因此C51方法放弃使用高斯分布等参数化方式表示$Z(x,a)$，而是使用了51个等间距的atoms描绘一个分布。实际上其与描绘分布的直方图差距不大，但是直方图更注重区间，而C51算法使用的comb form的方法更注重"点"。当然这个分布肯定是近似的结果，点的范围和个数取决于应用情况。实际上，假设分布$Z(x,a)$的上界为$V_{MAX}$，下界$V_{MIN}$，将最大值和最小值之间的区间均匀化为$N$个点，则每个点$\mathcal{z}_i$有：

$\left\{ z _ { i } = V _ { \min } + i \Delta z : 0 \leq i < N , \Delta z : = \frac { V _ { M A X } - V _ { M I N } } { N - 1 } \right\}$

例如，当$N=7$时，一共有7个atoms，其分布可能为：

现在我们来考虑C51的网络架构。在DQN中，网络的输入为$s$，输出为$|\mathcal{A}|$维的一个向量$(Q(s,a_1), Q(s,a_2),...,Q(s,a_m))$。向量的每一个元素代表对应动作状态对$s,a$的预测$Q$值。而在C51算法中，输出变成了一个矩阵。矩阵的行数为$|\mathcal{A}|$，每一行代表给定$s,a$的值分布。矩阵的列数为$N$，每一列代表值分布在给定atom上的概率。若atoms集合为$\left\{ z _ { 0 } , z _ { 1 } , \cdots , z _ { N - 1 } \right\}$，C51的输出应该为$(\left\{ p _ { 0 } ^ { a_1 } , p _ { 1 } ^ { a_1 } , \cdots , p _ { N - 1 } ^ { a_1 } \right\}, \left\{ p _ { 0 } ^ { a_2 } , p _ { 1 } ^ { a_2 } , \cdots , p _ { N - 1 } ^ { a_2 } \right\} \cdots , \left\{ p _ { 0 } ^ { a_m } , p _ { 1 } ^ { a_m } , \cdots , p _ { N - 1 } ^ { a_m } \right\})$。

为了保证神经网络输出的矩阵的每一行构成一个离散的分布，可以对神经网络的输出进行softmax操作。具体为：

$Z _ { \theta } ( x , a ) = z _ { i } \quad w . p . p _ { i } ( x , a ) : = \frac { e ^ { \theta _ { i } ( x , a ) } } { \sum _ { j } e ^ { \theta _ { j } ( x , a ) } }$

C51损失函数

网络结构定义完成后，我们需要设置损失函数，从而完成训练过程。在DQN中，由于网络输出为一个实数，因此我们可以使用均方误差作为损失函数。但是C51算法输出为一个分布，因此我们需要找到一个衡量分布之间距离的方法。通常我们使用Wasserstein Metric作为衡量分布距离的方法（具体解释看这篇文章，其实也看不太懂呜呜）。使用Wasserstein Metric可以保证Bellman算子在输出分布的情况下也满足$\gamma - contraction$的条件。但是，对于Q-learning实际用到的最优Bellman算子则没有这样的理论保证。而且使用SGD方法训练时是没有办法保持Wasserstein Metric的。因此，C51算法使用了KL散度启发式的衡量两个分布的距离。

投影Bellman更新

使用KL散度的离散形式计算可以很容易的算出两个分布之间的距离。但是在实际运算中还会产生一个问题。我们可以通过模拟整个过程找出问题。

首先我们从Buffer中采样一个$(s,a,r,s')$，根据q-learning函数使用的最优Bellman算子，有：

$\mathcal { T } Z ( x , a ) : \stackrel { D } { = } R ( x , a ) + \gamma P Z ( x , a )$

因此为了改善$Z(x,a)$，我们需要计算$R(x,a) + \gamma Z(x',a^*)$。在DQN中，$a^*$的选择是根据$a ^ { * } \leftarrow \arg \max _ { a } Q \left( x' , a \right)$标准的。在C51中我们沿用这个标准，并通过$Q(x,a) = \mathbb{E}[Z(x,a)]$来计算$Q(x,a)$。具体有：

$Q \left( s _ { t + 1 } , a \right) : = \sum _ { i } z _ { i } p _ { i } \left( s _ { t + 1 } , a \right)$

根据上述标准，我们可以从target network中获得$Z(x',a^*)$的结果，是一个$N$维的向量，我们设其为$p'$。那么$R(x,a) + \gamma Z(x',a^*)$分布就可以表示为:

有$p'_{0}$的概率取$R + \gamma z_0$

有$p'_{1}$的概率取$R + \gamma z_1$

......

有$p'_{N-1}$的概率取$R + \gamma z_{N-1}$

但我们检查$Z(x,a)$可以发现，其可以表示为：

有$p'_{0}$的概率取$z_0$

有$p'_{1}$的概率取$z_1$

......

有$p'_{N-1}$的概率取$z_{N-1}$

可以发现，$Z(x,a)$与$R(x,a) + \gamma Z(x',a^*)$的atoms没有对齐。C51的解决方法是，将$R + \gamma z_k$的概率分摊到$z_k$和$z_{k+1}$上。具体的分摊比例按照$R + \gamma z_k$到$z_k$与$z_{k+1}$的距离的反比即可。

以这篇文章中的例子，令$\hat { \mathcal { T } } z _ { 0 } = R + \gamma z_0$对应的概率为$p_0(x',a')$，将其投影到相邻的支点$z_0,z_1$上，如图：

则$z_0$分得的概率的大小为：

$[1 - \frac { \hat { \mathcal { T } } z _ { 0 } - z _ { 0 } } { z_1 - z_0 }]p _ { 0 } \left( x ^ { \prime } , a ^ { \prime } \right)$

经过上述分配，我们可以得到一个新的分布$\Phi { \mathcal { T } } Z ( x , a )$，损失函数可得为：

$D_{KL}(\Phi { \mathcal { T } } Z ( x , a ) || Z(x,a))$

C51具体算法流程如下：