这个分支的几篇paper倒是有点意思。我们知道,RL的很多问题中,一旦$\mathcal{P}(s|s, a), \mathcal{R}(s, a),\pi(s,a)$中有一个是non-deterministic,那么必然会使得$Q(s, a)$也应该是non-determinsitc的,而我们之前提到的算法中, DNN输出的$Q(s, a)$却是determinisitic的,那么训练中的target Q value必然会产生震荡。本节的几篇文章都是围绕着这个问题展开的。

A Distributional Perspective on Reinforcement Learning

C51的解决思路很简单。就是将$Q(s,a)$的值分布视为在$[-10, 10]$之间的51个均匀区间的categorial distribution,如此,则$Q(s,a)$的输出则为概率分布。那么这种情况下,Bellman Equation是如何实现的呢?

$$\begin{aligned}Q^\pi(s,a) &= r(s,a) + \gamma\sum_{s'\in\mathcal{S}}p(s'|s,a)\sum_{a'\in\mathcal{A}}\pi(s',a')Q^\pi(s',a')\\&= r(s, a) + \gamma \mathbb{E}_{s', a'}[Q^\pi(s',a')]\end{aligned}$$

如果$Q^\pi(s', a')$是一个分布,那么Bellman Equation的RHS: $r(s, a) + \gamma Q^\pi(s', a')$是对原分布的变形,如原文中的下图所示:

其对应的伪代码如下:

我们接下来一行一行解释以上code。首先,算法规定了$Q(s,a)$的取值范围, 比如$[-10, 10]$, 那么将这个区间平均分成50份,那么一共会有以下51个端点$\{-10, -9.6, -9.2, \dots, 9.6, 10\}=\{z_i\}$,文中被称为51个support $\{z_i\}$,这也是算法被简称为C51的由来(51 Categories)。Q Network的output dimension 由原来的$|\mathcal{A}|$变为$|\mathcal{A}|\times 51$,且对每一个action对应的$1\times 51$ array取softmax以代表categorical probability $\{p_i(s, a)\}$,如此,则$Q(s, a)=\sum_i z_i p_i(s, a)$。

我们以$r = 1, \gamma = 0.99$ 为例,$\{z_i'\}=\{\mathcal{T}z_i\} = \{[r + \gamma z_i]_{-10}^{10}\} = \{-8.9, -8.504, \dots, 9.712, 10, 10, 10\}$,这一步就是所谓的compute the projection of $\mathcal{T}z_i$ on to $z_i$。接下来,我们看看何为distribute probability of $\mathcal{T} z_i$。对于$z'_0 = -8.9$, 它在$z_2 = -9.2$与$z_3 = -8.8$之间,我们需要将$z'_0$ 对应的概率$p_0$ 分配到$z_2,z_3$上,因为$z_0'$到$z_2, z_3$的距离比为$(9.2-8.9):(8.9-8.8)=3:1$,所以将$p_0$ 的$0.25$分配到$z_2$上,$0.75$分配到$z_3$上,以此类推。然而GPU代码层面要实现以上逻辑会复杂的多,因为涉及到寻找一个batch的index的问题。

经过以上操作,Bellman Equation的RHS和LHS的categorical distribution即有相同的support, 那么两个分布之间的距离即可用cross entropy来衡量,并以此作为loss。

Distributional Reinforcement Learning with Quantile Regression

QR-DQN是在C51的基础上做了稍许改进,甚至和C51一枚硬币的两面。C51中,support是fixed and evenly distributed,Q Network输出的是support对应的probability;而在QRDQN中,probability是fixed and evenly distributed,Q Network输出的是support。

相比起C51,QR-DQN在code层面要简洁和简单很多。以作者的标准算法为例,Q Network的输出为$|\mathcal{A}| \times 200$的array,每个action对应的$1 \times 200$的array即代表quantile distribution的两百个evenly distributed quantiles $\{z_i\}$,如此则$Q(s, a) = N^{-1}\sum_{i=1}^N z_i$。

那么quantile distributions之间的loss又如何计算呢?文中的下图即是作者的思路。不同于C51中计算又相同support的categorical distribution的cross entropy,QR-DQN想要用quantile distribution去逼近return的真实distribution,计算的是两个distributions的cumulated distribution probability(CDF)之间的1-Wasserstein error。

在具体计算中,作者用到了Quantile Regression的概念。所以接下来我们简单介绍下何为Quantile Regression。设random variable $Y$的cumulative distribution function为$F_Y(y) = P(Y \leq y)$,我们设分为分位数$\tau \in [0, 1]$对应的分位点为$\bar y$,则有$P(Y \leq \bar y) = \tau$,亦或者$\bar y = q(\tau) = F_Y^{-1}(\tau) = \inf \{y: F_Y(y) \geq \tau \}$。

那么问题来了,对于给定的$F_Y(y)$,如何求分位数$\tau$对应的分位点$F_Y^{-1}(\tau)$呢?我们设$L_\tau(u) = \mathbb{E}[(Y-u)(\tau-\mathbb{I}_{Y\lt u})]$,则

$$\begin{aligned}q(\tau) &= \arg \min_u L_\tau(u) \\&= \arg\min_u \mathbb{E}[(Y-u)(\tau-\mathbb{I}_{Y\lt u})] \\&= \arg \min_u \{ (\tau - 1)\int_{-\infty}^u(y - u)dF_Y(y) + \tau \int_{u}^\infty (y-u)dF_Y(y) \}\end{aligned}$$

其中

$$\mathbb{I}_{Y \lt u} = \begin{cases}1, Y \lt u \\0, Y \geq u\end{cases}$$

我们对$L_\tau(u)$求导令其为零,则有

$$\begin{aligned}\frac{\partial L_\tau(u)}{\partial u} &= (1-\tau)\int_{-\infty}^{u}dF_Y(y) - \tau \int_{u}^{\infty}dF_Y(y) \\&= \int_{-\infty}^u dF_Y(y) - \tau \int_{-\infty}^{\infty}dF_Y(y) \\&= F_Y(u) -\tau \\&= 0\end{aligned}$$

即$F_Y(u) = \tau, u= F^{-1}(\tau)=q(\tau)$。所以我们最小化$L_\tau(u)$,即可求得$\tau$的分位点$q(\tau)$。

如果$F_Y(y)$是quantile distribution,那么quantile regression loss为:

$$L_\tau(u) = (\tau - 1)\sum_{y_i \lt u}(y_i - u) + \tau \sum_{y_i\geq q}(y_i-u) = \sum_i |\tau - \mathbb{I}_{y_i\lt u}||y_i-u|$$

下面举例说明两个quantile distribution之间的QR loss。假设我们有quantile distribution $\mathcal{Q}=\{-5, -2, 3, 5, 8\}$ ,其对应的分位数$\tau_Q=\{0.1, 0.3, 0.5, 0.7, 0.9\}$,另有quantile distribution $\mathcal{P} = \{-3, -1, 2, 4, 7\}$ 且$\tau_P = \tau_Q$,那么$\mathcal{P, Q}$之间的Quantile Regression Loss

$$\begin{aligned}L_\tau(\mathcal{P}) =& \frac{1}{5}(|0.1 - 1|\times|-3+5| +|0.3-0|\times|-3+2| + |0.5-0|\times|-3-3|\\& + |0.7-0|\times|-3-5| + |0.9-0|\times|-3-8| + \\& |0.1-1|\times|-1+5| + |0.3-1|\times|-1+2| + |0.5-0|\times|-1-3|\\& + |0.7-0|\times|-1-5| + |0.9-0|\times|-1-8|+ \\& \dots+ \\& |0.1-1|\times|7+5|+|0.3-1|\times|7+2|+|0.5-1|\times|7-3|\\& +|0.7-1|\times|7-5|+|0.9-0|\times|7-8|) \\=& 18.8\end{aligned}$$

其对应code为

import torch
p = torch.tensor([-3, -1, 2, 4, 7])
q = torch.tensor([-5, -2, 3, 5, 8])
tau = torch.tensor([0.1, 0.3, 0.5, 0.7, 0.9])
diff = p.view(-1, 1) - q.view(1, -1)

weight = tau - (diff > 0).float()
loss = diff.abs() * weight.abs()
loss = loss.sum(-1).mean(-1)

Implicit Quantile Networks for Distributional Reinforcement Learning

IQN又在QR-DQN的基础上做了改进。我们知道,QR-DQN的分位数$\tau$是fixed and evenly distributed,所以很自然的想到,有没有办法也让$\tau$是variable?IQN的思路是在network方面动手,如文中的下图所示。

IQN的网络输入,除了RL的state之外,还有一个随机采样的分位数$\tau$,这个分位数被嵌入到网络中,并与convolution layers出来的feature互动,输出对应的分位点$q(\tau)$。

import torch
import torch.nn as nn
import numpy as np

class Network(nn.Module):
  # ......
  def forward(self, states, N):
    B = states.size(0) # Batch Size
    x = self.convs(states).view(B, -1) # convolution layers
    
    taus = torch.rand(B, N, 1) # Sample taus from uniform distribution
		pis = np.pi * torch.arange(1, self.num_cosines + 1).to(x).view(1, 1, -1)
    cosine = pis.mul(taus).cos().view(B * N, self.num_cosines)
    
    tau_embed = self.cosine_emb(cosine).view(B, N, -1)
    state_embed = x.view(B, 1, -1)
    features = (tau_embed * state_embed).view(B * N, -1)
    
    q = self.dense(features).view(B, N, -1) # B x N x |A|
    return q, taus

我在上面大概写了下算法对应的python code。网络每forward一次,便输出一个$B \times N$的$\tau$,以及与其对应的$B\times N \times |\mathcal{A}|$的quantiles,其余部分与QR-DQN相同。很多人会疑惑cosine embed这层的操作,其实这完全是调参试出来的结果,可以在Appendix里面的Figure 5看到,因为加了这个效果更好。

文中花了大量篇幅讲distortion risk measure,作者本期望通过distortion function来调整policy的risk偏好,结果发现还是identity distortion function,即neutral policy的总体效果最好。这里我们也简单解释下何为distortion risk measure。首先,我们需要定义一个distortion function $D:[0, 1] \rightarrow [0, 1]$,且要求$D$是continuous and increasing, 如此则$D(0) = 0, D(1) = 1$。之前我们提到,对于random variable $Y$,其CDF为$F_Y(y) = P(Y \leq y)$,quantile function为$q(\tau) = F_Y^{-1}(\tau) = \inf \{y: F_Y(y) \geq \tau \}, \tau \in [0, 1]$,那么distortion risk measure为

$$\rho_D(Y) = \int_{0}^1 F_Y^{-1}(\tau)dD(\tau) = \int_{-\infty}^\infty y d(D \circ F(y))$$

而我们知道

$$\mathbb{E}[Y] = \int_{-\infty}^\infty y d(F(y)) = \int_{-\infty}^\infty yf_Y(y) dy$$

对比以上两式我们发现,本质上,distortion risk measure是在将CDF distort之后的expected 。我画出了文中提到的主要的distortion function $D(\tau)$,如下图

如果被distort的是normal distribution,那么distort之后的distribution如下

我们接着理解,为什么convex distortion function会给出risk seeking policy而concave distortion function会给出risk averse policy。我们看到,concave的$\text{Wang}(-0.75)$给出的distribution,相对于identity normal distribution而言向右移动了,也就是说,我们的policy更倾向于选higher estimated return的action(即便我们的estimation可能没那么好),而避免选择lower estimated return的action,这就是所谓的risk averse;反之,convex的$\text{Wang}(0.75)$是risk seeking,也是一样的道理。而这其实就是exploration和exploitation的问题,risk seeking policy更倾向于explore未知的领域,即便它可能带来lower return。

Fully Parameterised Quantile Function for

Distributional Reinforcement Learning**

FQF这篇文章又在IQN的基础上做了改进。我们知道,IQN的分位点$\tau$是随机采样的,FQF的想法是,$\tau$应该是根据输入的states产生的。