BAKU: An Efficient Transformer for Multi-Task Policy Learning

TL;DR

BAKU以约1000万参数融合多模态、动作分块与Transformer,在LIBERO-90达90%,真实机器人达91%。

cs.RO 🟡 进阶级 2024-06-12 24 次浏览
Siddhant Haldar Zhuoran Peng Lerrel Pinto
机器人学习 多任务策略 Transformer 模仿学习 动作分块

核心发现

方法论

BAKU将策略分为感知编码器、观测主干和动作头。图像使用FiLM条件化ResNet-18,机器人本体状态使用两层MLP,文本使用6层MiniLM;各模态投影到同一维度后输入因果Transformer解码器。模型预测长度为10的动作块,并用指数时间平滑;动作头可替换为MLP、GMM、BeT、VQ-BeT或Diffusion。

关键结果

  • 在129个模拟任务上,BAKU总体较RT-1和MT-ACT提升18个百分点;LIBERO-90(90任务)成功率0.90,较最佳基线提升36个百分点;Meta-World为0.79,DMC为0.70。
  • 真实xArm实验包含30个厨房操作任务、520条示范,平均每任务17条。BAKU使用MLP动作头达到86%成功率;改用VQ-BeT后达到91%,较最强基线提升35个百分点。
  • 消融显示,LIBERO中动作分块使成功率由0.76升至0.90;Transformer主干较MLP为0.90对0.81;114M模型在LIBERO仅0.19,提示大模型可能过拟合。

研究意义

BAKU回应了机器人数据昂贵、示范稀缺而多任务策略通常弱于单任务策略的长期难题。它表明,性能提升未必来自堆叠超大模型,而可能来自对观测融合、时间建模和动作生成方式的精确组合。对学术界而言,论文提供了清晰的模块化比较框架;对工业界而言,平均17条示范即可训练30任务统一策略,降低了真实机器人数据采集门槛。

技术贡献

核心贡献是模块化的约1000万参数架构:约2.1M用于感知编码器、6.5M用于观测主干、1.4M用于动作头。因果Transformer把多视角图像、本体状态和语言作为观测token,并用动作token输出控制信号;独立动作头允许直接接入VQ-BeT等生成模型。论文还证明,将动作块作为单一拼接向量预测,比逐时刻解码更有效。

新颖性

BAKU并非提出全新的Transformer或损失函数,而是首次系统地把FiLM视觉条件化、因果观测主干、动作分块、多模态动作头和时间平滑整合为面向低数据多任务机器人的简洁框架。相较RT-1的离散动作分类和MT-ACT的编码器—解码器设计,BAKU更强调可替换模块及经验验证。

局限性

  • 真实实验仅覆盖xArm 7、固定初始机器人姿态和有限物体位置,且每任务仅5次评估,尚不能代表开放家庭环境中的泛化能力。
  • 研究主要依赖离线行为克隆;面对示范分布之外的状态、遮挡、动态障碍和长时间误差累积,策略仍可能产生协变量偏移与失败。

未来方向

后续可扩大机器人、场景和对象分布,研究跨本体迁移、在线纠错与不确定性估计;也可结合更强视觉基础模型、扩散动作头和人类视频数据。作者已计划公开数据、训练代码和评测代码,这将便于复现及系统化比较。

AI 总览摘要

让一个机器人同时完成开门、取物、擦桌等任务,难点不只是动作复杂,而是每条示范都要通过真实世界采集。视觉语言领域可以从互联网获得海量数据,机器人却常被限制在少量昂贵示范中。传统多任务策略因此经常不如分别训练的单任务策略。

纽约大学团队提出BAKU,将问题拆成感知编码器、观测主干和动作头。它用FiLM条件化ResNet-18处理视觉,用MiniLM编码指令,用MLP处理本体状态,再由因果Transformer融合多模态和时间信息。模型不是只预测下一步,而是一次预测一段动作,并通过指数平滑形成更连贯的控制;动作头还可替换为VQ-BeT等多峰生成模型。

在LIBERO-90、Meta-World和DeepMind Control的129个模拟任务上,BAKU总体比RT-1和MT-ACT高18个百分点,LIBERO达到90%,提升36个百分点。真实xArm在30项厨房任务、平均每任务17条示范下,MLP版本成功率86%,VQ-BeT版本达到91%。结果说明,精心设计的中等规模模型能够有效利用稀缺数据。不过实验仍集中于单一机器人和受控环境,未来需要验证开放世界泛化、在线恢复和跨机器人迁移。

深度分析

研究背景

机器人多任务学习借鉴了RT-1、MT-ACT等Transformer策略,也吸收了GMM、BeT、VQ-BeT和Diffusion等动作生成方法。问题在于机器人数据无法像图文数据一样廉价扩张:示范需要时间、设备和操作者。已有大数据路线成本高,多任务策略仍常显著落后于单任务策略,因此需要更高数据效率的架构。

核心问题

给定离线专家轨迹,策略需根据图像、本体状态及任务目标预测动作。挑战包括多模态信息对齐、任务间共享与区分、动作的时间相关性,以及低数据条件下的分布外状态。单步回归容易产生抖动和误差累积;复杂动作头也可能在有限数据下难以训练。

核心创新

  • �� 用FiLM让语言调制ResNet-18视觉特征,形成任务相关表示。
  • �� 用因果Transformer观测主干,将各模态视为token并以动作token预测控制。
  • �� 采用动作分块:LIBERO和Meta-World预测未来10步,DMC预测3步,再做指数时间平滑。
  • �� 将动作头与观测编码解耦,可插入MLP、GMM、BeT、VQ-BeT和Diffusion。
  • �� 系统消融模型规模、目标形式、历史观测和多步损失。

方法详解

  • �� 输入:多摄像头RGB、机器人本体状态和文本或目标图像。
  • �� 编码:ResNet-18提取视觉特征,FiLM由任务语言调制;两层MLP编码状态;6层MiniLM编码文本。
  • �� 融合:所有特征映射到统一维度,组成观测token;因果Transformer decoder只访问当前及过去信息。
  • �� 预测:可学习动作token产生动作特征,动作头输出连续动作块。
  • �� 训练:使用行为克隆均方误差,基本形式为L=E[||a-π(o|g)||²];历史输入时可对所有时间步计算多步损失。
  • �� 控制:以10Hz部署,并对重叠动作块进行指数时间平均。

实验设计

模拟实验覆盖LIBERO-90的90任务、Meta-World的30任务和DMC的9个状态控制任务。示范数量分别为每任务50、35和500;图像分辨率为LIBERO 128×128、Meta-World 84×84,每任务评估10次。真实实验使用xArm 7、四个128×128 RGB视角、520条示范和30任务,平均17条;30Hz采集、10Hz部署,每任务5次。基线为RT-1和MT-ACT。

结果分析

BAKU在LIBERO-90、Meta-World、DMC分别达到0.90、0.79、0.70,而RT-1为0.16、0.65、0.66,MT-ACT为0.54、0.13、0.59。长时程LIBERO-10为0.86,真实五项长任务为0.84,均高于MT-ACT的0.68和0.64。真实30任务中,MLP头为86%,VQ-BeT为91%。历史观测若只监督末步会严重退化,多步监督可平均提升47%,但最终未改善整体策略。

应用场景

BAKU适合低示范成本的多任务桌面操作、厨房协作、仓储取放和实验室自动化。部署前需固定或校准摄像头、机器人状态接口、任务指令格式及动作空间,并准备少量覆盖各任务的遥操作示范。其模块化动作头便于根据场景在简单MLP与多峰VQ-BeT之间选择。

局限与展望

实验主要在模拟器和单一xArm厨房中完成,物体位置、初始姿态及评估次数有限;论文没有证明跨机器人、跨场景或真实开放世界的零样本能力。行为克隆仍受协变量偏移影响,遮挡、动态物体和长时程失败恢复尚未充分处理。114M模型在LIBERO的0.19也显示,扩大容量并不必然提升数据效率。

通俗解读 非专业人士也能看懂

把机器人想成一个只有少量培训视频的新厨师。它先看菜单,也就是任务指令;再看厨房里的多个摄像头,观察锅、瓶子和自己的手;还要知道手臂现在在哪里。BAKU像一个总协调员,把这些信息放到同一张工作台上,而不是让每个摄像头各自做决定。

普通方法常常只说“下一步移动一点”,像厨师每秒都重新猜动作,结果容易抖动。BAKU会一次规划接下来十个小动作,再把相邻计划平滑地混合起来,所以手臂更像连续完成“伸手—抓住—移动—放下”。如果同一任务存在多种合理做法,VQ-BeT动作头还能保留多种可能,而不是强行平均成一个奇怪动作。

它最重要的地方是节省培训材料。模拟实验中,LIBERO-90成功率达到90%;真实机器人平均每项任务只有17条示范,仍达到91%。但它还像在熟悉的厨房里考试:换机器人、换房间或突然有人挡住视线,表现仍需进一步验证。

简单解释 像给14岁少年讲一样

想象你在玩一个机器人游戏:任务可能是“打开烤箱”“把番茄酱放进冰箱”或“从架子拿杯子”。机器人不能只看一张照片,因为它还要知道自己的手臂位置,并理解文字任务。BAKU就像一个很会观察的队友,把多个摄像头、手臂状态和任务说明一起读懂。

它的聪明之处是不会每次只决定一个动作。比如拿杯子时,它会连续想好“伸手、靠近、抓住、抬起、收回”这一小串动作,然后边做边修正。这叫动作分块,但你可以把它理解成游戏里的连招。动作之间还会被平滑处理,所以机器人不会一顿一顿地抽动。

BAKU大约只有1000万参数,却在129个模拟任务上比RT-1和MT-ACT总体高18个百分点。在更难的LIBERO-90中,它达到90%成功率。真实xArm机器人用平均17条示范学习每个任务,VQ-BeT版本完成率达到91%,非常像只看少量攻略就学会一整套厨房关卡。

当然,它还不是万能机器人。实验多在一个xArm和受控厨房完成;如果房间改变、物体被遮住,或者动作失败后需要自己补救,它可能会迷糊。下一步研究要让它适应更多机器人和真实家庭环境。

术语表

Behavior Cloning(行为克隆)

从专家示范中直接学习观测到动作的映射。论文采用均方误差目标L=E[||a-π(o|g)||²]。

BAKU的基本训练范式,不使用环境奖励。

Action Chunking(动作分块)

一次预测连续多个未来动作,而非只预测下一步。它利用动作的时间相关性并减少控制抖动。

LIBERO和Meta-World预测10步,DMC预测3步。

FiLM Conditioning(FiLM条件化)

根据任务信息调制视觉特征,通常通过缩放和偏移中间表示实现。它使同一视觉输入可产生任务相关解释。

语言条件化ResNet-18视觉编码器。

VQ-BeT

结合向量量化与Behavior Transformer的多模态动作生成器。它能表达同一情境下多种合理动作。

真实机器人中将成功率从86%提高到91%。

Observation Trunk(观测主干)

融合各传感器编码结果并形成动作表示的网络部分。BAKU使用因果Transformer decoder。

接收视觉、状态、文本token并输出动作token特征。

开放问题 这项研究留下的未解疑问

  • 1 BAKU能否在未见过的房间、物体和机器人本体上迁移仍未知。现有实验主要是受控分布,需要跨域数据和系统化零样本评估。
  • 2 行为克隆遇到失败状态时缺少主动恢复机制。未来需结合在线交互、不确定性估计或安全规划,减少长时程误差累积。

应用场景

近期应用

多任务厨房机器人

研究团队或机器人公司可用少量遥操作示范训练取放、擦拭、开关门等统一策略。前提是提供多视角RGB、本体状态和清晰任务指令,并对动作空间及摄像头完成标定。

仓储与实验室操作

在物体种类有限但任务较多的工作站中,BAKU可统一处理抓取、搬运和放置。建议先使用MLP动作头建立基线,再以VQ-BeT处理存在多种操作路径的场景。

远期愿景

通用家庭机器人

若结合更大规模视觉语言数据、跨机器人训练和在线纠错,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/.

cs.RO