CS285 Part 3 - 策略梯度方法
本部分介绍策略梯度相关的方法。这是一种 Model-Free 的强化学习算法，即假设不能知道转移概率。
策略梯度的推导
回忆在强化学习中我们需要优化的目标是
$$
{J}(\theta) = \mathbb{E}_{\tau\sim p_\theta(\tau)}\left[\sum_{t}r(s_t, a_t)\right]
$$
其中 $\tau = \{(s_1, a_1, …, s_T, a_T)\}$ 是一个 trajectory，分布 $p$ 在上一节中有给出：
$$
p_{\theta}(\mathbf{s}_1, \mathbf{a}_1, …, \mathbf{s}_T, \mathbf{a}_T) = p(\mathbf{s}_1)\prod_{t = 1}^T \pi_{\theta}(\mathbf{a}_t | \mathbf{s}_t)p(\mathbf{s}_{t + 1} | \mathbf{s}_t, \mathbf{a}_t)
$$
我们尝试推导 $\nabla_\theta J(\theta)$。
$$
\begin{aligned}
\nabla_\theta J(\theta) &= \nabla_\theta\int p_{\theta}(\tau) r(\tau) \mathrm{d}\tau \\
&= \int p_\theta(\tau)\nabla_\theta \log p_\theta(\tau) r(\tau) \mathrm{d}\tau
\\
&= \mathbb{E}_{\tau\sim p(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_t r(a_{t}, s_{t})\right]
\end{aligned}
$$
注意 $\log p(s_{t + 1} | s_t, a_t)$ 和 $\theta$ 无关，所以求偏导之后没有这一项。此时我们可以用采样一个大小为 $N$ 的 batch 来估计这个期望：
$$
\boxed{\nabla_\theta J(\theta)\approx \frac{1}{N}\sum_{i=1}^N\sum_t \nabla_\theta \log \pi_\theta(a_{i, t} | s_{i, t})\sum_t r(a_{i, i}, s_{i, t})} \label{directpolicygrad}
$$
这就是策略梯度的计算方法。在本节末尾我们略微在机器学习的视角下谈一些这个式子的意义。
一方面，从直觉上来讲，奖励更大的策略将产生更大的梯度，参数确实将朝着奖励大的决策方向下降。
回忆在 imitation learning 中，我们本质上是在对专家策略做极大似然估计，专家策略的 likelihood 的对数就是
$$
\frac{1}{N}\sum_{i=1}^N\sum_{t} \log \pi_\theta(a_{i, t} | s_{i, t})
$$
而上面的式子积分之后得到的 pseudo-loss（不是真正优化的东西，但是其梯度正好是原来函数的梯度）是
$$
\frac{1}{N}\sum_{i=1}^N\sum_t \log \pi_\theta(a_{i, t} | s_{i, t})\sum_t r(a_{i, t}, s_{i, t})
$$
虽然这里策略是模型的策略而非专家策略，但是后面乘上了一个关于奖励的系数 $r$，这等价于在采样时，奖励更大的轨迹有更大概率被采样。于是策略梯度算法实际上是在最大化高奖励轨迹的似然。
此外，如果我们只有 observation 而没有完整的 state，需要注意的是得到观察 $o_t$ 的概率仅仅是在得到 $s_t$ 的概率上面乘一个与 $\theta$ 无关的条件概率，所以只知道 observation 时需要优化的式子是
$$
\begin{aligned}
\nabla_\theta J(\theta)&\approx \frac{1}{N}\sum_{i=1}^N\sum_t \nabla_\theta \log \pi_\theta(a_{i, t} | o_{i, t}) r(\tau)
\end{aligned}
$$
注意最后这个求和在 policy gradient 当中只是一条轨迹的 reward（并已经在模拟后给出），所以它是怎么算出来的毫不重要（到底是 $r(s, a)$ 还是 $r(o, a)$）。因此即便只知道 observation，我们的梯度仍然是正确的。
算法的问题
在实际操作中，我们发现用梯度下降方法去做 RL 会很困难，这里列举一些具体原因。
高方差
Remark. 这部分网上很多东西很没有道理，课程中给出的例子看起来也很没有道理，我们直接按着自己的思路来写。
注意在 $\ref{directpolicygrad}$ 中我们用采样的方法来估计梯度，这样采样自然得到的是一个无偏估计。然而在 MDP 中采样，方差几乎是随时间累计的，这就是策略梯度的高方差问题，它导致估计的准确性是很难 bound 住的。我们期望找一些方法来在不增加采样量的前提下降低采样的方差。
Baseline Trick
引理 1. 给奖励函数加上一个常数 $b$ 将不会改变梯度的期望。即
$$
\nabla_\theta \mathbb{E}_{\tau \sim p_\theta(\tau)}\left[\sum_{t} r(s_t, a_t) - b\right] = \nabla_\theta \mathbb{E}_{\tau \sim p_\theta(\tau)}\left[\sum_{t} r(s_t, a_t)\right]
$$
但是此举会影响梯度的方差，简便起见记 $\nabla_\theta \log \pi_\theta(a_{t} | s_{t}) = g(\tau)$：
$$
\begin{aligned}
&\mathrm{Var}\left[g(\tau)^2(r(\tau) - b)\right] \\
=& \mathbb{E}[g(\tau)^2(r(\tau) - b)^2] - \mathbb{E}[\text{Grad}]^2 \\
=& \mathrm{Var}[\text{Grad}] - 2b\mathbb{E}[g(\tau)^2r(\tau)] + b^2\mathbb{E}[g(\tau)^2]
\end{aligned}
$$
取
$$
b = \frac{\mathbb{E}[g(\tau)^2r(\tau)]}{\mathbb{E}[g(\tau)^2]}
$$
立即得到最优的方差。这实际上就是奖励的一个加权平均作为 baseline，策略向优于 baseline 的方向移动。当然直接取
$$
b = \frac 1N \sum_{i = 1}^n r(\tau)
$$
也可以跑的非常好。
Reward-to-go Trick
重新审视 $\ref{directpolicygrad}$，发现这里我们相当于给每一步操作都乘上了它所在的 trajectory 的全部奖励，直觉上来说我们只需要乘这一步以后的奖励。不妨做如下假设：
假设（Causality Assumption）. 对于 $t < t’$，$t’$ 时刻的操作和 $t$ 时刻的奖励、操作都独立。换言之，我们的 Agent 只需要考虑行动的后果。
我们对梯度的式子进行一些变形：
$$
\begin{aligned}
\nabla_{\theta}J(\theta) &= \mathbb{E}_{\tau\sim p(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_t r(a_{t}, s_{t})\right] \\
&= \mathbb{E}_{\tau\sim p(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_{t’ \geq t} r(a_{t’}, s_{t’})\right] \\
&+ \color{orange}\mathbb{E}_{\tau\sim p(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_{t’ < t} r(a_{t’}, s_{t’})\right]
\end{aligned}
$$
现在来分析橙色部分具体是什么。
$$
\begin{aligned}
&\mathbb{E}_{\tau\sim p(\tau)}\left[\sum_{t’} r(a_{t’}, s_{t’})\sum_{t > t’} \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\right] \\
=~& \sum_{t’}\sum_{t > t’} \mathbb{E}[r(a_{t’}, s_{t’})\nabla_\theta\log\pi_\theta(a_t | s_t)] \\
=~& \sum_{t’} \mathbb{E}[r(a_{t’}, s_{t’})]\sum_{t > t’} \color{blue}\mathbb{E}[\nabla_\theta\log\pi_\theta(a_t | s_t)]
\end{aligned}
$$
末尾蓝色的部分非常眼熟，它正是
$$
1 = \int p(\tau) \mathrm{d}\tau
$$
的梯度，注意 $\pi_\theta(a_t | s_t)$ 无非是 $p(\tau)$ 的一个边缘分布。
于是橙色的部分是 $0$，计算梯度只需要未来的奖励：
$$
\nabla_{\theta}J(\theta) = \mathbb{E}_{\tau\sim p(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_{t’ \geq t} r(a_{t’}, s_{t’})\right]
$$
直觉上，这种方法通过缩小了变量的值域来减小了方差。因此也被归类在解决“高方差”这一小节之下。实际上这个方法的可能更加利好后面的 actor-critic 模型的推导。
如果用采样来计算，需要优化的式子即为
$$
\boxed{\frac 1N \sum_{i=1}^n\sum_t \nabla_\theta \log \pi_\theta(a_{i, t} | s_{i, t})\sum_{t’ \geq t} r(a_{i, t’}, s_{i, t’})} \label{reward-to-go-policy-grad}
$$
折扣因子
通过让每一时间步获得的 Reward 都乘上 $\gamma$ 来减小方差。下面的式子用来折扣普通的策略梯度：
$$
\nabla_\theta J(\theta)\approx \frac{1}{N}\sum_{i=1}^N\sum_t \nabla_\theta \log \pi_\theta(a_{i, t} | s_{i, t})\sum_t {\color{red}{\gamma^{t - 1}}}r(a_{i, i}, s_{i, t})
$$
而带上 causality assumption 的版本则是
$$
\frac 1N \sum_{i=1}^n\sum_t \nabla_\theta \log \pi_\theta(a_{i, t} | s_{i, t})\sum_{t’ \geq t} {\color{red}{\gamma^{t’ - t}}}r(a_{i, t’}, s_{i, t’})
$$
采样效率低
注意我们要优化的式子中，优化一轮就需要采样一轮，开销非常大。现在考虑是否有办法在原来的分布上进行采样。
引理 2（Importance Sampling）.
$$
\mathbb{E}_{x\sim p(x)}[f(x)] = \mathbb{E}_{x\sim q(x)} \left[\frac{p(x)}{q(x)}f(x)\right]
$$
假设我们此前从一个旧版本的神经网络 $\bar\pi$ 中采样了一组 trajectory，那么我们需要计算
$$
\frac{\pi_\theta(\tau)}{\bar{\pi}(\tau)} = \frac{p(s_1)\prod p(s_{t + 1} | s_t, a_t){\pi}_\theta(a_t | s_t)}{p(s_1) \prod p(s_{t + 1} | s_t, a_t)\bar{\pi}(a_t | s_t)} = \frac{\prod {\pi}_\theta(a_t | s_t)}{\prod \bar{\pi}(a_t | s_t)}
$$
套上 Importance Sampling 得到优化目标函数
$$
J(\theta) = \mathbb{E}_{\tau \sim \bar\pi(\tau)}\left[\frac{\pi_\theta(\tau)}{\bar\pi(\tau)}r(\tau)\right]
$$
所有和 Model 相关的东西都消掉了，因此此方法是可行的。那么简单求导知道：
$$
\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \bar\pi(\tau)}\left[\prod_t \frac{\pi_{\theta}(a_t | s_t)}{\bar\pi(a_t | s_t)}\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\sum_t r(a_{t}, s_{t})\right]
$$
如果带上 causality assumption，有
$$
\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \bar\pi(\tau)}\left[\sum_t \nabla_\theta \log \pi_\theta(a_{t} | s_{t})\prod_{t’ \leq t} \frac{\pi_{\theta}(a_{t’} | s_{t’})}{\bar\pi(a_{t’} | s_{t’})}\sum_{t’’ \geq t} r(a_{t’’}, s_{t’’})\right]
$$
Remark. 这个就有点不知道是什么道理了。