FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance¶
会议: NeurIPS 2026
arXiv: 2605.30749
项目主页: https://flag-rl.github.io/
领域: 强化学习
关键词: 最大熵强化学习、流匹配策略、重要性采样、隐变量增强MDP、高维连续控制
一句话总结¶
针对生成式连续策略在最大熵强化学习(MaxEnt-RL)中面临的全局重要性采样权重坍塌以及穿透常微分方程的时间反向传播(BPTT)不稳定困境,FLAG 提出将状态空间提升至隐变量增强 MDP,利用流匹配确定性映射构建以锚点动作为中心的局部高斯策略,通过交叉熵代理目标和期望最大化(EM)实现局部重要性采样监督,在高维复杂控制任务与极限动作评估预算下实现稳定高效的策略优化。
研究背景与动机¶
最大熵强化学习(MaxEnt-RL)通过在期望累积回报中引入策略熵正则化项,赋予智能体强大的多模态行为探索能力与鲁棒性。传统实用算法如软行动者-评论家(SAC)受限于可解析优化的需求,通常将策略强制约束为单模态高斯分布,这在本质上抹杀了复杂决策环境中探索多最优解的能力。为了突破表达力瓶颈,近期工作尝试引入扩散模型(Diffusion)与连续归一化流匹配(Flow Matching)等生成模型作为强表达力策略架构。
然而,将生成式策略深度融入在线强化学习面临两难困境。一类方法依赖穿越生成轨迹全过程的时间反向传播(BPTT)计算动作梯度(如 DACER、DIME),随着常微分方程(ODE)求解步数增加,敏感性方程沿局部雅可比矩阵累积极易导致梯度消失或爆炸,计算极耗内存且高度数值脆弱;另一类方法沿袭 MPO 路线,采用期望最大化(EM)将策略更新转化为目标分布上的监督学习。但由于真实最大熵目标分布难以直接采样,该范式依赖全动作空间的重要性采样(Global IS)。在高维动作空间中,全局提议分布与目标分布之间的支撑集严重错配,导致重要性权重急剧坍塌为零,使监督信号高度稀疏,无法有效追踪在线 RL 中非平稳的目标模式。
本文的切入角度是:与其在广袤的全动作空间中进行盲目且稀疏的全局重要性采样,不如借助流匹配模型确定性连续变换的几何特性,将提议分布与目标分布联合局部化在流生成的锚点动作近邻区域,从而在局部受限区域内实现致密且鲁棒的局部重要性采样。核心 idea:将原始 MDP 提升为融合流潜在向量的隐变量增强 MDP,引入可解析评估的交叉熵作为代理最大熵目标,并推导具备单调提升理论保证的局部 EM 更新,将局部重要性采样优化的目标动作作为确定性指引蒸馏回流匹配网络中。
方法详解¶
整体框架¶
FLAG 围绕“以流匹配锚点为中心的局部高斯策略构建 \(\to\) 提升至隐变量增强 MDP(\(\mathbf{z}\)-MDP)\(\to\) 局部期望最大化(EM)求解改进目标动作 \(\to\) 条件流匹配(CFM)监督蒸馏回流策略”这一闭环管线展开。整体架构规避了穿透常微分方程的高阶梯度计算,并将困难的全局高维采样退化为锚点周围的高效局部扰动加权。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["环境状态 s 与先验潜变量 z ~ p_z 输入"] --> B["流生成锚点动作与局部高斯策略构建<br/>通过流ODE映射生成锚点 T_theta(s,z) 并叠加各向同性高斯扰动"]
B --> C["提升至隐变量增强 MDP 与交叉熵代理目标<br/>在 (s,z) 状态空间评估 Q 函数与基于 Hutchinson 的交叉熵"]
C --> D["局部重要性采样与期望最大化更新<br/>在锚点近邻执行 E-step 局部 IS,通过一阶矩匹配得出目标动作 mu*"]
D --> E["目标动作缓存与条件流匹配监督蒸馏<br/>存入 Guidance Buffer 并通过 CFM 目标监督更新流向量场"]
E --> F["输出表达力提升且具备多模态探索能力的全局流策略"]
关键设计¶
1. 局部高斯策略参数化与隐变量增强 MDP:收敛重要性采样支撑集
为了彻底摆脱全动作空间全局采样导致的维度灾难,FLAG 首先利用参数化流变换 \(T_\theta(s, z) = a^1 = z + \int_0^1 v_\theta(a^\tau, \tau, s) d\tau\) 将先验分布 \(z \sim p_z\) 映射为状态条件下的锚点动作。随后,以该确定性锚点为均值中心,叠加方差为 \(\Sigma = \sigma^2 I\) 的辅助局部高斯分布,定义局部策略 \(\hat{\pi}(a \mid s, z; \theta) = \mathcal{N}(a; T_\theta(s, z), \Sigma)\)。对隐变量 \(z\) 求边际即可无损复原全局复合策略 \(\pi(a \mid s; \theta) = \int p_z(z) \hat{\pi}(a \mid s, z; \theta) dz\)。
为了给局部策略赋予坚实的序贯决策理论基础,论文定义了隐变量增强 MDP(\(\mathbf{z}\)-MDP)\(\hat{\mathcal{M}}\),其增广状态为 \(\hat{s} = (s, z) \in \mathcal{S} \times \mathcal{Z}\),状态转移核满足 \(\hat{p}(\hat{s}' \mid \hat{s}, a) = p(s' \mid s, a) p_z(z')\),即每步从先验中独立刷新隐变量。论文在理论上严格证明了边际一致性(Corollary 4.2)与 Q 函数一致性(Corollary 4.4):即在 \(\mathbf{z}\)-MDP 中任何关于可测函数的期望均精确等价于原始 MDP 中的边际期望,且两者共享完全相同的动作价值函数 \(Q^{\hat{\pi}}(\hat{s}, a) = Q^{\pi}(s, a)\)。这一构造在数学上保证了在增广空间中直接对局部策略 \(\hat{\pi}\) 进行优化,等价于在全局尺度上提升整体复合策略 \(\pi\)。
2. 交叉熵代理目标与理论一致性:破解不可解析熵计算瓶颈
在标准 MaxEnt-RL 中,复合策略 \(\pi\) 包含对隐变量 \(z\) 的边际积分,其对数密度 \(\log \pi(a \mid s)\) 不存在闭式解,若直接采样估计熵项将导致极其高昂的计算代价和方差。FLAG 另辟蹊径,采用复合策略 \(\pi\) 与基流策略 \(\tilde{\pi}_\theta\) 之间的交叉熵作为代理熵: $\(\mathcal{H}(\pi(\cdot \mid s), \tilde{\pi}_\theta(\cdot \mid s)) = \mathbb{E}_{z \sim p_z} \left[ \mathcal{H}(\hat{\pi}(\cdot \mid \hat{s}), \tilde{\pi}_\theta(\cdot \mid s)) \right] = \mathbb{E}_{z \sim p_z, a \sim \hat{\pi}} \left[ -\log \tilde{\pi}_\theta(a \mid s) \right]\)$ 其中基流策略的瞬时对数密度 \(\log \tilde{\pi}_\theta(a \mid s)\) 可借助连续时间变化率方程与 Hutchinson 随机迹估计器以无偏蒙特卡洛形式高效求得。
论文进一步证明了该替代的数学有效性(Proposition 4.3):当高斯局部协方差迹满足 \(\text{tr}(\Sigma) \ll 1\) 时,复合策略与基流策略之间的总变差(TV)距离与 KL 散度分别满足 \(\mathcal{D}_{\text{TV}}(\pi, \tilde{\pi}) = \mathcal{O}(\sqrt{\text{tr}(\Sigma)})\) 以及 \(\mathcal{D}_{\text{KL}}(\pi \parallel \tilde{\pi}) = \mathcal{O}(\text{tr}(\Sigma)^2)\)。这意味着代理交叉熵与真实微分熵之间的差距为二阶小量 \(\mathcal{H}(\pi, \tilde{\pi}) = \mathcal{H}(\pi) + \mathcal{O}(\text{tr}(\Sigma)^2)\)。因此,在训练中随着局部高斯方差衰减,最大化该代理目标在理论上严格渐近逼近原始最大熵强化学习。
3. 局部重要性采样与一阶矩匹配:零阶梯度规避 BPTT
在 \(\mathbf{z}\)-MDP 下,第 \(k\) 步迭代时的能量函数定义为 \(f_{\hat{s}, k}(a) = Q^{\pi_k}(s, a) - \alpha \log \tilde{\pi}_{\theta_k}(a \mid s)\)。EM 框架的 E 步给出受限 KL 约束下的最优非参数变分目标分布 \(q_k(a \mid \hat{s}) \propto \hat{\pi}(a \mid \hat{s}; \theta_k) \exp(f_{\hat{s}, k}(a) / \lambda)\)。在 M 步中,为了将变分分布投影回参数化策略,由于局部策略 \(\hat{\pi}\) 具有各向同性高斯分布形式且方差固定,最小化 \(D_{\text{KL}}(q_k \parallel \hat{\pi})\) 精确退化为匹配其一阶矩(均值): $\(\theta_{k+1} \in \arg\min_\theta \mathbb{E}_{(s, z) \sim \hat{\rho}_{\hat{\pi}}} \| T_\theta(s, z) - \mu_k^*(\hat{s}) \|^2, \quad \mu_k^*(\hat{s}) \triangleq \mathbb{E}_{q_k}[a]\)$ 为了以极低计算开销求解 \(\mu_k^*(\hat{s})\),算法在局部高斯提议分布上仅采样 \(N\) 个局部扰动 \(\delta_i \sim \mathcal{N}(0, \Sigma_k)\),通过自归一化重要性采样(SNIS)进行加权聚合: $\(\mu_k^*(\hat{s}) \approx \sum_{i=1}^N \bar{w}_i (T_{\theta_k}(s, z) + \delta_i), \quad \bar{w}_i = \text{softmax} \left( \frac{f_{\hat{s}, k}(T_{\theta_k}(s, z) + \delta_i)}{\lambda} \right)\)$ 从理论深度上看,论文指出目标均值偏移量 \(\mu_k^*(\hat{s}) - T_{\theta_k}(s, z) = \Sigma_k \nabla_\mu \log \mathbb{E}[\exp(f / \lambda)]\) 基于 Stein 引理精确等价于能量函数 \(f_{\hat{s}, k}\) 关于动作动作梯度的一阶 Log-Sum-Exp 平滑零阶估计。它在不需要对 Critic 网络进行反向求导、不需要反向穿透流 ODE 求解器的情况下,获得了与 SAC 相同的策略改进方向,从而在根本上杜绝了数值不稳定性。
4. 条件流匹配蒸馏与单调提升保障:协方差调度与引导缓存协同
在获得改进的局部目标动作 \(\mu_k^*(\hat{s})\) 后,FLAG 并未通过强化学习价值梯度直接拉扯网络权重,而是将其视为监督引导标志(Latent Augmented Guidance),调用标准条件流匹配(CFM)回归损失对流速度场进行监督蒸馏: $\(\mathcal{L}(\theta; s, z) = \mathbb{E}_{\tau \sim \mathcal{U}[0, 1]} \left[ \| u_\theta^\tau(s, (1-\tau) z + \tau \mu^*) - (\mu^* - z) \|^2 \right]\)$ 论文证明了单调性能提升定理(Theorem 4.5):单步策略提升量下界满足 \(\mathcal{J}_{k+1} - \mathcal{J}_k \ge \lambda \beta \mathcal{G}_k - \alpha C_\Sigma \sigma_k^2 \beta \mathcal{G}_k - \lambda \epsilon_k^{\text{proj}} + \mathcal{O}(\beta^2)\)。其中第一项为冻结奖励下的 MPO 提升量,第二项为由于参数更新引发交叉熵奖励漂移的惩罚项(正比于局部方差 \(\sigma_k^2\)),第三项为 CFM 网络拟合目标动作的投影误差 \(\epsilon_k^{\text{proj}}\)。
为了确保提升量始终严格非负,FLAG 引入了两个工程核心机制: - 协方差退火调度(Covariance Scheduling):将高斯扰动标准差 \(\sigma_k\) 从初始值 \(\sigma_{\text{init}}\) 逐步退火至微小值 \(\sigma_{\text{final}}\)。在训练初期保持适当探索宽度,后期将奖励漂移项压缩至零,收敛理论保证。 - 引导缓存区(Guidance Buffer):将近期计算的改进动作标签 \(\mu_k^*(\hat{s})\) 存入一个滚动轻量级缓存中并在多次离线批次更新中复用。这大幅平滑了非平稳漂移,使得 CFM 回归的监督信号更密集,将经验投影误差 \(\epsilon_k^{\text{proj}}\) 压低在可控范围内。
实验关键数据¶
主实验¶
论文在极具挑战性的高维连续控制基准 DMC Dog(包含 Dog-stand、Dog-walk、Dog-trot、Dog-run 四个任务,动作维度高且高度欠驱动)以及高接触物理交互模拟基准 MyoSuite(肌肉骨骼控制,包含 Reach-Hard、Obj-Hold-Hard、Key-Turn-Hard、Pen-Twirl-Hard)上进行了系统评测。所有方法统一采用 CrossQ 评论家网络架构与分布式 Critic 以保证比较的公平性,评估指标报告 100 万环境交互步数(1M steps)下的四分位均值(IQM)及 95% 置信区间。
| 基准环境 | 算法类型 | 具体算法 | 任务 1 | 任务 2 | 任务 3 | 任务 4 | 平均性能 (Avg.) |
|---|---|---|---|---|---|---|---|
| DMC Dog | 高斯策略 | CrossQ (SAC基线) | 0.957 ± 0.121 | 0.882 ± 0.079 | 0.890 ± 0.016 | 0.467 ± 0.080 | 0.799 |
| 动作梯度 | DIPO | 0.906 ± 0.030 | 0.570 ± 0.164 | 0.404 ± 0.080 | 0.264 ± 0.017 | 0.536 | |
| 动作梯度 | QSM | 0.061 ± 0.131 | 0.137 ± 0.144 | 0.142 ± 0.053 | 0.064 ± 0.055 | 0.101 | |
| BPTT 架构 | DACERv2 | 0.923 ± 0.031 | 0.726 ± 0.258 | 0.367 ± 0.226 | 0.090 ± 0.108 | 0.526 | |
| BPTT 架构 | DIME | 0.968 ± 0.017 | 0.944 ± 0.024 | 0.886 ± 0.050 | 0.629 ± 0.056 | 0.857 | |
| BPTT 架构 | FlowRL | 0.964 ± 0.014 | 0.910 ± 0.011 | 0.870 ± 0.018 | 0.452 ± 0.036 | 0.799 | |
| 目标匹配 | FLAG (本文) | 0.974 ± 0.016 | 0.950 ± 0.011 | 0.912 ± 0.044 | 0.680 ± 0.062 | 0.879 | |
| MyoSuite | 高斯策略 | CrossQ (SAC基线) | 0.600 ± 0.217 | 1.000 ± 0.100 | 0.900 ± 0.150 | 0.100 ± 0.133 | 0.650 |
| 动作梯度 | DIPO | 0.600 ± 0.219 | 0.340 ± 0.201 | 0.560 ± 0.408 | 0.760 ± 0.174 | 0.565 | |
| 动作梯度 | QSM | 0.089 ± 0.069 | 0.000 ± 0.000 | 0.209 ± 0.153 | 0.026 ± 0.048 | 0.081 | |
| BPTT 架构 | DACERv2 | 0.800 ± 0.107 | 0.925 ± 0.157 | 0.911 ± 0.100 | 0.800 ± 0.129 | 0.859 | |
| BPTT 架构 | DIME | 0.900 ± 0.167 | 1.000 ± 0.017 | 1.000 ± 0.000 | 0.700 ± 0.233 | 0.900 | |
| BPTT 架构 | FlowRL | 0.000 ± 0.008 | 0.017 ± 0.042 | 0.000 ± 0.000 | 0.050 ± 0.067 | 0.017 | |
| 目标匹配 | FLAG (本文) | 0.930 ± 0.078 | 1.000 ± 0.000 | 0.930 ± 0.142 | 0.870 ± 0.174 | 0.933 |
注:DMC Dog 任务列依次为 Dog-stand、Dog-walk、Dog-trot、Dog-run;MyoSuite 任务列依次为 Reach-Hard、Obj-Hold-Hard、Key-Turn-Hard、Pen-Twirl-Hard。加粗数值为各任务第一名。
消融实验:采样预算与核心组件敏感性¶
论文首先评估了在严苛的高维控制任务 DMC Dog-trot 与 Dog-run 中,策略更新时局部重要性采样样本预算 \(N\) 对学习稳定性的影响;其次在 Dog-run 任务上深入消融了引导缓存区大小与局部高斯方差调度对最终回报的影响。
| 局部采样数 \(N\) | Dog-trot 归一化回报 (1k) | Dog-run 归一化回报 (1k) | 引导缓存大小 | 归一化回报 (1k) | 协方差调度 \(\sigma_{\text{init}}(\sigma_{\text{final}})\) | 归一化回报 (1k) |
|---|---|---|---|---|---|---|
| \(N = 2\) | 0.620 ± 0.072 | 0.818 ± 0.095 | 0 (无缓存) | 0.601 | \(10^{-1} (10^{-1})\) | 0.614 |
| \(N = 4\) | 0.639 ± 0.061 | 0.827 ± 0.106 | 10.24k (默认) | 0.680 | \(10^{-1} (10^{-3})\) | 0.732 |
| \(N = 8\) | 0.680 ± 0.066 | 0.912 ± 0.047 | 51.2k | 0.670 | \(10^{-2} (10^{-3})\) (默认) | 0.680 |
| \(N = 16\) | 0.687 ± 0.073 | 0.874 ± 0.087 | 102.4k | 0.618 | \(10^{-2} (10^{-4})\) | 0.597 |
| \(N = 32\) | 0.641 ± 0.075 | 0.879 ± 0.057 | 204.8k | 0.589 | \(10^{-2} (10^{-6})\) | 0.269 |
关键发现¶
- 对严苛样本预算的极强鲁棒性:全局提议采样基线(如 SDAC、DPMD、QVPO)在 \(N \le 8\) 时在多目标导航和高维控制中彻底失效,必须配合昂贵的 Best-of-32 候选重抽样;而 FLAG 在极其有限的 \(N=2\) 极端样本约束下,在 Dog-run 上仍能取得 0.818 的优异表现,充分证明了局部化采样在消除重要性权重坍塌方面的理论威力。
- 超越穿透式 BPTT 方法的优越性:对比依赖 BPTT 优化复杂逆向扩散轨迹的 DACERv2 与 DIME,FLAG 完全规避了多步反向传播带来的数值发散与巨大显存开销,在 DMC Dog(0.879 vs 0.857)与 MyoSuite(0.933 vs 0.900)的综合性能上均确立了清晰的 SOTA 优势。
- 协方差适度退火的决定性作用:当方差恒定偏大时(如 \(10^{-1}\)),交叉熵奖励漂移惩罚使渐近性能受阻;而若过度剧烈退火(如降至 \(10^{-6}\)),局部提议区域在训练过早阶段发生收缩坍缩,无法向更优模式迁移,导致最终回报暴跌至 0.269。平缓的对数退火与理论单调提升下界的设计完美呼应。
- 引导缓存区权衡监督密度与陈旧偏差:完全移除 Guidance Buffer 会使 CFM 投影误差放大(回报跌落至 0.601);而当缓存过大(如 204.8k)时,大量陈旧过时的历史改进动作与当前动态不匹配,同样会损伤策略表现,10k~50k 的中等缓冲容量最为理想。
亮点与洞察¶
- 从全空间随机搜索到确定性流局部引导:巧妙利用了流匹配模型先验潜变量到输出动作的确定性双射映射,将原本无法直接求解的全局测度变换规约到高斯锚点周围的局部摄动,从理论上攻克了高维强化学习中重要性采样退化的死结。
- 无梯度反向传播的“软策略改进”实现:通过 Stein 引理建立了自归一化一阶矩匹配与 SAC 动作梯度的一阶对齐,证明了通过加权扰动聚合即可获得零阶、平滑且抗噪声的动作更新方向,彻底甩掉了 BPTT 的沉重显存负担与数值不稳定枷锁。
- 完备自洽的理论收敛闭环:理论分析层层递进,从 \(\mathbf{z}\)-MDP 的测度一致性、交叉熵二阶误差控制,一直贯穿至单步策略改善的单调不等式,将算法层面的协方差退火和目标缓存直接与理论界中的误差项绑定。
局限与展望¶
- 局部高斯头与流模型的启发式组合:当前架构是在确定性流模型的输出末端显式拼接一个各向同性高斯头,虽然结构简单直观,但并非最内生、最优雅的形式。正如作者在结论中指出的,未来探索通过常微分方程到随机微分方程(ODE-to-SDE)的自然映射或伴随匹配(Adjoint Matching)来原生地赋予流模型随机性,将是更完备的理论方向。
- 超高维复杂多峰环境下的局部极小陷阱:局部重要性采样高度依赖锚点动作在全局空间的初始分布覆盖。在动作空间极大且奖励极度稀疏的极端任务中,若初始流模型生成的锚点均落在次优局部凹坑内,过小的局部方差可能会加剧对局部极值的过拟合,需要结合自适应方差或全局混合提议机制进行探索补偿。
相关工作与启发¶
- vs DIME / DACERv2 (Celik et al., 2025; Wang et al., 2025):DIME 和 DACER 均采用穿透生成轨迹的 BPTT 计算策略梯度。FLAG 采用局部 EM 框架将优化转化为纯前向监督蒸馏,完全免除了微分方程反向积分,在训练速度、显存占用和抗 Critic 噪声扰动上具备压倒性优势。
- vs RSM / MaxEntDP / QVPO (Ma et al., 2025; Dong et al., 2025; Ding et al., 2024):此类基于目标匹配的方法依赖全局重要性采样,在高维动作空间中面临严峻的支撑集消失与权重爆炸难题。FLAG 通过构建 \(\mathbf{z}\)-MDP 将采样严格限制在以流锚点为中心的微小局部邻域,使得在仅需 \(N=2 \sim 8\) 个样本的极端预算下即可稳定优化。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立基于潜变量增强 MDP 的流策略局部重要性采样与监督指导范式。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 DMC Dog、MyoSuite 和 MuJoCo 等高维控制全集,消融深度直击理论误差项。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导完备自洽,从问题发现、理论证明到算法闭环浑然一体。
- 价值: ⭐⭐⭐⭐⭐ 为解决生成式策略在强化学习中的计算与扩展性瓶颈提供了极具启发性的通用新方案。