强化学习中的策略梯度与baseline

一、策略梯度中的 baseline

本小结的主要内容是做数学推导得到带 baseline 的策略梯度的公式。

策略梯度方法常用 baseline 来降低方差，可以让收敛更快。

1.1 baseline 的性质

baseline 指的是一个函数 b ，它不依赖于动作 $A$ ，可以得到公式：

$$\mathbb{E}_{A \sim \pi} \left[ b \cdot \frac{\partial \log \pi(A|s; \theta)}{\partial \theta} \right] = 0 \tag{1.1}$$

证明过程如下：

$$\begin{aligned} \mathbb{E}_{A \sim \pi} \left[ b \cdot \frac{\partial \log \pi(A|s; \theta)}{\partial \theta} \right] \\ & （b 不依赖与动作 A） \\ & = b \cdot \mathbb{E}_{A \sim \pi} \left[ \frac{\partial \log \pi(A|s; \theta)}{\partial \theta} \right] \\ & （假设动作是离散的，展开期望） \\ & = b \cdot \sum_{a \in A}\pi(a|s; \theta) \frac{\partial \log \pi(a|s; \theta)}{\partial \theta} \\ & （求导的链式法则） \\ & = b \cdot \sum_{a \in A}\pi(a|s; \theta) \left[ \frac{1}{\pi(a|s; \theta)} \cdot \frac{\partial \pi(a|s; \theta)}{\partial \theta} \right] \\ & = b \cdot \sum_{a \in A} \frac{\partial \pi(a|s; \theta)}{\partial \theta} \\ & （连加与求导的对象不同） \\ & = b \cdot \frac{\partial \sum_{a \in A} \pi(a|s; \theta)}{\partial \theta} \\ & （所有动作的概率和为1） \\ & = b \cdot \frac{\partial 1}{\partial \theta} \\ & （常数的导数为0） \\ & = 0 \\ \end{aligned}$$

1.2 带 baseline 的策略梯度公式

在基于策略的强化学习中推导了策略梯度的公式：

$$\frac{\partial V(s; \theta)}{\partial \theta} = \mathbb{E}_{A \sim \pi(\cdot | s; \theta)} \left[ \frac{\partial \log \pi(A | s; \theta)}{\partial \theta} \cdot Q_{\pi}(s, A) \right]$$

结合上面推导得到公式1和策略梯度公式，可以得到 带 baseline 的策略梯度公式 ：

$$\begin{aligned} \frac{\partial V_{\pi}(s)}{\partial \theta} & = \mathbb{E}_{A \sim \pi} \left[ \frac{\partial \log \pi(A | s; \theta)}{\partial \theta} \cdot Q_{\pi}(s, A) \right] \\ & = \mathbb{E}_{A \sim \pi} \left[ \frac{\partial \log \pi(A | s; \theta)}{\partial \theta} \cdot Q_{\pi}(s, A) \right] - 0 \\ & = \mathbb{E}_{A \sim \pi} \left[ \frac{\partial \log \pi(A | s; \theta)}{\partial \theta} \cdot Q_{\pi}(s, A) \right] - \mathbb{E}_{A \sim \pi} \left[ b \cdot \frac{\partial \log \pi(A|s; \theta)}{\partial \theta} \right] \\ & = \mathbb{E}_{A \sim \pi} \left[ \frac{\partial \log \pi(A | s; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s, A) - b \right) \right] \end{aligned}$$

即，如果 baseline 函数 b 不依赖于动作 $A_t$，下面公式成立：

$$\frac{\partial V_{\pi}(s_t)}{\partial \theta} = \mathbb{E}_{A_t \sim \pi} \left[ \frac{\partial \log \pi(A_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, A_t) - b \right) \right] \tag{1.2}$$

从公式中我们可以看出：baseline 不会影响策略梯度的正确性，不管函数 b 是什么，得到的期望都是一样的。

1.3 蒙特卡洛近似

将 带 baseline 的策略梯度公式 中求期望的部分记作 $g(A_t)$，即：

$$g(A_t) = \frac{\partial \log \pi(A_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, A_t) - b \right)$$

1. 从策略 $\pi(\cdot | s; \theta)$ 中抽样出动作 $a_t$，然后计算 $g(a_t)$ 。计算得到的 $g(a_t)$ 就是 $\mathbb{E}_{A_t \sim \pi} \left[ g(A_t) \right]$ 的蒙特卡洛近似。
2. $g(a_t)$ 是策略梯度的一个无偏估计，这是因为 $g(A_t)$ 关于 $A_t$ 的期望等于策略梯度。$g(a_t)$ 其实是一个随机梯度，它是对策略梯度的蒙特卡洛近似。实际训练策略网络的时候，用的都是随机梯度 $g(a_t)$ ，而不是用真正的策略梯度。 $$\mathbb{E}_{A_t \sim \pi} \left[ g(A_t) \right] = \frac{\partial V_{\pi}(s_t)}{\partial \theta}$$
3. 使用随机梯度上升更新策略网络的参数 $\theta$ $$\theta \leftarrow \theta + \beta \cdot g(a_t)$$

从 带 baseline 的策略梯度公式 看到，虽然使用不同的 baseline 函数 b 不会改变期望 $\mathbb{E}_{A_t \sim \pi} \left[ g(A_t) \right]$ 的值，但是会影响 $g(a_t)$ 的值。如果选择的baseline 函数 b 很好，接近 $Q_{\pi}(s_t, A_t)$，随机策略梯度 $g(a_t)$ 的方差就会变小，就能让算法收敛更快。这也是为什么要使用 baseline 函数 b 的原因。

1.4 baseline 的选择

1. b=0

   相当于不适用 baseline，得到的是标准的策略梯度。

   $$\frac{\partial V_{\pi}(s_t)}{\partial \theta} = \mathbb{E}_{A_t \sim \pi} \left[ \frac{\partial \log \pi(A_t | s_t; \theta)}{\partial \theta} \cdot Q_{\pi}(s_t, A_t) \right]$$

2. $b=V_{\pi}(s_t)$
   因为 $V_{\pi}(s_t)$ 只依赖于当前状态的 $s_t$，状态 $s_t$ 是先于 $A_t$ 被观测到的，所以 $s_t$ 不依赖于 $A_t$，可以作为 baseline 。

   而选择使用 $V_{\pi}(s_t)$ 作为 baseline 的原因在于，$V_{\pi}(s_t)$ 的值接近 $Q_{\pi}(s_t, A_t)$，能够有效的降低方差。

   $$V_{\pi}(s_t) = \mathbb{E}_{A_t}[Q_{\pi}(s_t, A_t)]$$

二、REINFORCE with Baseline

2.1 近似带 baseline 的策略梯度

在1.4 baseline 的选择中，证明了可以使用 $V_{\pi}(s_t)$ 来作为 Baseline：

$$\frac{\partial V_{\pi}(s_t)}{\partial \theta} = \mathbb{E}_{A_t \sim \pi} \left[ \frac{\partial \log \pi(A_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, A_t) - V_{\pi}(s_t) \right) \right]$$

在使用这个公式训练时，需要做总共3步近似：

1. 使用随机策略梯度来近似真正的策略梯度的期望（蒙特卡洛近似）
   记 $g(A_t)$ 为：

   $$g(A_t) = \frac{\partial \log \pi(A_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, A_t) - V_{\pi}(s_t) \right)$$

   从策略 $\pi(\cdot | s; \theta)$ 中抽样出动作 $a_t$，然后使用随机策略梯度 $g(a_t)$ 的值来代替真正的策略梯度的期望 $\mathbb{E}_{A_t \sim \pi} \left[ g(A_t) \right]$ 的值。

2. 使用折扣回报来近似 $Q_{\pi}(s_t, a_t)$ 的值（再一次蒙特卡洛近似）
   经过了上一步的近似后，我们得到：

   $$g(a_t) = \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, a_t) - V_{\pi}(s_t) \right)$$

   得到的等式中，$Q_{\pi}(s_t, a_t)$ 和 $V_{\pi}(s_t)$ 依然是未知的。

   使用策略 $\pi(\cdot | s; \theta)$ 控制智能体玩一局游戏到结束，得到轨迹：

   $$s_t, a_t, r_t, s_{t+1}, a_{t+1}, r_{t+1}, \cdots, s_n, a_n, r_n$$

   计算折扣回报：

   $$u_t = \sum_{i=t}^n \gamma^{i-t} r_i$$

   折扣回报 $u_t$ 就是 $Q_{\pi}(s_t, a_t)$ 的无偏估计。REINFORCE 算法就是使用折扣回报 $u_t$ 来近似 $Q_{\pi}(s_t, a_t)$ 的值。

3. 使用神经网络来近似 $V_{\pi}(s_t)$（函数近似）

   使用神经网络 $v(s; \mathbf{w})$ 来近似 $V_{\pi}(s_t)$ 的值。

经过3步近似后，可以得到 近似后的带 baseline 的策略梯度 ：

$$\frac{\partial V_{\pi}(s_t)}{\partial \theta} \approx g(a_t) \approx \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( u_t - v(s_t; \mathbf{w}) \right)$$

2.2 算法训练流程

1. 使用策略 $\pi(\cdot | s; \theta)$ 控制智能体玩一局游戏到结束，得到轨迹：

   $$s_t, a_t, r_t, s_{t+1}, a_{t+1}, r_{t+1}, \cdots, s_n, a_n, r_n$$

2. 计算折扣回报 $u_t$ 和误差 $\delta_t$

   $$u_t = \sum_{i=t}^n \gamma^{i-t} r_i$$ $$\delta_t = v(s_t; \mathbf{w}) - u_t（注意这里是v-q而不是q-v）$$

3. 更新策略网络的参数 $\theta$

   $$\theta \leftarrow \theta - \beta \cdot \delta_t \cdot \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta}$$

4. 更新价值网络的参数 $\mathbf{w}$

   $$\mathbf{w} \leftarrow \mathbf{w} - \alpha \cdot \delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$$

   其中 $\delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$ 是误差 $\delta_t$ 的均方差关于参数 $\mathbf{w}$ 的导数

   $$\frac{\partial \delta_t ^2 / 2}{\partial \mathbf{w}} = \delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$$

5. 对于第1步中的每个时刻，重复第2-4步。

三、Advantage Actor-Critic（A2C）算法

把 baseline 用到 Actor-Critic Method 上得到的算法就是A2C算法。

3.1 神经网络

1. 策略网络（Actor）
   与原始的 Actor-Critic Method 中的策略网络（Actor）相同，依然是使用神经网络 $\pi(a|s; \theta)$ 来近似策略函数 $\pi({a|s})$，其中 $\theta$ 是神经网络的参数。

2. 价值网络（Critic）
   价值网络使用的是状态价值网络 $V_{\pi}(s)$ 而不是动作价值网络 $Q_{\pi}(s, a)$。使用神经网络 $v(s; \mathbf{w})$ 来近似策略函数 $V_{\pi}(s)$，其中 $\mathbf{w}$ 是神经网络的参数。

   动作价值网络 $Q_{\pi}(s, a)$ 依赖于状态 $s$ 和 动作 $a$ ，而状态价值网络 $V_{\pi}(s)$ 仅依赖于状态 $s$，不依赖于动作 $a$，所以状态价值网络 $V_{\pi}(s)$ 比动作价值网络 $Q_{\pi}(s, a)$ 更好训练。

3.2 训练过程

1. 获取到一组 transition $(s_t, a_t, r_t, s_{t+1})$
2. 计算 TD target $$y_t = r_t + \gamma \cdot v(s_{t+1}; \mathbf{w})$$
3. 计算 TD error $$\delta_t = v(s_t; \mathbf{w}) - y_t$$
4. 更新策略网络的参数 $\theta$ $$\theta \leftarrow \theta - \beta \cdot \delta_t \cdot \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta}$$
5. 更新价值网络的参数 $\mathbf{w}$ $$\mathbf{w} \leftarrow \mathbf{w} - \alpha \cdot \delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$$

3.3 算法推导

3.3.1 价值函数性质的推导

在强化学习中时序差分学习推导出了公式：

$$Q_{\pi}(s_t, a_t) = \mathbb{E}[R_t + \gamma \cdot Q_{\pi}(S_{t+1}, A_{t+1})], \text{ for all }\pi$$

所以：

$$\begin{aligned} Q_{\pi}(s_t, a_t) & = \mathbb{E}[R_t + \gamma \cdot Q_{\pi}(S_{t+1}, A_{t+1})] \\ & = \mathbb{E}_{S_{t+1}, A_{t+1}}[R_t + \gamma \cdot Q_{\pi}(S_{t+1}, A_{t+1})] \\ & （R_t 与 A_{t+1} 无关） \\ & = \mathbb{E}_{S_{t+1}} \left[ R_t + \gamma \cdot \mathbb{E}_{A_{t+1}}[Q_{\pi}(S_{t+1}, A_{t+1})] \right] \\ & （状态价值函数的定义） \\ & = \mathbb{E}_{S_{t+1}}[R_t + \gamma \cdot V_{\pi}(S_{t+1})] \\ \end{aligned}$$

即：

$$Q_{\pi}(s_t, a_t) = \mathbb{E}_{S_{t+1}}[R_t + \gamma \cdot V_{\pi}(S_{t+1})] \tag{3.3.1a}$$

另，根据状态价值函数的定义：

$$\begin{aligned} V_{\pi}(s_t) & = \mathbb{E}_{A_t}[Q_{\pi}(s_t, A_t)] \\ & （使用公式3.3.1a） \\ & = \mathbb{E}_{A_t} \left[ \mathbb{E}_{S_{t+1}}[R_t + \gamma \cdot V_{\pi}(S_{t+1})] \right] \\ & = \mathbb{E}_{A_t,S_{t+1}}[R_t + \gamma \cdot V_{\pi}(S_{t+1})] \\ \end{aligned}$$

即：

$$V_{\pi}(s_t) = \mathbb{E}_{A_t,S_{t+1}}[R_t + \gamma \cdot V_{\pi}(S_{t+1})] \tag{3.3.1b}$$

3.3.2 价值函数性质的蒙特卡洛近似

假设我们获取到一组 transition $(s_t, a_t, r_t, s_{t+1})$

1. 公式3.3.1a 可以近似为

   $$Q_{\pi}(s_t, a_t) \approx r_t + \gamma \cdot V_{\pi}(s_{t+1}) \tag{3.3.2a}$$

   后面会使用这个公式来近似策略梯度，A2C算法的关键也在于这个公式。

2. 公式3.3.1b 可以近似为

   $$V_{\pi}(s_t) \approx r_t + \gamma \cdot V_{\pi}(s_{t+1}) \tag{3.3.2b}$$

   使用TD算法训练价值网络 $V_{\pi}(s_t)$ 时，TD target 就是由这个公式来的。

3.3.3 训练策略网络

本文前面定义了随机梯度 $g(a_t)$

$$g(a_t) = \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( Q_{\pi}(s_t, a_t) - V_{\pi}(s_t) \right)$$

$g(a_t)$ 中包含了 $Q_{\pi}(s_t, a_t) - V_{\pi}(s_t)$，而 $Q_{\pi}(s_t, a_t) - V_{\pi}(s_t)$ 正好是优势函数（Advantage function）的定义，所以算法名中包含了 Advantage 。

结合 $g(a_t)$ 的定义和 公式3.3.2a，有：

$$g(a_t) \approx \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( r_t + \gamma \cdot V_{\pi}(s_{t+1}) - V_{\pi}(s_t) \right)$$

使用神经网络 $v(s;\mathbf{w})$ 近似状态价值函数 $V_{\pi}(s)$，得到：

$$g(a_t) \approx \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( r_t + \gamma \cdot v(s_{t+1};\mathbf{w}) - v(s_t;\mathbf{w}) \right)$$

将公式中的 $r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 记作 $y_t$。

$y_t$ 其实就是 TD target，后面更新价值网络时也会用到。

既然计算出了近似的策略梯度，就可以使用梯度上升来训练策略网络了：

$$\theta \leftarrow \theta + \beta \cdot \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot (y_t - v(s_t; \mathbf{w}))$$

3.3.4 训练价值网络

将 公式3.3.2b 中的 $V_{\pi}(S)$ 使用神经网络 $v(S; \mathbf{w})$ 近似，得到：

$$v(s_t; \mathbf{w}) \approx r_t + \gamma \cdot v(s_{t+1}; \mathbf{w})$$

将 $r_t + \gamma \cdot v(s_{t+1}; \mathbf{w})$ 记作 $y_t$。使用TD算法训练时， $y_t$ 其实就是 TD target 。

所以 TD error 就是：

$$\delta_t = v(s_t; \mathbf{w}) - y_t$$

梯度：

$$\frac{\partial \delta_t ^2 / 2}{\partial \mathbf{w}} = \delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$$

最后使用梯度下降方法来更新网络参数 $\mathbf{w}$：

$$\mathbf{w} \leftarrow \mathbf{w} - \alpha \cdot \delta_t \cdot \frac{\partial v(s_t; \mathbf{w})}{\partial \mathbf{w}}$$

3.3.5 对算法的直观理解

回顾对策略梯度的近似：

$$g(a_t) \approx \frac{\partial \log \pi(a_t | s_t; \theta)}{\partial \theta} \cdot \left( r_t + \gamma \cdot v(s_{t+1};\mathbf{w}) - v(s_t;\mathbf{w}) \right)$$

这个公式用来更新策略网络 $\pi(a|s; \theta)$。

公式中的 $r_t + \gamma \cdot v(s_{t+1};\mathbf{w}) - v(s_t;\mathbf{w})$ 是状态价值网络 $v(s;\mathbf{w})$ 做出的评价。它评价了动作 $a_t$ 的好坏，可以指导策略网络 $\pi(a|s; \theta)$ 做改进，所以状态价值网络 $v(s;\mathbf{w})$ 被称为 Critic 。

但是 $r_t + \gamma \cdot v(s_{t+1};\mathbf{w}) - v(s_t;\mathbf{w})$ 中并未包含动作 $a_t$，又是如何评价动作 $a_t$ 的好坏呢？

1. $v(s_t;\mathbf{w})$ 是状态价值网络在 $t$ 时刻基于状态 $s_t$ 对状态 $s_t$ 的折扣回报 $\mathbb{E}[U_t | s_t]$ 的 预测值。预测值可以评价状态 $s_t$ 的好坏，状态 $s_t$ 越好，这个值就越大。

2. $r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 是状态价值网络在 $t+1$ 时刻基于状态 $s_t$ 和状态 $s_{t+1}$ 对状态 $s_t$ 的折扣回报 $\mathbb{E}[U_t | s_t, s_{t+1}]$ 的 预测值。

3. $v(s_t;\mathbf{w})$ 和 $r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 的 相同点：

   • 都是状态价值网络对状态 $s_t$ 的折扣回报的预估值。
   • 都能够评价状态 $s_t$ 的好坏。

4. $v(s_t;\mathbf{w})$ 和 $r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 的 区别：

   • $v(s_t;\mathbf{w})$ 是在 $t$ 时刻做出的预测，这个预测是在执行动作 $a_t$ 之前做出的，所以这个预测与动作 $a_t$ 无关。
   • $r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 是在 $t+1$ 时刻做出的预测，这个预测受动作 $a_t$ 的影响，动作 $a_t$ 越好，新的状态 $s_{t+1}$ 就越好，$r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 的值也就越大。
   • 如果动作 $a_t$ 越好，$r_t + \gamma \cdot v(s_{t+1};\mathbf{w})$ 的值就比 $v(s_t;\mathbf{w})$ 更大，它们的差值就是正数，反之亦然。所以两者的差可以反映动作 $a_t$ 带来的优势，两者的差就叫做 Advantage 。

---

参考：策略梯度与baseline