BAKU: An Efficient Transformer for Multi-Task Policy Learning

📄 arXiv: 2406.07539v2 📥 PDF

作者: Siddhant Haldar, Zhuoran Peng, Lerrel Pinto

分类: cs.RO

发布日期: 2024-06-11 (更新: 2024-07-16)

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

BAKU：一种高效的Transformer架构，用于多任务策略学习

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多任务学习 机器人策略 Transformer 模仿学习 数据效率

📋 核心要点

  1. 现有方法在多任务机器人策略学习中面临数据效率低的挑战，尤其是在真实机器人任务中获取大量数据成本高昂。
  2. BAKU通过结合观察主干、动作分块、多感官观察和动作头等技术，构建了一个高效的Transformer架构。
  3. 实验结果表明，BAKU在多个模拟和真实机器人任务中显著优于现有方法，尤其是在数据有限的情况下。

📝 摘要（中文）

训练能够解决各种任务的通用智能体极具挑战性，通常需要大量的专家演示数据集。这在机器人领域尤其突出，因为每个数据点都需要在现实世界中实际执行动作。因此，迫切需要能够有效利用现有训练数据的架构。本文提出了BAKU，一种简单的Transformer架构，能够有效学习多任务机器人策略。BAKU建立在离线模仿学习的最新进展之上，巧妙地结合了观察主干、动作分块、多感官观察和动作头，从而显著优于先前的工作。在LIBERO、Meta-World套件和Deepmind Control套件的129个模拟任务上的实验表明，BAKU的总体性能比RT-1和MT-ACT提高了18%，在更困难的LIBERO基准测试中提高了36%。在30个真实世界的操作任务中，平均每个任务仅需17个演示，BAKU的成功率达到了91%。

🔬 方法详解

问题定义：论文旨在解决多任务机器人策略学习中的数据效率问题。现有方法，如RT-1和MT-ACT，在需要大量数据才能训练出通用智能体，这在机器人领域尤其困难，因为真实世界的数据采集成本很高。因此，如何利用有限的数据训练出高性能的多任务策略是本研究的核心问题。

核心思路：BAKU的核心思路是通过精心设计的Transformer架构，更有效地利用现有的训练数据。该架构结合了多种技术，包括观察主干（observation trunks）用于提取环境信息，动作分块（action chunking）用于处理连续动作，多感官观察（multi-sensory observations）用于融合不同传感器信息，以及动作头（action heads）用于预测最终的动作。通过这些技术的结合，BAKU能够更好地学习任务之间的共性，从而提高数据效率。

技术框架：BAKU的整体架构是一个基于Transformer的编码器-解码器结构。编码器接收来自环境的观察信息，包括视觉、触觉等多种传感器数据。观察主干负责从这些原始数据中提取特征。然后，这些特征被输入到Transformer编码器中进行处理。解码器接收编码器的输出，并预测一系列动作。动作分块技术将连续动作序列分成多个块，从而更好地处理时间依赖性。最后，动作头将解码器的输出转换为实际的机器人动作。

关键创新：BAKU的关键创新在于其架构的精心设计，它巧妙地结合了多种现有技术，并针对多任务机器人策略学习进行了优化。与现有方法相比，BAKU更加注重数据效率，能够在数据有限的情况下取得更好的性能。此外，BAKU的多感官观察能力使其能够更好地理解环境，从而提高策略的鲁棒性。

关键设计：BAKU的关键设计包括：1) 观察主干的具体结构，例如使用卷积神经网络提取视觉特征；2) 动作分块的大小，需要根据任务的特点进行调整；3) 多感官观察的融合方式，例如使用注意力机制对不同传感器数据进行加权；4) 损失函数的设计，例如使用模仿学习损失来鼓励智能体模仿专家行为。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

BAKU在129个模拟任务上取得了显著的性能提升，相比RT-1和MT-ACT，总体性能提高了18%，在更具挑战性的LIBERO基准测试中提高了36%。更重要的是，在30个真实世界的操作任务中，平均每个任务仅需17个演示，BAKU的成功率达到了91%。这些实验结果表明，BAKU在数据效率和泛化能力方面具有显著优势。

🎯 应用场景

BAKU具有广泛的应用前景，可用于各种机器人任务，如物体抓取、装配、导航等。该研究的实际价值在于降低了机器人学习的成本，使得在数据有限的情况下也能训练出高性能的机器人策略。未来，BAKU可以应用于工业自动化、家庭服务、医疗康复等领域，提高机器人的智能化水平和服务能力。

📄 摘要（原文）

Training generalist agents capable of solving diverse tasks is challenging, often requiring large datasets of expert demonstrations. This is particularly problematic in robotics, where each data point requires physical execution of actions in the real world. Thus, there is a pressing need for architectures that can effectively leverage the available training data. In this work, we present BAKU, a simple transformer architecture that enables efficient learning of multi-task robot policies. BAKU builds upon recent advancements in offline imitation learning and meticulously combines observation trunks, action chunking, multi-sensory observations, and action heads to substantially improve upon prior work. Our experiments on 129 simulated tasks across LIBERO, Meta-World suite, and the Deepmind Control suite exhibit an overall 18% absolute improvement over RT-1 and MT-ACT, with a 36% improvement on the harder LIBERO benchmark. On 30 real-world manipulation tasks, given an average of just 17 demonstrations per task, BAKU achieves a 91% success rate. Videos of the robot are best viewed at https://baku-robot.github.io/.