Awesome World Models / OpenReview

Offline Transition Modeling via Contrastive Energy Learning

作者：作者信息待提取 · 来源：General Approaches to World Models · 发布日期：日期待提取 · 分类：8. World Models in the context of Model-Based RL

三维视觉先验基于模型的强化学习

精读笔记

Problem Setting

这篇论文处理的是 offline setting 下 transition model 的可靠外推问题，特别是模型被用于 OPE 或 offline MBRL rollout 时，目标策略诱导的 state-action 分布会偏离行为数据。困难不只是 one-step supervised learning，而是 transition error 在 rollout 中自激放大，并被 policy optimization 主动 exploit。

以前方法的主要瓶颈在于：标准 FTM 通常把 s,a 映射到 next-state Gaussian mean/variance，本质上是平滑回归器。对于 MuJoCo/real robotics 中常见的接触、摩擦、碰撞、阈值效应，真实 transition 可能有 discontinuity 或 large curvature。FTM 为了拟合这种局部尖锐结构，需要在函数输出上形成很大梯度；在离线数据稀疏处，这种梯度既难学，也容易产生错误插值。关键矛盾是：offline MBRL 最需要模型在低密度区域稳健，但 forward regression 的归纳偏置恰好会在这些区域做最危险的平滑泛化。

Motivation

作者不是简单认为“模型容量不够”，而是认为已有 transition modeling 的学习范式本身有偏。换 Transformer、autoregressive model 或更大 ensemble，主要是在扩大表达能力，但仍然需要模型显式输出 next state；这没有解决 smooth regressor 与非光滑 dynamics 的冲突。

核心观察来自 implicit behavior cloning / energy-based regression：如果不直接预测目标，而是学习输入与候选输出之间的 compatibility score，再通过 argmin/search 得到输出，那么最终输入到输出的映射可以是非连续的，即使 score function 相对平滑。这正好对应 transition modeling 中的跳变/接触问题。论文要补的缺口是：offline transition model 不应只看作条件均值回归，而应看作条件结构化预测或条件分布能量建模。

Core Idea

ETM 的核心思想是把 transition probability 隐式写成 pθ(s'|s,a) ∝ exp(-Eθ(s,a,s'))。模型学习的不是下一个状态本身，而是一个条件能量地形：真实 next state 应该落在低能量区域，错误 candidates 应该被推高。推理时通过 Langevin MCMC 或直接能量最小化得到 next state。

本质区别在于，FTM 学的是 fθ(s,a)≈s'，输出随输入连续变化；ETM 学的是 Eθ(s,a,s')，最终预测是 argmin_{s'} Eθ(s,a,s')。argmin 操作可以把平滑能量面转换成非平滑决策边界，因此更适合表示 contact-like 或 threshold-like transition。这个 inductive bias 更接近 discriminative structured prediction：不是要求模型在输出空间平滑插值，而是要求它在候选输出之间排序正确。

这也重新组织了信息流：FTM 的每个样本只提供一个回归目标；ETM 的每个样本通过负样本对能量面多个位置施加约束。因此训练信号从“点预测误差”变成“局部能量排序”，理论上更能约束 rollout 中会遇到的邻域。

Method

1. Energy-based transition parameterization：用 Eθ(s,a,s') 代替 Gaussian mean/variance。它解决的是 FTM 输出分布受限与平滑回归偏置问题。核心变化是 transition 被表示为条件能量面，而非显式函数图像。

2. Contrastive InfoNCE training：给定真实 transition，把真实 s' 与若干 negative s'_j 做分类式区分。它解决的是 EBM likelihood 中 partition function 不可 tractable 的问题，同时让训练显式比较正负候选。作者进一步把 InfoNCE 解释成对 MLE/CD 的 robust surrogate：当负样本分布与真实模型分布存在 KL 误差时，InfoNCE 类似在优化一个更保守的上界。这个推导是论文中比较有价值的理论连接，但仍依赖负样本足够好和数量足够。

3. Langevin negative sampling / inference：负样本不是纯随机，而是通过当前能量模型的 Langevin dynamics 生成，推理也通过迭代搜索低能量点。它解决的是 energy distribution 无法直接采样的问题，同时引入 test-time compute。这里的核心变化是模型预测不再是一次前向传播，而是“模型 + 优化过程”的组合。

4. Gradient penalty：用于限制能量面对 s' 的梯度，避免训练点附近能量过尖导致 Langevin 采样退化。这个机制更像 EBM 训练稳定性的必要工程组件，不是主要 conceptual novelty。

Key Insight / Why It Works

这篇最重要的 insight 是：transition modeling 的失败并不总是 epistemic uncertainty 或 capacity 问题，很多时候是输出参数化和真实 dynamics 几何不匹配。FTM 在非光滑动力学上会被迫做平滑插值；这种错误在 offline rollout 中尤其致命，因为模型需要预测的恰是训练分布边缘甚至外部的状态。

ETM 可能有效的真正原因有三层。第一是 better inductive bias：energy + argmin 允许非连续 input-output map，避免了 FTM 必须显式拟合 steep gradient。这个应该是核心贡献。第二是 contrastive shaping：负样本让模型不仅降低真实点能量，还提高附近错误点能量，因此学习到的不是单点均值，而是局部可区分的 transition manifold。第三是 test-time compute：Langevin/search 在推理时实际进行了局部优化，这部分能力不能算纯模型泛化；它把一部分预测难度推迟到 inference-time optimization。

需要直接指出：offline RL 分数提升不一定完全来自 energy modeling。EMPO 使用 ensemble、uncertainty penalty、SAC、rollout length 和任务级 penalty tuning，增益来源混合。更干净的证据是 OPE 中 ETM vs FTM 的 direct rollout 对比，而不是最终 policy score。论文声称的 OOD generalization 也要谨慎理解：D4RL 不同 level 数据集共享同一 simulator dynamics，OOD 更多是 policy/state distribution shift，不是跨物理系统泛化。

这不是 retrieval，也不是 memory reuse；也不是 scaling 论文。它的核心更接近“更合适的 structured prediction inductive bias + contrastive energy learning + test-time optimization”。如果未来用更强 sampler 或 amortized sampler 替代 Langevin，可能保留主要收益并降低成本。

Relation To Prior Work

最接近的技术谱系不是传统 world model，而是 energy-based conditional regression / implicit behavior cloning。Florence et al. 的 implicit BC 已经展示了 EBM 对非光滑动作映射和多模态输出的优势；这篇把同一思想迁移到 transition dynamics，并将其嵌入 offline OPE/RL pipeline。

相对 MOPO/COMBO/MOBILE 等 offline MBRL，ETM 的不同点不在 conservatism，而在 model class。MOPO 系列默认 forward probabilistic dynamics，然后处理模型误差；ETM 试图从源头改变误差结构。相对 Trajectory Transformer / autoregressive dynamics，它不是用更大序列模型吸收复杂性，而是改变输出预测方式。相对 diffusion planner / conditional generative model，ETM 更原始、更轻量，但采样/搜索同样是核心；文中和 diffusion transition/world model 的边界讨论不够。

看似新的部分如 InfoNCE 训练 EBM、Langevin sampling、gradient penalty都不是新技术；实质创新在于把它们组织成 offline transition model，并给出 InfoNCE-as-robust-MLE 的解释，再用 OPE 这个对 model fidelity 敏感的任务验证。真正新增的信息是：在 state-based offline control 中，energy transition model 的 OOD rollout error 可能系统性优于 FTM。

Dataset / Evaluation

实验覆盖了 didactic discontinuity、D4RL MuJoCo/Adroit OPE、D4RL offline RL，以及补充 NeoRL。任务范围主要是低维 state-based continuous control，没有视觉输入、没有真实机器人、没有跨 simulator 或跨物理参数泛化。因此它验证的是“在常见 offline RL benchmark 的 state dynamics 上，ETM 比 FTM 更适合 rollout”，不是广义 real-world world model。

OPE 是这篇最有说服力的 evaluation，因为 direct model rollout 的 value error 对 transition model 质量非常敏感，且与 model-free OPE baseline 对比能体现实际效用。didactic example 也很好地支撑了 discontinuity claim，但它是人工构造。offline RL 结果展示了方法可接入现有 MBRL pipeline，但不完全隔离变量；policy optimization 的提升可能受 conservative penalty、ensemble variance、rollout schedule 和调参影响。

benchmark 是否真正验证核心 claim？部分验证。它较好验证了“ETM 在相同 offline dataset 上比 FTM 有更好的 rollout/OPE 表现”；但没有充分验证“真实世界接触动力学中的 discontinuity 是主要收益来源”，因为 MuJoCo tasks 中哪些状态区域对应非光滑接触、ETM 是否在那里特别好，文中未做细粒度诊断。

Limitation

最根本的限制是 ETM 把 forward prediction 的归纳偏置问题转移成了 energy landscape optimization 问题。要成立，需要：负样本覆盖到有判别价值的错误 next states；Langevin 能在有限步内找到合理低能量区域；energy function 不在训练点附近形成不可采样的尖峰；offline 数据至少覆盖真实 transition manifold 的关键局部结构。这些前提在低维 MuJoCo 中相对可满足，在高维视觉状态或真实机器人中未必成立。

scalability 上限明显。对每个 transition 都要迭代 Langevin/search，长 horizon rollout 会放大计算成本。论文报告约 5 倍 inference time 和 40% learning time 增加，但这还只是低维 state tasks；如果 s' 是图像、点云或 latent token，直接 energy over output 会困难得多。除非引入 amortized sampler、latent energy 或 diffusion-like proposal，否则很难成为通用 world model backbone。

泛化也不能过度解读。D4RL level 间的 OOD 仍共享动力学和状态空间，很多所谓外推可能是 benchmark 内分布重组。核心能力可能主要来自更合理的局部排序和 test-time search，而不是学到了可迁移物理规律。对于严重超出数据覆盖的状态-action，ETM 没有理论保证；低能量区域甚至可能产生更隐蔽的 hallucination。

另一个未充分说明的问题是多模态 transition。EBM 理论上适合多模态，但 D4RL/MuJoCo dynamics 多数近似 deterministic；因此实验并没有强验证 flexible distribution prediction 的优势。论文中 distributional flexibility 的 claim 比 discontinuity/inductive-bias claim 更弱。

Takeaway

  • 1. 对 offline MBRL，transition model 的输出参数化可能比架构规模更关键；FTM 的平滑回归偏置是一个真实且被低估的问题。
  • 2. Energy-based conditional prediction 的价值在于把非光滑映射转化为能量排序 + test-time optimization，这个 insight 可迁移到 action prediction、contact dynamics、hybrid systems、constraint satisfaction 型 world models。
  • 3. OPE 是检验 dynamics model 的好场景，因为它比最终 offline RL score 更少受 policy optimizer 和 conservatism 影响；未来 world model paper 应更多使用这种隔离式评估。
  • 4. 下一步真正值得做的不是继续堆 benchmark 分数，而是解决 amortized inference、高维 observation、multi-modal stochastic dynamics，以及对能量面 OOD failure 的可检测性。

一句话总结

这篇论文把 offline transition modeling 从显式平滑 forward regression 推向 contrastive energy-based structured prediction，真正贡献是用更合适的归纳偏置和 test-time energy search 缓解非光滑/OOD rollout 误差。