BAKU: An Efficient Transformer for Multi-Task Policy Learning

We thank you for your positive response to our paper and we are glad that you like the paper writing and the results. We provide clarifications to your questions below.

“No variance or confidence intervals reported”: We have included the results for multiple seeds with standard deviation in Table 1 and Table 2 of the PDF attached to the global rebuttal. We will update these results in the final version of the paper.

“Limited number of baselines”: We would like to provide some clarification here. Our 2 primary baselines, RT-1 and MT-ACT, were selected as they propose different transformer-based architectures for offline imitation learning. Other popular methods include modifying the learning objectives to allow better modeling of the training data. These include GMMs, Behavior Transformers, and Diffusion Policy. We include the comparisons with these objectives in Table 2 showing that BAKU is a policy architecture that can incorporate the latest advances in policy learning objectives while learning from a smaller number of demonstrations and with a small model size.

“... how performance varies across data scales …”: We thank you for bringing up this point. We have included a data efficiency analysis in the global rebuttal with the results provided in Table 4 and Table 5 of the attached PDF. We observe that at each level of data availability, BAKU shows a significantly higher success rate than MT-ACT and RT-1.

We hope we have addressed your questions in the above clarifications. Kindly let us know if you have additional questions and we would be more than happy to discuss further.

We thank you for your positive response to our paper and we are glad that you found the paper clear and understandable. We provide clarifications to your questions below.

“... more extensive experiments … real world”: We include a new experiment studying the performance of BAKU on long-horizon tasks in the simulated LIBERO-10 benchmark and our real-world multi-task kitchen environment. We have included this in the global rebuttal. To summarize, we observe a similar trend where BAKU significantly outperforms our strongest baseline, MT-ACT, on these long-horizon tasks, achieving on average 19% higher success rate. We would also like to refer to line 171 in the paper where we mention that we perform 150 real-world evaluation trials per method. Doing this for the 4 methods in Table 1 brings the total count to 600 real-world trials which is very time-consuming. Hence, we restrict ourselves to using the simulated benchmarks for our ablation study. Moreover, from Table 1, LIBERO-90 is the most similar simulated benchmark to our real-world setup, with both using multiple camera views and robot proprioception as input. In this identical setting, we observe that the performance order of the baselines and BAKU remains the same with BAKU outperforming the baselines by a significant margin. We thank you for bringing up this important point and will include this explanation in the final version of the paper for more clarity.

“... doesn't include any methodological innovations, instead combining existing approaches”: You are correct that we try to find the right design choices for multi-task policy learning in BAKU. We have addressed this concern in the global rebuttal. To summarize, we observe that small design choices in the architecture lead to massive improvements in multi-task policy performance. To the best of our knowledge, no prior work has looked into this. Appendix D further elaborates on the differences between BAKU and state-of-the-art prior work.

“Is RT-1 fine-tuned on the tasks, trained from scratch, or are the numbers its zero-shot performance?”: All baselines in the paper and BAKU are trained from scratch and evaluated on the tasks they are trained on with randomized object initialization during evaluations. We primarily focus on the effect of architecture design choices on policy performance in this work and agree that understanding the zero-shot performance of multi-task policies is an interesting future direction.

“Are the models trained using relative or absolute actions?”: The policies on LIBERO, Metaworld, and the real-world setup are trained on relative actions. We stick to the base configuration of the benchmarks that we evaluate. We will include this detail in the final version of the paper.

“MLP action heads … relative actions … diffusion policies … absolute actions”: Thank you for this suggestion. Recent work [1] shows that even for diffusion policy, delta or relative actions show a better performance than absolute actions. We will be happy to add a comparison across different action spaces in the final version of the paper.

“... good to have some numbers for what the inference speed tradeoffs …”: Thank you for the feedback. We agree that providing performance metrics is important for potential users. Accordingly, we have evaluated the inference time tradeoffs between using a shared vision encoder and separate view-specific encoders. We tested scenarios with 2 camera views, matching the setup in LIBERO, and 4 camera views like the real-world xArm environment. A shared encoder allows batching frames from all views into a single forward pass, while separate encoders require individual passes. The table below shows inference times for a single forward pass with BAKU:

We see a modest 3.94% increase when going from 2 to 4 views with a shared encoder, but a much larger 39% increase with separate encoders. Further, our visual encoders have 1.4M parameters in a 10M parameter model. Separate encoders would result in 11.4M and 14.2M parameter models (for 2 and 4 camera views respectively), increasing the size by 14% and 42% respectively. We thank you for suggesting this study and we will include this in the final version of the paper.

We hope we have been able to address your questions in the above clarifications. Kindly let us know if you have additional questions or concerns that stand between us and a higher recommendation for acceptance.

[1] Chi, Cheng, et al. "Universal manipulation interface: In-the-wild robot teaching without in-the-wild robots." arXiv preprint arXiv:2402.10329 (2024).

We thank you for your comments about the paper and we are glad that you found the paper well organized and the experiments extensive. We provide clarifications to your questions below.

“... novelty of the algorithm … limited”: We have addressed the novelty of the algorithm in the global response. To summarize, we observe that small design choices in the architecture lead to massive improvements in multi-task policy performance. To the best of our knowledge, no prior work has looked into these design decisions. We have further elaborated the differences between BAKU and state-of-the-art prior work in Appendix D.

“... it is essential to clarify what makes it efficient”: Thank you for pointing this out. The efficiency in this work pertains to both your points - a small model size (resulting in high frequency execution on a real robot) and a high success rate achieved with a small number of demonstrations (also evident from Table 4 and Table 5 in the PDF attached to global rebuttal). The explanation provided in Appendix D compares BAKU with current state-of-the-art policy architectures such as MT-ACT and RT-1 and highlights the source of this performance improvement. Given the importance of this concern, we will move Appendix D to the main paper in the final version to address any concerns about novelty.

“... unclear why it achieves such improvements”: We would again like to refer you to Appendix D which highlights the differences between BAKU and RT-1/MT-ACT. You are correct that BAKU integrates useful ideas from several prior works and we show, through an extensive ablation study (Table 2), the effect of each of these components on the policy performance.

“... 10M model performs the best, even surpassing the 114M model”: We have included an explanation for this result in line 199 (Section 4.3). In recent years, we have witnessed that larger models tend to require more data [1, 2]. However, when it comes to robotics, collecting large amounts of data can be challenging. For instance, the LIBERO benchmark only has 50 demonstrations available per task. We suspect that operating in this low data regime leads the 114M parameter model to overfit the training data, resulting in poor performance.

We hope we have been able to address your questions in the above clarifications. Kindly let us know if you have additional questions or concerns that stand between us and a higher score.

[1] Hoffmann, Jordan, et al. "Training compute-optimal large language models." Proceedings of the 36th International Conference on Neural Information Processing Systems. 2022.

[2] Hernandez, Danny, et al. "Scaling laws for transfer." arXiv preprint arXiv:2102.01293 (2021).

We thank you for your insightful comments about the paper and we are glad that you found the paper easy to understand and the experiments extensive. We provide clarifications to your questions below.

“... combination of existing architecture/design choices”: You are correct that we try to find the right design choices for multi-task policy learning in BAKU. We have addressed this concern in the global rebuttal. To summarize, we observe that small design choices in the architecture lead to massive improvements in multi-task policy performance. To the best of our knowledge, no prior work has looked into this. We have further elaborated the differences between BAKU and state-of-the-art prior work in Appendix D.

“... code is not provided along with this submission”: Thank you for pointing this out. We completely agree that making the code public is essential for such a project and have added an anonymized version of the code to our anonymous webpage. Following the updated instructions, we have also privately shared the code with the AC.

“... did not consider the time efficiency and data efficiency for multi-task policy learning …”: We agree with the reviewer that these are very important factors for robot learning, especially when learning tasks on a real robot. Accordingly, we have included a data efficiency study in the global rebuttal and would refer the reviewer to Table 4 and Table 5 in the PDF attached to the global rebuttal. To summarize, we observe that at each level of data availability, BAKU shows a significantly higher success rate than MT-ACT and RT-1.

“RVT-1 on LIBERO:” We believe that there has been a confusion here. We compare BAKU with RT-1 which is another policy architecture that uses transformers. As you mentioned, RVT assumes the possibility of keyframe extraction so we do not compare BAKU against this line of work.

“... only focuses on the specific area of robotics manipulation”: In addition to robotic manipulation tasks in the LIBERO and Metaworld benchmarks and our real-world experiments, we have also included locomotion tasks in the DeepMind Control suite (DMC) in Table 1 and Table 2 to highlight the usefulness of the proposed architecture in domains other than manipulation.

We hope we have been able to address your questions in the above clarifications. Kindly let us know if you have additional questions or concerns that stand between us and a higher score.

Thanks for the authors' response. I appreciate the code sharing and the detailed clarification. But I'm still concerned about the novelty of this work. Overall, I slightly increase my score.

We thank the reviewer for the detailed comments about the work. We are glad that you found the study systematic and would like to provide clarifications to the questions below:

“The statistical significance of the experiments is questionable”: We have conducted additional experiments and report the mean and standard deviation across 3 different seeds for the simulated environments in Table 1 and Table 2 of the PDF attached to the global rebuttal. We find the performance of the model is quite robust to seeds and hence significant.

Initial state selection for LIBERO and Metaworld: For both simulated benchmarks, we reset the environment to a random state and evaluate on the sampled state. We have done this across 3 different seeds for each method to avoid a favorable seed value from impacting the results.

Lower performance of MT-ACT compared to other works: Thank you for bringing these works to our notice. We believe there might be some confusion here. Both PRISE and QueST use 100 demonstrations for each task on Metaworld as compared to the 35 demonstrations per task used by us. Further, both these papers also run MT-ACT on LIBERO-90 using the same 50 demonstrations per task provided by the benchmark. The multitask MT-ACT performance on LIBERO-90 reported by PRISE is 46.6% (Figure 6) and by QueST is also 46.6% (Figure 2a). Under the same setting, we report a multitask MT-ACT performance of 54% (Table 1). Hence, we believe that the small dataset size used in BAKU is the cause of lower performance in Metaworld.

“... key differences between BAKU and BC-ResNet-Transformer from LIBERO …”: We would like to note that the LIBERO paper primarily focuses on lifelong learning and provides a comparison with multi-task learning performance (48% success in Table 2 and Figure 5 of the paper). Though we primarily evaluate BAKU in the multi-task learning setting, we want to point out that BAKU is also compatible with lifelong learning. There are some key differences between BAKU and the BC-ResNet-Transformer (BC-ResNet-T) architecture used in LIBERO. While BC-ResNet-T uses a Gaussian mixture model (GMM) for its action head, the vanilla version of BAKU uses a simpler deterministic head with action chunking, which significantly improves performance (Figure 1b and Table 2). Additionally, BAKU shows that state-of-the-art multimodal policy classes like VQ-BeT and Diffusion Policy can be unified under a common framework. Further, BC-ResNet-T incorporates a history of observations. We show in Table 2 that using history without a multi-step loss can hurt performance and could be a potential reason behind the low performance of BC-ResNet-T. Hence, BAKU, with its simple design choices, results in a significant performance boost on the same LIBERO tasks as compared to BC-ResNet-T.

“Line 45 claims randomized object initialization but line 170 mentions initialization from fixed set during evaluation”: Thank you for pointing this out. We meant that the objects are initialized at random locations and we use the same set of randomized locations for all baselines to keep the comparisons fair. We will make this clearer in the final version of the paper.

We hope we have been able to address your questions in the above clarifications. Kindly let us know if you have additional questions or concerns that stand between us and a higher score.