PaperlayerView exact revision on arXiv

Source-linked AI summary

BAKU: An Efficient Transformer for Multi-Task Policy Learning

Siddhant Haldar, Zhuoran Peng, Lerrel Pinto

arXiv:2406.07539v2cs.RO

TL;DR

Robot generalist policies are difficult to train when demonstrations require costly physical interaction and multi-task policies may underperform single-task ones. BAKU combines multimodal encoders, a transformer observation trunk, and an action head for multi-task policy learning. It reports an 18% absolute improvement over prior multi-task methods across 129 simulated tasks.

  • Problem

    Robot demonstrations require physical interaction, while multi-task policies trained on collected data can perform below independently trained single-task policies.

    Source: §1 · p. 1

  • Method

    BAKU combines multimodal sensory encoders, a transformer observation trunk, and an action head, with action chunking for multi-task policy learning.

    Source: §1 · p. 1, §3 · p. 3, §3.3 · p. 4

  • Results

    18% absolute improvement over prior state-of-the-art multi-task algorithms across 129 simulated tasks, including a 36% absolute improvement on LIBERO.

    Source: §1 · p. 2

  • Takeaways & Limitations

    With an average of 17 demonstrations per task, BAKU achieved 91% average success across 30 real-world tasks.

    Source: §1 · p. 2

    The study focuses on one skill at a time, and BAKU struggled with some precise real-world manipulation tasks.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

1 Introduction

Robotics makes generalist policy learning difficult because collecting demonstrations requires costly physical interaction, while multi-task policies trained on available data often lag single-task policies. BAKU combines a transformer architecture with multi-task learning and reports gains across simulated and real-world tasks.

Motivation

  • Collecting robot data requires physical interaction, making it more laborious in time and cost than amassing large vision and language datasets.Multi-task policies trained on large demonstration datasets can still perform below independently trained single-task policies.Source: §1 · p. 1

BAKU

  • BAKU combines a multimodal transformer encoder with temporal context, a task-conditioned visual encoder, and an action head separate from the observation trunk.The architecture is presented as a combination of existing ideas suited to multi-task decision making.Source: §1 · pp. 1–2

Main findings

  • 18% absolute improvement over prior state-of-the-art multi-task algorithms across 129 simulated tasks, with a 36% absolute improvement on LIBERO and 91% average success across 30 real-world tasks.The real-world result used an average of 17 demonstrations per task.Source: §1 · p. 2

Resources

  • The authors state that datasets, training code, and evaluation code will be made publicly available.Videos of trained policies are also linked from the paper.Source: §1 · p. 2

2 Background

The paper frames imitation learning around learning from expert-policy trajectories, using observation-based demonstrations as a feasible source of data for robotics. Its multi-task behavior-cloning formulation conditions action prediction on a goal.

Imitation Learning

  • The paper uses observation-based trajectories from expert or near-expert demonstrators, a setting described as feasible for real-world data collection.The trajectories contain observations and actions.Source: §2 · pp. 2–3

Multi-task Behavior Cloning

  • Behavior cloning trains a policy to mimic demonstrated actions from observations, with the stated fixed-variance Gaussian formulation cast as regression.The single-task objective is given in Equation 1.Source: §2 · p. 3
  • For multi-task behavior cloning, action prediction is conditioned on a goal variable represented as a text task description or goal image.The paper uses the same behavior-cloning formulation with this goal conditioning.Source: §2 · p. 3

3 BAKU

BAKU organizes multi-task policy learning into sensory encoders, an observation trunk, and an action head. The design supports multiple sensory modalities and action-generation choices, with action chunking and temporal smoothing in the described policy.

3 BAKU

  • BAKU divides multi-task policy learning into sensory encoders, an observation trunk, and an action head, and studies architectural choices through ablations.The encoders map raw inputs to features, the trunk combines them, and the head predicts actions.Source: §3 · p. 3

3.1 Sensory Encoders

  • Vision, proprioception, and task instructions are encoded with a FiLM-conditioned ResNet-18, a two-layer MLP, and a six-layer MiniLM text encoder, respectively.The representations are projected to a shared dimensionality for combination.Source: §3.1 · pp. 3–4

3.2 Observation Trunk

  • The observation trunk either concatenates encoded inputs for an MLP or treats each encoded input as a token in a causal transformer decoder.The transformer uses action tokens to predict actions from past observations.Source: §3.2 · p. 4

3.3 Action Head

  • The action head maps trunk features to actions; BAKU tests five head types and predicts each action chunk as one concatenated vector with exponential temporal averaging.The authors report that this chunking simplification improves performance and that temporal averaging produces smoother behavior and counters covariate shift often seen in low-data imitation learning.Source: §3.3 · p. 4

3.4 Putting it all together

  • The described base model uses a modified FiLM-conditioned ResNet-18, an MLP proprioception encoder, a pretrained text encoder, a causal transformer decoder, and an MLP action head.It uses the current observation as input; the model has approximately 10M parameters in total.Source: §3.4 · p. 4

4 Experiments

Experiments compare BAKU with RT-1 and MT-ACT across simulated benchmarks and real-world robot tasks, then assess long-horizon performance and individual design choices.

4.1 How well does BAKU work for multi-task learning?

  • 36% on LIBERO-90, 14% on Meta-World, and 4% on DMC are BAKU’s reported gains over the strongest baseline.The comparison covers 90 LIBERO-90 tasks, 30 Meta-World tasks, and 9 DMC tasks.Source: §4.1 · pp. 5–6, Table 1 · p. 6

4.2 How does BAKU perform on real-world tasks?

  • 91% success with a VQ-BeT head, versus 86% with an MLP head, was achieved across 30 real-world manipulation tasks.The xArm evaluations used five runs per task, with a fixed robot pose and fixed object positions.Source: §4.2 · p. 6

4.3 How does BAKU perform on long-horizon tasks?

  • 19% higher average success than MT-ACT was reported on long-horizon tasks in LIBERO-10 and the real-world kitchen.The real-world tasks chained two shorter tasks sequentially and used an average of 19 demonstrations per task.Source: §4.3 · p. 6, Table 2 · p. 6

4.4 What design decisions affect multi-task policy learning?

  • A transformer trunk outperformed an MLP trunk by 9% on LIBERO-90; 4.4M, 10M, and 31M models performed similarly, while 114M underperformed on LIBERO-90.The authors suspect overfitting may explain the largest model’s result and selected the 10M model as the smallest with the best performance on two of three simulated benchmarks.Source: §4.4 · p. 8
  • On simulated benchmarks, the MLP head performed as well as or better than multimodal heads; on the real robot, VQ-BeT reached 91%, 5% above MLP.VQ-BeT was the best-performing multimodal variant in simulation.Source: §4.4 · p. 8
  • Removing action chunking dropped LIBERO-90 performance by 14%, made no perceptible difference on Meta-World, and raised DMC performance by 4%.The experiments used chunks of 10 future actions for LIBERO-90 and Meta-World, and 3 for DMC.Source: §4.4 · p. 9
  • Multi-step action prediction over observation history improved single-step-loss performance by 47% on average, but history did not noticeably improve overall policy performance.Naively using history with loss only at the last step significantly degraded performance; the final architecture uses only the most recent observation.Source: §4.4 · p. 9
  • Text instructions, goal images, and intermediate goal images had similar performance, while FiLM-conditioned vision performed equally well or better than unconditional vision on image-based benchmarks.Goal images were sampled from demonstrations; intermediate goals used a future frame. The FiLM encoder was used for image-based experiments.Source: §4.4 · p. 9

5 Related Work

Robotics has advanced in single-task policy learning, but multi-task performance still trails single-task policies. BAKU combines ideas from prior work in a simpler architecture for multi-task control.

5 Related Work

  • Despite progress in single-task robotic policy learning, a gap remains between single-task and multi-task policy performance [44] [57].The literature includes behavior cloning, inverse reinforcement learning, multimodal action generation, and long-horizon imitation learning.Source: §5 · p. 9
  • Multi-task robotics spans grasping, language understanding, goal-reaching, and diverse datasets, with transformer policies increasingly used for navigation, locomotion, and manipulation.Some policies condition on goal images or videos, and some co-train with visual question answering or image captioning.Source: §5 · pp. 9–10
  • Multi-task policy architectures can be complicated, and the importance of individual components is sometimes unclear; BAKU combines prior ideas in a simpler model.The paper characterizes high-capacity transformer control policies as increasingly common in robotics.Source: §5 · p. 10

6 Conclusion and Limitations

BAKU is presented as an improved multi-task policy architecture, with appendices describing FiLM conditioning, action-head options, temporal smoothing, and training settings.

6 Conclusion and Limitations

  • BAKU improves multi-task policy learning across simulated and real-world domains, but the paper identifies limits in precise skills, skill chaining, and analysis of generalization as task diversity grows.The real-world policy struggled on some precise manipulation tasks; it focuses on one skill at a time and does not analyze how generalization changes with increasing task count.Source: §6 · p. 10

A.1 FiLM Conditioning

  • FiLM applies feature-wise scaling and shifting to network activations using a conditioning input such as text instructions or observations.The paper describes this as adapting network computation and output based on the instructions.Source: §A.1 · p. 16

A.2 Action Heads

  • A separate action prediction module lets BAKU use different action-generation methods, including MLP, GMM, BeT, VQ-BeT, and diffusion heads.The described variants range from mixture-based multi-modal sampling and clustered action prediction to iterative denoising; diffusion inference is slower due to denoising iterations.Source: §A.2 · p. 16

A.3 Temporal smoothing over action chunking

  • Temporal ensembling combines overlapping action-chunk predictions with exponential weights, supporting smooth and precise motion without extra training cost.The policy is queried each timestep; ensembling adds inference-time computation, and both chunking and ensembling were found important for precise, smooth motion.Source: §A.3 · pp. 16–17

A.4 Hyperparameters

  • Training settings vary by benchmark, with the reported runs using 600k steps for LIBERO and Meta-World, 2M for DM Control, and 200k for xArm.On a single NVIDIA RTX A4000, reported training times were around 10.5, 8, 26, and 6 hours, respectively; the appendix also gives baseline implementation settings.Source: §A.4 · p. 17

B Simulation Tasks

The simulation evaluation covers LIBERO-90, Meta-World, and DM Control, using benchmark-provided or expert-generated demonstrations.

B Simulation Tasks

  • BAKU is evaluated on LIBERO-90, Meta-World, and DM Control; the study uses 90 LIBERO tasks, 30 Meta-World tasks, and 9 DM Control tasks.The LIBERO-90 dataset provides demonstrations for all 90 tasks; Meta-World and DM Control demonstrations come from expert RL agents, and only tasks with available expert data are included.Source: B Simulation Tasks · p. 17

C Robot Tasks

The real-world evaluation uses 30 tasks in a multi-task kitchen environment, with deployment controlled at separate high- and low-level rates.

C Robot Tasks

  • BAKU is evaluated on 30 tasks in the real-world multi-task kitchen environment, with deployment rollouts supplied for individual tasks and long-horizon tasks.The paper points to Figures 5–9 for task rollouts and Figure 10 for long-horizon rollouts.Source: C Robot Tasks · p. 17
  • Policies run at 10Hz through a high-level controller, while a low-level Minimum-Jerk Controller runs at 100Hz to facilitate smooth robot motion.Source: C Robot Tasks · p. 17

D Baselines

The section contrasts BAKU with MT-ACT and RT-1, highlighting architectural differences and reporting that BAKU outperforms RT-1 in the experiments. It also includes rollout figures for complex manipulation tasks.

D Baselines

  • MT-ACT encodes camera views, proprioception, and instructions separately, then uses a transformer encoder-decoder to predict action chunks.A CVAE style variable conditions its encoder; at inference, that variable is set to zero, producing unimodal behavior.Source: D Baselines · pp. 17–19
  • Unlike MT-ACT, BAKU uses a decoder-only transformer to predict action features for past observations and supports multimodal behavior at inference.The passage links multimodal action heads to improved real-world performance.Source: D Baselines · pp. 17–19
  • Figures 5–8 show BAKU policy rollouts in complex manipulation tasks.Table 5 lists tasks in Meta-World and DM Control; Table 4 lists hyperparameters.Source: Table 4 · p. 18, Table 5 · p. 19, Figure 5 · p. 20
  • BAKU’s design choices lead to significant performance improvements over RT-1 in the experiments.RT-1 discretizes continuous actions and uses a Token Learner to reduce visual tokens; BAKU uses final vision features and predicts continuous actions.Source: D Baselines · pp. 19–21

E.1 Real-World Task-wise Results

This section presents task-wise real-world results for BAKU’s multi-task kitchen evaluation, including long-horizon tasks and rollout figures. The supplied passages report the demonstration count and identify the tables and figures, but do not provide task-level success values.

E.1 Real-World Task-wise Results

  • The real-world evaluation covers 30 kitchen tasks, using an average of 17 demonstrations per task and 520 demonstrations overall.Table 6 gives task-wise performance.Source: §E.1 · p. 23, Table 6 · p. 25
  • Figure 9 shows BAKU policy rollouts in complex real-world manipulation tasks.The caption describes the rollouts as showing BAKU’s capability in those tasks.Source: Figure 9 · p. 24
  • Table 7 presents task-wise performance for the real-world long-horizon tasks.Figure 10 shows BAKU rollouts on long-horizon manipulation tasks.Source: Figure 10 · p. 24, §E.1 · p. 25

E.2 Additional Analysis

The additional analysis compares shared and separate vision encoders and alternative observation-trunk inputs. The broader-impacts passage emphasizes data efficiency in robotics, where expert demonstrations are costly and time-consuming to collect.

E.2 Additional Analysis

  • Separate vision encoders improve LIBERO-90 performance by 2%, while increasing parameters by 15% per added camera view.BAKU uses a shared encoder to keep the model compact and assist faster inference, particularly relevant to the four camera views in the real-world experiments.Source: §E.2 · p. 25, Table 8 · p. 26
  • Separate observation tokens improve LIBERO-90 performance by 3% over concatenating inputs, with no performance difference on single-input Meta-World and DMC.The comparison concerns two camera views and proprioceptive state in LIBERO-90.Source: §E.2 · p. 26, Table 8 · p. 26

F Broader Impacts

  • BAKU aims to make generalist robot-policy training more data-efficient, reducing reliance on costly, time-consuming expert demonstrations.Robotics data collection requires physical engagement with the world.Source: F Broader Impacts · p. 26

Related papers

Loading 2406.07539v2…