TL;DR

扩散规划器（Diffusion Planners）在离线强化学习中表现强劲，但常常会选中那些“看起来高分但物理上无法执行”的幻觉轨迹。本文提出的 SAGE 是一种即插即用的推理端算法，它通过自监督学习的潜空间预测一致性作为“能量”信号，将不可行的动作序列拒之门外。在保持生成模型和评价模型不变的前提下，SAGE 在 Locomotion, Navigation 和 Manipulation 任务上全面刷新了性能基准。

背景定位：价值导向的代价

在离线 RL 领域，将控制问题转化为轨迹生成问题已成为主流。扩散模型因其强大的多模态分布建模能力，能够生成多样化的候选路径。然而，现有的规划逻辑存在一个致命弱点：价值函数（Critic）往往只关心终点或长期累积回报，却忽视了第一步动作是否能基于当前状态物理实现。

当一个计划的初始段落（Prefix）与真实环境动力学不匹配时，Agent 就会 commit 到一个错误的动作上，导致执行过程极其脆弱。

核心直觉：可行性与价值的解耦

作者认为，可行性（Feasibility） 应该是一个独立于 价值（Value） 的信号。SAGE 的核心思想是建立一个“安检门”：它不参与扩散模型的复杂训练，而是在推理阶段，对生成的每个候选方案进行物理一致性审查。

方法论：两阶段自监督学习

SAGE 的技术架构主要分为两个离线训练阶段：

1. 特征表示学习 (JEPA Stage)：学习一个状态编码器。通过屏蔽状态（Masking）并预测未来的潜空间嵌入，确保 Latent Space 能够捕捉到环境中关键的动力学特征，避免直接重建像素或原始状态带来的噪声。
2. 动作条件预测 (AC Predictor Stage)：在冻结的 Latent Space 中，训练一个 Transformer 预测器。它根据当前潜状态和动作，预测下一时刻的潜状态。如果扩散模型生成的轨迹在这一步预测误差很大，说明该轨迹在离线数据分布中从未出现，物理上极大概率是不可行的。

上图展示了 SAGE 的流程：(a) 训练 JEPA 编码器；(b) 训练动作条件预测器；(c) 推理时通过能量函数进行重排序。

推理端重排序 (Inference Gating)

在测试阶段，扩散模型生成 $C$ 个候选轨迹。SAGE 对每个轨迹的前 $K$ 个步骤计算潜空间一致性能量 $E$。

选择策略非常直观：

• 硬过滤：保留能量最低（最可靠）的前 $P$ 部分候选。
• 软惩罚：在剩余候选里，结合价值评分 $J$ 和能量惩罚项 $\lambda E$ 选择最优动作。

实验分析：抑制“幻觉”轨迹

在 Maze2D 的对比实验中，传统的 MCSS 采样有时会生成“穿墙”轨迹，因为价值函数可能误判穿墙后的状态为高价值。而 SAGE 通过能量评分，精准地识别并抑制了这些穿墙样本。

左侧为标准 MCSS，可见大量轨迹穿过墙壁；右侧为 SAGE，轨迹严格遵守地图约束且保持了多样性。

关键战绩

• Locomotion (MuJoCo)：在 9 个子任务中，SAGE 相比强基线 DV* 实现了 1.5% 的稳健提升。
• Navigation (AntMaze)：在稀疏奖励环境下表现亮眼，均分从 81.6 提升至 84.5。
• 计算开销：由于只需对极短的 Prefix 进行 Latent 级别前向计算，推理耗时仅增加了 6.8%，远低于引入额外扩散引导方法的开销。

深度洞察

SAGE 的成功揭示了离线 RL 中的一个通用事实：生成模型的表达能力（Expression）往往超前于物理一致性约束（Consistency）。

通过消融实验（Ablation Study）可以发现：

• Prefix 长度 $K$ 存在甜点位。过短（如 K=1）可能无法捕捉时序特征，过长则会因预测误差累计而误伤长线计划。
• 自监督的目标：相比于在原始状态空间做回归（Ridge Regression），在 JEPA Latent Space 中进行比对具有更高的辨别力（AUROC 0.98 vs 0.88）。

总结与展望

SAGE 为扩散规划器的落地提供了一个低门槛、高性能的增强方案。它不需要重新训练笨重的生成模型，仅需一个轻量级的预测器即可补齐可行性的短板。未来的研究可以探索如何将这种能量信号动态地融合进扩散模型的反向去噪过程，以实现更高效的引导生成。

关键词：Diffusion Planner, JEPA, Action Gating, Offline RL, Feasibility Energy.