BAKU: An Efficient Transformer for Multi-Task Policy Learning

TL;DR

BAKU is a transformer-based architecture that significantly improves multi-task robot policy learning efficiency, achieving high success rates with minimal demonstrations across simulated and real-world tasks.

Contribution

The paper introduces BAKU, a novel transformer architecture that effectively leverages offline imitation learning techniques for multi-task robotic policy learning.

Findings

18% improvement over RT-1 and MT-ACT on simulated benchmarks

36% improvement on LIBERO benchmark

91% success rate on real-world tasks with few demonstrations

Abstract

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/.