Research map / Reinforcement learning / Imitation learning
BAKU: An Efficient Transformer for Multi-Task Policy Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents BAKU, a ~10M-parameter decoder-only transformer for multi-task robot policy learning that combines FiLM-conditioned ResNet-18 encoders, observation tokens, action chunking with temporal ensembling, and swappable action heads, outperforming RT-1 and MT-ACT on 129 simulated tasks and 30 real xArm tasks.
Problem
Multi-task robot policy learning is data-hungry and multi-task policies often underperform single-task ones, so architectures that use limited demonstration data efficiently are needed.
Approach
A simple decoder-only transformer over separate modality tokens (vision, proprioception, language/goal) with an independent, swappable action head that predicts a concatenated action chunk, combined with FiLM conditioning and exponential temporal ensembling; design choices are systematically ablated.
Key results
- BAKU achieves 0.90 success on LIBERO-90 vs 0.54 for MT-ACT and 0.16 for RT-1, a 36% absolute gain over the strongest baseline.
- On 30 real xArm tasks (~17 demos/task), BAKU with VQ-BeT head reaches 91% success vs 56% for MT-ACT and 37% for RT-1.
- Across 129 simulated tasks, BAKU shows an overall 18% absolute improvement over RT-1 and MT-ACT, with 0.79 on Meta-World and 0.70 on DMC.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BAKU instead of the strongest baseline | Multi-task policy learning · LIBERO-90 | success rate: 0.54 -> 0.9 | Helped |
| Use BAKU instead of the strongest baseline | Multi-task policy learning · Meta-World | success rate: 0.65 -> 0.79 | Helped |
| Use BAKU instead of the strongest baseline | Multi-task policy learning · DMC | success rate: 0.66 -> 0.7 | Helped |
| Use BAKU instead of the strongest baseline | Multi-task policy learning · Real Robot (20 tasks) | success rate: 0.56 -> 0.86 | Helped |
| Use BAKU instead of MT-ACT | Long-horizon policy learning · LIBERO-10 | success rate: 0.68 -> 0.86 | Helped |
| Use BAKU instead of MT-ACT | Long-horizon policy learning · Real Robot (5 tasks) | success rate: 0.64 -> 0.84 | Helped |
| Replace the MLP observation trunk with a transformer (ablation) | BAKU · LIBERO-90 | success rate: 0.81 -> 0.90 | Helped |
| Increase model size to 114M parameters (ablation) | BAKU · LIBERO-90 | success rate: 0.9 -> 0.19 | Hurt |
| Replace the MLP action head with GMM (ablation) | BAKU · LIBERO-90 | success rate: 0.90 -> 0.84 | Hurt |
| Replace the MLP action head with VQ-BeT (ablation) | BAKU · Real Robot (20 tasks) | success rate: 0.86 -> 0.91 | Helped |
| Add action chunking (ablation) | BAKU · LIBERO-90 | success rate: 0.76 -> 0.90 | Helped |
| Add action chunking (ablation) | BAKU · DMC | success rate: 0.74 -> 0.70 | Hurt |
Surprising findings
- The 114M-parameter variant severely underperforms on LIBERO-90 (0.19), which the authors attribute to overfitting.
- A simple MLP action head performs as well or better than multimodal heads (GMM, BeT, VQ-BeT, diffusion) on simulated benchmarks.
Limitations
- Struggles on precise manipulation tasks such as opening an oven door or placing a tea bottle in the fridge.
- Only single-skill tasks are studied; chaining multiple skills is not addressed.
- Does not analyze generalization benefits as the number of tasks grows.
Experimental setup
Datasets: LIBERO-90 Meta-World demonstrations (expert-generated) DM Control demonstrations (expert-generated) real xArm kitchen demonstrations
Benchmarks: LIBERO-90 LIBERO-10 Meta-World DM Control real xArm kitchen tasks
Metrics: success rate
Baselines: RT-1 MT-ACT
Builds on: Action Chunking Transformer Behavior Transformer VQ-BeT FiLM MT-ACT
Closest papers by meaning
- Active Fine-Tuning of Multi-Task Policies (ICML 2025 · similarity 0.87)
- Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers (NeurIPS 2024 · similarity 0.86)
- Policy Decorator: Model-Agnostic Online Refinement for Large Policy Model (ICLR 2025 · similarity 0.86)
- RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation (ICLR 2025 · similarity 0.86)
- Quantization-Free Autoregressive Action Transformer (NeurIPS 2025 · similarity 0.86)
- Predictive Inverse Dynamics Models are Scalable Learners for Robotic Manipulation (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.