一种局部优化TRPO的思路

UNOboros

已于 2024-07-24 15:56:00 修改

阅读量243

点赞数 4

文章标签：人工智能算法机器学习

于 2024-07-24 15:48:40 首次发布

本文链接：https://blog.csdn.net/UNOboros/article/details/140660547

版权

之前我曾谈到过TRPO的一项重要成果就是从理论上确立了策略梯度一定会保持单调递增的条件：

$max_\theta \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim\pi_{\theta}(s)}[Q_{\theta'}(s,a)]\\ s.t.\ max_s\text{KL}[\pi_{\theta'}(\cdot|s)),\pi_\theta(\cdot|s))]<\delta$

这个优化问题很难求解，因为它涉及到需要搜索全部状态空间去找到KL的最大值，这几乎是不可能的。所以TRPO需要退而求其次，用期望代替求最大值，然后再用二阶泰勒展开去近似求解。这是一种全局的优化思想，因为它侧重于在整体上（求最值也好，取期望也好）约束新策略和旧策略的距离。

另一个思路是局部优化每一个状态，也就是说对于每一个给定的s, 我都去算一下在这个s下最优策略在同等的距离约束下是什么样子的：

$max_\theta \mathbb{E}_{a\sim\pi_{\theta}(s)}[Q_{\theta'}(s,a)]\\ s.t.\text{KL}[\pi_{\theta'}(\cdot|s)),\pi_\theta(\cdot|s))]<\delta$

可见这个局部优化问题比全局优化问题更加困难，因为我们要在每一个状态上去求解一遍约束优化，这并不比原来的TRPO问题要简单，我们依然需要用泰勒展开去近似，而且更困难的是我们怎么才能把每个单独算出来的局部策略统合成一个完整唯一的策略？这里就用到了非参数化技巧：我们不去优化策略的参数，而是直接优化策略函数本身。或者说不对函数的自变量求梯度，而是直接对函数本身求梯度

$max_{\pi_\theta} \mathbb{E}_{a\sim\pi_{\theta}(s)}[Q_{\theta'}(s,a)]\\ s.t.\text{KL}[\pi_{\theta'}(\cdot|s)),\pi_\theta(\cdot|s))]<\delta$

这个问题可以直接用Lagrangian Multiplier去求解：

$L= \mathbb{E}_{a\sim\pi_{\theta}(s)}[Q_{\theta'}(s,a)]+\lambda(\delta-\text{KL}[\pi_{\theta'},\pi_\theta])\\ \frac{\partial L}{\partial \pi_{\theta}}=Q_{\theta'}(s,a)+\lambda(log\pi_\theta-log\pi_{\theta'}+1)\\ \frac{\partial L}{\partial \pi_{\theta}}=0\Rightarrow \pi^*_{\theta}=\pi_{\theta'}exp(\frac{1}{\lambda}Q_{\theta'}(s,a))$

注意到策略是概率，所以我们要归一化处理一下：

$\pi^*_{\theta}=\frac{1}{Z(s)}\pi_{\theta'}exp(\frac{1}{\lambda}Q_{\theta'}(s,a))\\ Z(s)=\sum_a\pi_{\theta'}exp(\frac{1}{\lambda}Q_{\theta'}(s,a))=\mathbb{E}_{a\sim \pi_{\theta'}(\cdot|s)}[exp(\frac{1}{\lambda}Q_{\theta'}(s,a)]$

这样，非参数化技巧使得我们可以计算出针对每个状态的最优策略，这个最优策略可以表示成当前策略的某种加权。最后的问题就是，我们怎么把这个最优策略参数化，也就是把每一个 $\pi^*_{\theta}(s,a)$ 整合成唯一的一个新策略？

一种直观的解法就是我让新策略和每一个 $\pi^*_{\theta}(s,a)$ 的KL距离都尽量接近，也就是求解另一个优化问题：

$min_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}[\text{KL}[\pi^*(\cdot|s)),\pi_\theta(\cdot|s))]]\\ =min_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim \pi^*}[-log\pi_\theta(a|s)]\\ =min_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim \pi_{\theta'}}[-\frac{\pi^*(a|s)}{\pi_{\theta'}(a|s)}log\pi_\theta(a|s)]\\ =min_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim \pi_{\theta'}}[-\frac{\frac{1}{Z(s)}\pi_{\theta'}exp(\frac{1}{\lambda}Q_{\theta'}(s,a))}{\pi_{\theta'}(a|s)}log\pi_\theta(a|s)]\\ =min_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim \pi_{\theta'}}[-\frac{1}{Z(s)}log\pi_\theta(a|s)exp(\frac{1}{\lambda}Q_{\theta'}(s,a))] \\=max_{\theta} \mathbb{E}_{s\sim\rho_{\theta'}(s)}\mathbb{E}_{a\sim \pi_{\theta'}}[\frac{1}{Z(s)}log\pi_\theta(a|s)exp(\frac{1}{\lambda}Q_{\theta'}(s,a))]$

这个目标函数相对于TRPO或者PPO的有很好的性质：它只依赖于从先前的策略上去采样，同时又不会显式的去使用先前的策略。虽然TRPO和PPO也可以从先前的策略上去采样，但是却必须依赖于重要性采样，不可避免的就要用到先前策略。这个性质使得这种局部优化的方法特别适合于不知道先前策略的情况，也就是offline RL。这个方法唯一的问题就是对Z(s)的估计会很复杂。

另一种很特殊的解法只出现在训练大语言模型这个特定场景下。通过局部最优策略去反解奖励模型，然后用preference model去消掉Z(s)项，这就是当今最火的DPO的核心思想了。