👾 xlinker's blog

策略梯度算法

Nov 20, 202418 min read

我去这也太全了 Policy Gradient Algorithms | Lil’Log

直接优化策略函数（Policy Gradient）

（参考Part 3: Intro to Policy Optimization — Spinning Up documentation，Vanilla Policy Gradient — Spinning Up documentation）

调整分布以最大化奖励

首先考虑一个最简单的情况，是我做动作的概率，由参数控制，x是动作，R(x)是单个动作的奖励。我现在想要调整概率分布以最大化期望奖励。

最简单的办法就是选择让R(x)最大的x，但是一般来说很难，因为x可能的空间会很大，比如说可以是一条轨迹，可以是一段token序列。所以一个最简单的想法就是求梯度。
那怎么样对分布求梯度呢？推导如下

Misplaced &\nabla_\theta J(\theta)&=\nabla_\theta \underset{x\sim P_\theta}{\mathrm{E}} \left[R(x)\right] \\ &= \nabla_{\theta}\int_{x}R(x)P_{\theta}(x) \\ &= \int_{x}R(x) \nabla_{\theta}P_{\theta}(x) \\ &= \int_{x}R(x) \nabla_{\theta}\log P_{\theta}(x) P_{\theta}(x)\\ &= \underset{x\sim P_\theta}{\mathrm{E}} [ R(x) \nabla_{\theta}\log P_{\theta}(x) ]\\ \end{aligned}$$ 在倒数第三步的时候用了一个小技巧，因为$\nabla_{\theta} \log P_{\theta}(x) = \frac{\nabla_{\theta}P(x)}{P_{\theta}(x)}$，所以把$\nabla_{\theta}P(x)$替换成$\nabla_{\theta} \log P_{\theta}(x)P_{\theta}(x)$，这样原来的求和又变成了求期望的形式。 这就是说，想要让期望奖励变大，就是要对$\nabla_{\theta}\log P_{\theta}(x)$按照奖励大小进行加权。 假如所有动作奖励一样会怎样呢？假设R(x)=1,这时$J(\theta)=1$，不管$P_\theta$分布如何改变，都不会影响到期望奖励。把$R(x)=1$带入上面的推导过程，发现$\underset{x\sim P_\theta}{\mathrm{E}} [\nabla_{\theta}\log P_{\theta}(x) ]=0$，即当每个动作奖励一样的时候，所有动作梯度的权重都一样，因此梯度的期望就是0. 这个结论非常常用，也叫做Expected Grad-Log-Prob Lemma **EGLP Lemma.** Suppose that $P_\theta$ is a parameterized probability distribution over a random variable, $x$. Then: $$\mathrm{E}_{x\sim P_\theta}\left[\nabla_\theta \log P_\theta(x)\right]=0.$$ 同时我们也发现，奖励的绝对大小并不重要，它可以任意的加减一个常数，因为总会抵消掉。 另外，每更新一次参数，都需要从最新的$P_\theta$里采样$x$来估计梯度，这是这种方法的一个很大的问题。只要想通过梯度来调整分布最大化奖励，就总会面临这种问题。ppo就是对新旧$P_\theta$之间的距离做了一定的约束，让我们可以用旧的采样$x$（$x\sim P_\theta$）来估计新策略$P_\theta'$的梯度。 ### 调整轨迹分布以最大化奖励 接下来我们让x是策略$\pi_\theta$和环境交互所展开的一条轨迹$\tau$。$J(\theta)=\mathrm{E}_{\tau\sim P_\theta}\left[R(\tau)\right]$ 某个轨迹序列$\tau=(s_0,a_0,...,s_{T+1})$，它出现的概率为： $$P(\tau|\theta)=\rho_0(s_0)\prod_{t=0}^TP(s_{t+1}|s_t,a_t)\pi_\theta(a_t|s_t).$$ 其中$P(s_{t+1}|s_t,a_t)$是环境的转移概率，$\pi_\theta(a_t|s_t)$是我们能控制的策略。参照上面的推导，想要求$\nabla_\theta J(\theta)$，我们需要$\nabla_{\theta} \log P(\tau|\theta)$ 对概率两边取对数 $$\log P(\tau|\theta)=\log\rho_0(s_0)+\sum_{t=0}^T\bigg(\log P(s_{t+1}|s_t,a_t)+\log\pi_\theta(a_t|s_t)\bigg).$$ 然后对$\theta$求梯度得到 $$\begin{aligned} \nabla_{\theta}\log P(\tau|\theta)& =\nabla_{\theta}\log\rho_{0}(s_{0})+\sum_{t=0}^{T}\left(\nabla_{\theta}\log P(s_{t+1}|s_{t},a_{t})+\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})\right) \\ &=\sum_{t=0}^T\nabla_\theta\log\pi_\theta(a_t|s_t). \end{aligned}$$ 因此，$\nabla_{\theta}J(\theta) =\underset{\tau\sim\pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(\tau)\right]$ 意思就是说对每个动作，按照其所在轨迹获得的奖励进行加权。其中对轨迹的奖励，可以完全来自于轨迹的结束。也就是只在游戏结束给奖励。但也可以在游戏途中给，这样的话$R(\tau)=\sum_{t=0}^TR(s_{t},a_{t},s_{t+1})$ 代入得到：$\nabla_{\theta}J(\theta) =\underset{\tau\sim\pi_{\theta}}{\mathrm{E}}\left[ \left( \sum_{t=0}^{T}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})\right) \left(\sum_{t=0}^TR(s_{t},a_{t},s_{t+1})\right) \right]$ ### 简化梯度的权重 #### 简化梯度的权重中过去奖励 **reward-to-go** 但是这不太符合直觉，因为每个动作的好坏，与达到该动作之前的奖励无关，应该由该动作后果的平均好坏来评估。通过推导，可以得到下面的梯度估计公式： $$\nabla_\theta J(\theta)=\underset{\tau\sim\pi_\theta}{\mathrm{E}}\left[\sum_{t=0}^T\nabla_\theta\log\pi_\theta(a_t|s_t)\sum_{t^{\prime}=t}^TR(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})\right].$$ 定义$\hat{R}_t\doteq\sum_{t'=t}^TR(s_{t'},a_{t'},s_{t'+1})$，叫做reward-to-go，就是从$t$往后的所有奖励。 **简化梯度权重中过去奖励的推导** 详细推导 [Extra Material — Spinning Up documentation](https://spinningup.openai.com/en/latest/spinningup/extra_pg_proof1.html) 下面我按自己的理解写一下。 $$\begin{aligned} \nabla_{\theta}J(\pi_{\theta})& =\underset{\tau\sim\pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(\tau)\right] \\ &=\underset{\tau\sim\pi_{\theta}}{\operatorname*{E}}\left[\sum_{t=0}^{T}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})\sum_{t^{\prime}=0}^{T}R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})\right] \\ &=\sum\limits_{t=0}^{T}\sum\limits_{t'=0}^{T}\mathop{\mathrm{E}}\limits_{\tau\sim\pi_{\theta}}\left[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t'},a_{t'},s_{t'+1})\right], \end{aligned}$$ 接下来要说明 $\mathop{\mathrm{E}}_{\tau\sim\pi_{\theta}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})]$ ，对$t'<t$时为0. 反之则不为0 我们先把轨迹中和期望里面随机变量相关的东西提出来，然后对期望做拆分。因为条件概率${s_t,a_t|s_{t'},a_{t'},s_{t'+1}}$中$s_{t'},a_{t'},s_{t'+1}$是确定的，所以可以把奖励那一项提出来。 $$\begin{aligned} & \mathop{\mathrm{E}}_{\tau\sim\pi_{\theta}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})] \\ =& \mathop{\mathrm{E}}_{s_t,a_t,s_{t'},a_{t'},s_{t'+1}\sim\pi_{\theta}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}[\mathop{\mathrm{E}}_{s_t,a_t|s_{t'},a_{t'},s_{t'+1}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})]] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}\left[ R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1}) \mathop{\mathrm{E}}_{s_t,a_t|s_{t'},a_{t'},s_{t'+1}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})]\right] \\ \end{aligned}$$ 在$t'<t$时，由于决策的马尔可夫性质（只与上一时刻的状态有关）$a_t\sim\pi_\theta(\cdot|s_t)$，因此根据上面的EGLP引理 $\mathop{\mathrm{E}}_{s_t,a_t|s_{t'},a_{t'},s_{t'+1}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})]=0$ 在$t'\ge t$时，由于动作已经做完了，${a_t|s_t,s_{t'},a_{t'},s_{t'+1}}$. 也就是说，知道未来的某个状态会影响我现在动作的条件概率。因此无法使用EGLP引理，所以 $\mathop{\mathrm{E}}_{s_t,a_t|s_{t'},a_{t'},s_{t'+1}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})]\ne 0$ 因此每个$\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})$ 的权重可以改写成reward-to-go的形式。 #### 简化梯度权重中未来的奖励 更进一步，考虑一个不常见的情况，假如已知未来的奖励与当前动作无关呢？是不是每个$\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})$ 的权重都可以简化成自己的奖励了呢？ 未来奖励与当前动作无关，意味着$P({s_t,a_t,R(s_{t'},a_{t'},s_{t'+1})})=P({R(s_{t'},a_{t'},s_{t'+1})|s_t,a_t})$. 即$P({s_t,a_t,R(s_{t'},a_{t'},s_{t'+1})})=P(s_t,a_t)$ 参考上面的推导，可以把${s_t,a_t|s_{t'},a_{t'},s_{t'+1}}$换成${s_t,a_t|R(s_{t'},a_{t'},s_{t'+1})}$ $$\begin{aligned} & \mathop{\mathrm{E}}_{\tau\sim\pi_{\theta}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})] \\ =& \mathop{\mathrm{E}}_{s_t,a_t,s_{t'},a_{t'},s_{t'+1}\sim\pi_{\theta}}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}[\mathop{\mathrm{E}}_{s_t,a_t|R(s_{t'},a_{t'},s_{t'+1})}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})]] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}\left[ R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1}) \mathop{\mathrm{E}}_{s_t,a_t|R(s_{t'},a_{t'},s_{t'+1})}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})]\right] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}\left[ R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1}) \mathop{\mathrm{E}}_{s_t}~ \mathop{\mathrm{E}}_{P(a_t|s_t,R(s_{t'},a_{t'},s_{t'+1}))}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})]\right] \\ =& \mathop{\mathrm{E}}_{s_{t'},a_{t'},s_{t'+1}}\left[ R(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1}) \mathop{\mathrm{E}}_{s_t}~ \mathop{\mathrm{E}}_{\pi_\theta(a_t|s_t)} \left[\frac{P(a_t|s_t,R(s_{t'},a_{t'},s_{t'+1}))}{\pi_\theta(a_t|s_t)}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t}) \right]\right] \\ \end{aligned}$$ 如果满足$\frac{P(a_t|s_t,R(s_{t'},a_{t'},s_{t'+1}))}{\pi_\theta(a_t|s_t)}=1$，即使在$t'>t$时，这个奖励也可以忽略. 所以 $$\nabla_\theta J(\theta)=\underset{\tau\sim\pi_\theta}{\mathrm{E}}\left[\sum_{t=0}^T\nabla_\theta\log\pi_\theta(a_t|s_t)R(s_{t},a_{t},s_{t+1})\right].$$ 这相当于是认为每个动作即时获得奖励，因此不必考虑以后的事情。 比如在某个游戏中，我们已知之后是否获得奖励与之前的动作关系不大，或者说在我们设计奖励的时候，仅仅是想奖励当前动作本身，而不想对过去行为进行奖励。 举个例子，我们玩多臂赌博机的时候，如果赌博机里面不保存状态，就是上面说的这种情况，此时可以忽略后面的奖励。 而如果机器里面保存状态，按下按钮的奖励和之前行为有关的话，那就需要考虑后面的奖励了。而由于没人知道里面的状态转移逻辑，就只能不断尝试，统计之后的奖励，尝试估计优势函数。 如果真实情况下我们已知未来奖励和现在关系不大，就可以采用折中的办法，适当减小GAE中的$\gamma$，加强对未来奖励的惩罚力度。因为知道未来的奖励和现在关系不大，所以减少$\gamma$的同时，也不用担心梯度估计会过度不准。这样的好处就是能获得方差更小的优势函数估计。 如果游戏中有多种类别的奖励，那么可能对它们分别采用不同的折扣系数$\gamma$会比较合理。 但问题是这样可能需要训练多个critic（价值网络）来估计不同折扣系数的奖励对应的价值。 #### 降低reward-to-go的方差 $$\nabla_\theta J(\theta)=\underset{\tau\sim\pi_\theta}{\mathrm{E}}\left[\sum_{t=0}^T\nabla_\theta\log\pi_\theta(a_t|s_t)\sum_{t^{\prime}=t}^TR(s_{t^{\prime}},a_{t^{\prime}},s_{t^{\prime}+1})\right].$$ 上面的式子还可以进一步简化，reward-to-go可以替换成$Q(s_t,a_t)$ 一个感性的理解是，我采样了成千上万条轨迹，那么肯定有几条轨迹在相同的状态$s_t$下采用了相同的动作$a_t$. 这些项对应的梯度也是一样的，都是$\nabla_\theta\log\pi_\theta(a_t|s_t)$，只是它们后续获得的奖励不同。我们将它们后续获得的奖励合并到一起，就是$(s_t,a_t)$之后平均获得的奖励，就是Q函数的定义。 #### 减去梯度权重中的$b(s_t)$ todo #### 使用critic估计梯度的权重 [[GAE：广义优势函数估计]] --- 这个梯度更新公式其实就是在说，根据当前策略随机尝试一个动作，如果平均结果比预期好，那么就增加这个动作的概率，如果平均结果比预期差，那么就减少这个动作的概率。 它告诉了我们应该怎样估计目标函数J(x)的梯度。同时告诉了我们怎样去信用分配（credit assignment）。信用分配问题指的是，在一串决策后获得了高奖励，我应该如何知道究竟是哪一个动作让我获得了这么高的奖励呢？策略梯度的推导告诉我们，好的动作优势函数A(s,a)高。所以只需要多采样，采样采的够多了，你自然可以通过准确估计优势函数来处理奖励分配问题。优势函数的估计的误差是由方差和偏差造成的，你可以通过优势函数估计（GAE）来平衡它们。 但如果奖励过于稀疏，最后还是需要很大的采样开销才能有效的学习，所以需要合理设计奖励来引导。 一般来说，在有监督学习的神经网络中，是在训练集数据上计算并最小化损失函数。损失函数是对固定的训练集数据求期望。神经网络的参数只影响期望里面的东西。 但是policy gradient的目标函数不同，这个目标函数里是对轨迹的分布求期望奖励。神经网络输出的策略会影响轨迹分布。这就决定了每一步梯度更新都需要在当前策略上采样轨迹和奖励数据，进而计算当前策略的梯度。 [Deep Reinforcement Learning: Pong from Pixels](http://karpathy.github.io/2016/05/31/rl/) 这篇2016年的博客以策略梯度算法为例讲了强化学习，很有启发性。其中最有意思的地方是它强调了强化学习可以控制一个不可微的黑箱。 简单策略梯度的公式虽然简单，但每更新一步就需要采样大量轨迹，这样数据利用效率太低了。于是[John Schulman](http://joschu.net/publications.html)提出了TRPO和更简洁的PPO以提升轨迹数据的利用效率。 V函数和Q函数都和使用什么策略有关，因此随着策略函数更新，需要用新轨迹数据来同步更新V函数的值。 ## PPO 我们回顾一下刚刚最开始的的推导。考虑一个多臂老虎机，$P_\theta$是我做动作的概率，由参数$\theta$控制，x是动作，R(x)是单个动作的奖励。我现在想要调整概率分布$P_\theta$以最大化期望奖励$J(\theta)$。 $J(\theta)=\mathrm{E}_{x\sim P_\theta}\left[R(x)\right]$ 想要让最大化$J(\theta)$，梯度按照下面的式子来估计 $$\begin{aligned} \nabla_\theta J(\theta)&=\nabla_\theta \underset{x\sim P_\theta}{\mathrm{E}} \left[R(x)\right] \\ &= \nabla_{\theta}\int_{x}R(x)P_{\theta}(x) \\ &= \int_{x}R(x) \nabla_{\theta}P_{\theta}(x) \\ &= \int_{x}R(x) \nabla_{\theta}\log P_{\theta}(x) P_{\theta}(x)\\ &= \underset{x\sim P_\theta}{\mathrm{E}} [ R(x) \nabla_{\theta}\log P_{\theta}(x) ]\\ \end{aligned}$$ 也就是说，每次估计梯度时，从$P_\theta$中采样动作$x$，以拿到的奖励为权重，对对应动作的梯度加权。这里的问题是每次更新策略都需要从新策略$P_\theta$中进行采样，能不能从旧策略中进行采样估计新策略的梯度呢？应该是可以的，使用拒绝采样的技巧就可以。上式中倒数第二步可以变成 $$\begin{aligned} &= \int_{x}R(x) \nabla_{\theta}\log P_{\theta}(x) P_{\theta}(x)\\ &= \int_{x}R(x) \nabla_{\theta}\log P_{\theta}(x) \frac{P_{\theta}(x)}{P_{\theta'}(x)} P_{\theta'}(x)\\ &= \underset{x\sim P_\theta}{\mathrm{E}} [ R(x) \nabla_{\theta}\log P_{\theta}(x) \frac{P_{\theta}(x)}{P_{\theta'}(x)}]\\ &= \underset{x\sim P_\theta}{\mathrm{E}} [ R(x) \frac{\nabla_{\theta}P_{\theta}(x)}{P_{\theta'}(x)}]\\ \end{aligned}$$ 也就是说，只要能直接算出新旧策略对应动作的概率，直接乘一个比例系数就可以准确估计新动作概率的梯度了。最终，策略将收敛到奖励最高的动作上。一切都是那么完美~ 然而一般的强化学习解决的是序列决策问题，每个动作的"权重"，或者说优势函数，都与当前策略有关！为了确保旧策略采样出来的轨迹，估计出来的优势函数/奖励在新策略上也适用，我们不得不限制采样轨迹的策略与当前策略之间的距离！这就不是简单的乘一个拒绝采样系数$\frac{P_{\theta}(x)}{P_{\theta'}(x)}$就可以解决的事了，必须实际去采样才行（ 这就是轨迹必须on-policy的根本原因 不过虽然优势函数和策略有关，reward-to-go就只和轨迹有关了。而语言模型的ppo恰好一般使用reward-to-go - baseline TRPO [Trust Region Policy Optimization — Spinning Up documentation](https://spinningup.openai.com/en/latest/algorithms/trpo.html#background) 策略梯度直接对优化目标$J(\pi)$进行优化，但是每训练一步都需要依据当前策略采样大量数据。也就是要求训练数据on-policy。 TRPO希望提升数据利用效率，即希望旧策略采样出来的轨迹数据能在新策略上多用几次，旧的轨迹数据能多训练几步。但是最多训练几步呢？作者对$\pi_\text{old}$和$\pi_\theta$之间的距离做了限制，这样的限制保证了代理目标函数的提升等于实际目标函数的提升。也就是我可以安全的在这个条件下做优化。但是这个算法的问题在于，实际训练时需要算输出相对于模型参数的黑塞矩阵。训练效率太低。 PPO [Proximal Policy Optimization — Spinning Up documentation](https://spinningup.openai.com/en/latest/algorithms/ppo.html) PPO利用TRPO的思想，把距离限制改成了启发式的约束。下面简单描述一下 训练两个网络，策略网络$\pi(s,a)$和价值网络$V(s)$。策略网络采样得到动作序列，使用[Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438) (GAE)估计优势函数A(s,a)，利用估计到的优势函数调整策略。比如A(s,a1)>0，那么就提高策略做出动作a1的概率，即提高$\pi(s,a1)$。反之降低概率。 但是A(s,a1)可能被错误的估计，比如说过大，这时就可能激励$\pi(s,a1)$错误的提升$a1$的概率。PPO就以$\pi_\text{old}(s,a1)$为参照，约束$\pi(s,a1)/\pi_\text{old}(s,a1)$的偏离不超过一定比例。实验发现这种办法简单好用。 另一种启发式的办法是增加旧策略和新策略之间KL散度的正则化项，并动态调整系数。

Graph View

Backlinks