BAKU: An Efficient Transformer for Multi-Task Policy Learning
BAKU combines multimodal conditioning and action chunking, reaching 90% on LIBERO-90 and 91% on real xArm tasks.
Key Findings
Methodology
BAKU factorizes policy learning into sensory encoders, an observation trunk, and an action head. Images are encoded by a FiLM-conditioned ResNet-18, proprioception by a two-layer MLP, and language by a six-layer MiniLM. Projected modality features become tokens for a causal Transformer decoder. The policy predicts action chunks—10 steps for image-based manipulation and 3 for DMC—then applies exponential temporal averaging. The head is replaceable: MLP, GMM, BeT, VQ-BeT, or diffusion.
Key Results
- Across 129 simulated tasks, BAKU improves overall performance by 18 absolute percentage points over RT-1 and MT-ACT. It reaches 0.90 on LIBERO-90, a 36-point gain over prior work, with 0.79 on Meta-World and 0.70 on DMC.
- On a real xArm 7, 520 demonstrations cover 30 kitchen manipulation tasks, averaging 17 demonstrations per task. The MLP-head BAKU reaches 86% success; replacing it with VQ-BeT increases success to 91%, 35 points above the strongest baseline.
- Ablations show that action chunking raises LIBERO performance from 0.76 to 0.90, while a Transformer trunk reaches 0.90 versus 0.81 for an MLP trunk. A 114M model achieves only 0.19 on LIBERO-90, suggesting overfitting.
Significance
BAKU addresses a central robotics bottleneck: demonstrations are expensive because every data point requires physical interaction, yet multitask policies often underperform individually trained policies. Its results suggest that data efficiency can improve through disciplined architectural composition rather than simply scaling parameter count. For research, the paper supplies a modular framework for isolating design choices. For industry, 91% success with only about 17 demonstrations per task indicates a plausible path toward lower-cost deployment of versatile manipulation systems.
Technical Contribution
The approximately 10M-parameter design allocates about 2.1M parameters to sensory encoders, 6.5M to the observation trunk, and 1.4M to the action head. A causal Transformer treats multimodal observations as tokens and uses action tokens for control prediction. Separating the head enables direct retrofitting with modern action generators such as VQ-BeT. A further engineering contribution is predicting an entire chunk as one concatenated vector, rather than decoding each timestep independently, followed by exponential smoothing.
Novelty
BAKU does not introduce a new Transformer objective or a new generative action model. Its novelty is the systematic integration of FiLM-conditioned vision, causal multimodal fusion, action chunking, temporal smoothing, and interchangeable action heads into a compact low-data multitask policy. Compared with RT-1’s discretized action classification and MT-ACT’s encoder-decoder design, BAKU emphasizes modularity and controlled ablations.
Limitations
- Real-world evidence is limited to one xArm 7, a kitchen setup, fixed robot initialization, constrained object positions, and five evaluation runs per task. This does not establish robust generalization to homes or unseen embodiments.
- Offline behavior cloning remains vulnerable to covariate shift. Occlusion, dynamically changing objects, out-of-distribution states, and recovery after mistakes are not comprehensively addressed.
Future Work
Future work should test cross-robot and cross-environment transfer, online correction, uncertainty estimation, and safer recovery. Stronger visual representations, human-video pretraining, and diffusion-based heads may extend coverage of diverse behaviors. The authors’ planned release of datasets, training code, and evaluation code should enable reproducibility and more systematic comparisons.
AI Executive Summary
A robot that must open an oven, fetch a glass, wipe a board, and store food faces a problem unlike internet-scale vision or language learning: demonstrations require costly physical execution. Existing multitask policies often remain weaker than separate single-task systems, encouraging an expensive cycle of collecting ever more expert data.
BAKU, from New York University, addresses this bottleneck with a compact modular Transformer. A FiLM-conditioned ResNet-18 processes images, MiniLM encodes instructions, and an MLP encodes proprioception. A causal Transformer fuses these modality tokens, while a separate action head predicts a short sequence of future controls rather than one isolated action. Exponential temporal averaging then smooths overlapping chunks; VQ-BeT can replace the basic MLP head when multiple behaviors are useful.
Across 129 simulated tasks, BAKU improves overall performance over RT-1 and MT-ACT by 18 absolute points, reaching 90% on the difficult LIBERO-90 benchmark. On a real xArm robot, it achieves 86% with the MLP head and 91% with VQ-BeT across 30 kitchen tasks, using only 520 demonstrations—about 17 per task. The evidence is promising but bounded: experiments use one robot and controlled layouts, leaving open-world generalization, failure recovery, and cross-embodiment transfer for future work.
Deep Analysis
Background
Recent robot learning has adopted Transformer policies such as RT-1 and MT-ACT, alongside multimodal action generators including GMMs, BeT, VQ-BeT, and diffusion policies. Yet robot data cannot be collected as cheaply as web-scale image-text data. Teleoperation is costly, and multitask policies frequently lag behind single-task specialists. The remaining challenge is to extract more capability from limited expert trajectories rather than relying only on dataset expansion.
Core Problem
Given offline expert trajectories, a policy must map images, proprioceptive state, and a task goal to continuous actions. It must align heterogeneous modalities, share useful structure across tasks while preserving task distinctions, model temporal dependencies, and remain stable under distribution shift. Single-step regression can produce jerky behavior and compound errors; expressive multimodal heads may also be difficult to fit when demonstrations are scarce.
Innovation
- �� FiLM-conditioned ResNet-18 makes visual features task-specific through language modulation.
- �� A causal Transformer trunk represents modalities as tokens and predicts through action tokens.
- �� Action chunking predicts 10 future steps for LIBERO/Meta-World and 3 for DMC, followed by exponential smoothing.
- �� A decoupled action head supports MLP, GMM, BeT, VQ-BeT, and diffusion models.
- �� Controlled ablations examine trunk type, model size, goal modality, history, FiLM, and supervision design.
Methodology
- �� Input: multi-view RGB observations, robot proprioception, and text or goal-image instructions.
- �� Encoding: ResNet-18 extracts vision features and FiLM conditions them on language; a two-layer MLP encodes state; six-layer MiniLM encodes text.
- �� Projection: modality representations are mapped to a common dimensionality and assembled as observation tokens.
- �� Fusion: a causal Transformer decoder processes only available current or historical observations; learnable action tokens produce action features.
- �� Learning: behavior cloning minimizes L=E[||a−π(o|g)||²]. With history, actions can be supervised at every timestep for richer gradients.
- �� Control: the deployed policy runs at 10 Hz and exponentially averages overlapping action chunks.
Experiments
Simulation covers 90 LIBERO-90 tasks, 30 Meta-World tasks, and 9 state-based DeepMind Control locomotion tasks. Demonstrations per task are 50, 35, and 500 respectively; image sizes are 128×128 for LIBERO and 84×84 for Meta-World, with 10 rollouts per task. Real experiments use an xArm 7, four 128×128 RGB views, 520 demonstrations across 30 tasks, 30-Hz collection, and 10-Hz deployment. Baselines are RT-1 and MT-ACT; ablations compare five action heads and multiple architectural choices.
Results
BAKU obtains 0.90, 0.79, and 0.70 on LIBERO-90, Meta-World, and DMC, versus RT-1’s 0.16, 0.65, and 0.66 and MT-ACT’s 0.54, 0.13, and 0.59. On long-horizon LIBERO-10 and five real chained tasks, it scores 0.86 and 0.84 versus MT-ACT’s 0.68 and 0.64. Real-task success rises from 86% with MLP to 91% with VQ-BeT. Naive history with last-step loss harms performance, whereas multi-step supervision improves it by 47% on average without improving the final no-history policy.
Applications
BAKU is suitable for low-demonstration desktop manipulation, kitchen assistance, warehouse pick-and-place, and laboratory workcells. Deployment requires calibrated cameras, proprioceptive access, a defined action space, and task instructions, plus a small but representative teleoperation dataset. Its modular head allows simple MLP control for narrow behavior distributions and VQ-BeT when several valid action modes exist.
Limitations & Outlook
The evidence is concentrated in simulation and one controlled xArm kitchen. The paper does not demonstrate zero-shot transfer across embodiments, homes, object distributions, or severe visual disturbances. Offline behavior cloning remains exposed to covariate shift and lacks a strong mechanism for detecting and recovering from failure. The 114M model’s 0.19 LIBERO score also shows that scaling capacity can hurt under limited data. Future work should combine broader data, online adaptation, uncertainty-aware control, and safety constraints.
Plain Language Accessible to non-experts
Imagine a new cook learning an entire kitchen from a small set of videos. The menu tells the cook what dish to prepare; several cameras show the ingredients and tools; a small body sensor reports where the cook’s hands are. BAKU is like one head chef who combines all these clues before deciding what to do, instead of letting each camera make a separate guess.
A weaker system might choose only the next tiny movement: reach a little, stop, look again, and guess again. That can make a robot shake. BAKU plans a short sequence—reach, grasp, lift, move, release—then blends overlapping plans so motion remains smooth. When several actions could work, VQ-BeT keeps several possibilities instead of averaging them into an awkward compromise.
The striking point is how little teaching data it needs. It reaches 90% on LIBERO-90 and 91% on a real robot using about 17 demonstrations per task. Still, this is like passing an exam in a familiar kitchen. A different room, a hidden object, or a failed grasp may require abilities the paper has not yet tested.
ELI14 Explained like you're 14
Think of a robot as a teammate in a game with lots of missions: open the oven, grab a glass, put ketchup in the fridge, or wipe a table. It cannot rely on one picture. It needs the mission text, several camera views, and information about where its own arm is. BAKU is a smart teammate that reads all of those clues together.
Its coolest trick is planning a combo instead of pressing one button at a time. For picking up a cup, it predicts a short chain: move closer, reach, close the gripper, lift, and pull back. Then it keeps checking the world while acting. This makes movement smoother, like a game character performing a practiced combo rather than freezing after every move!
BAKU has about 10 million parameters and was tested on 129 simulated missions. It scored 90% on the challenging LIBERO-90 benchmark and beat RT-1 and MT-ACT overall by 18 percentage points. With a real xArm robot, only about 17 demonstrations per mission were used on average, yet the VQ-BeT version succeeded 91% of the time.
But it is not a magic robot. Most tests used one robot and a controlled kitchen. What if someone blocks the camera, the object moves, or the robot drops something? The next challenge is teaching BAKU to handle surprises, recover from mistakes, and work in many different homes and robots.
Glossary
Behavior Cloning
A method that learns to imitate expert actions directly from demonstrations. In BAKU, the objective is mean-squared action regression, L=E[||a−π(o|g)||²].
It is the paper’s offline imitation-learning framework.
Action Chunking
Predicting a short sequence of future actions rather than only the next action. It exploits temporal correlation and can reduce jitter and compounding errors.
BAKU predicts 10-step chunks for manipulation and 3-step chunks for DMC.
FiLM Conditioning
Feature-wise modulation that scales and shifts intermediate neural representations according to conditioning information. Plainly, language changes how visual evidence is interpreted.
Task instructions condition the ResNet-18 visual encoder.
VQ-BeT
A Vector-Quantized Behavior Transformer that models multiple plausible action modes. It is useful when one observation admits several valid behaviors.
It raises real-robot success from 86% to 91%.
Observation Trunk
The network component that combines encoded sensory inputs into a representation for action prediction. BAKU uses a causal Transformer decoder as its trunk.
Vision, state, and language tokens are fused before the action head.
Open Questions Unanswered questions from this research
- 1 Can BAKU transfer to unseen homes, objects, camera layouts, and robot embodiments? Current experiments use controlled distributions, so cross-domain and zero-shot tests are still required.
- 2 How should a policy recover after an error? Offline cloning offers limited information about failure states; online interaction, uncertainty estimates, and safety-aware recovery may be necessary.
Applications
Immediate Applications
Multitask kitchen robots
Robot developers can collect a small teleoperation set for opening, wiping, fetching, and placing tasks, then train one policy. Required ingredients include calibrated RGB cameras, proprioception, a fixed action interface, and clear language instructions.
Warehouse and laboratory workcells
In structured stations with repeated objects but many operations, BAKU can unify picking, transport, and placement. An MLP head is a practical baseline; VQ-BeT is preferable when several valid motion styles exist.
Long-term Vision
General-purpose domestic robots
With cross-embodiment data, stronger visual-language representations, and online correction, BAKU-like modular policies could support many household tasks. Major obstacles remain safety, open-world generalization, and reliable long-horizon recovery.
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/.