Policy Gradients • CoderXL's Blog

本笔记聚焦于在线的强化学习。

数学#

我们希望最大化策略的累积奖励值的期望。

符号化地，就是：

\theta^* = \arg \max_{\theta} ~ \underbrace{\mathbb{E}_{\tau \sim p_{\theta}(\tau)}[r(\tau)]}_{J(\theta)}

其中 $\theta$ 代表模型的参数，也即我们采取的策略。

当我们要优化（此处是最大化）某个值的时候，我们会使用梯度下降。但是此处的 $J(\theta)$ 无法梯度下降的！因为在期望的计算过程中，有一个关键步骤是不可微分的，那就是从 $p_\theta(\tau)$ 中采样一个 $\tau$ . 这一步随机采样截断了反向传播路径，使得期望无法对 $\theta$ 求梯度！

因此，我们需要改造其数学形式。首先将期望写成积分形式，然后使用梯度和积分的可调换性：

\begin{aligned} \nabla_\theta J(\theta)&=\nabla_\theta \int p_\theta(\tau)r(\tau)\mathrm{d}\tau\\ &=\int r(\tau) \mathrm{d}\tau ~ \nabla_\theta ~ p_\theta(\tau)\\ \end{aligned}

然后，使用对数拆分技巧（ $\mathrm{d}x = x~\mathrm{d}\log x$ ），并重新还原成期望形式：

\begin{aligned} \nabla_\theta J(\theta)&=\int r(\tau) \mathrm{d}\tau ~ p_\theta(\tau) \nabla_\theta \log p_\theta(\tau)\\ &=\mathbb{E}_{\tau \sim p_\theta(\tau)}[r(\tau)\nabla_\theta\log p_\theta(\tau)] \end{aligned}

到这一步，我们发现，我们的梯度不再套在期望的外面，而进入了期望的里面。从期望的梯度变为了梯度的期望，现在可以求得梯度了：我们通过采样数据集，为每个样本计算相应梯度，然后通过上式计算期望，能够得到一个合理的整体梯度，从而进行梯度下降。

同时，展开 $p_\theta(\tau)$ 可得：

p_\theta(\tau)=p(s_1) \prod_t \pi_\theta(a_t \mid s_t)p(s_{t+1} \mid a_t, s_t)

因此上式中的 $\log p_\theta(\tau)$ 实际上可以展开为：

\log p_\theta(\tau)=\log p(s_1) + \sum_t \log \pi_\theta(a_t \mid s_t) + \sum_t p(s_{t+1} \mid a_t, s_t)

有由于是对 $\theta$ 取梯度，因此有：

\nabla_\theta \log p_\theta(\tau) = \nabla_\theta \sum_t \log \pi_\theta(a_t \mid s_t) = \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)

因此：

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}[r(\tau)\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)]

此外，当 reward 函数可以分配到每一步的时候，有 $\displaystyle{r(\tau) = \sum_t r(s_t, a_t)}$ ，上式便还可以写成：

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)]~[\sum_t r(s_t, a_t)]

这基本上就是我们推导的最终版本，具有良好的性质。我们可以每次对当前 $\theta$ 采样多个 $\tau$ ，分别计算 $\displaystyle{[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)]~[\sum_t r(s_t, a_t)]}$ ，然后取平均，并认为随着采样变多，我们对真实梯度的估计就足够好。

同时，我们发现此处 $\log \pi_\theta(a_t \mid s_t)$ 的计算和交叉熵是十分类似的。事实上，在监督学习（或者模仿学习）和强化学习里面，都是在求最大似然估计，而最大似然估计的损失函数就是交叉熵。
区别是，在监督学习（或模仿学习）里，我们始终希望交叉熵最大，因此损失函数是交叉熵的负数；
在强化学习里，我们通过一个 $\tau$ 的奖励函数 $r(\tau)$ 作为权重，来衡量这个样本的交叉熵是否值得被增大、增大多少。

从抽象的角度看，强化学习其实就是数学化的 Trial and Error.

然而，观察刚刚推导出的 $\nabla_\theta J(\theta)$ ，会发现一个问题：由于我们的每一步“交叉熵”都会乘以整个轨迹的 reward 值，这导致一旦模型过往出现失误，未来不论做什么，即使是补救措施，也会被惩罚。但这是不合理的，因为后发生的 action 无法影响已经产生的 reward，因此也不应该用这些 reward 来奖励/惩罚后发生的 action. 为此，我们要让每一步的预测都只受到未来的 reward 的影响。

因此可以将公式改成：

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'=t}^T r(s_{t'}, a_{t'})]

这一步其实在数学上没有改变 $J(\theta)$ 的定义与取值，只是换成了方差更小的估计方式。

因为这相当于手动去掉了“期望一定是 $0$ ”的项。
设想有一个复杂的期望，其中一项是均值为 $0$ 方差为 $1$ 的正态分布。这一项的期望一定是 $0$ ，但如果通过多次采样求平均来隐式地计算，则方差很大，可能无法让这一项收敛到 $0$ 。因此不如通过数学推导，直接证明其期望为 $0$ ，然后将其从期望公式中消除。

下面推导新形式和原形式的期望是等价的（注意，对于单个 $\tau$ ，二者计算结果不会相等，只是期望相等）：

将原来的形式中的 $r$ 分配进去，并拆分为过去和未来：

\begin{aligned} \nabla_\theta J(\theta) &= \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)]~[\sum_t r(s_{t'}, a_{t'})]\\ &= \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'} r(s_{t'}, a_{t'})]\\ &=\mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'})] + \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t' \ge t} r(s_{t'}, a_{t'})] \end{aligned}

后一项就是新的形式。下面证明前一项在数学上是 $0$ ：

\begin{aligned} & \mathbb{E}_{\tau \sim p_\theta(\tau)}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'})]\\ =& \sum_t \mathbb{E}_{\tau \sim p_\theta(\tau)} [\nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'})] \end{aligned}

取其中一个时间点 $t$ ，由于 $\tau=(s_1, a_1, \dots, s_t, a_t, \dots, s_T, a_T)$ ，可以取 $h_t=(s_1, a_1, \dots, s_t)$ ，那么上式方括号内的取值是 $h_t$ 和 $a_t$ 的函数。再结合全期望公式可得：

\begin{aligned} & \mathbb{E}_{\tau \sim p_\theta(\tau)} [\nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'})]\\ =& \mathbb{E}_{h_t}[\mathbb{E}_{a_t}[\nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'}) \mid h_t]] \end{aligned}

由于 $a_t$ 和 $r(s_{t'}, a_{t'})$ 是独立变量，因此可以化为：

\begin{align} & \mathbb{E}_{h_t}[\mathbb{E}_{a_t}[\nabla_\theta \log \pi_\theta(a_t \mid s_t)\sum_{t'<t} r(s_{t'}, a_{t'}) \mid h_t]]\\ =& \mathbb{E}_{h_t}[\sum_{t'<t} r(s_{t'}, a_{t'}) ~ \mathbb{E}_{a_t}[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \mid h_t]] \tag{\#} \end{align}

下面，最终、最关键的一步：

\begin{aligned} & \mathbb{E}_{a_t}[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \mid h_t]\\ =& \sum_{a_t} \pi_\theta (a_t \mid s_t) \nabla_\theta \log \pi_\theta(a_t \mid s_t)\\ =& \sum_{a_t} \nabla_\theta \pi_\theta(a_t \mid s_t)\\ =& \nabla_\theta \sum_{a_t} \pi_\theta(a_t \mid s_t)\\ =& \nabla_\theta ~ 1\\ =& ~ 0 \end{aligned}

或者，连续动作空间的情形：

\begin{aligned} & \mathbb{E}_{a_t}[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \mid h_t]\\ =& \int \pi_\theta(a_t \mid s_t) \nabla_\theta \log \pi_\theta(a_t \mid s_t) ~ \mathrm{d}a_t\\ =& \int \nabla_\theta \pi_\theta(a_t \mid s_t) ~ \mathrm{d}a_t\\ =& \nabla_\theta \int \pi_\theta(a_t \mid s_t) ~ \mathrm{d}a_t\\ =& \nabla_\theta ~ 1\\ =& ~ 0 \end{aligned}

由此得出 $(\#)$ 式等于 $0$ ，得证。

有了上面的经验，我们发现 reward 还有更多优化空间。比如对于一个 batch 的所有 $\tau$ ，我们对 $r(\tau)$ 取平均得到 baseline，然后将每个 $r(\tau)$ 减去这个均值。这样做符合直觉：可以抑制低于 baseline 的行为，促进高于 baseline 的行为。而事实上，可以通过与上一节类似的方法证明减去 baseline 完全没有影响理想情况下 $J(\theta)$ 的取值，只是优化了对它的估计。

因此，我们最终的公式版本就会是：

\begin{aligned} \nabla_\theta J(\theta) &\approx {1\over N}\sum_{i=1}^N ~ \sum_{t=1}^T \left[ \nabla_\theta \log \pi_\theta(a_{i,t} \mid s_{i,t}) \left(\left( \sum_{t'=t}^T r(s_{i,t'}, a_{i,t'}) \right) -b \right)\right]\\ &= \nabla_\theta ~ {1\over N}\sum_{i=1}^N ~ \sum_{t=1}^T \left[ \log \pi_\theta(a_{i,t} \mid s_{i,t}) \left(\left( \sum_{t'=t}^T r(s_{i,t'}, a_{i,t'}) \right) -b \right)\right]\\ \end{aligned}

可以看到，我们通过采样有限个样本去近似期望；同时我们在最外层计算梯度，这样只需要反向传播一次，性能更好。

实践#

通常，当奖励较密集（Dense）的时候，可以使用随机初始化的 Policy 开始训练；但如果奖励较为稀疏（Sparse），则建议先进行模仿学习，再转移到在线的强化学习。