Temporal Difference Flows

📄 arXiv: 2503.09817v1 📥 PDF

作者: Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Rémi Munos, Alessandro Lazaric, Ahmed Touati

分类: cs.LG, cs.AI, stat.ML

发布日期: 2025-03-12

💡 一句话要点

提出TD-Flow，通过概率路径上的贝尔曼方程和流匹配技术，学习长时域精确的几何视界模型。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 时间差分学习 几何视界模型 流匹配 长时域预测 贝尔曼方程 概率路径 世界模型

📋 核心要点

  1. 现有世界模型在长时域预测中误差累积严重，几何视界模型(GHMs)虽能避免，但训练时自举预测导致性能受限。
  2. TD-Flow利用概率路径上的贝尔曼方程和流匹配技术，降低训练时的梯度方差，从而学习更精确的GHMs。
  3. 实验表明，TD-Flow在多种任务上显著提升了GHMs的预测精度和规划性能，尤其在长时域任务中优势明显。

📝 摘要（中文）

预测模型是智能体进行推理和规划的基础。一种常见策略是学习世界模型并在推理时逐步展开，但小误差会迅速累积。几何视界模型(GHMs)提供了一种引人注目的替代方案，它直接预测未来状态，避免累积推理误差。GHMs可以通过时间差分(TD)学习的生成式模拟方便地学习，但现有方法受到训练时自举预测的负面影响，难以生成高质量的长时域预测。本文提出了时间差分流(TD-Flow)，它利用概率路径上新颖的贝尔曼方程结构以及流匹配技术，以超过现有方法5倍的时域长度学习精确的GHMs。理论上，我们建立了一个新的收敛结果，并将TD-Flow的有效性主要归因于训练期间降低的梯度方差。我们进一步表明，类似的论点可以扩展到基于扩散的方法。在实验上，我们在各种领域验证了TD-Flow在生成指标和下游任务（包括策略评估）上的有效性。此外，将TD-Flow与最近的行为基础模型集成，以进行基于预训练策略的规划，证明了显著的性能提升，突显了其在长时域决策中的前景。

🔬 方法详解

问题定义：论文旨在解决长时域预测中，现有几何视界模型（GHMs）训练不稳定、预测精度不足的问题。现有方法依赖于自举（bootstrapping）预测，导致训练过程中误差累积，尤其在长时域预测中表现更差。这限制了GHMs在复杂决策任务中的应用。

核心思路：TD-Flow的核心思路是将时间差分学习（TD learning）与流匹配（flow matching）技术相结合，在概率路径上求解贝尔曼方程。通过这种方式，TD-Flow能够更稳定地训练GHMs，并生成更准确的长时域预测。关键在于利用流匹配来定义概率路径，并在此路径上应用TD学习，从而降低梯度方差，提高训练效率。

技术框架：TD-Flow的整体框架包括以下几个主要部分： 1. 概率路径定义：使用流匹配技术定义从初始状态到未来状态的概率路径。 2. 贝尔曼方程求解：在定义的概率路径上，求解一个新颖的贝尔曼方程，该方程描述了状态值函数在时间上的演变。 3. 模型训练：使用时间差分学习方法，基于贝尔曼方程训练几何视界模型，使其能够准确预测未来状态。 4. 策略规划：将训练好的几何视界模型应用于策略规划任务，以优化智能体的行为。

关键创新：TD-Flow的关键创新在于： 1. 概率路径上的贝尔曼方程：提出了一个在概率路径上定义的贝尔曼方程，这与传统的贝尔曼方程不同，更适合于生成式模型的训练。 2. 流匹配与TD学习的结合：巧妙地将流匹配技术与时间差分学习相结合，利用流匹配来定义概率路径，并在此路径上应用TD学习，从而降低了训练的梯度方差。 3. 长时域预测能力：通过降低梯度方差，TD-Flow能够训练出更稳定的GHMs，从而实现更准确的长时域预测。

关键设计：TD-Flow的关键设计包括： 1. 流匹配网络结构：使用了特定的神经网络结构来实现流匹配，例如连续归一化流（Continuous Normalizing Flows, CNF）。 2. 损失函数：设计了基于贝尔曼方程的时间差分损失函数，用于训练几何视界模型。该损失函数旨在最小化预测值与真实值之间的差异。 3. 训练策略：采用了特定的训练策略，例如梯度裁剪（gradient clipping）和学习率衰减（learning rate decay），以进一步提高训练的稳定性。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，TD-Flow在多个领域都取得了显著的性能提升。例如，在某些任务中，TD-Flow能够将预测时域长度提高到现有方法的5倍以上。此外，TD-Flow在策略评估和规划任务中也表现出色，与现有的行为基础模型集成后，能够进一步提高智能体的性能。

🎯 应用场景

TD-Flow在机器人控制、自动驾驶、游戏AI等领域具有广泛的应用前景。它可以用于训练智能体进行长时域规划，例如，机器人需要在复杂环境中完成一系列动作才能达到目标，自动驾驶系统需要预测车辆在未来一段时间内的行为，游戏AI需要制定长期的战略。TD-Flow能够提高这些智能体的决策能力和适应性。

📄 摘要（原文）

Predictive models of the future are fundamental for an agent's ability to reason and plan. A common strategy learns a world model and unrolls it step-by-step at inference, where small errors can rapidly compound. Geometric Horizon Models (GHMs) offer a compelling alternative by directly making predictions of future states, avoiding cumulative inference errors. While GHMs can be conveniently learned by a generative analog to temporal difference (TD) learning, existing methods are negatively affected by bootstrapping predictions at train time and struggle to generate high-quality predictions at long horizons. This paper introduces Temporal Difference Flows (TD-Flow), which leverages the structure of a novel Bellman equation on probability paths alongside flow-matching techniques to learn accurate GHMs at over 5x the horizon length of prior methods. Theoretically, we establish a new convergence result and primarily attribute TD-Flow's efficacy to reduced gradient variance during training. We further show that similar arguments can be extended to diffusion-based methods. Empirically, we validate TD-Flow across a diverse set of domains on both generative metrics and downstream tasks including policy evaluation. Moreover, integrating TD-Flow with recent behavior foundation models for planning over pre-trained policies demonstrates substantial performance gains, underscoring its promise for long-horizon decision-making.