FPO——流匹配策略梯度:避开复杂的对数似然计算,通过「最大化基于CFM损失计算优势加权比率」做策略优化,兼容PPO-CLIP

FPO——流匹配策略梯度:避开复杂的对数似然计算,通过「最大化基于CFM损失计算优势加权比率」做策略优化,兼容PPO-CLIP
前言在解读π0.6『详见此文《π∗0.6——通过RL框架RECAP微调流式VLA π0.6先基于示教数据做离线RL预训练再SFT(即IL微调)最后在线RL后训练(与环境自主交互从经验数据中学习且必要时人工干预)》』的过程中便有意再写一篇文章解读下π*0.6的第82篇参考文献『82-Flow matching policy gradients』即FPO如此本文便来了第一部分 流策略优化(Flow Policy OptimizationFPO)1.1 引言与相关工作1.1.1 引言首先根据此文《强化学习极简入门通俗理解MDP、DP MC TD和Q学习、策略梯度、PPO》可知在RL中目标是训练一个智能体使其在与环境的交互中通过最大化累积奖励获取一套最优的行为策略以下是策略梯度公式这里的关键是即在策略下采取动作的对数似然(log-likelihood)是优势函数简单理解就是动作比当前状态下的平均动作好多少策略梯度公式公式告诉我们 优化的方向就是增加那些带来正优势的动作的对数似然其次在机器人控制等连续动作空间任务中策略通常被建模为一个高斯分布。也就是说策略网络会输出一个均值和一个标准差然后从中采样得到动作高斯分布的优点是 易于实现和理解且对数似然有解析解计算快缺点是如此文所说的“ 高斯分布有一个天生的缺陷它是单峰的unimodal好比你开车到一个十字路口左转和右转都能顺利到达目的地并且耗时差不多。一个优秀的策略应该能告诉你左转和右转都是好选择。但一个高斯策略的「导航仪」做不到这一点。经过训练它可能会发现左转的期望回报是 10右转是 9.9于是它会坚定地认为 左转是唯一最优解并不断强化这个信念它无法表达 左转和右转都很好 这种多峰值multimodal 的决策”总之在许多复杂任务中这种多峰值决策能力至关重要。比如人形机器人为了保持平衡可能有多种不同的手臂摆放姿势。一个只能认准一种姿势的高斯策略显然不够灵活和鲁棒流策略优化Flow Policy Optimization, FPO这是一种用于优化基于流的生成模型的策略梯度算法。FPO将策略优化重新表述为最大化一个基于『条件流匹配(conditional flow matching简称CFM)目标函数[13-Flowmatching for generative modeling]』计算得到的优势加权比值其避开了通常与流模型相关的复杂似然计算而是在策略梯度中使用流匹配损失作为对数似然的替代项这使得目标能够直接与提升高奖励动作的证据下界(ELBO)对齐与之前将去噪过程重构为马尔可夫决策过程(MDP)的方法不同后者将训练绑定到特定的采样方法并延长了奖励归因的时序跨度而FPO 在 rollout期间将采样过程视为黑盒1.1.2 相关工作首先对于基于策略的强化学习而言在该领域中通过对参数化策略进行优化以使其在给定环境中获得最大累积奖励通常这一问题通过策略梯度方法进行求解这类方法通过将动作的对数概率与观测到的奖励或优势相加权从而避开了对环境奖励函数可微性的需求15-Policy gradientmethods for reinforcement learning with function approximation16-imple statistical gradient-following algorithms for connectionist reinforce-ment learning17-A natural policy gradient.18-Natural actor–critic19-Trustregion policy optimization20-Proximalpolicy optimization algorithms即PPO21-Asynchronous methods for deep reinforcementlearning22-Sample efficient actor–critic with experience replay23-Deepseekmath策略梯度方法在普通的连续控制任务24-Benchmarking deepreinforcement learning for continuous control25-Open rl benchmark: Comprehensive tracked experiments for reinforcementlearning机器人运动26-Learning to walk in minutesusing massively parallel deep reinforcement learning.27-Curiosity-driven learning of joint locomotion and manipulation tasks.28-Sym-metry considerations for learning task symmetric robot policies29-Visual imitationenables contextual humanoid control以及操作控制30-Solving rubik’s cubewith a robot hand31-system for general in-hand object re-orientation32-General in-hand object rotation with vision and touch33-From simpleto complex skills: The case of in-hand object reorientation中处于核心地位。这些方法也越来越多地被用于在预训练生成模型中搜索和改进先验分布且在与人类偏好进行对齐34-Training language models to follow instructions with human feedback即instructGPT35-Deep reinforcement learning from human prefere2017年6月OpenAI联合DeepMind首次正式提出的Deep Reinforcement Learning from Human Preferences即基于人类偏好的深度强化学习简称RLHF和利用可验证奖励提升推理能力等方面36-Deepseek-r137-Mistral-AI此类方法也已被证明行之有效在本研究中作者提出了一种用于训练基于流的生成策略(如扩散模型)的简单算法该算法基于策略梯度框架通过利用近期流匹配(flow matching)[13-Flowmatching for generative modeling]的研究成果训练的策略能够表示比连续控制强化学习中最常用的对角高斯分布[26–29,32,33]更为丰富的分布同时仍然兼容标准的演员-评论家(actor-critic)训练技术其次对于扩散模型而言扩散模型是用于建模复杂连续分布的强大工具已在广泛领域取得了显著成功。这些模型已成为图像生成38-DDPM39-DDIM40-High-resolution image synthesis with latent diffusion models即SD的奠基论文41-Generative modeling by estimating gradients of the data distribution视频生成42-Video diffusion models43-Make-a-video:Text-to-video generation without text-video data44-Imagen video: High definition video generation with diffusion models4-Video generation models as world simula音频生成7-Audioldm: Text-to-audio generation with latent diffusion models45-Grad-tts: A diffusion probabilistic model for text-to-speech46-Wavegrad 2: Iterative refinement for text-to-speech synthes8-Diffwave: A versatilediffusion model for audio synthesis以及最近的机器人动作生成9-Diffusion policy: Visuomotor policy learning via action diffusion47-π048-Gr00t n1的主流方法在这些应用中扩散模型的目标是从感兴趣的数据分布中进行采样这些分布可能来源于互联网爬取或通过人类远程操作收集Flow matching [13] 简化并推广了扩散模型框架。它学习一个将样本从易处理的先验分布传输到目标数据分布的矢量场条件流匹配(CFM)目标训练模型去对被高斯噪声扰动的数据进行去噪给定数据和噪声CFM 目标可以表达为(定义为公式1)其中表示在流动步骤下部分加噪的样本它是由超参数和定义的噪声与数据的插值是模型对原始数据速度的估计是条件流该模型还可以将去噪样本x 或噪声分量ϵ 作为优化目标进行估计而不是速度所学习到的速度场是一个连续映射通过ODE 积分将样本从简单且易处理的分布如高斯噪声迁移到训练数据分布总之通过流模型直接优化似然性是可行的但这需要进行散度估计 [49-The superposition of diffusion models using the itô density estimator]并且在计算上非常昂贵因此流匹配通过上述简单的去噪损失(即条件流匹配损失)来优化似然的变分下界在本研究中作者直接在策略梯度框架中利用流匹配这种方法能够基于奖励来训练扩散模型而无需进行计算成本高昂的似然计算最后对于扩散策略基于扩散的策略在机器人和决策制定应用中已展现出良好的前景50-Diffusion policy51-Is conditional generative modeling all you need for decision making?47-π0目前大多数方法通过行为克隆52-Planning with diffusionfor flexible behavior synthesis9-Diffusion policy: Visuomotor policy learning via action diffusion训练这些模型即在没有奖励反馈的情况下对策略进行监督以模仿专家轨迹受扩散模型和流模型强大生成能力的启发已有多项工作探索使用RL对扩散模型进行微调尤其是在文本到图像生成等领域 [53–55]Psenka 等人[56-Learning a diffusion modelpolicy from rewards via q-score matching]最近的研究探索了通过 Q-score 匹配对扩散策略进行离策略训练尽管离策略强化学习持续取得进展57-Fasttd3: Simple, fast, and capable reinforcement learning for humanoid control58-Addressing function approximation error inactor-critic methods但当前实际应用中主要还是以顺策略方法为主————例如DDPO54-Training diffusionmodels with reinforcement learningDPPO59-Diffusion policy policy optimization以及 Flow-GRPO55-Flow-grpo: Training flow matching models via online rl等方法采用顺策略策略梯度通过将初始噪声值视为环境观测把去噪过程建模为马尔可夫决策过程并在每一步都将其作为高斯策略用 PPO 进行训练本文要介绍的方法与之不同直接将条件流匹配(CFM)的目标集成到类似 PPO 的框架中同时保持标准扩散正向与反向过程的结构由于 FPO 将流匹配作为其基础原语因此在训练和推理阶段对于采样方法的选择是无关的这与行为克隆中的流匹配是一致的1.2 先给结论流匹配策略梯度与流策略优化1.2.1 策略梯度与PPORL的目标是学习一个策略以最大化在给定环境中的期望回报。在每一次在线RL的迭代中策略会被执行以收集每个环境时间步的观测、动作和奖励三元组的批次这些轨迹可以用于策略梯度目标函数[15-Policy gradientmethods for reinforcement learning with function approximation]以增加导致更高奖励的动作的概率其中是根据回合奖励和学习到的价值函数估算的优势 [60-High-dimensional continuous control using generalized advantage estimation]标准的策略梯度仅在当前策略参数附近局部有效而若进行大幅度的更新可能会导致策略崩溃或学习过程不稳定为了解决这一问题PPO [20] 通过对似然比进行裁剪引入了信赖域机制其中是可调节的阈值表示当前动作概率与旧动作概率的比值(定义为公式4)PPO因其稳定性、简洁性和较高性能是在策略强化学习中的流行选择『如果对PPO还不熟悉的可以参见此文《强化学习极简入门通俗理解MDP、DP MC TD和Q学习、策略梯度、PPO》的4.4节』然而与标准策略梯度方法类似它需要对采样动作计算精确的似然概率对于简单的高斯或分类动作空间这些量是可计算的但对于流匹配和扩散模型来说估算这些量的计算成本极高几乎难以实现1.2.2 流策略优化定义条件流匹配损失函数流策略优化(FPO)是一种用于以流模型表示策略的在线强化学习算法。在实际应用中它与高斯PPO主要有两个关键区别在 rollout 过程中流模型通过一系列学习得到的变换将随机噪声转化为动作[13-Flow matching for generative model]从而使策略表达能力远超标准PPO此外在更新策略时FPO 用流匹配损失替代了原有的高斯似然损失直观来说FPO 的目标是引导策略的概率流向高回报行为。作者不再计算似然值而是利用标准的流匹配损失来构造一个简单的比率估计(定义为公式6)————正如下文将要讨论的这可以通过优化证据下界(即Evidence Lower Bound, ELBO也称变分下界)推导出来进一步而言对于给定的动作和观察对是每个样本条件流匹配损失『而这个便是由公式1所定义的』的一个估计该估计的定义为(下述三个式子分别为公式7、公式8、公式9)其中用表示流动时间步用表示环境时间步。作者在中包含了两个时间步这表示在回合时间下、噪声水平为时的动作——依据公式1且作者在和之间使用了相同的和样本在公式6中FPO的比率估计可以直接替代PPO的似然比『PO’s ratio estimate in Equation 6 serves as a drop-in replacement for the PPO likelihoodratio』因此FPO继承了对优势估计算法如GAE [60] 和 GRPO [23]的兼容性在不影响一般性的前提下它同样兼容基于噪声[38] 或干净动作[1] 估计的流和扩散实现这些实现可以通过重加权实现与[61, 13] 在数学上的等价性(作者将在下文FPO比率的推导中利用这一特性)1.3 再做证明——先证变分下界ELBO与对数似然等价再证流匹配损失与ELBO等价1.3.1 FPO代理目标优化变分下界(ELBO)作为对数似然的替代目标本节摘要如此文所说首先在变分推断中经常用证据下界(Evidence Lower Bound, ELBO)来近似真实的对数似然即可以认为总之在基于流的模型中精确的似然估计在计算上非常昂贵。因此通常采用优化证据下界(ELBO)作为对数似然的替代目标(定义为公式10)其中是ELBO 与真实对数似然之间的KL 间隙πθ 是通过从流模型中采样获得的分布。流匹配和扩散模型都使用条件流匹配损失『即一种简单的MSE 去噪目标[62, 13]』来优化ELBOFPO 比率(公式11)利用了流模型可以通过ELBO 目标进行训练的这一事实即具体而言作者计算当前策略与旧策略下ELBO 的比值(定义为公式11)对该比率进行分解可以得到真实似然比『见公式4』的一个缩放变体(定义为公式12)在这里该比值被分解为标准似然比和包含KL间隙的逆向修正项因此最大化该比值能够提升模型的似然度同时减少KL间隙——这两者对于策略优化都是有益的前者促进策略倾向于选择具有正优势的动作而后者则收紧了对真实对数似然的近似1.3.2 利用流匹配估算FPO比率最小化流匹配损失 最大化 ELBO(从而也等价最大化似然)本节摘要可用损失函数近似 ELBO研究表明流匹配的训练损失即条件流匹配损失(Conditional Flow Matching Loss)与 ELBO 之间有如下关系衡量的是模型预测的「流动方向」和真实「流动方向」之间的均方误差(MSE)这个公式告诉我们最小化流匹配损失就等价于最大化 ELBO从而间接等价于最大化真实似然作者直接使用流匹配目标来估计FPO 比率这源自Kingma 和Gao [63-UnderstandingDiffusion Objectives as theELBOwith Simple Data Augmentation本文第二部分 会更详细介绍该工作] 所建立的加权去噪损失与ELBO 之间的关系是流匹配和去噪扩散损失的一个更通用的形式它将模型参数化为预测即对模型输入中真实噪声的估计对于一个干净动作加权去噪损失的形式为其中是加权的选择表示噪声水平时的log-SNR作者用对时间步和噪声的蒙特卡洛抽样来估计该值权重选择将条件流匹配损失和标准扩散损失作为更一般族的特例纳入其中作者在这里关注恒定权重情况(扩散调度)这带来了最简单的理论联系类似的结果也适用于许多常用的调度方式包括最优传输和方差保持调度[13]。详情请参见补充材料总之对于扩散调度[63-Understanding Diffusion Objectives as theELBOwith Simple Data Augmentation] 证明如下其中是关于的常数。从几何上看最小化使流动更多地指向最小化也最大化了ELBO『即公式10』因此的似然性也会增大所以流向某一特定动作会使其更有可能被采样这一直观理解与策略梯度目标自然一致希望增加高优势动作的概率。通过将流动重定向到这些动作(即最小化它们的扩散损失)使它们在学习到的策略下更有可能被采样利用该关系作者将 FPO 比率(公式 11)用流匹配目标来表示(定义为公式16)这就是 FPO 的核心技巧它用一个基于损失函数差值的指数替代了那个难以计算的似然比打个比方我们想知道新配方 A 比旧配方 B 做出的蛋糕好吃多少似然比。直接测量「好吃程度」很难。但可以测量制作过程的「费力程度」CFM 损失如果对于同一个蛋糕新配方 A 比旧配方 B 省力得多损失更小则就推断 A 做出的蛋糕更好吃似然更高这个「省力程度」的差值就被 FPO 用来近似「好吃程度」的比率其中如公式7 所示——可以通过对次采样的平均来估计。且作者发现样本数量是调控学习效率的一个有用超参数。虽然在实际操作中作者只使用少量采样但该估计器在极限情况下能够还原精确的FPO 比率一个可能的担忧是较小的值会带来偏差仅由一个对估算得出的比值仅在期望意义上这是对真实比率的一个上界可以通过Jensen不等式证明如下为了理解这种上偏差可以利用对数导数技巧对 FPO 梯度进行分解由于梯度算子与期望运算可交换右侧的梯度项是无偏的(定义为公式20)换句话说即使在比值被最坏情况高估的情况下梯度估计在方向上也是无偏的————作者的实验结果与这一结论一致虽然增加样本数量有帮助但在第4.2 节中通过实证观察到即使在的情况下FPO 也可以通过训练优于高斯PPO算法1详细介绍了FPO在该数学框架下的实际实现第一步定义策略策略网络不再输出和而是变成一个流模型输入状态、加噪的动作和时间步输出速度第二步与环境交互Rollout当需要一个动作时先从高斯分布中采样一个初始噪声然后用一个 ODE 求解器(如欧拉法)在策略的引导下从出发积分几步得到最终的动作用这个动作与环境交互收集数据(s, a, r, …)第三步策略更新Update在 PPO 的更新循环中对于每个采集到的动作计算它在旧策略下的 CFM 损失和在新策略下的 CFM 损失而CFM 损失是通过对动作随机加噪几次并求平均来估计的用计算似然比然后代入到PPO中的clip损失函数中最后像往常一样最大化这个目标函数来更新策略网络1.3.3 去噪MDP对比现有的基于扩散模型的策略内强化学习算法 [54,59,55]将去噪过程本身重新表述为马尔可夫决策过程MDP。这些方法通过将采样链中的每一步视为一个独立的动作且每一步都被参数化为高斯策略步骤从而绕过了流模型的似然计算而FPO 针对这些方法存在的一些局限性进行了改进首先去噪 MDP 会将时间跨度乘以去噪步数通常为 10-50 步这增加了信用分配的难度其次这些 MDP 在计算似然时并不考虑初始噪声样本。相反这些噪声值被视为环境的观测值[59]——这显著增加了学习问题的维度最后去噪 MDP 方法由于其构建方式仅限于随机采样过程相比之下FPO采用流匹配因此继承了标准流/扩散模型对于采样器选择的灵活性。这包括快速的确定性采样器、高阶积分和可以任意选择采样步数也许更为重要的是FPO 的实现更为简单因为它不需要自定义采样器或额外环境步数的概念第二部分(选读) 将扩散目标理解为带有简单数据增强的 ELBO扩散模型可以被解释为深度变分自编码器(VAE)[Kingma和Welling,2013; Rezende等,2014]的一种特殊情形其推断模型和生成模型有特定的选择与VAE类似最初的扩散模型[Sohl-Dickstein等,2015]也是通过最大化数据对数似然的变分下界(也称为证据下界简称ELBO)进行优化的Variational Diffusion ModelsVDM[Kingma等,2021]以及[Song等,2021a]展示了如何利用ELBO目标优化连续时间的扩散模型并在图像密度估计算法基准上实现了最先进的似然性能然而就样本质量指标如FID分数而言最佳结果是通过其他目标实现的例如去噪得分匹配目标[Song and Ermon,2019]或简单的噪声预测目标[Ho等,2020]这些如今广泛使用的目标函数表面上看与传统流行的最大似然目标和ELBO目标有很大不同但通过本文的分析作者揭示出所有当前最先进扩散模型中采用的训练目标实际上都与ELBO目标密切相关比如在下文中作者展示出如果加权函数是时间的单调函数则加权损失等价于在数据增强(即高斯噪声扰动)下最大化ELBO。这一结论适用于例如[Salimans 和 Ho, 2022]提出的v-预测损失以及[Lipman 等人,2022]提出的最优传输路径上的流匹配// 待更