AI Research 技术调研

BAKU: An Efficient Transformer for Multi-Task Policy Learning

14分钟阅读

一句话定位

NYU（Siddhant Haldar、Zhuoran Peng、Lerrel Pinto）提出的多任务模仿学习策略架构 BAKU：把「感知编码器 + 观测 trunk（因果 Transformer）+ 独立动作头」拆成三段可插拔模块，系统消融多任务策略学习中常见的设计选择（trunk 类型、模型规模、动作头、动作分块、观测历史、目标模态、FiLM），拼出一个约 10M 参数的极简架构；在 LIBERO-90 / Meta-World / DMC 共 129 个仿真任务上比 RT-1、MT-ACT 平均高 18%（LIBERO-90 上高 36%），在 xArm 真机 30 个任务、人均仅 17 条演示下做到 91% 成功率。NeurIPS 2024。

背景与定位

多任务模仿学习长期存在「多任务策略性能明显低于同等规模的单任务专家策略」的差距（论文引用 gradient surgery、actor-mimic、policy distillation 等工作佐证这一现象）。既有代表性架构如 rt-1（离散动作 token + FiLM ResNet + TokenLearner 压缩视觉 token）与 roboagent-mt-act（MT-ACT，ACT 的多任务扩展，CVAE 编码器-解码器架构）各自把观测编码、时序建模、动作生成耦合在一个整体设计里，导致「究竟哪个组件真正贡献了性能」难以厘清。

BAKU 的定位不是提出全新算法，而是做一次系统性架构解耦与消融：把 sensory encoder、observation trunk、action head 拆成三段独立可换的模块，逐项消融组件选择，最终拼出一个同时吸收 transformer trunk + FiLM 条件化视觉编码器（类似 octo 的做法）与动作分块（action chunking，源自 MT-ACT / ACT）优点、又能即插即用当代动作生成头（GMM / BeT / VQ-BeT / diffusion）的极简架构。评测基准覆盖操作任务基准 libero（LIBERO-90 / LIBERO-10）与 meta-world，以及 DeepMind Control Suite 的状态输入运动控制任务。

模型架构

三段式设计：感知编码器 → 观测 trunk → 动作头（Figure 2）。

  • 感知编码器：视觉用 FiLM-conditioned ResNet-18（LIBERO 官方提供的改版）；本体感受（proprioception）用两层 MLP；语言指令用 Sentence-Transformers 提供的 6 层 MiniLM。各模态编码后经额外 MLP 投影到同一维度。多相机视角下共用同一个视觉编码器（而非每视角独立编码器）。
  • 观测 trunk（causal transformer decoder）：每个模态的编码作为一个 observation token，附加一个可学习 action token，一起送入基于 minGPT 实现的因果 Transformer 解码器（8 层、4 头）；因果 mask 保证动作预测只依赖过去观测。论文也实现了 MLP trunk 变体做对比（消融见下）。最终架构不使用观测历史，只用当前时刻观测。
  • 动作头（可插拔，五种变体）：base 模型用两层 MLP 头；另评测 GMM（两层，5 个模式，Softplus 激活）、BeT（Behavior Transformer，64 个 k-means 动作聚类 + focal loss 离散头 + L2 offset 头）、VQ-BeT（两层 residual VQ，codebook size 16、latent dim 256）、Diffusion（两层 transformer-based diffusion 头）。
  • 动作分块（action chunking）+ 指数时间平滑：不同于 MT-ACT/ACT 逐时间步解码，BAKU 把一个 chunk 的动作预测为单个拼接向量。推理时每步都查询策略（而非每 k 步查询一次），对重叠的历史 chunk 预测做指数加权时间集成 $w_i = \exp(-m \cdot i)$ 以获得平滑动作，该集成只增加推理时计算量、不增加训练成本。
  • FiLM 条件化：文本指令编码后作为条件向量，通过 $\text{FiLM}(x) = \gamma(z) \odot x + \beta(z)$ 对视觉编码器做逐通道仿射调制，让视觉表征随任务变化。

参数量（base MLP 头配置）：感知编码器 ≈2.1M，观测 trunk ≈6.5M，动作头 ≈1.4M，总计 ≈10M。隐藏维度 256；图像输入 128×128（LIBERO-90、xArm）或 84×84（Meta-World）。

模型规模消融覆盖 4.4M / 10M / 31M / 114M 四档（同时改变 trunk 与动作头大小），114M 在 LIBERO-90 上从 0.90 崩到 0.19（怀疑过拟合），最终选用 10M 档。

数据

  • LIBERO-90（libero）：90 个操作任务，官方提供的 50 条演示/任务，输入为第三人称 + 夹爪相机两个视角 + 本体感受，图像 128×128。
  • Meta-World（meta-world）：30 个操作任务，35 条演示/任务，来自用演示引导强化学习（demonstration-guided RL）训练出的专家策略，仅用第三人称视角，图像 84×84。
  • DeepMind Control Suite (DMC)：9 个运动控制任务，500 条演示/任务，来自 DrQ-v2 训练的专家，纯状态输入（无视觉）。
  • xArm 真机（Ufactory xArm 7 + 夹爪，多任务厨房环境）：30 个操作任务，共 520 条演示，人均 17 条；4 个相机视角（含固定在夹爪上的第一人称相机）；动作空间为末端执行器位姿 + 夹爪状态；演示通过 VR 遥操作系统（Open Teach）以 30Hz 采集，策略以 10Hz 闭环部署。
  • 长程任务：LIBERO-10（10 个任务，50 条演示/任务）；真机长程任务 5 个（每个由两个短程技能顺序拼接，人均 19 条演示）。
  • 目标模态消融数据：文本指令（任务元数据自带）、目标图像（随机采样一条演示的最后一帧）、中间目标图像（同一演示中未来第 k 步的帧，LIBERO-90 取 k=50、Meta-World 取 k=30）。
  • 未使用大规模跨本体预训练数据混合——各基准分别从零训练，强调的是「数据效率」而非「数据规模」。

训练方法

  • 目标函数：标准多任务行为克隆（behavior cloning），对专家动作做 MLE 回归（L2 损失），按文本或图像目标条件化。
  • 优化器：Adam，学习率 1e-4（各基准通用）。
  • Batch size：64（LIBERO-90、Meta-World、xArm），128（DM Control）。
  • 训练步数：LIBERO-90 / Meta-World 各 600k 步；DM Control 2M 步；xArm 200k 步。
  • 动作分块 + 指数时间平滑联合使用（借鉴 ACT/MT-ACT），论文强调二者对真机上产生精确且平滑的运动都是必要的，且平滑集成不增加训练成本。
  • 无 RL、无蒸馏——各基准均为单阶段纯 BC 训练。

Infra（训练 / 推理工程）

  • 训练硬件：单张 NVIDIA RTX A4000 GPU（未披露多卡/集群并行）——刻意保持极低算力需求，是论文「efficient」主张的一部分。
  • 单卡训练耗时（Appendix A.4）：
    • LIBERO-90：600k 步、batch 64、2 相机视角 + 本体感受 → 约 10.5 小时
    • Meta-World：600k 步、batch 64、1 相机视角 → 约 8 小时
    • DM Control：2M 步、batch 128、状态输入 → 约 26 小时
    • xArm 真机：200k 步、batch 64、4 相机视角 + 本体感受 → 约 6 小时
  • 推理：真机闭环控制策略查询频率 10Hz，另有一个独立的低层 Minimum-Jerk 控制器以 100Hz 运行以平滑运动；论文未披露 GPU 推理延迟 / FPS 等具体数字。
  • 精度 / 并行策略：未披露。

评测 benchmark

Table 1（仿真每任务 10 次 rollout，真机每任务 5 次评测；成功率）：

MethodLIBERO-90 (90 tasks)Meta-World (30 tasks)DMC (9 tasks)Real Robot
RT-10.160.650.660.37
MT-ACT0.540.130.590.56
BAKU（MLP 头）0.900.790.700.86
BAKU w/ VQ-BeT0.900.780.700.91

（原文 Table 1 表头把 Real Robot 列标注为「20 tasks」，但正文 4.2 节明确写「30 manipulation tasks… 5 次评测/任务，共 150 次评测/方法」，附录 E.1 的逐任务表也列出 30 个真机任务——判断为原文表头笔误，本页按正文与附录一致的 30 任务口径记录。）

Table 2（长程任务成功率）：

MethodLIBERO-10 (10 tasks)Real Robot 长程 (5 tasks)
MT-ACT0.680.64
BAKU0.860.84

Table 3 消融（LIBERO-90 / Meta-World / DMC 成功率）：

  • 观测 trunk：MLP 0.81/0.78/0.68 vs Transformer 0.90/0.79/0.70（Transformer 在 LIBERO-90 上领先 9%）
  • 模型规模：4.4M→0.85/0.78/0.68，10M→0.90/0.79/0.70，31M→0.87/0.81/0.70，114M→0.19/0.81/0.68（114M 在 LIBERO-90 上崩溃）
  • 动作头：MLP 0.90/0.79/0.70，GMM 0.84/0.65/0.67，BeT 0.89/0.78/0.60，VQ-BeT 0.90/0.78/0.70，Diffusion 0.89/0.45/0.61（仿真上 MLP 与 VQ-BeT 相当或更优；真机上 VQ-BeT 反超 MLP 5%，见 Table 1）
  • 动作分块：不用 → 0.76/0.78/0.74（DMC 上反而更好）；用 → 0.90/0.79/0.70（LIBERO-90 提升 14%，Meta-World 几乎无差，DMC 下降 4%）
  • 观测历史：无历史 0.90/0.79/0.70；历史 + 仅末步损失 0.54/0.08/0.37（严重退化）；历史 + 多步损失 0.90/0.82/0.68（多步损失比仅末步损失平均高 47%，但仍未超过无历史基线）
  • 目标模态（LIBERO-90/Meta-World，DMC 不适用）：文本 0.90/0.79，目标图像 0.88/0.81，中间目标图像 0.91/0.80——三者大致相当
  • FiLM：关闭 0.87/0.79（DMC 不适用）；开启 0.90/0.79——持平或更优，最终采用

Appendix E.2 补充消融（LIBERO-90 / Meta-World / DMC）：

  • 独立 vs 共享视觉编码器（仅 LIBERO-90 两视角场景）：共享 0.90，独立 0.92（+2%，但每多一个视角参数量增加 15%——真机 4 视角场景下代价更大，故 BAKU 采用共享编码器以保持模型紧凑、加快推理）
  • 观测 trunk 输入格式：分离 token（各模态独立送入 trunk）0.90/0.79/0.70 vs 单一拼接向量 0.87/0.79/0.70（LIBERO-90 上分离 token 领先 3%，单一输入源的 Meta-World/DMC 无差异）

Baseline 实现细节：RT-1（论文自行复现，动作均匀离散化分箱，FiLM-ResNet-18 + TokenLearner 把 k×k×512 特征图压成 8 个 512 维 token）；MT-ACT（开源实现，默认超参，各基准分别调节动作分块长度）。

创新点与影响

  • 把「感知编码器 / 观测 trunk / 动作头」解耦为三段独立可换的模块，对多任务策略学习中的常见设计选择（trunk 类型、模型规模、动作头、分块、历史、目标模态、FiLM）做了系统性受控消融，把此前分散在 RT-1、MT-ACT、BeT、VQ-BeT、Diffusion Policy 等工作里的最佳实践统一到一个框架内直接对比。
  • 若干反直觉发现：观测历史本身不提升性能（多步损失只是修复了朴素历史用法的坑，最终架构反而只用当前观测）；模型越大不代表越好（114M 在 LIBERO-90 上从 0.90 崩到 0.19，疑似过拟合）；简单 MLP 动作头在仿真数据上不输甚至优于更复杂的多模态头，只有部署到真机时多模态头（VQ-BeT）才体现优势（+5%）。
  • 极简、低算力路线：整套方法在单张 RTX A4000 上数小时到一天内即可训完各基准，不依赖大规模跨本体预训练数据，仍在 129 个仿真任务上对 RT-1/MT-ACT 平均领先 18%（LIBERO-90 领先 36%），真机 30 任务人均 17 条演示做到 91% 成功率。
  • 作者自陈局限：(a) 真机上在少数需要精细操作的任务（如开/关烤箱门、把茶瓶放进冰箱）表现不佳，提示跨难度任务共享数据可能损害精细技能，需要专门方法处理不同复杂度任务的统一策略学习；(b) 目前只处理单技能执行，未研究把多个技能链接起来完成长程任务的算法；(c) 只研究了策略架构本身，未分析任务数量增长是否会带来涌现式的泛化收益。

原始链接

一手源存档（sources/）

关系图谱

反向链接