基于Bayesian Flow Networks的离线轨迹规划

# 用于离线轨迹规划的贝叶斯流网络
来源：https://arxiv.org/html/2608.25163
Ludvig Killingberg，Helge Langseth  
所属机构：挪威科技大学

###### 摘要

离线强化学习（\(RL\)）利用静态数据集学习决策策略，而无需实时环境交互。尽管近期基于序列模型的方法依赖于连续扩散模型进行轨迹合成，但将这些方法应用于离线规划任务需要采用分类变量形式，而非标准的高斯构造。我们提出了BFN-RL，一个基于贝叶斯流网络（BFNs）的离线强化学习统一生成建模框架。BFN-RL通过迭代更新分布参数（而非含噪数据实例）的方式，在单一概率框架内原生地建模离散与连续轨迹空间。该分类规划器生成未来状态序列，并通过学习得到的逆动力学模型将连续生成的状态转化为动作。在离线规划和连续控制任务上的评估表明，BFN-RL能够跨越分类与连续状态空间生成有效轨迹。我们的研究结果确立了BFNs作为一种通用生成基础框架在跨数据模态的离线轨迹规划中的价值。

## 1 引言

离线强化学习（\(RL\)）（Sutton and Barto, 2018 (https://arxiv.org/html/2608.25163#bib.bib18)；Levine et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib17)）是一种强大的范式，它利用静态的、预先收集的数据集来学习有效的决策策略，而无需实时环境交互。通过消除在线探索带来的安全隐患，离线强化学习特别适用于自动驾驶和医疗决策等高风险领域。近年来，将离线强化学习框架构建为条件序列建模任务（Janner et al., 2021 (https://arxiv.org/html/2608.25163#bib.bib12)；Chen et al., 2021 (https://arxiv.org/html/2608.25163#bib.bib15)）已成为一种有前景的替代方案，以替代传统的基于价值的方法，后者常常在处理离分布的状态-动作对时产生价值过估计问题（Agarwal et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib16)；Levine et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib17)）。通过将轨迹生成视为条件生成建模任务，基于序列的智能体能够捕获长时间跨度上的复杂时间依赖性，从而合成高回报轨迹（Janner et al., 2022 (https://arxiv.org/html/2608.25163#bib.bib3)；Ajay et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib2)）。

尽管在实证上取得了成功，当前的序列建模方法几乎完全依赖高斯去噪扩散概率模型（DDPMs）（Ho et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib4)；Janner et al., 2022 (https://arxiv.org/html/2608.25163#bib.bib3)；Ajay et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib2)）。虽然扩散模型在连续控制领域表现出色，但离散的状态和动作需要分类变量形式，而非标准的高斯构造（Austin et al., 2021 (https://arxiv.org/html/2608.25163#bib.bib20)；Lou et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib19)）。这促使我们研究一种能原生处理两种数据类型的生成框架。

为了解决这一局限性，我们提出了*BFN-RL*，一个基于贝叶斯流网络（BFNs）（Graves et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib11)）的离线强化学习框架。不同于扩散模型通过迭代去噪损坏的数据实例来工作，BFNs通过由连续时间参数流驱动的贝叶斯推断来更新输入分布的参数（Graves et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib11)）。这种表述产生了一种统一的生成范式，能够在同一概率框架内原生地处理分类、连续和离散化变量。

我们的主要贡献总结如下：

- •统一生成规划框架：我们引入了BFN-RL，建立了一种用于离线序列建模的参数流范式，该范式能够在离散和连续的状态-动作空间上无缝运行（Graves et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib11)）。
- •分类轨迹规划：我们将分类BFN状态序列模型与学习得到的逆动力学模型相结合，得到一个离散规划器。
- •跨领域评估：我们证明了相同的BFN规划方法在离散和连续控制问题中均是可行的。

## 2 预备知识

### 2.1 强化学习

强化学习（\(RL\)）是一个在环境中学习决策的框架（Sutton and Barto, 2018 (https://arxiv.org/html/2608.25163#bib.bib18)）。与环境的交互被建模为马尔可夫决策过程（\(MDP\)），它是一个元组 \(S,A,P,R,γ\) (\(\mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma\))，其中 \(S\) (\(\mathcal{S}\)) 是状态空间，\(A\) (\(\mathcal{A}\)) 是动作空间，\(P\) (\(\mathcal{P}\)) 是转移函数，\(R\) (\(\mathcal{R}\)) 是奖励函数，\(γ\) (\(\gamma\)) 是折扣因子。在一个环境中，智能体在状态 \(s\) (\(s \in \mathcal{S}\)) 下执行动作 \(a\) (\(a \in \mathcal{A}\))，下一个状态 \(s'\) (\(s' \in \mathcal{S}\)) 从 \(P(s,a)\) (\(\mathcal{P}(s,a)\)) 中采样得到，即仅依赖于当前状态和动作，而不依赖于先前状态和动作的历史。换句话说，该领域遵循马尔可夫性质。设 \(r_t\) 表示在时间 \(t\) 获得的奖励，并令 \(R_t = \sum_{i=0}^{\infty} \gamma^i r_{t+i}\) 是从时间 \(t\) 开始获得的折扣累积回报。现在，强化学习的一个目标是学习一个策略 \(\pi\) (\(\pi: \mathcal{S} \rightarrow \mathcal{A}\))，最大化期望回报 \(\mathbb{E}[R_t]\)，其中期望取自转移函数和随机策略 \(\pi\) 定义的不确定性。

探索-利用权衡是强化学习中的一个基本挑战，通常与在线学习场景相关，即智能体迭代地与环境交互以学习最优策略。探索涉及采样动作以收集关于环境的信息，有可能发现更好的策略；而利用则涉及利用已知信息以最大化期望回报。在线强化学习的大部分研究致力于在探索和利用之间取得平衡，设计算法以有效地驾驭这种权衡，从而收敛到最优或近最优策略。

### 2.2 离线强化学习

在离线强化学习领域，主要目标是从固定的数据集中学习有效的策略，而无需在线交互（Levine et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib17)）。在这种探索方面固有缺失的情况下，重点转向如何有效利用可用数据集来优化策略。传统上，强化学习关注估计稳态策略或单步模型，利用马尔可夫性质将问题在时间上进行分解。然而，将标准的强化学习方法应用于离线设置是具有挑战性的，因为依赖于估计价值函数的方法常常会过度估计离分布状态-动作对的价值。已经提出了各种方法来解决这个问题，包括将策略约束在接近数据分布的范围内（Peters et al., 2010 (https://arxiv.org/html/2608.25163#bib.bib10)）或使用保守的价值函数（Kumar et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib9)）。另一方面，我们的解决方案是生成一系列步骤，这些步骤将根据强化学习智能体的目标进行条件生成。当我们从序列建模的视角看待强化学习时，一个有趣的视角出现了。与其将其视为一个专门的领域，我们可以将强化学习视为一个通用的序列建模问题。这种观点的关键在于生成一系列导致高奖励序列的动作。早期的工作通过根据回报来调节模型，从而在在线设置中生成具有高回报的轨迹来解决这个问题（Ajay et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib2)；Janner et al., 2021 (https://arxiv.org/html/2608.25163#bib.bib12)）。采用这种观点，我们可以简化设计决策，并省略离线强化学习算法中常见的许多组件。这种方法不仅在长时域动力学预测、模仿学习、目标条件强化学习和离线强化学习等多种任务中展现出灵活性，而且在稀疏奖励、长时域场景中产生了最先进的规划器（Janner et al., 2022 (https://arxiv.org/html/2608.25163#bib.bib3)；Ajay et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib2)）。

### 2.3 去噪扩散概率模型

由于该领域的当前最新技术（Janner et al., 2022 (https://arxiv.org/html/2608.25163#bib.bib3)；Ajay et al., 2023 (https://arxiv.org/html/2608.25163#bib.bib2)）依赖于扩散模型作为生成模型，我们将简要介绍去噪扩散概率模型（\(DDPMs\)）（Ho et al., 2020 (https://arxiv.org/html/2608.25163#bib.bib4)）。这也将作为我们接下来在第2.5节 (https://arxiv.org/html/2608.25163#S2.SS5) 讨论贝叶斯流网络的背景。DDPMs是一种受非平衡热力学启发的生成模型。该模型由一个*前向过程*定义，该过程缓慢地向数据添加高斯噪声，以及其逆过程，该过程通过学习迭代地去噪来重建样本。扩散模型主要用于图像生成，但在视频生成和3D模型等其他领域也显示出最先进的性能（Ho et al., 2022 (https://arxiv.org/html/2608.25163#bib.bib5)；Luo and Hu, 2021 (https://arxiv.org/html/2608.25163#bib.bib6)）。

给定数据 \(\mathbf{x}_0 \sim q(\mathbf{x})\)，我们定义*前向过程*以产生一系列含噪样本 \(\mathbf{x}_1, \dots, \mathbf{x}_K\)：
\[
q(\mathbf{x}_k | \mathbf{x}_{k-1}) = \mathcal{N}(\mathbf{x}_k; \sqrt{1-\beta_k}\mathbf{x}_{k-1}, \beta_k \mathbf{I}).
\]
其中 \(\{\beta_i \in (0,1)\}_{i=1}^K\) 是一个精心选择的方差调度。前向过程的一个良好性质是，我们可以直接在任意步骤 \(i\) 采样 \(\mathbf{x}_k\)，因为分布 \(q(\mathbf{x}_k | \mathbf{x}_0)\) 可以使用不相关正态分布随机变量之和仍为正态分布的性质推导出来。令 \(a_k = 1 - \beta_k\) 和 \(\bar{a}_k = \prod_{i=1}^k a_i\)，则
\[
q(\mathbf{x}_k | \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_k; \sqrt{\bar{a}_k}\mathbf{x}_0, (1-\bar{a}_k)\mathbf{I}).
\]
另请注意
\[
q(\mathbf{x}_{k-1} | \mathbf{x}_k, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{k-1}; \tilde{\boldsymbol{\mu}}(\mathbf{x}_k, \mathbf{x}_0), \tilde{\beta}_k \mathbf{I}), \quad (1)
\]
其中
\[
\tilde{\boldsymbol{\mu}}(\mathbf{x}_k, \mathbf{x}_0) = \frac{\sqrt{\bar{a}_{k-1}}\beta_k}{1-\bar{a}_k}\mathbf{x}_0 + \frac{\sqrt{a_k}(1-\bar{a}_{k-1})}{1-\bar{a}_k}\mathbf{x}_k, \qquad \tilde{\beta}_k = \beta_k\frac{1-\bar{a}_{k-1}}{1-\bar{a}_k}. \quad (2)
\]
虽然前向过程创建了数据的噪声表示，但*逆过程*旨在通过建模并随后从 \(q(\mathbf{x}_{k-1} | \mathbf{x}_k)\) 中采样，从而从噪声中迭代地重建样本。设 \(p_\theta(\mathbf{x}_{k-1} | \mathbf{x}_k)\) 是 \(q(\mathbf{x}_{k-1} | \mathbf{x}_k)\) 的参数化近似。这意味着我们定义了一个具有可训练参数 \(\theta\) 的神经网络模型，该模型输出 \(\boldsymbol{\mu}_\theta(\mathbf{x}_k, k)\) 和 \(\boldsymbol{\Sigma}_\theta(\mathbf{x}_k, k)\)，使得
\[
p_\theta(\mathbf{x}_{k-1} | \mathbf{x}_k) = \mathcal{N}(\mathbf{x}_{k-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_k, k), \boldsymbol{\Sigma}_\theta(\mathbf{x}_k, k)).
\]
Ho et al. (2020) (https://arxiv.org/html/2608.25163#bib.bib4) 选择将方差项 \(\boldsymbol{\Sigma}_\theta(\mathbf{x}_k, k)\) 固定为常数 \(\sigma_k^2 = \tilde{\beta}_k\)，见公式 (2 (https://arxiv.org/html/2608.25163#S2.E2))。因此我们只关注 \(\boldsymbol{\mu}_\theta(\mathbf{x}_k, k)\) 是如何估计的。首先，考虑恒等式
\[
\tilde{\boldsymbol{\mu}}_k(\mathbf{x}_k, \mathbf{x}_0) = \frac{1}{\sqrt{a_k}}\left(\mathbf{x}_k - \frac{1-a_k}{\sqrt{1-\bar{a}_k}}\boldsymbol{\epsilon}_k\right),
\]
其中 \(\boldsymbol{\epsilon}_k \sim \mathcal{N}(0, \mathbf{I})\)；参见公式 (1 (https://arxiv.org/html/2608.25163#S2.E1)) 和 (2 (https://arxiv.org/html/2608.25163#S2.E2))。由于在训练期间 \(\mathbf{x}_k\) 是已知的，我们可以选择预测 \(\boldsymbol{\epsilon}_k\)，而不是直接预测 \(\tilde{\boldsymbol{\mu}}_k\)。在经验上，这显示了更好的结果。让我们定义 \(\boldsymbol{\epsilon}_\theta(\mathbf{x}, k)\) 为预测噪声 \(\boldsymbol{\epsilon}_k\) 的模型。这意味着我们可以定义
\[
\boldsymbol{\mu}_\theta(\mathbf{x}_k, k) = \frac{1}{\sqrt{a_k}}\left(\mathbf{x}_k - \frac{1-a_k}{\sqrt{1-\bar{a}_k}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_k, k)\right).
\]
Ho et al. (2020) (https://arxiv.org/html/2608.25163#bib.bib4) 推导出以下损失函数以最小化 \(\boldsymbol{\mu}_\theta\) 与 \(\tilde{\boldsymbol{\mu}}\) 之间的差异：
\[
L(\theta) = \mathbb{E}_{k\sim[1,K], \mathbf{x}_0, \boldsymbol{\epsilon}_k}\left[\frac{\beta_k^2}{2\sigma_k^2 a_k (1-\bar{a}_k)}\left\|\boldsymbol{\epsilon}_k - \boldsymbol{\epsilon}_\theta(\mathbf{x}_k, k)\right\|^2\right].
\]
他们还提出了以下简化的损失函数，事实证明能给出更好的经验结果：
\[
L(\theta) = \mathbb{E}_{k\sim[1,K], \mathbf{x}_0, \boldsymbol{\epsilon}_k}\left\|\boldsymbol{\epsilon}_k - \boldsymbol{\epsilon}_\theta(\mathbf{x}_k, k)\right\|^2.
\]

### 2.4 引导扩散

我们将讨论扩散模型对变量进行条件化的三种方式。第一种，*分类器引导扩散*（Dhariwal and Nichol, 2021 (https://arxiv.org/html/2608.25163#bib.bib7)），其出发点是...