BAKU: An Efficient Transformer for Multi-Task Policy Learning

Summary

The paper combines a series of architectural modifications made by prior work in offline imitation learning and shows that the new, composite architecture performs well in multi-task simulated settings and in data scarce, real-world robotics scenarios. Comprehensive ablations show the impact of each design choice made in selecting the final architecture.

Strengths

- Strong results relative to the chosen baselines and in both simulated and real-world collections of tasks. - Thorough investigation of different viable architecture designs. - Comprehensive ablations. - Excellent written communication.

Weaknesses

- No variance or confidence intervals reported for results on simulated environments. - Limited number of baselines, with RT-1 being somewhat outdated. - It would be useful to see how performance varies across data scales for the different architectures (incl. ablations and baselines).

Questions

- Why have you only considered two offline imitation learning baselines? - Why do you not report variance in your results?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The paper appropriately addresses its limitations.

Summary

This paper presents a model for multi-task behavior cloning. The paper analysis several different existing architectural options across three simulated environments to find the best options. The best model is then tested in the real world across a series of tasks on real robots.

Strengths

The paper tests on a large number of simulated and real world tasks. The paper is very clear and understandable.

Weaknesses

Most of the conclusions about architectural choices are based on experiments in simulation. However, the real world results show that the ordering in the simulation doesn't always stay the same when the architecture is used in the real world. Since we want these systems to work in the real world, it would be good to have more extensive experiments about when the ordering between architectures remains the same in the real world. This is also important for when the simulation experiments had inconclusive results, or when minimal performance improvements in simulation were sacrificed for increased inference speed. The paper doesn't include any methodological innovations, instead combining existing approaches. While this is ok, it makes having a strong evaluation even more critical.

Questions

Is RT-1 fine-tuned on the tasks, trained from scratch, or are the numbers its zero-shot performance? Are the models trained using relative or absolute actions? In my understanding, MLP action heads tend to work best with relative actions, while diffusion policies work best with absolute actions. It is important to compare the various models under the most favorable conditions for each one, so it would be good to include experiments with and without absolute actions for each action head. There are a few instances where choices that decrease success rate were chosen because of inference speed or model parameter count, such as choosing not to have shared vision encoders. It would be good to have some numbers for what the inference speed tradeoffs were because other users of the model may want to make different choices.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The paper discusses its limitations.

Summary

This work comprehensively investigates the architecture for multi-task policy learning on robotics manipulation tasks. Considering imitation learning with a limited number of demonstrations, the authors study the design choice of the network architecture, including encoders for observation with multi-modality, observations chunk, and action head. This work presents experiments on diverse benchmarks in the simulator and real-world robot manipulation tasks. According to the detailed ablation study of design choices on the components, this work selects the architecture design with a performance around 20% improvement over the baselines.

Strengths

The presentation is clear and easy to understand The authors conducted extensive experiments and solid evidence to support their proposed architecture It is impressive to achieve >90% real-world performance with a limited number of demonstrations on the real robot.

Weaknesses

The contribution of this work is mostly the empirical study of the combination of existing architecture/design choices. It does not propose any new architecture/design choices Since the main contribution is empirical study rather than a new idea, it is important to share the code so that other researchers can replicate the experiments and benefit from this work for experiments on other tasks. However, the code is not provided along with this submission. This work did not consider the time efficiency and data efficiency for multi-task policy learning when comparing different architectures and baselines. But these factors are also very important for robot learning. It will be much better to also discuss these factors.

Questions

I’m concerned about the baseline implementation, especially RVT-1 on LIBERO and real-robot tasks, since it seems much poorer than RVT-1 paper (though RVT-1 shows experiments on RLBench and a different set of real-robot tasks.) About keyframe extraction, which is an important part of the success of RVT-1, how do you define keyframes on LIBERO and real-robot data? About action prediction, did you use the same discretization way as presented in RVT-1 paper? Is it possible to study BAKU on RLBench, so we can have a more convincing comparison? Do you have any intuitive explanation about why the baseline fails? RVT-1 also works well on real-world tasks in their paper. Their task set also consists of general pick&place tasks. Why does its success rate drop a lot in the task set in this submission? Only the architecture difference from BAKU contributes to the performance gap? This work will be more convincing if the baselines can be discussed more in detail.

Rating

5

Confidence

4

Soundness

2

Presentation

4

Contribution

2

Limitations

This work studies the architecture design with detailed ablative study, which is definitely informative and useful for other researchers in the area. However, it does not propose any new ideas/insights. So it is more like a technical report or system engineering paper, rather than a research paper. Its contribution only focuses on the specific area of robotics manipulation, and the evidence or conclusions seem not generalizable/beneficial for any other AI area. Not sure whether NeurIPS is the proper venue. Perhaps it is more appropriate to be presented in a robotics conference/journal.

Summary

The paper presents a novel transformer architecture, BAKU, designed for efficient multi-task policy learning in robotics. BAKU consists of three components: sensory encoders, an observation trunk, and an action head. According to the experimental results, BAKU is showcased to be effective in both simulated and real-world tasks, offering a promising solution to the challenge of training generalist policies for physical robots.

Strengths

1. The paper is well organized, and the writing is clear and well-understood. The author clearly describes the algorithm BAKU, including the objective of the optimization process and the details of the algorithm implementation. 2. The authors conduct extensive experiments compared to other baseline methods in both simulated and real-world tasks, as well as detailed ablation studies. The results show that BAKU outperforms other baselines by a large margin and can accomplish tasks using few demonstrations with high data utilization and efficiency.

Weaknesses

1. The novelty of the algorithm presented in this paper is limited, as all components of the BAKU model are derived from other papers. For instance, RT-1 employs a transformer architecture and FiLM layers, while ACT utilizes action chunking techniques. The authors should discuss in greater detail how this algorithm differs from these previous works. 2. As indicated in the paper's title, "Efficient Transformer," it is essential to clarify what makes it efficient. Does efficiency pertain to the model's size, or is it the high success rate achieved with a small number of demonstrations? Additionally, which component is critical for enhancing efficiency? 3. Even though the experimental results indicate that this algorithm significantly outperforms RT-1 and MTACT, it remains unclear why it achieves such improvements. From a design perspective, BAKU does not propose a novel network architecture or optimization objective compared to RT-1 and MTACT; it rather appears to integrate elements from both. 4. In Table 2, I noticed an intriguing phenomenon: the 10M model performs the best, even surpassing the 114M model. What could be the reason behind this outcome?

Questions

Details are listed above.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

YES

Summary

This paper presents a multitask BC model for robot learning. They systematically compare several components of recently proposed BC models and combines their best aspects into a common architecture. Key design choices include: observation trunk, model sizes, action heads, goal modalities, action chunking, observation history etc. Experimental comparison with RT-1, MT-ACT and BAKU-variants on 3 simulated benchmarks and on 30 real world tasks validates their choices and design methodology.

Strengths

This paper presents a very systematic ablation study across various model components proposed in prior works. The final model is simple, it debunks some overcomplicated/compute-expensive choices like GMM/diffusion heads, observation history or FiLM etc. and highlights importance of others like Transformer trunk, action chunking etc.

Weaknesses

1. The statistical significance of the experiments is questionable. For probabilistic approaches (GMM/diffusion) especially, a single unfavorable random seed could significantly skew the performance results for a baseline run. 2. Both LIBERO and MetaWorld environments offer 50 different initial states for each task, yet the reported results are based on only 10 rollouts. Many methods are known to be sensitive to out-of-distribution (OOD) initial states, but the paper lacks details on how initial states were selected for both LIBERO and MetaWorld experiments. This omission severely undermines the reliability and generalizability of the presented results. Furthermore, for such low number of rollouts, it becomes crucial to conduct evaluations across multiple random seeds to ensure the robustness of the findings. 3. The results reported for MT-ACT on MetaWorld appear to be notably lower than those presented in other recent works, such as [PRISE] (fig:15|80%) and [QueST] (fig:3|90%). This discrepancy raises questions about the MT-ACT implementation and evaluation in this work. 4. This work does not offer many new insights to the community. For example usefulness of probabilistic action heads over deterministic ones for multimodal data have been shown in [Robomimic] paper. Action chunking helps in libero, but the chunk size is a more sensitive parameter which has not been studied. Transformers over MLP and model size overfitting is a common knowledge. 5. The proposed model is very similar to BC-ResNet-Transformer from [Libero] paper, but still performs much better. This is not addressed in the paper. 6. Line 45 claims randomized object initialization but line 170 mentions initialization from fixed set during evaluation. This claim should be removed from main contributions if later is the case. [PRISE]: https://arxiv.org/abs/2402.10450 [QueST]: https://arxiv.org/abs/2407.15840v2 [Libero]: https://arxiv.org/abs/2306.03310

Questions

1. How can reported results be considered reliable given single training seed and limited evaluation rollouts? 2. Could you describe the initial state selection process for LIBERO and MetaWorld experiments? 3. Why are MT-ACT MetaWorld results lower compared to those reported in other aforementioned works? 4. What key differences between BAKU and BC-ResNet-Transformer from LIBERO contribute to BAKU's superior performance?

Rating

6

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

Limitations are adequately addressed.