AVID:使视频传播模型适应世界模型
24年11月来自微软的论文“AVID: Adapting Video Diffusion Models To World Models”。
大规模生成模型在许多领域取得了显著的成功。然而,对于机器人等顺序决策问题,动作标记数据通常很少,因此扩大决策基础模型仍然是一个挑战。一个潜在的解决方案,是利用广泛可用的未标记视频来训练模拟动作后果的世界模型。如果世界模型准确,则可以用它来优化下游任务中的决策。图像-到-视频的扩散模型,已经能够生成高度逼真的合成视频。然而,这些模型不是动作条件的,最强大的模型是闭源的,这意味着它们无法进行微调。在这项工作中,将预训练的视频扩散模型调整为动作条件的世界模型,而无需访问预训练模型的参数。该方法 AVID,在一个小型特定域的动作标记视频数据集上训练一个适配器。AVID 使用学习掩码来修改预训练模型的中间输出并生成准确的动作条件视频。在视频游戏和现实世界的机器人数据上对 AVID 进行评估,表明它在扩散模型自适应方面优于现有的基线。
做决策的扩散。许多方法利用扩散模型来生成动作(Pearce,2023;Janner,2022)。然而,工作的重点是生成用于世界建模的动作条件合成数据(Ha & Schmidhuber,2018),这些数据可在下游用于规划(Hafner,2021)。SynTHER(Lu,2024)采用无条件扩散模型为强化学习(RL)生成合成数据。其他方法(Alonso,2024;Yang,2024a;Rigter,2024;Hu,2023;Zhang,2024)利用扩散模型来训练动作条件世界模型。所有这些作品都是从头开始训练扩散模型。虽然自然语言适合描述高级动作,但对于建模低层动作的目标来说却是不够的 (McCarthy,2024)。
世界模型的视频预训练。为了解决动作标记的现实世界数据的稀缺问题,一些研究调查了使用未标记视频来提高世界模型训练的效率。Seo (2022) 和 Wu (2024) 预训练一个自回归视频预测模型,随后使用动作标记数据对其进行了微调。Mendonca (2023) 通过定义两个具身之间共享的高级动作空间,从人类视频中开发出一个世界模型。另一种方法涉及从视频中学习潜在动作(Bruce,2024;Schmidt & Jiang,2024)。
为扩散模型添加可控性。分类器指导(Dhariwal & Nichol,2021)使用单独的分类器为预训练的扩散模型添加条件。无分类器指导(Ho & Salimans,2021)实现更强的经验性能,但需要使用条件信号对模型进行训练,因此不适合事后应用于预训练模型。
最近的进展包括 ControlNet(Zhang,2023a)和 T2I-Adapter(Mou,2024),它们通过冻结原始 UNet 并向预训练网络注入附加信号,将条件控制引入预训练扩散模型。ControlNet-XS(Zavadski,2023)通过增加预训练网络和适配器网络之间的交互来扩展这一概念。Li(2023)将额外的可训练层合并到冻结的预训练 UNet 中。这些技术已被用于将语言(Xing,2024)、光流(Hu & Xu,2023)和深度图(Chen,2023b)等控制引入视频模型。但是,它们不适用于环境,因为它们需要访问原始模型的权重,假设这些权重是无法访问的。
文本反转 (Gal,2023) 和空文本反转 (Mokady,2023) 优化文本嵌入,以增加文本条件扩散模型的可控性。然而,这需要通过预训练模型进行反向传播,这在设置中是不可行的。级联扩散模型 (Ho,2022a;b) 训练一系列扩散模型,每个模型都以先前扩散模型的输出为条件。然而,这些工作的重点是提高每一步的空间和时间分辨率,而我们的重点是结合一个新的条件信号,即动作,这是预训练模型中没有的。
组合生成模型。组合生成模型中生成模型是概率组合的 (Liu,2022;Nie,2021;Du,2020)。 RoboDreamer 将这一想法应用于世界建模,将语言条件分解为多个组件,生成多个预测,然后将它们组合起来(Zhou,2024)。Yang(2024b)将其输出与独立训练的域特定扩散模型相结合,修改预训练视频扩散模型的风格。
本文解决利用预训练视频扩散模型生成动作条件预测而无需访问预训练模型参数的问题。受到 Yang (2024b) 最近研究的启发,假设只能访问预训练扩散模型的噪声预测。AVID,是一个域特定的适配器,它以动作为条件并修改预训练模型的噪声预测以生成准确的动作条件预测。为了训练适配器,假设可以访问特定域的动作标记视频数据集。
在视频扩散模型中,每个数据点都是一个视频,x = [x0,x1 ,… ,xT−1 ],其中 xτ 表示视频帧。请注意,使用上标索引 xτ 表示时间步骤,使用下标索引 x/i 表示扩散过程中的步骤。假设可以访问预训练的图像-到-视频扩散模型 ε/pre,该模型是在网络规模数据上训练的。给定初始图像 x0,图像-到-视频模型 ε/pre 生成合成视频 ˆ x = [x0,xˆ1,…,^xˆT −1]。
将每个视频视为由部分可观马尔可夫决策过程 (POMDP)(Kaelbling,1998)生成的一系列观测,相应的动作序列 a = [a0,a1,…,aT−1]。对于每个域,假设可以访问带有动作标记的视频数据集 D = {(x, a), . . .}。给定一个新的初始图像 x^0 和动作序列 a,目标是生成一个合成视频 xˆ,该视频准确描述动作的真实实现。
为了利用预训练模型 ε/pre,需要向模型添加动作条件,因为没有动作就无法生成准确的视频。在以前的研究中,已经很好地探索向预训练扩散模型添加新条件信号的方法(Zhang,2023a;Mou,2024;Mokady,2023)。然而,大多数现有研究都假设可以访问预训练模型的参数。这项工作假设无法访问预训练模型的参数。
AVID 并不试图组成两个独立训练的模型,而是使用预训练模型的输出来训练直接优化去噪损失的适配器,如图所示。

AVID 是一种新的扩散模型自适应方法,不需要访问预训练模型。AVID 的动机是,虽然预训练的图像-到-视频模型可以生成逼真的视频,但它们无法生成与给定动作序列相关的准确视频。为了实现准确性,必须引导预训练模型生成正确的动作序列。然而,正如实验表明的那样,诸如无分类器引导之类的技术在这种设置下表现不佳。AVID 通过训练轻量级适配器来调整预训练模型的输出,以实现准确的动作条件视频预测,从而解决了这个问题。
AVID 训练一个适配器,将预训练模型的输出作为输入。AVID 学习生成掩码,并使用此掩码将预训练模型的输出与适配器的输出相结合。最终的组合输出,用于计算标准去噪损失,并训练适配器参数以优化此损失。
用来自动作标记数据集的样本 (x, a) 来训练 AVID 适配器。对于预训练的图像到视频模型 ε/pre,假设不必访问它的参数 θ/pre,但可以使用该模型运行推理以获得它的噪声预测。为此,将一个嘈杂的视频 x/i、初始图像 x^0 和扩散步骤 i 输入到预训练的图像-到-视频模型中,以获得其噪声预测 ε/pre(x/i, i, x^0)。预训练模型 εpre 的参数未修改。
训练的适配器是一个 3D UNet(Ho,2022c;Ronneberger,2015),由一系列具有残差连接的时空块组成。每个时空块由 3D 卷积、空间注意和时间注意组成(Ho,2022c)。 UNet 将张量作为输入。输入包括噪声视频 x/i、预训练模型的输出 ε/pre (x/i , i, x^0 ) 和在时间维度上重复 T 次的初始图像 x^0。这三个输入按通道连接以创建输入张量。
适配器还以扩散步骤 i 和动作序列 a 为条件。对于噪声视频 x/i,根据学习的嵌入表,嵌入扩散时间步长,以获得扩散步骤嵌入 e/i。对于噪声视频 x/i 的每个时间步长 τ,嵌入相应的动作 aτ 以使用离散动作的嵌入表或连续动作的线性层来计算动作嵌入 eτ/a。在每个块中,这两个嵌入被连接起来并由 MLP 处理以计算第 τ 帧的比例和平移参数 γτ 和 β^τ。这些参数在每次 3D 卷积之后缩放和移动第 τ 帧的特征图(Perez,2018)。
为了将正确的动作条件注入预训练模型的预测中,适配器需要删除预训练模型预测的错误动作并添加正确的动作。为了实现这一点,适配器输出一个张量,该张量由一个掩码 m 组成,该掩码由一个 S 形的层限制在 0 和 1 之间,以及来自适配器 ε/adapt 的噪声预测。然后使用掩码将来自预训练模型和适配器的噪声预测结合起来。
当 ε/pre 为潜扩散模型时,假设也可以在相应的编码器和解码器上运行推理。对于每个训练示例,首先对视频进行编码:z = enc(x),将噪声添加到视频的这个潜表示中以产生嘈杂的潜 z/i。管道的其余部分以相同的方式进行,只是初始图像 x0 被对应于第一帧的潜 z^0 替换,并且嘈杂视频 x/i 被嘈杂的潜在 z/i 替换。适配器损失预测添加到潜在样本中的噪声,然后解码采样的潜向量,以生成最终视频。
Procgen 预训练模型是一个像素空间图像-到-视频扩散模型(Ho et al., 2022c),该模型使用从 Procgen(Cobbe et al., 2020)的 16 个程序生成游戏中 15 个中采样的视频进行训练,但不包括第 16 个游戏 Coinrun。自适应方法是使用从 Coinrun 采样的动作标记数据集进行训练的。在每个时间步,离散动作是 15 个可能的键盘输入之一,并且模型在 10 帧的序列上进行训练。每种自适应方法在单个 A100 GPU 上的训练时间限制为 3 天。测试三种数据集大小,Coin-run100k/500k/2.5M,并在保留的初始帧和动作测试集上评估模型。
RT1 + DynamiCrafter 在此基准测试中,预训练模型是 DynamiCrafter (Xing et al., 2023),它目前是 VBench 图像-到-视频排行榜 (Huang et al., 2024) 中表现最好的图像-到-视频模型之一。DynamiCrafter 是一种潜图像-到-视频扩散模型,它使用来自 Stable Diffusion (Rombach et al., 2022) 的自动编码器和在网络规模视频数据上训练的 1.4B 参数 3D UNet。由于 DynamiCrafter 是一种潜扩散模型,假设可以在编码器和解码器上运行推理。对于带动作标记的数据集,用由真实世界的机器人视频组成的 RT1 数据集 (Brohan et al., 2022)。每一步的动作都是一个 7 维连续向量,对应于末端执行器的移动和旋转以及夹持器的打开或关闭。模型在 16 帧的轨迹上进行训练。每种自适应方法在 4× A100 GPU 上的训练时间限制为 7 天,并使用一组保留的真值轨迹测试集进行评估。
更多推荐
所有评论(0)