【强化学习论文】柔性行为综合的扩散规划

• 文献题目：Planning with Diffusion for Flexible Behavior Synthesis
• 发表时间：2022

摘要

• 基于模型的强化学习方法通常只将学习用于估计近似动力学模型，将其余的决策工作卸载给经典的轨迹优化器。虽然概念上很简单，但这种组合有许多经验上的缺点，这表明学习模型可能不太适合标准轨迹优化。在本文中，我们考虑了将尽可能多的轨迹优化管道折叠到建模问题中会是什么样子，以便从模型中采样和用它进行规划变得几乎相同。我们的技术方法的核心在于通过迭代去噪轨迹来计划的扩散概率模型。我们展示了分类器引导的采样和图像修补如何被重新解释为连贯的规划策略，探索了基于扩散的规划方法的不同寻常和有用的属性，并证明了我们的框架在强调长期决策和测试时间灵活性的控制设置中的有效性。

引言

• 使用学习模型进行规划是用于强化学习和数据驱动决策的概念简单框架。它的吸引力来自于只在最成熟和最有效的地方使用学习技术:在相当于监督学习问题的未知环境动态的近似。之后，可以将学习到的模型插入经典轨迹优化例程(Tassa等人，2012;Posa等人，2014;Kelly, 2017)，这些词在最初的语境中同样很好理解。
• 然而，这种组合很少像描述的那样工作。由于强大的轨迹优化器利用学习模型，由该过程生成的计划通常看起来更像对抗的例子，而不是最优轨迹(Talvitie, 2014;Ke等人，2018)。因此，当代基于模型的强化学习算法通常更多地继承自无模型方法，如价值函数和策略梯度(Wang等人，2019)，而不是从轨迹优化工具箱中继承。那些依赖于在线规划的方法倾向于使用简单的无梯度轨迹优化例程，如随机射击(Nagabandi等人，2018)或交叉熵法(Botev等人，2013;Chua等人，2018)来避免上述问题。
• 在这项工作中，我们提出了一种数据驱动的轨迹优化的替代方法。==核心思想是训练一个直接服从轨迹优化的模型，从某种意义上说，从模型中采样和用它进行规划几乎是相同的。==这一目标要求改变模型的设计方式。由于学习动力学模型通常意味着环境动力学的代理，因此通常通过根据潜在的因果过程构建模型来实现改进(Bapst等人，2019)。相反，我们考虑的是如何设计一个与将要使用它的计划问题相一致的模型。例如，由于模型最终将用于计划，所以动作分布与状态动态一样重要，长视距精度比单步误差更重要。另一方面，该模型应该对奖励函数保持不可知，以便它可以用于多个任务，包括那些在训练过程中看不到的任务。最后，模型的设计应使其计划(而不仅仅是预测)随着经验的积累而改进，并能抵抗标准的基于射击的规划算法的短视失败模式。
• 我们将这一想法实例化为轨迹级扩散概率模型(Sohl-Dickstein等人，2015;Ho等人，2020)称为扩散器，如图2所示。而标准的基于模型的计划技术是在时间上自回归地向前预测，而Diffuser则同时预测计划的所有时间步长。扩散模型的迭代采样过程导致了灵活的调节，允许辅助向导修改采样程序，以恢复具有高回报或满足一组约束的轨迹。这种数据驱动的轨迹优化公式有几个吸引人的特性:
• 世界范围的可伸缩性 扩散器是针对其生成轨迹的准确性而不是其单步误差进行训练的，因此它不会受到单步动力学模型的复合滚出误差的影响，并且在长期规划范围内更优雅地缩放。
• 组合性任务 奖励函数提供了在抽样计划时使用的辅助梯度，允许通过添加它们的梯度同时组合多个奖励，从而实现一种直接的计划方式。
• 时间组合性 扩散器通过迭代改进局部一致性生成全局连贯轨迹，允许它通过将分布中的子序列拼接在一起来泛化到新的轨迹。
• 有效的非贪婪规划 通过模糊模型和规划器之间的界限，改进模型预测的训练过程也具有提高其规划能力的效果。 这种设计产生了一个学习规划器，可以解决许多传统规划方法难以解决的长期、稀疏回报问题。
• 这项工作的核心贡献是为轨迹数据设计的去噪扩散模型和用于行为合成的相关概率框架。 虽然与基于深度模型的强化学习中常规使用的模型类型相比是非常规的，但我们证明 Diffuser 具有许多有用的属性，并且在需要长期推理和测试时间灵活性的控制设置中特别有效。

背景

• 我们的规划方法是对过去使用轨迹优化的行为综合工作的基于学习的模拟（Witkin & Kass，1988 年；Tassa 等人，2012 年）。 在本节中，我们简要介绍了轨迹优化所考虑的问题设置以及我们为该问题采用的生成模型类别。

问题设置

• 考虑一个由离散时间动态 $s_{t+1}=f(s_t,a_t)$ 控制的系统，在状态 $s_t$ 给定一个动作 $a_t$。 轨迹优化指的是找到一系列动作 $a_{0:T}^{\ast }$ 最大化（或最小化）目标 $J$ 分解为每个时间步长的奖励（或成本）$r(s_t,a_t)$：
• 
• 其中 $T$ 是规划范围。 我们使用缩写 $\tau =(s_0,a_0,s_1,a_1,..,s_T,a_T)$ 来指代交错状态和动作的轨迹，并使用 $J(\tau )$ 来表示该轨迹的目标值。

扩散概率模型

• 扩散概率模型（Sohl-Dickstein 等人，2015 年；Ho 等人，2020 年）将数据生成过程视为迭代去噪过程 $p_{\theta }({\tau }^{i-1}|{\tau }^i)$。 这种去噪是前向扩散过程 $q({\tau }^i|{\tau }^{i-1})$ 的逆过程，前向扩散过程 $q({\tau }^i|{\tau }^{i-1})$ 通过添加噪声慢慢破坏数据中的结构。模型引起的数据分布由下式给出：
  
• 其中 $p({\tau }^N)$ 是标准高斯先验，${\tau }^0$ 表示（无噪声）数据。 通过最小化反向过程的负对数似然的变分界限来优化参数 $\theta$：${\theta }^{\ast }=argmi{n}_{\theta }-E_{{\tau }^0}[log{p}_{\theta }({\tau }^0)]$。 逆向过程通常被参数化为具有固定时间步长相关协方差的高斯分布：
  
• 正向过程 $q({\tau }^i|{\tau }^{i-1})$ 通常是预先指定的。
• 符号。 在这项工作中有两个“时间”在起作用：扩散过程的时间和规划问题的时间。 我们使用上标（未指定时为 $i$）表示扩散时间步长，使用下标（未指定时为 $t$）表示规划时间步长。 例如，$s_t^0$ 指的是无噪声轨迹中的第 $t$ 个状态。 当上下文明确时，省略无噪声量的上标：$\tau ={\tau }^0$。 我们通过将轨迹 $\tau$ 中的第 $t$ 个状态（或动作）称为 ${\tau }_{s_t}$（或 ${\tau }_{a_t}$）来稍微重载符号。

扩散规划

• 使用轨迹优化技术的一个主要障碍是它们需要了解环境动力学 $f$。 大多数基于学习的方法都试图通过训练近似动力学模型并将其插入常规规划程序来克服这一障碍。 然而，学习模型通常不太适合设计时考虑到真实模型的规划算法类型，导致规划者通过寻找对抗性示例来利用学习模型。
• 我们建议在建模和规划之间建立更紧密的耦合。 我们没有在经典规划器的上下文中使用学习模型，而是将尽可能多的规划过程纳入生成建模框架，这样规划就变得几乎与抽样相同。 我们使用轨迹的扩散模型 $p_{\theta }(\tau )$ 来执行此操作。 扩散模型的迭代去噪过程有助于通过从以下形式的扰动分布中采样来进行灵活调节：
  
• 函数 $h(\tau )$ 可以包含有关先验证据（例如观察历史）、期望结果（例如要达到的目标）或要优化的一般函数（例如奖励或成本）的信息。 在这种扰动分布中执行推理可以看作是第 2.1 节中提出的轨迹优化问题的概率模拟，因为它需要找到在 $p_{\theta }(\tau )$ 下物理现实和在 $h(\tau )$ 下的高回报（或约束满足）的轨迹 。 因为动力学信息与扰动分布 $h(\tau )$ 分离，所以单个扩散模型 $p_{\theta }(\tau )$ 可重复用于同一环境中的多个任务。
• 在本节中，我们将描述 Diffuser，这是一种为学习轨迹优化而设计的扩散模型。 然后，我们讨论使用 Diffuser 进行规划的两个具体实例，实现为分类器引导采样和图像修复的强化学习对应物。

轨迹规划的生成模型

• 时间顺序。 模糊轨迹模型采样和规划之间的界限会产生一个不寻常的约束：我们不能再按时间顺序自回归预测状态。 考虑目标条件推理 p(s1 | s0, sT )； 下一个状态 s1 取决于未来状态和先前状态。 这个例子是一个更普遍的原则的例子：虽然动态预测是因果关系，从现在由过去决定的意义上说，决策和控制可以是反因果的，从某种意义上说，现在的决定取决于未来。 因为我们不能使用时间自回归排序，所以我们设计 Diffuser 来同时预测计划的所有时间步长。
• 时间局部性。 尽管不是自回归或马尔可夫模型，但 Diffuser 具有一种松散的时间局部性形式。 在图 2 中，我们描绘了由单个时间卷积组成的扩散模型的依赖图。 给定预测的接受域仅包含过去和未来的附近时间步长。 结果，去噪过程的每一步只能根据轨迹的局部一致性进行预测。 然而，通过将许多这些去噪步骤组合在一起，局部一致性可以推动全局一致性。
• 在一般的强化学习环境中，为了编写动态规划递归的目的，对未来的条件化源于对未来最优性的假设。 具体而言，这表现为行动分布 $logp(a_t|s_t,O_{t:T})$ 中的未来最优变量 $O_{t:T}$ (Levine, 2018)。
  
• 图 2. 扩散器通过对由可变数量的状态-动作对组成的二维数组进行迭代去噪来对计划进行采样。 一个小的感受野限制模型只在单个去噪步骤中强制执行局部一致性。 通过将许多去噪步骤组合在一起，局部一致性可以推动采样计划的全局一致性。 可选的指导函数 $J$ 可用于将计划偏向那些优化测试时间目标或满足一组约束的计划。
• 轨迹表示。 扩散器是为规划而设计的轨迹模型，这意味着从模型中得出的控制器的有效性与状态预测的质量一样重要。 因此，轨迹中的状态和动作是联合预测的； 出于预测的目的，动作只是状态的附加维度。 具体来说，我们将 Diffuser 的输入（和输出）表示为二维数组：
  
• 规划范围的每个时间步长一列。
• 架构。 我们现在有了指定扩散器架构所需的成分：(1) 整个轨迹应该以非自回归方式预测，(2) 去噪过程的每个步骤都应该是时间局部的，以及(3) 轨迹表示应允许沿一个维度（规划范围）的等变性，但不允许沿另一个维度（状态和动作特征）。 我们使用由重复（时间）卷积残差块组成的模型来满足这些标准。 整体架构类似于在基于图像的扩散模型中取得成功的 U-Net 类型，但二维空间卷积被一维时间卷积取代（图 A1）。 因为模型是完全卷积的，所以预测的范围不是由模型架构决定的，而是由输入维度决定的； 如果需要，它可以在计划期间动态更改。
  
• 图 A1。 Diffuser 具有 U-Net 架构，其残差块由时间卷积、组归一化和 Mish 非线性组成。
• 训练。 我们使用 Diffuser 参数化轨迹去噪过程的学习梯度 ${\epsilon }_{\theta }({\tau }^i,i)$，从中可以以封闭形式求解平均 ${\mu }_{\theta }$ (Ho et al., 2020)。 我们使用简化的目标来训练 $\epsilon$-模型，由下式给出：
  
• 其中 $i\sim U\{1,2,...,N\}$ 是扩散时间步长，$\epsilon \sim N(0,I)$ 是噪声目标，${\tau }^i$ 是被噪声 $\epsilon$ 破坏的轨迹 ${\tau }^0$。 逆过程协方差 ${\Sigma }^i$ 遵循 Nichol & Dhariwal (2021) 的余弦时间表。

强化学习作为引导抽样

• 为了用 Diffuser 解决强化学习问题，我们必须引入奖励的概念。 我们呼吁控制作为推理图形模型（Levine，2018）这样做。 令 $O_t$ 为二元随机变量，表示轨迹的时间步长 $t$ 的最优性，其中 $p(O_t=1)=exp(r(s_t,a_t))$。 我们可以通过在等式 1 中设置 $h(\tau )=p(O_{1:T}|\tau )$ 从一组最优轨迹中采样：
  
• 我们已经将强化学习问题换成了条件抽样之一。 值得庆幸的是，之前已经有很多关于使用扩散模型进行条件抽样的工作。 虽然很难从该分布中精确采样，但当 $p(O_{1:T}|{\tau }^i)$ 足够平滑时，反向扩散过程转换可以近似为高斯分布（Sohl-Dickstein 等人，2015 年）：

• 其中 $\mu ,\Sigma$ 是原始逆过程转换的参数 $p_{\theta }({\tau }^{i-1}|{\tau }^i)$ 和
  
• 这种关系提供了分类器引导采样（用于生成类别条件图像（Dhariwal 和 Nichol，2021））与强化学习问题设置之间的直接转换。 我们首先在所有可用轨迹数据的状态和动作上训练扩散模型 $p_{\theta }(\tau )$。 然后我们训练一个单独的模型 $J_{\phi }$ 来预测轨迹样本 ${\tau }^i$ 的累积奖励。 $J_{\phi }$ 的梯度用于通过根据等式 3 修改反向过程的均值 $\mu$ 来指导轨迹采样过程。采样轨迹 $\tau \sim p(\tau |O_{1:T}=1)$ 的第一个动作可以在环境中执行，之后规划过程在标准后退水平控制循环中再次开始。 算法 1 给出了引导规划方法的伪代码。
  

目标条件强化学习作为修复

• 一些规划问题更自然地表现为约束满足而不是奖励最大化。 在这些设置中，目标是生成满足一组约束的任何可行轨迹，例如在目标位置终止。 求助于等式 2 描述的轨迹的二维数组表示，此设置可以转化为修复问题，其中状态和动作约束类似于图像中观察到的像素（Sohl-Dickstein 等人，2015 年）。 阵列中所有未观察到的位置必须由扩散模型以与观察到的约束一致的方式填充。
• 此任务所需的扰动函数是观测值的 Dirac delta 和其他地方的常量。 具体来说，如果 $c_t$ 是时间步 $t$ 的状态约束，则
  
• 动作约束的定义是相同的。 在实践中，这可以通过从未受干扰的逆向过程 ${\tau }^{i-1}\sim p_{\theta }({\tau }^{i-1}|{\tau }^i)$ 中采样并在所有扩散时间步 $i\in \{0,1,...，N\}$。
• 即使是奖励最大化问题也需要通过修复进行调节，因为所有采样轨迹都应从当前状态开始。 算法 1 中的第 10 行描述了这种调节。

扩散规划器的属性

• 我们讨论了 Diffuser 的一些重要属性，重点关注那些不同于标准动力学模型或非自回归轨迹预测不常见的属性。

• 图 3.（扩散规划器的属性）（a）学习的长期规划：扩散器的学习规划过程不会受到射击算法常见的近视失败模式的影响，并且能够在奖励稀少的情况下进行长期规划。 (b) 时间组合性：即使模型不是马尔可夫模型，它也会通过迭代改进来生成轨迹以达到局部一致性。 因此，它展示了通常与马尔可夫模型相关的泛化类型，能够将训练数据中的轨迹片段拼接在一起以生成新颖的计划。 © 可变长度计划：尽管是轨迹级模型，但 Diffuser 的计划范围并不由其架构决定。 通过改变输入噪声的维数，可以在训练后更新地平线。 (d) 任务组合性：扩散器可以与新的奖励函数组合，以规划训练期间未见过的任务。 在所有子图中， 表示起始状态并表示目标状态。
• 学会了长远规划。 单步模型通常用作地面实况环境动态 $f$ 的代理，因此不依赖于任何特定的规划算法。 相比之下，算法 1 中的规划例程与扩散模型的特定可供性密切相关。 因为我们的规划方法几乎与抽样相同（唯一的区别是受扰动函数 h(τ) 的指导），扩散器作为长期预测器的有效性直接转化为有效的长期规划。 我们在图 3a 中展示了在目标达成设置中学习规划的好处，表明扩散器能够在稀疏奖励设置类型中生成可行的轨迹，而基于射击的方法众所周知会遇到困难。 我们在第 5.1 节中探讨了此问题设置的更多定量版本。
• 时间组合性。 单步模型通常使用马尔可夫属性来激发，允许它们组合分布内转换以推广到分布外轨迹。 由于 Diffuser 通过迭代改进局部一致性（第 3.1 节）生成全局一致的轨迹，它还可以以新颖的方式将熟悉的子序列拼接在一起。 在图 3b 中，我们在仅沿直线行进的轨迹上训练扩散器，并表明它可以通过在交点处组合轨迹来泛化为 V 形轨迹。
• 可变长度计划。 因为我们的模型在其预测的范围维度上是完全卷积的，所以它的规划范围不是由架构选择指定的。 相反，它由初始化去噪过程的输入噪声 ${\tau }^N\sim N(0,I)$ 的大小决定，允许可变长度计划（图 3c）
• 任务组合性。 虽然扩散器包含有关环境动态和行为的信息，但它独立于奖励函数。 由于该模型作为可能未来的先验，因此可以通过对应于不同奖励的相对轻量级扰动函数 $h(\tau )$（甚至多个扰动的组合）来指导规划。 我们通过计划在扩散模型训练期间看不到的新奖励函数来证明这一点（图 3d）。

实验

相关工作

• 深度生成建模的进展最近进入了基于模型的强化学习，多行工作探索参数化为卷积 U 网络（Kaiser 等人，2020 年）、随机循环网络（Ke 等人，2018 年； Hafner 等人，2021a；Ha 和 Schmidhuber，2018 年）、矢量量化自编码器（Hafner 等人，2021b；Ozair 等人，2021）、神经 ODE（Du 等人，2020a）、归一化流（Rhinehart 等人 al., 2020; Janner et al., 2020)，生成对抗网络 (Eysenbach et al., 2021)，基于能量的模型 (EBMs; Du et al. 2019)，图神经网络 (Sanchez-Gonzalez et al., 2018)、神经辐射场 (Li et al., 2021) 和 Transformers (Janner et al., 2021; Chen et al., 2021a)。 这些调查通常假设模型和规划器之间存在抽象障碍。 具体来说，学习的作用被降级为近似环境动力学； 一旦学习完成，该模型可能会被插入到各种规划（Botev 等人，2013 年；Williams 等人，2015 年）或政策优化（Sutton，1990 年；Wang 等人，2019 年）算法中，因为规划器的形式并不强烈依赖于模型的形式。 我们的目标是通过设计一个相互训练的模型和规划算法来打破这种抽象障碍，从而产生一个采样和规划几乎相同的非自回归轨迹级模型。
• 许多平行的工作线研究了如何以不同的方式打破模型学习和规划之间的抽象障碍。 方法包括训练用于奖励预测的自回归潜在空间模型（Tamar 等人，2016 年；Oh 等人，2017 年；Schrittwieser 等人，2019 年）； 按状态值权衡模型训练目标（Farahmand 等人，2017 年）； 并将搭配技术应用于学习的单步能量。 相比之下，我们的方法计划通过同时生成轨迹的所有时间步长，而不是自回归，并使用辅助引导函数调节采样轨迹。
• 扩散模型已成为一类很有前途的生成模型，它将数据生成过程制定为迭代去噪过程（Sohl-Dickstein 等人，2015 年；Ho 等人，2020 年）。 去噪过程可以看作是参数化数据分布的梯度 (Song & Ermon, 2019)，将扩散模型连接到分数匹配 (Hyvarinen ¨, 2005) 和 EBM (LeCun et al., 2006; Du & Mordatch, 2019; Nijkamp 等人，2019 年；Grathwohl 等人，2020 年）。 基于梯度的迭代采样有助于实现灵活的调节（Dhariwal & Nichol，2021）和组合性（Du 等人，2020b），我们使用它们从异构数据集中恢复有效行为，并计划在训练期间看不到的奖励函数。 虽然已经开发了扩散模型来生成图像（Song 等人，2021 年）、波形（Chen 等人，2021c）、3D 形状（Zhou 等人，2021 年）和文本（Austin 等人，2021 年） )，据我们所知，它们以前从未在强化学习或决策制定的背景下使用过。

结论

• 我们提出了 Diffuser，一种用于轨迹数据的去噪扩散模型。 使用 Diffuser 进行规划与从中采样几乎相同，不同之处仅在于添加了用于指导采样的辅助扰动函数。 学习到的基于扩散的规划程序具有许多有用的特性，包括对稀疏奖励的优雅处理、无需重新训练即可规划新奖励的能力，以及这允许它通过将分布内的子序列拼接在一起来产生分布外的轨迹。 我们的结果指向一类新的基于扩散的规划程序，用于基于深度模型的强化学习。