Awesome World Models / OpenReview

Temporal Difference Flows

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

基于模型的强化学习

精读笔记

Problem Setting

Temporal Difference Flows 处理的是 model-based RL 语境下的长 horizon future-state generative prediction：给定离线 transition 数据和目标 policy，学习 mπ(·|s,a)，即 discounted future state occupancy 的生成模型。它不是一步 dynamics model，也不是普通 offline RL policy learning；它关心的是能否直接生成远未来状态分布，并用它做 value estimation / GPI。

真正难点是 train-time bootstrapping。GHM 避免了传统 world model 在 inference 时逐步 rollout 的 compounding error，但它必须通过 Bellman fixed point 学 successor measure。已有 TD-GHM 在每轮用当前模型生成 pseudo target，然后拟合这个 target；当 γ 接近 1 时，bootstrap 项权重接近 1，训练几乎完全由模型自身采样驱动，采样方差和路径不匹配会主导学习。以前方法卡住的不是表达能力本身，而是“用生成模型拟合另一个由自己生成的长程分布”这个训练过程不稳定。

这个任务的关键矛盾是：直接预测长未来可以省掉 rollout error，但长未来分布又无法直接从 off-policy 数据低方差采样；TD 能提供 off-policy fixed point，但 TD target 本身是自举生成分布，越长 horizon 越噪。

Motivation

已有路线不够的地方很明确。Monte Carlo 采样 successor measure 需要 on-policy 长轨迹或 importance sampling，离线 off-policy 下方差不可接受；Bellman TD-GHM 可离线训练，但 naive bootstrapping 会在长 horizon 下退化；GAN/VAE/normalizing-flow 只是换了生成模型形式，没有改变 TD target 的噪声结构。

作者的关键观察是：flow matching / diffusion 的训练对象不是只有 endpoint distribution，而是整条从 noise 到 data 的 probability path，且 path 上的 velocity / score 是可以作为监督信号的。Bellman operator 又是分布的线性 mixture，因此可以把 Bellman backup 作用到 path 上，而不是只作用到 t=1 的 endpoint。缺口就是：现有 GHM 只在“分布终点”做 TD，没有利用生成过程内部结构；现有 flow/diffusion 又假设能直接采样真实 target distribution，不能直接处理 successor measure 的 TD fixed point。

Core Idea

核心思想是把 TD backup 从 successor-measure endpoint 推广到生成模型的 probability path。对 Bellman target 的两个组成部分——one-step transition 和 bootstrapped future distribution——分别构造 path，然后在每个 t 上学习 mixture vector field。这样，GHM 的训练目标不再只是“拟合一个自举出来的未来状态样本”，而是“拟合 Bellman mixture 对应的生成动力学”。

最重要的版本是 TD2：bootstrap 部分不再先采样旧模型的 endpoint X1 再构造一条新的 conditional path，而是直接沿旧模型 ODE 到当前 t，使用旧 vector field 作为监督。这相当于对生成过程做 TD-style self-distillation。它引入的 inductive bias 是 path consistency：successor Bellman fixed point 不仅约束最终 occupancy，也约束从 base noise 到 occupancy 的传输过程。相比 prior，它重新组织了 bootstrap 信息流，把高方差 endpoint pseudo-label 转成低方差 local velocity / score matching target，因此更适合 γ→1 的长 horizon。

Method

方法层面应抓住三个机制。

1. Path-level Bellman operator：论文定义 Bt m = (1−γ)Pt + γPπm，其中 Pt 是 one-step transition 通过 conditional path 扩展到中间 t 后的分布。它解决的问题是：普通 Bellman operator 只告诉 t=1 的 successor measure，而 flow/diffusion 训练需要 t∈[0,1] 的监督。核心变化是把 TD fixed point 扩展成整条 probability path 的 fixed point，并证明 Bt 仍是 γ-contraction。

2. TD-CFM / TD-CFM(C)：TD-CFM 是直接的 endpoint bootstrap CFM；TD-CFM(C) 进一步把生成 bootstrap sample 的 noise X0 和 endpoint X1 绑定起来，避免“同一个 target 但路径随便接”的 crossing path 问题。它们解决的是如何把 CFM 放进 TD-GHM，但真正必要性在于暴露了 naive endpoint pseudo-data 的方差来源。Coupled 版本说明 coupling 能显著改善 transport geometry，但还不是最干净的解法。

3. TD2-CFM / TD2-DD：TD2 在 bootstrap 项上直接回归旧模型在同一中间态的 velocity / score。它解决的是 bootstrapped pseudo target 的条件方差：给定 Xt 时，naive CFM 仍要对隐藏 endpoint X1 的 conditional velocity 做采样估计；TD2 用旧 marginal vector field 相当于对这个条件期望做了 Rao-Blackwellization。核心变化是从 sample-level bootstrapping 变成 field-level bootstrapping。

Diffusion 扩展是同一思想的 score-matching 版本：bootstrap 项直接对齐旧 score，而不是采样到 data 再重新 noising。它证明这个机制不是 flow matching 特例，而是适用于 iterative generative models 的 TD path backup 模式。

Key Insight / Why It Works

这篇最重要的 insight 是：长 horizon GHM 的瓶颈不是生成模型不够强，而是 bootstrapped generative TD target 的梯度方差太大；TD2 的主要贡献是把这个方差消掉一部分。理论上，TD-CFM 和 TD2-CFM 在理想精确优化下收敛到同一个 successor-measure path fixed point，因此性能差异不应来自 fixed point 更好，而应来自有限样本 / SGD 下的 gradient variance。论文的 Theorem 2/3 正好说明 naive TD-CFM 比 TD2 多出一个 γ² 加权的 conditional covariance 项；γ 越接近 1，这个项越致命。这是全文最有解释力的结果。

从技术判断看，TD2 本质上是对 bootstrap 生成目标做 Rao-Blackwellization / self-distillation：不采样 endpoint 后再估计 conditional velocity，而直接用旧模型的 marginal velocity / score 作为 target。它不是 retrieval，不是 memory，更不是 test-time compute 的胜利；它主要是 better inductive bias + lower-variance bootstrapped supervision。Coupled TD-CFM 的作用更像 transport geometry 修正：当 straight paths 不交叉时，它可以接近 TD2；但一旦换 curved path，TD-CFM(C) 明显退化，说明它依赖路径几何巧合，而 TD2 更稳。

哪些部分可能只是辅助：UNet-style 架构、EMA target network、ODE solver、DDIM step 数等更多是 engineering；flow 比 diffusion 好的结果也不应过度解读，可能来自 diffusion 在高噪声长 horizon target 上训练更难，而非 flow matching 本质上更适合 RL。GAN/VAE baseline 崩得厉害并不意外，不能单独证明 TD2 的优势；真正有说服力的是 TD-CFM vs TD2-CFM 的 controlled comparison 和 curved path ablation。

需要警惕的是，planning 结果可能部分来自在一个预训练 policy family 上做 value-based reranking，而不是“形成了可组合的长期世界模型”。GPI 使用 reward expectation over generated successor states，本质更像 occupancy-conditioned policy retrieval / selection。它确实需要 long-horizon occupancy 质量，但不等价于 learned model 能做 arbitrary counterfactual planning。

Relation To Prior Work

这篇处在 successor representation / successor measure、γ-model / GHM、flow/diffusion generative modeling、model-based RL 的交叉处。最接近的是 Janner et al. 的 γ-models、Thakoor et al. 的 GHM、以及 Bellman diffusion / TD diffusion 类工作。和普通 world model 的本质差异是它不学习 local transition kernel 并 rollout，而是直接学习 discounted occupancy；和 successor features / FB representation 的差异是它保留了完整 future-state distribution，而不是只预测 reward-linear features 或 value inner product。

看似新的部分中，使用 flow matching / diffusion 本身不是创新；用 TD fixed point 训练 successor measure 也不是创新；用 target network / EMA 更不是创新。实质创新是把 Bellman equation 写到 probability path / vector field 层，并发现 TD2-style objective 是低方差的 path-level bootstrap。这个新增信息不是“更强生成模型”，而是“生成过程内部可以承载 TD backup”。

它属于一条更大的技术谱系：把 RL 的 Bellman self-consistency 从 scalar value 扩展到 distribution / representation / generative process。TD2-CFM 可以被看作 successor-measure 版本的 generative self-distillation，也可以看作对 flow matching conditional objective 的 Bellman-structured Rao-Blackwellization。

Dataset / Evaluation

实验覆盖了 Maze、Walker、Cheetah、Quadruped 等 DMControl/ExoRL 风格任务，包含 single-policy GHM 评估和 policy-conditioned GHM + GPI。任务范围对连续控制和低维状态足够有代表性，但不是视觉世界模型，也没有真实机器人或强随机环境。数据来自 10M exploratory transitions，覆盖性较强，这对 off-policy successor learning 很关键。

评估指标包括 EMD、NLL、value-function MSE 和 GPI return。它们基本能验证论文的核心 claim：TD2 在长 effective horizon 下比 naive bootstrapped generative GHM 稳定，且生成的 successor measure 对 value estimation / policy improvement 有用。尤其 horizon scaling 和 curved-path ablation 支持了“variance reduction 而非单纯 capacity”的归因。

明显 limitation 是 benchmark 仍偏 clean：状态低维、模拟环境、目标 policy/family 与数据分布之间可能有足够 overlap；value evaluation 的 reward 是环境已知函数，GPI 也是在已有 policy family 上选 policy。它验证的是 discounted occupancy modeling，不是通用 long-horizon reasoning。没有充分展示 out-of-distribution policy、partial observability、图像观测、真实部署下是否仍成立。

Limitation

核心前提有四个。第一，Bellman backup 依赖 transition data 对目标 policy 后继状态的覆盖；如果 π(s′) 落在数据支持外，TD2 只是更稳定地传播错误。第二，方法假设 future-state distribution 可由 Euclidean flow/diffusion 有效建模；复杂 observation 或 discrete/structured state 需要额外设计。第三，TD2 减少 variance，不消除 approximation bias；旧模型 path 如果偏了，field-level self-distillation 可能会把 bias 固化。第四，planning 仍依赖大量采样估计 Qπw，计算上不轻。

scalability 上限主要在 sampling cost 和 policy-conditioned distribution complexity。对于 GPI，每个 state 要对许多 candidate policies 采样 successor states，本质是 test-time generative evaluation；这比一步模型 rollout 稳，但不一定便宜。论文提到 consistency / self-distillation 作为未来方向，说明当前采样成本是实质瓶颈。

泛化是否真实存在仍不清楚。多 policy 实验条件在 FB 产生的 policy manifold 上训练和测试，可能存在较强 manifold overlap；所谓 planning gain 更像在候选 policy set 上做 occupancy-based reranking，而非开放式规划。核心能力可能主要来自数据覆盖 + policy family 限制 + 低方差 GHM，而不是一般意义的世界模型推理。

此外，flow 优于 diffusion 的归因文中未充分说明。可能是 diffusion noising 加剧了 noisy TD target，也可能是实现和采样步数差异。GAN/VAE baseline 弱也不能说明太多，因为它们本来就不适合高维长 horizon bootstrapped distribution。

Takeaway

  • 1. 对 long-horizon GHM 来说，关键不是换一个更强 generator，而是控制 bootstrapped generative TD 的方差；TD2 给出了一个很干净的 path-level 解法。
  • 2. Bellman consistency 可以作用在生成过程的中间表征上，而不只是 endpoint distribution / scalar value。
  • 这一 insight 可迁移到 diffusion policy evaluation、occupancy modeling、imitation learning、distributional RL 乃至 model-based planning。
  • 3. Coupling 和 path geometry 很重要，但最稳的做法是直接复用旧模型的 marginal velocity / score；endpoint pseudo-data self-training 在 γ→1 时天然脆弱。

一句话总结

Temporal Difference Flows 是把 successor-measure TD learning 从 endpoint distribution 提升到 flow/diffusion probability path 的方法，真正贡献是用 field-level bootstrapping 降低长 horizon GHM 的训练方差，而不是简单把更强生成模型套到 model-based RL 上。