PPO in the Fisher-Rao geometry¶
会议: NeurIPS 2026
arXiv: 2506.03757
代码: https://github.com/deterministicdavid/sb3-contrib-with-frppo
领域: 强化学习 / Reinforcement Learning
关键词: 近端策略优化, Fisher-Rao几何, 镜像下降, 单调策略提升, 理论收敛界
一句话总结¶
针对传统 PPO 启发式截断缺乏单调提升和收敛保证的理论鸿沟,本文在 Fisher-Rao 几何下推导出基于平方密度差的紧致代用目标下界,提出了具有单调改进和无维度依赖次线性收敛保障的 FR-PPO 算法,且可无缝作为 PPO 损失函数的即插即用替代。
研究背景与动机¶
策略梯度(Policy Gradient, PG)方法在强化学习中占据着基石地位,但在更新策略时选择合适的步长极其敏感,微小的步长偏差都可能导致性能雪崩。为了约束更新幅度,信任域策略优化(TRPO)在策略间引入平均 KL 散度约束,从理论上确保了单调性能提升,然而其依赖共轭梯度与二阶回溯线性搜索,在面对高维连续控制和大规模网络时计算开销高昂。近端策略优化(PPO)通过概率比率截断(clipping)机制将硬约束转换为无约束的一阶优化,凭借简洁高效的实现成为了工业界与学界的默认基线。
然而,PPO 截断机制的理论基础长期模糊不清。PPO 截断代用目标的直观动机是惩罚新旧策略间的总变差(Total Variation, TV)距离,但由于截断算子与 TV 距离本质上是非光滑的,且 TV 距离不构成 Bregman 散度,无法直接利用一阶最优性条件或镜像下降的三点不等式建立收敛性理论。既有研究甚至表明,PPO 在实际训练中频繁突破预设的比率边界;理论分析要么只能在特殊神经网络与参数化下给出弱收敛界,要么完全丧失单调提升保证。
本文的切入视角是重新审视代用目标底层的几何框架:TV 惩罚的非光滑性可以通过几何提升(lifting)在 Fisher-Rao 几何中化解。核心 idea:将策略概率测度提升到平方密度空间,利用 Fisher-Rao 距离构造自然诱导 Bregman 散度的紧致二次 TV 下界,提出具有单调策略提升与 \(\mathcal{O}(1/n)\) 收敛保障的 FR-PPO 算法,并在实践中化为可直接替代 PPO 截断损失的光滑二次惩罚项。
方法详解¶
整体框架¶
FR-PPO 的理论推导与算法流程从 MDP 性能差异引理(Performance Difference Lemma)出发。首先,借助新旧策略占用测度之间的相对集中性假设,在旧策略占用测度下建立包含积分 \(\mathrm{TV}^2\) 项的严格下界;接着利用 Cauchy-Schwarz 不等式将状态级 \(\mathrm{TV}^2\) 控制在平方密度的 Fisher-Rao 距离(即 \(\chi^2\) 散度诱导的 Bregman 散度)之下,从而将非光滑的策略比率截断转化为 Fisher-Rao 几何下的策略镜像下降。
整个算法保持了与标准 PPO 完全一致的数据采样与架构设计,仅需在策略更新阶段用光滑的 Fisher-Rao 惩罚损失替换原始的非光滑截断损失。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["环境交互与轨迹收集<br/>基于当前策略采样批次状态动作对"] --> B["广义优势估计 GAE<br/>计算标量优势估计值"]
B --> C["Fisher-Rao 代用下界构建<br/>将 TV 距离提升至平方密度 FR 几何"]
C --> D["投影自然策略梯度与镜像更新<br/>基于 Bregman 三点不等式优化参数"]
D --> E["无缝替换策略损失函数<br/>执行一阶 Adam 优化更新网络参数"]
关键设计¶
1. Fisher-Rao 几何代用下界构建:将状态极值松弛为占用测度期望
TRPO 的经典性能下界依赖于全状态空间上 TV 或 KL 散度的最大值 \(\max_{s\in S} \mathrm{TV}^2(\pi'(\cdot|s), \pi(\cdot|s))\),这在连续状态空间下无法通过采样精确优化。本文首先引入相对集中性假设,证明了辅助状态分布与基准占用测度之间的 Radon-Nikodym 导数具有上界 \(C_\rho\)。在此基础上推导出以期望形式给出的紧致下界: $\((V_{\pi'} - V_\pi)(\rho) \ge \left\langle \frac{\mathrm{d}\pi'}{\mathrm{d}\pi} A_\pi, \pi \right\rangle_{d_\rho^\pi} - \frac{4\gamma \|r\|_\infty}{(1-\gamma)^3} (1+C_\rho) \int_S \mathrm{TV}^2(\pi'(\cdot|s), \pi(\cdot|s)) d_\rho^\pi(\mathrm{d}s)\)$ 为克服 \(\mathrm{TV}^2\) 非光滑且非 Bregman 散度的缺陷,作者将概率密度提升至其平方根密度所在的 \(L^2\) 球面上,并利用不等式 \(\mathrm{TV}^2(P, Q) \le \frac{1}{16} \mathrm{FR}^2(P^2, Q^2)\),将价值函数差异严格绑定至平方密度间的 Fisher-Rao 距离,该距离与 \(\chi^2\) 散度诱导的 Bregman 散度 \(D_h\) 严格等价。这一转换使得价值函数呈现出相对光滑性(Relative Smoothness),为凸优化理论中的三点引理铺平了道路。
2. 投影自然策略梯度与镜像更新:绕过参数光滑性假设建立收敛界
在直接参数化设定下,FR-PPO 的迭代格式定义为对优势线性化项与 Fisher-Rao 正则项的最优权衡: $\(\pi^{n+1} = \operatorname{argmax}_{\pi \in \mathcal{P}_\lambda(A|S)} \left[ \left\langle \frac{\mathrm{d}\pi}{\mathrm{d}\pi^n} A_{\pi^n}, \pi^n \right\rangle_{d_\rho^{\pi^n}} - \frac{1}{8\tau} \int_S \operatorname{FR}^2(\pi(\cdot|s)^2, \pi^n(\cdot|s)^2) d_\rho^{\pi^n}(\mathrm{d}s) \right]\)$ 通过验证上述优化问题在点态上的强凹性与凸集投影性质,作者首先证明了任意正惩罚参数 \(\tau > 0\) 下的严格单调提升 \(V_{n+1}(\rho) \ge V_n(\rho)\),彻底解决了 PPO 启发式截断无法保证单调上升的理论痛点。更进一步,作者引入希尔伯特特征空间下的线性策略参数化,提出投影自然策略梯度(Proj-NPG)及其实际逼近形式 AProj-NPG。不同于现有方法对参数级策略光滑性强依赖所导致的步长误差累积,本文证明了 Proj-NPG 对应着策略空间在 Bregman 散度度量下的精确变分字符,借助 Bregman 三点不等式建立了无维度依赖的 \(\mathcal{O}(1/n)\) 次线性全局收敛速率,将既有工作在兼容函数逼近下的 \(\mathcal{O}(1/\sqrt{n})\) 上界大幅提升。
3. 即插即用策略损失函数设计:单步经验期望的平滑替换
为了使精妙的几何理论能够直接落地到现代深度 RL 算法管线中,作者推导了理论更新在参数化策略下的样本无偏估计。在任意给定的旧策略 \(\pi_{\theta_{\text{old}}}\) 与优势估计 \(\hat{A}\) 下,传统 PPO 的截断项被平滑替代,对应的经验目标函数 \(L_{\text{FR-PPO}}(\theta)\) 写为: $\(L_{\text{FR-PPO}}(\theta) = \mathbb{E}_{s \sim d_\rho^{\pi_{\text{old}}}, a \sim \pi_{\text{old}}} \left[ \frac{\pi_\theta(a|s)}{\pi_{\theta_{\text{old}}}(a|s)} \hat{A}(s,a) - \frac{1}{2\tau} \left| \frac{\pi_\theta(a|s)}{\pi_{\theta_{\text{old}}}(a|s)} - 1 \right|^2 \frac{\mathrm{d}\pi_{\theta_{\text{old}}}}{\mathrm{d}\lambda}(a|s) \right]\)$ 该目标函数仅用策略概率比率的二次惩罚项平滑替代了 PPO 带有折点的 \(\operatorname{clip}_{1-\varepsilon}^{1+\varepsilon}\),计算复杂度完全相同,可以直接复用 Stable-Baselines3 等现成框架中的 Rollout 采样、GAE 优势计算与网络反向传播结构,无需任何高阶近似或回溯线搜索。
损失函数 / 训练策略¶
在具体的神经网络实现中,演员网络 \(\pi_\theta\) 与评论家网络 \(V_\varphi\) 联合优化。总体损失函数定义为标量形式: $\(L(\theta, \varphi) = - L^{\text{pol}}(\theta) + L^{\text{val}}(\varphi)\)$ 其中策略损失 \(L^{\text{pol}}(\theta)\) 在每个大小为 \(|B|\) 的小批量样本上计算: $\(L^{\text{pol}}(\theta) = \frac{1}{|B|} \sum_{(s_i, a_i) \in B} \left[ r_i(\theta) \hat{A}_i - \frac{1}{2\tau} |r_i(\theta) - 1|^2 \pi_{\theta_{\text{old}}}(a_i|s_i) \right]\)$ 其中比率 \(r_i(\theta) = \frac{\pi_\theta(a_i|s_i)}{\pi_{\theta_{\text{old}}}(a_i|s_i)}\),价值网络则通过时序差分回归目标进行更新:\(L^{\text{val}}(\varphi) = \frac{1}{|B|} \sum_{(s_i, R_i) \in B} (V_\varphi(s_i) - R_i)^2\)。训练采用标准 Adam 优化器,超参数设置沿用标准 PPO 的基线配置,重点通过调节罚常数 \(\tau \in \{0.05, 0.1\}\) 来控制信任域约束强弱。
实验关键数据¶
主实验¶
论文在 MuJoCo 连续控制基准(5 个环境)与 Atari 离散控制基准(6 个环境)上对比了 PPO 与 FR-PPO。下表展示了综合归一化回报与两算法在胜出环境数量(包含最大值胜出与总和胜出)上的统计对比。
| 评估领域 | 归一化回报 (PPO) | 归一化回报 (FR-PPO) | PPO 获胜环境数 (Max, Sum) | FR-PPO 获胜环境数 (Max, Sum) | 结论分析 |
|---|---|---|---|---|---|
| MuJoCo (5 环境) | 5.4 | 4.4 | (3, 2) | (2, 3) | 连续控制领域表现具备高竞争力,Sum 维度更优 |
| Atari (6 环境) | 5.8 | 5.6 | (3, 3) | (3, 3) | 离散控制领域总体回报相当,各胜出 3 个环境 |
消融实验¶
论文详细测试了不同超参数配置下的表现对比(PPO 考察截断系数 \(\varepsilon \in \{0.1, 0.2\}\),FR-PPO 考察惩罚系数 \(\tau \in \{0.1, 0.05\}\))。下表展示了 Atari 与 MuJoCo 代表性环境的平均归一化得分数据(基于 12 个随机种子独立训练、每个策略评估 40 次并进行极值归一化)。
| 测试环境 (Atari / MuJoCo) | PPO (\(\varepsilon=0.1\)) | PPO (\(\varepsilon=0.2\)) | PPO (Sum) | FR-PPO (\(\tau=0.1\)) | FR-PPO (\(\tau=0.05\)) | FR-PPO (Sum) | 表现分析 |
|---|---|---|---|---|---|---|---|
| BeamRider-v4 (Atari) | 0.65 | 0.24 | 0.89 | 0.44 | 0.52 | 0.96 | FR-PPO 在总和及更强约束下更稳健 |
| Breakout-v4 (Atari) | 0.35 | 0.28 | 0.63 | 0.39 | 0.21 | 0.60 | FR-PPO 在 \(\tau=0.1\) 取得单项最高 0.39 |
| KungFuMaster-v4 (Atari) | 0.47 | 0.48 | 0.95 | 0.46 | 0.31 | 0.77 | PPO 保持微弱优势 |
| MsPacman-v4 (Atari) | 0.44 | 0.64 | 1.08 | 0.35 | 0.35 | 0.70 | PPO 在较大截断下得分更高 |
| Pong-v4 (Atari) | 0.89 | 0.75 | 1.64 | 0.85 | 0.90 | 1.75 | FR-PPO 总分 1.75 领先 |
| SpaceInvaders-v4 (Atari) | 0.28 | 0.35 | 0.63 | 0.28 | 0.54 | 0.82 | FR-PPO 在 \(\tau=0.05\) 取得 0.54 大幅提升 |
| HalfCheetah-v5 (MuJoCo) | 0.48 | 0.49 | 0.97 | 0.63 | 0.67 | 1.30 | FR-PPO 在两种 \(\tau\) 下均显著超越 PPO |
| Hopper-v5 (MuJoCo) | 0.73 | 0.62 | 1.35 | 0.43 | 0.40 | 0.83 | PPO 表现更佳 |
| Reacher-v5 (MuJoCo) | 0.54 | 0.72 | 1.26 | 0.68 | 0.70 | 1.38 | FR-PPO 总和 1.38 领先 |
| Swimmer-v5 (MuJoCo) | 0.30 | 0.27 | 0.57 | 0.15 | 0.51 | 0.66 | FR-PPO 在 \(\tau=0.05\) 逆势大幅提升至 0.51 |
| Walker2d-v5 (MuJoCo) | 0.71 | 0.55 | 1.26 | 0.10 | 0.11 | 0.21 | FR-PPO 出现明显未对齐掉点 |
关键发现¶
- 在 Atari 基准中,FR-PPO 与 PPO 各自赢得 3 个环境的最大值和总和表现,且在 Pong 和 SpaceInvaders 上展现出更稳定的策略提升。
- 在 MuJoCo 基准中,FR-PPO 在 HalfCheetah、Reacher、Swimmer 上表现出明显超越 PPO 的收敛速度与稳定性,证明了光滑比率惩罚在复杂动力学环境中的有效性。
- 超参数敏感性分析表明,惩罚强度参数 \(\tau\) 对 FR-PPO 的影响比截断参数 \(\varepsilon\) 对 PPO 的影响更为敏感(例如在 Swimmer 中,\(\tau=0.05\) 的表现相比 \(\tau=0.1\) 从 0.15 跃升至 0.51),这意味着实际工程落地中需要更精细地调优 \(\tau\) 或引入自适应机制。
亮点与洞察¶
- 首次确立了不依赖 KL 散度的 PPO 风格算法严格理论闭环:通过将 TV 范数巧妙升阶到平方密度空间上的 Fisher-Rao 几何,化解了非光滑截断导致梯度理论失效的顽疾。
- 建立了首个无状态/动作维度依赖的 \(\mathcal{O}(1/n)\) 次线性收敛界:借助 Bregman 三点不等式与变分等距性质,摆脱了参数空间曲率估计的误差累积。
- 理论成果到工业落地的转化路径极其优雅:复杂的测度空间投影推导最终落地为一个极其简单的二次概率比率惩罚项,能够 100% 复用已有 PPO 代码库。
局限与展望¶
- 超参数 \(\tau\) 的选取对不同环境动态敏感,目前缺乏像 TRPO 中自适应阈值那样动态调整 \(\tau\) 的理论驱动算法。
- 经验公式中的参考测度导数项 \(\frac{\mathrm{d}\pi_{\theta_{\text{old}}}}{\mathrm{d}\lambda}(a|s)\) 在高维连续动作空间且非均匀参考测度时可能会增加采样方差,当前实验主要基于标准分布做了近似简化。
- 未来工作可探索将 Fisher-Rao 几何拓展至大语言模型对齐(如 RLHF / PPO-RLVR)场景,研究在长文本生成中消除概率比率异常剧烈波动的潜力。
相关工作与启发¶
- vs TRPO: TRPO 依赖于全状态空间上严格的 KL 散度约束,需要二阶海森矩阵向量积与回溯线性搜索,计算复杂度高;FR-PPO 采用 Fisher-Rao 几何下的等价 Bregman 正则化,转化为纯一阶优化算法,理论上有相同的单调提升保证而计算代价减半。
- vs PPO-clip: PPO-clip 采用硬截断启发式约束,虽然计算极其简洁,但缺乏单调提升证明且实际比率频繁越界;FR-PPO 将经验截断替换为平滑的二次差惩罚,保留了 PPO 简易结构的同时提供了首个收敛界保证。
- vs Neural-PPO-clip: 既有工作在单隐层 ReLU 网络假设下仅能获得 \(\mathcal{O}(1/\sqrt{n})\) 的极小迭代收敛速率;FR-PPO 在兼容线性参数化下通过变分特征表示将速率提升到了 \(\mathcal{O}(1/n)\)。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用 Fisher-Rao 几何与平方密度 Bregman 散度刻画 PPO 截断机制,视角极其深刻。
- 实验充分度: ⭐⭐⭐⭐ 覆盖了标准离散 Atari 与连续 MuJoCo 经典基准,但缺乏超大规模复杂场景(如机器人灵巧手或大模型对齐)的实证验证。
- 写作质量: ⭐⭐⭐⭐⭐ 数学证明严密自洽,从测度论基础到具体算法损失函数的映射逻辑极其清晰顺畅。
- 价值: ⭐⭐⭐⭐⭐ 为强化学习中最流行的基线算法之一补全了缺失已久的坚实理论基石。