【AI论文】Dita：扩展扩散变换器以构建通用视觉-语言-动作策略

摘要：尽管近期在多样化机器人数据集上训练的视觉-语言-动作模型在有限领域内数据上展现出了良好的泛化能力，但它们依赖紧凑的动作头来预测离散或连续动作，这限制了其对异构动作空间的适应性。我们提出了Dita，这是一个可扩展的框架，它利用Transformer架构，通过统一的多模态扩散过程直接对连续动作序列进行去噪。与先前通过浅层网络在融合嵌入上进行去噪条件化的方法不同，Dita采用了上下文条件化——这使得去噪后的动作能够与来自历史观察的原始视觉标记进行精细对齐。这种设计显式地建模了动作变化和环境细节。通过扩展扩散动作去噪器以及利用Transformer的可扩展性，Dita能够有效地整合跨不同相机视角、观察场景、任务和动作空间的跨实体数据集。这种协同作用增强了模型对各种变化的鲁棒性，并促进了长时程任务的成功执行。在广泛基准上的评估表明，Dita在模拟中达到了最先进或相当的性能。值得注意的是，Dita仅通过10次微调，使用第三人称相机输入，就能在现实世界中实现对环境变化和复杂长时程任务的鲁棒适应。该架构为通用机器人策略学习建立了一个多功能、轻量级且开源的基准。项目页面：https://robodita.github.io。Huggingface链接：Paper page，论文链接：2503.19757

研究背景和目的

研究背景

传统的机器人学习策略通常依赖于为特定机器人和任务收集的大规模数据，然而，为通用任务收集数据既耗时又成本高昂，这主要是由于现实世界机器人硬件的固有局限性。目前，自然语言处理和计算机视觉领域的基础模型，通过在大规模、多样化且任务无关的数据集上进行预训练，已经展示出了在解决下游任务方面的卓越效果，无论是通过零样本学习还是仅使用少量任务特定样本。这一成就表明，一个预训练于异构机器人数据并通过最小监督进行微调的通用机器人策略，对于实现视觉-语言-动作（VLA）模型的真正泛化至关重要。然而，跨广泛的跨实体数据集（涵盖不同的传感器、动作空间、任务、相机视角和环境）训练这样的策略仍然是一个开放性的挑战。

近年来，一些研究直接将视觉观察和语言指令映射到动作上，使用扩展的VLA模型进行导航或操作，从而在新颖环境中展示了零样本或少样本泛化能力。尽管这些VLA模型在从大规模跨实体数据集中学习机器人策略方面展现出了巨大潜力，但数据集内机器人配置的内在多样性仍然限制了其泛化能力。

扩散策略在单任务模仿学习范式下的机器人策略学习中表现出了可靠的性能。然而，现有的基于扩散的操控策略主要依赖于U-Net架构或用于单任务的浅层交叉注意力网络，这限制了它们在多模态应用中的可扩展性。此外，尽管一些最新的通用模型采用了视觉语言模型（VLM）嵌入与紧凑的多层感知机（MLP）扩散器相结合的方式，但它们可能无法有效地捕捉历史观察中的细微变化，这对于动作预测至关重要。

研究目的

本研究旨在开发一个可扩展的框架，以解决现有视觉-语言-动作模型在泛化和适应性方面的局限性。具体来说，我们希望通过利用Transformer架构的强大能力，构建一个能够直接从跨实体数据集中学习，并在新颖环境和任务中展现出强大泛化和适应能力的通用机器人策略。这一目标对于推动机器人技术的进一步发展，特别是在复杂、多变环境下的自主操作和任务执行方面，具有重要意义。

研究方法

模型架构

我们提出的Dita模型是一个可扩展的扩散变换器（DiT），它采用上下文条件化机制与因果变换器相结合，直接对连续动作序列进行去噪。Dita模型仅使用语言指令和第三人称相机图像作为输入。语言指令通过冻结的CLIP模型进行标记化，而图像观察则首先通过DINOv2模型提取图像块特征。为了提高计算效率，我们还采用了Q-Former与FiLM条件化来根据指令上下文选择图像特征。

在模型设计中，我们采用了DiT结构，该结构对动作标记块而不是单个动作标记进行去噪。这是通过因果变换器使用上下文条件化方法直接对图像观察和指令标记进行条件化来实现的。具体来说，我们将语言标记、图像特征和时间戳嵌入在序列开头进行连接，将带噪声的动作与指令标记一起处理。这种设计保留了Transformer网络的可扩展性，并允许去噪直接基于图像块进行条件化，从而使模型能够捕捉历史观察中动作的细微变化。

训练目标

去噪网络E_θ(c_lang, c_obs, t, x_t)是在因果变换器上构建的，其中c_obs表示图像观察，c_lang表示语言指令，t ∈ {1, 2, ..., T_train}表示总去噪步骤T_train内的时间戳索引。在训练过程中，我们在每个时间戳t处采样一个高斯噪声向量x_t ~ N(0, I)，并将其添加到动作a中以形成带噪声的动作标记^a。网络E_θ被训练为预测噪声向量^x_t，其中t是随机采样的。Dita的优化目标是最小化噪声向量x_t和预测噪声向量^x_t之间的均方误差（MSE）损失。

预训练数据和实现细节

我们选择OXE数据集对Dita模型进行预训练。预训练采用DDPM扩散目标，去噪步骤为1000步，而在零样本评估时采用DDIM，评估步骤为20步以加速推理。网络通过AdamW优化器优化100,000步，学习率为1e-4（对于因果变换器和Q-Former）和1e-5（对于DINOv2）。训练在32个NVIDIA A100 GPU上进行，批量大小为8192，每个GPU分配256个样本。

研究结果

模拟实验

我们在四个模拟基准上评估了预训练的Dita模型：SimplerEnv、LIBERO、CALVIN和ManiSkill2。在SimplerEnv上，Dita在零样本评估中展现出了强大的泛化性能，特别是在处理各种类型的变化（如背景、纹理、对象、空间位置等）时表现尤为突出。在LIBERO上，Dita在大多数子数据集上超越了基线方法，平均成功率提高了近6%，特别是在长时程任务（LIBERO-LONG）上实现了10%的改进。在CALVIN上，Dita实现了与仅使用单个静态RGB相机作为感知输入的方法相当的性能，特别是在处理长时程任务时表现出色。在ManiSkill2上，Dita在相机视角泛化基准上展示了优越的性能，特别是在处理复杂任务时表现尤为突出。

真实机器人实验

我们通过10次微调评估了Dita模型在未见过的机器人环境中执行复杂、长时程、多模态任务的能力。实验结果表明，Dita在两步任务上实现了63.8%的成功率，其中第二阶段贡献了近一半的成功率。与基线方法相比，Dita在所有复杂任务上均表现出色，特别是在处理需要极端3D旋转的任务时优势明显。此外，Dita还展示了在不同环境变化（如背景变化、非目标对象排列和照明条件）下的鲁棒性。

研究局限

尽管Dita模型在模拟和真实机器人实验中均表现出了强大的性能，但仍存在一些局限性。首先，Dita模型的规模相对较小（334M参数），这可能限制了其在处理极其复杂和多样化的任务时的泛化能力。通过增加模型规模，我们有望进一步提高Dita的性能和泛化能力。其次，Dita目前仅使用第三人称相机输入，这可能限制了其在某些需要更精细感知能力的任务中的表现。未来工作可以探索如何集成其他输入模态（如腕部相机图像、目标图像预测、机器人状态、触觉反馈等）以进一步提高模型的性能。

未来研究方向

扩展模型规模

增加Dita模型的规模是提高其性能和泛化能力的一个直接途径。通过采用更大的Transformer架构和更多的训练数据，我们有望训练出具有更强泛化能力的通用机器人策略。

集成多模态输入

除了第三人称相机输入外，未来工作还可以探索如何集成其他多模态输入（如腕部相机图像、目标图像预测、机器人状态、触觉反馈等）以进一步提高Dita的性能。这些额外的输入模态可以提供更丰富的感知信息，有助于模型更好地理解环境和任务需求。

增强鲁棒性和适应性

尽管Dita已经展示出了在不同环境变化下的鲁棒性，但未来工作仍可以进一步探索如何增强其鲁棒性和适应性。例如，可以通过引入更复杂的噪声模型和更先进的正则化技术来提高模型的泛化能力；同时，也可以通过采用元学习或在线学习等方法来使模型能够更好地适应新环境和任务。

应用于更复杂的任务和环境

目前，Dita主要在模拟环境和相对简单的真实机器人任务上进行了评估。未来工作可以探索如何将其应用于更复杂的任务和环境（如户外环境、动态环境、人机协作任务等）中。这将有助于进一步验证Dita的通用性和实用性，并推动其在现实世界中的应用。

综上所述，本研究提出了一种可扩展的扩散变换器框架Dita，用于构建通用视觉-语言-动作策略。通过利用Transformer架构的强大能力和统一的多模态扩散过程，Dita在模拟和真实机器人实验中均表现出了强大的性能和泛化能力。然而，仍存在一些局限性需要进一步探索和改进。未来工作可以围绕扩展模型规模、集成多模态输入、增强鲁棒性和适应性以及应用于更复杂的任务和环境等方面展开。