BAKU: An Efficient Transformer for Multi-Task Policy Learning
Authors: Siddhant Haldar
Zhuoran Peng
Lerrel Pinto
Venue: NeurIPS 2024 poster (2024)
Keywords: Robot learning
Imitation Learning
Multitask Learning
Primary area: robotics
TL;DR: We present BAKU, a simple architecture for multi-task policy learning that provides highly efficient training, particularly in data-scarce problems such as robotics.

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 baku-robot.github.io.

BAKU: An Efficient Transformer for Multi-Task Policy Learning
Siddhant Haldar
Zhuoran Peng
Lerrel Pinto
New York University
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 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 36% improvement on the harder LIBERO benchmark. On 30 real-world manipulation tasks, given an average of just 17 demonstrations per task, BAKU achieves 91% success rate. Videos of the robot are best viewed baku-robot github.o..........
1 Introduction
Learning generalist policies that can solve multiple tasks is a long standing problem in decision making and robotics. While significant advances have been made computer vision [4, 59] and natural language processing [2, 53, 69], algorithms that can effectively do so for physical agents are far behind. A key reason for this is the scale of available data. While large-scale datasets in vision and language can readily be amassed from the Internet, robotics presents unique challenge. Given its interactive nature, data acquisition requires physical engagement with the world, making robot data considerably more laborious to obtain in terms of both time and financial costs.
prominent approach for training multi-task policies to bite the bullet and collect large amounts of data, often by contracting teleoperators [37, 6, 61]. However, policies trained on such data are quite inefficient, often achieving performance far below independently trained single-task policies [75, 44, 57]. The current best answer to solve this problem is unfortunately to collect even more demonstration data from experts.
In this work, we present BAKU, a simple architecture for multi-task policy learning that provides highly efficient training, particularly in data-scarce problems such as robotics. BAKU builds upon recent work in multitask learning [5, 6] and has three key features. First, a transformer encoder that fuses information from multiple modalities like vision and language while incorporating temporal context. Second, FiLM- SLLL--oondiiooeed [46] visual encoder helps the model learn task-specific repre- sentations by adapting the encoder to the task. Third, an action prediction head that separated from the observational encoding trunk, enabling BAKU to be easily retrofitted with state-of-the-ee action
Correspondence to: siddhanthaldarGe nyu.edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Instruction
Close door of the
Baku K
Sensory Encoders
Observation Trunk
Action Head
Time
Camera Views
a
Robot State
End effectee pose
(a) Overview
MEFACT
Open the doer of over
Petch glass from the mads
cream chense In basket
Wipe board towel
Transformer BAKU)
0.60
Action Chunking
yoghert fiigs
knifc segarizer
History
0.87
FiLM Conditioning
Success rate
(b) Performance on LIBERO-90
(c) Real-world tasks (8 of 30 shown)
Figure I: (a) We present BAKU, simple transformer architecture learning multi-task policies across a diverse range of tasks. BAKU encodes inputs from different modalities using modality-specinf encoders. The encoded representations are merged in observation trunk before predicting actions through an action head. (b) We develop a unified policy for 90 tasks in the LIBERO-90 benchmark, discussing design choices that impact multi-task performance. (c) On our xArm robot, BAKU can learn single multi-task policy for 30 tasks with an average of 17 demonstrations collected per task.
generation models [39, 60, 31, 10]. The novelty of BAKU hence lies in carefully combining these ideas to produce a new transformer architecture particularly suited for multitask decision making.
To demonstrate the effectiveness of BAKU, we run extensive experiments on 129 simulated tasks across LIBERO [34]. Meta- World [76], and DeepMind Control [67], and 30 robotic manipulation tasks xArm robot (see Fig. 1). Our main findings are summarized below:
1. BAKU exhibits an overall 18% absolute performance improvement over prior state-of-the-arr multi-task learning algorithms on 129 tasks across simulated environment suites (Section 4.1). BAKU sets ofate-ofthheart performance on LIBERO with 90% average success rate, 36% absolute improvement over prior work (Table 1).
real-world tasks, with an average of demonstrations per task, BAKU achieves an average success rate of 91 % across 30 diverse tasks in multi-task kitchen environment, with randomized object initialization. This outperforms prior state-of-the-an algorithms by 35% (Section 4.2).
Through ablation analysis, study the importance of each component BAKU (Section 4.4), particularly the role of action chunking [77] and a multimodal action head in boosting performance, especially our real-worle experiments.
All of our datasets, and training and evaluation code will be made publicly available. Videos of trained policies can be seen here: baku-robot.githue
2 Background
Imitation Learning: The goal of imitation learning is learn behavior policy n given access to either the expert policy xt or trajectories derived from the expert policy T While there are multitude settings with differing levels access expert [68], this work operates the setting where the agent only has access to observation-base trajectories, i.c. T {(o,, a )
Here

---

Observation Trunk
Action Head
Ton Encesdr
Image Ercode
Iooge Ercoder
Ieagee Encoder
Imiaee Encoder
Sieee Enceskr
Aation Token
Multi-stop Action Prediction
Clise the lfcrof
10 Hz Control
T
6
S
Figure 2: Overview of BAKU, broken down into modality-specific sensory encoders, an observation trunk, and an action head predicting a chunk of actions. BAKU takes input observations from multiple camera views Or robot proprioceptive state and a task instruction T and enables performing closed-loop control 10Hz in our real world experiments on the xArm.
N and T denote the number of trajectory rollouts and episode timesteps respectively. We choose this specific setting since obtaining observations and actions from expert near-expert demonstrators is feasible in real-world settings [77, 24] and falls in line with recent work in this area [77, 11, 31, 10].
Multi-task Behavior Cloning: Behavior Cloning (BC) corresponds solving the maximum likelihood problem shown in Eq. 1. Here T refers to expert demonstrations. When parameterized by normal distribution with fixed variance, the objective can be framed as regression problem where, given observations d", π needs to output a".
C BC ( - (o'))
(1)
After training, it enables BC to mimic the actions corresponding the observations seen in the demonstrations. In multi-task settings, use the same formulation for BC but condition the action prediction on a goal variable g". Thus, the loss function for multi-task BC becomes the following.
BC Ei .n y BC (o')g))
(2)
In this work, we represent goals as either text description of the task [6, 5] or goal image [11, 72].
3 BAKU
The design of multi-task learning algorithms involves numerous decisions regarding model architec- ture and component selection. This often results in complex architectures where the importance of individual components is sometimes unclear. In this work, we perform a systematic and thorough ablation study across the various multi-task learning architectures proposed by prior works [6. 5, 63] and introduce BAKU, simple architecture for multi-task policy learning. To facilitate our analysis, we divide the overall model architecture into three main components: sensory encoders, an observa- tion trunk, and an action head. Sensory encoders process raw sensor inputs from different modalities into useful feature representations. The observation trunk combines the encoded information from the different modalities. Finally, the action head utilizes the combined information to predict actions. Below, we describe these three components in detail, with additional algorithmic details provided in Appendix A.
3.1 Sensory Encoders
In the real-world, robots encounter diverse data modalities, including vision, depth feedback, propri- oceptive feedback, and task instructions in various forms such as text, goal images, or task videos. In BAKU, we focus on vision, robot proprioception, and text or goal image based task instructions. For vision, ResNet-18 [19] visual encoder to process images of the scene, enhanced with

---

FiLM [46] layer to integrate task-specific information. Robot proprioception data is processed through two-layer multilayer perception (MLP) encoder. For text, we use a 6-layer version of MiniLM [73] provided in Sentence Transformers [54]. We project the representations obtained from all modalities to the same dimensionality through additional MILP layers, to facilitate combining the encoded information. We have included description of FiLM conditioning in Appendix A.1.
3.2 Observation Trunk
The encoded inputs from all sensory modalities combined in the observation trunk. We explore two variants of the trunk network:
Multilayer Perceptron (MLP) The encoded inputs are concatenated into a single feature vector and passed through multilayer perceptron. When using history of observations, the inputs corresponding to all time steps are concatenated.
Transformer Each encoded input is treated as an observation token and passed through a trans- former decoder network [70]. A learnable action token is appended to the list of observation tokens and used to predict the action. When using a historical observation, separate action token is added for each time step to enable predicting actions for all time steps. A causal mask is applied to the transformer ensure that actions are predicted solely based on past observations.
Both variants output action feature vectors (corresponding the action tokens for a transformer), which are then passed through an action head to predict actions.
3.3 Action Head
The final component of our architecture is the action head, an action prediction module that takes as input the action feature vectors obtained from the observation trunk and predicts the corresponding actions. An independent action prediction module enables to unify several state-of-the- action generation models within the same framework. We experiment with five action head variants: vanilla MLP, Gaussian Mixture Model (GMM) [36], Behavior Transformer (Be'T) [60]. Vector-Quantin Behavior Transformer (VQ-BeT) [31], and diffusion policy [45, 10, 55J. Considering the temporal correlation in robot movements, we follow prior work [77, 5] and include action chunking with exponential temporal averaging to produce smoother behaviors and counteract the covariate shift often seen in low-data imitation learning scenarios. In contrast to previous works [77, 5] that decode actions for each time step separately, we predict the action chunk as a single concatenated vector. We find that this simplification improves performance (see Table 1). More details about each action head variant has been provided in Appendix A.2 along with details about the exponential temporal smoothing technique in Appendix A.3.
3.4 Putting it all together
Our proposed architecture is depicted in Figure 2. Through extensive experimentation (see Section 4), our final architecture includes modified FiLM-conditioned ResNet- 18 vision encoder (provided with the LIBERO benchmark [34]), an MLP encoder for robot proprioception, and a pre-trained text encoder for task instructions. For environments with multiple camera views, we use common visual encoder across all views. We provide only the current observation as an input and the observation trunk uses causal transformer decoder architecture [29]. The base version of our model uses an MLP action head with action chunking and temporal smoothing to produce smoother motions. We also experiment with multimodal variants of action heads and have provided the results in Section 4.4.
The parameter counts are approximately 2. M. for the sensory encoders, 6.5M for the observation trunk, and 1.4M for the action head, bringing the total model size to approximately 10M parameters.
4 Experiments
Our experiments are designed to answer the following questions: (a) How well does BAKU work for multi-task learning? (b) How does BAKU perform real-world tasks? (c) How does BAKU perform

---

(a) LIBERO-90
(b) Metaworld
(C) DM Control
Figure 3: BAKU is evaluated on 3 simulated benchmarks LIBERO, Meta- World, and DM Control.
on long- horizon tasks? (d) What design decisions affect multi-task policy learning? Additional results and analysis have been provided in Appendix D.
Simulation Tasks: We experiment with 90 manipulation tasks from the LIBERO-90 bench- mark [34], 30 manipulation tasks from Meta-World suite [76], and locomotion tasks from DeepMind Control Suite (DMC) [67]. Figure 3 depicts the simulated environments. For LIBERO-90 50 demonstrations per task provided with the benchmark and use images from third-person and gripper camera views, well as robot proprioception demonstrations per task from an expert policy trained with demonstration guided reinforcement learning [17, 18], using only the
input. For Meta-World, we obtain 35
view as input. We images of size 128 x 128 for LIBERO-90 and 84 84 for Meta- Meta-World. For DMC, we train based locomotion policies using 500 demonstrations per task obtained from experts trained with DrQ-v2 [74]. All evaluations conducted using 10 policy rollouts per task. More details about the simulated environments can be found in Appendix B.
Robot Tasks: Our real-world experiments are performed on a Ufactory xArm 7 robot with an xArm Gripper in multi-task kitchen environment. The policies are trained on RGB images of size 128 x 128 obtained from four different camera views, including an egocentric camera attached to the robot gripper. The action space comprises the robot end effector pose and the gripper state. We collect a total of 520 demonstrations across 30 tasks, averaging 17 demonstrations per task. The demonstrations were collected using VR-based teleoperation system [24] a 30Hz frequency. The learned policies are deployed at 10Hz. Figure 4 shows selected tasks from our real-world environment. More details about each task and robot control can be found in Appendix C.
Strong Baselines: BAKU,
In this section, provide detailed explanation of our baselines in relation to
MT-ACT [5]: Multi-task Action-Chunking Transformer (MT-ACT) state-of-the-art trans- former encoder-decoder architecture for learning multi-task policies. MT-ACT extends Action- Chunking Transformer (ACT) [77] to multi-task setting. MT-ACT takes input observations from multiple camera views, robot proprioception, and task instructions. Fach input modality passes through dedicated encoders. The encoded observations are then fused transformer encoder, the output of which conditions transformer decoder to predict chunks of future actions. Each predicted action corresponds a position embedding input to the decoder. In ACT and MT- ACT, a conditional variational autoencoder (CVAE) is used to learn multimodal style variable which conditions the encoder to deal with multimodal action distributions. During inference, the style variable is set to zero, leading to unimodal behavior. In contrast, BAKU uses a decoder-only transformer architecture that directly predicts action features corresponding to past observations. This enables us to (1) leverage recent advances in multimodal action generation by plugging in several unimodal and multimodal heads for action prediction, and (2) incorporate a history of observations to predict actions for each time step in the history (see Section 4.4 for results on the
5

---

Table 1: Performance of multi-task policies learned using BAKU on 3 simulated benchmarks LIBERO-90, Meta- Meta-World, and DM Control and real xArm robot. We observe that BAKU significantly outperforms prior work on both simulated and real world tasks.
LIBERO-90
Meta- World
DMC
Real Robot
Method
(90 tasks)
(30 tasks)
(9 tasks)
(20 tasks)
RT-1
0.16
0.65
0.66
0.37
MTACT
0.54
0.13
0.59
0,56
BAKU (Ours)
0.9
0.79
0.7
0.86
BAKU w/ VQ-BeT (Ours)
0.9
0.78
0.7
0.91
use of observation history). Further, using a multimodal action head enables BAKU to exhibit multimodal behavior during inference, improving wol-w performance (Table 1).
RT-1 [6]: RT-1 is transformer- based multi-task policy learning architecture that models actions discrete classes by uniformly discretizing them into bins. RT-1 uses a FiLLM-conditioned vision encoder (ResNet-18 in our implementation), but instead of directly using the final 512-dimensional representation, it splits intermediate feature map of size k k 512 into k-2 tokens of 512 dimensions each. These tokens are passed through a Token Learner [58] module to reduce them 8 tokens per image. The reduced number of tokens is then passed through decoder-only transformer architecture to predict discrete action. In contrast, BAKU directly uses the final 512-dimensional representation from the vision encoder, without summarizing tokens via a token learner. Additionally. BAKU predicts a continuous action through an unimodal or multimodal action head. Based on our experiments (Table 1), we observe that these design choices in BAKU lead to significant improvements in performance over RT-1.
4.1 How well does BAKU work for multi-task learning?
We evaluate the multi-task performance of BAKU on 90 tasks from the LIBERO-90 benchmark, 30 tasks from Meta-World, and tasks from DMC. Table compares the performance BAKU with our baselines, RT-1 1 [6] and MT-ACT [5]. BAKU outperforms the strongest baseline by 36% and 14% on LIBERO-90 and Meta- Meta-World respectively. demonstrating more effective multi-task learning on complex manipulation tasks. On the simpler DMC locomotion tasks, BAKU outperforms the strongest baseline by 4%. Overall, these results suggest that BAKU more effectively leverages relationships between tasks to achieve superior multi-task learning performance compared to prior methods.
4.2 How does BAKU perform on real-world tasks?
We evaluate BAKU 30 manipulation tasks in real-world kitchen environment, comparing with MT-ACT and RT-1, During evaluations, the xArm was always initialized at the same pose and the objects being manipulated were placed in fixed set of positions for all methods. We conducted 5 evaluation runs per task, totaling 150 evaluation runs per method. Table includes our real-world results. We observe that BAKU achieves 86% success rate across all tasks, outperforming the strongest baseline by 30%. Replacing the MI.P action head with a multimodal VQ-BeTT [31] head further improves the success rate to 91 %. outperforming the strongest baseline by 35%. Figure shows real-world rollout trajectories for selected task set with all tasks included in Appendix C. Appendix D.I provides the task-wise performance for each method. Overall, these results indicate BAKU'S promise for deploying multi-task policies on real-world robotic systems.
4.3 How does BAKU perform on long-horizon tasks?
We also evaluate BAKU on long-horizon tasks in the simulated LIBERO-10 benchmark and our real-world multi-task kitchen environment. Table 2 provides the results on 10 tasks in LIBERO- 10 and 5 long-horizon tasks in the real kitchen environment, each composed of two shorter tasks chained sequentially. We use 50 demonstrations per task for LIBERO-10 10 and an average of 19 demonstrations per task for the real robot. We observe that BAKU significantly outperforms our strongest baseline, MT-ACT, on these long horizon tasks, achieving on average 19% higher success rate. This highlights BAKU'S ability to learn policies that effectively plan and execute sequenees of actions over

---

Open oven door: Open the door the oven.
Fetch glass from rack: Fetch the glass from the lower rack.
Put cheese in basket: Pick up the block of cream cheese and put in the basket.
Wipe towel: Wipe the cutting board with towel.
röe
Lift white plate from rack: Lift the white plate kept on the rack.
Put ketchup bottle inside fridge: Pick up the bottle of tomato ketchup and put inside the fridge.
von
SOS
Fetch yoghurt bottle from fridge door: Take the bottle of yoghurt out from the door of the fridge
Fetch knife from organizer: Fetch knife from organizer placed on the kitchen counter.
Figure 4: Real-world policy rollouts showing BAKU'S capability in complex manipulation tasks.

---

Table 2: Performance of multi-task policies learned using BAKU on long-horizon tasks in the LIBERO-10 simulated benchmark and a real xArm robot. We observe that BAKU significantly outperforms prior work on both simulated and real world tasks.
LIBERO-10
Real Robot
Method
(10 tasks)
(5 tasks)
MT-ACT
0.68
0.64
BAKU (Ours)
0.86
0.84
extended time horizons. Real world rollouts of these long-horizon tasks have been included in Appendix C with the task-wise performance and demonstration details in Appendix D.I.
4.4 What design decisions affect multi-task policy learning?
As described in Section 3, our multi-task policy architecture consists of three main components: sensory encoders, an observation trunk, and an action head. In this section, we analyze the design choices within each component and their effect overall multi-task performance. We consider BAKU with an MLP action head (described Section 3,4) as our base model. For ablations, we vary only a single property at time while keeping all other aspects identical. This experimental setup allows us to clearly isolate the impact of individual design decisions. We examine different observation trunks, model sizes, action heads, goal modalities, and the use of action chunking, observation history, and task conditioning through FiLM [46]. The results of our ablation study are provided in Table 3 with more analysis in Appendix D. The results provide insights into which components and properties most important for effective multi-task learning with BAKU.
Effect of Observation Trunk: We experiment with two trunk types: an MLP and a transformer architecture. In Table 3, we observe slight performance dip when using an MLP trunk Meta- World and DMC. For LIBERO-90, our most complex simulated benchmark, an MLP trunk resulted in 9% lower success rate than transformer trunk. This highlights the efficacy of transformers for modeling complex relationships between observations from multiple sensing modalities and actions.
Effect of Model Size: We study the effect of model size on multi-task performance by evaluating configurations with 4.4M, 10M, 31M, and 14M parameters. For each variant, we vary the size of the observation trunk and the action head while keeping the sensory encoders constant. The results in Table 3 show that the 4.4M, 10M, and 3 31M parameter models achieve similar performance across benchmarks. Surprisingly. the largest 1 114 14M parameter model severely underperforms on the harder LIBERO-90 benchmark. We suspect this poor performance may have been due to overfitting on the training data with larger capacity. Based on these results, we use the 10M parameter model for BAKU since it is the smallest model with the best performance on of the 3 simulated benchmarks.
Effect of Action Head: We compare the performance of BAKU retrofitted with five different action heads: MLP, GMM [36]. BeT [60]. VQ-BeT [31], and diffusion [10]. Having independent action head enables us to extend these state-of-the-art action generation models to multi-task settings. Table 3 shows that simulated benchmarks, simple MLP head performs just well better than other multimodal action heads. Among the multimodal variants, VQ-BeT achieves the best performance. As result, also evaluate BAKU with VQ-BeT action head in real-world setup (Table 1). On the real robot, having multimodal action head proves advantageous with the VQ-BeT head achieving a success rate of 91 %. 5% improvement over an MLP action head. Hence, our experiments demonstrate that while multimodal heads may provide benefits for real-world deployment, simple MLP head can perform well on simulated data with limited behavioral diversity.
Effect of Action Chunking: We study the effect of action chunking on multi-task policy perfor- mance. For the image-based LIBERO-90 and Meta-Wont benchmarks, we predict a chunk of 10 future actions. For the locomotion tasks in DMC, we predicted 3 future actions. Based on the results in Table 3, we observe the largest difference on LIBERO-90, where removing action chunking and instead predicting single action led to a 14% drop in performance. In contrast, there is no perceptible difference in performance
Woold with and without chunking. For the locomotion domains in
8

---

Table 3: Study of design decisions for BAKU that affects multi-task performance.
Category
Variant
LIBERO-90
Meta- World
DMC
Observation Trunk
MLP
0.81
0.78
0.68
Transformer
0.90
0.79
0.70
Model Size
4.4M
0.85
0.78
0.68
10M
0.9
0.79
0.70
31M
0.87
0.81
0.70
114M
0.19
0.81
0.68
Action Head
MLP
0.90
0.79
0.70
GMM
0.84
0.65
0.67
BeT
0.89
0.78
0.60
VQ-BeT
0.90
0.78
0.70
Diffusion
0.89
0.45
0.61
Action Chunking
X
0.76
0.78
0.74
0.90
0.79
0.70
Observation History
X
0.90
0.79
0.70
With last-step loss
0.54
0.08
0.37
With multi-step loss
0.90
0.82
0.68
Goal Modality
Text
0.90
0.79
N/A
Goal Image
0.88
0.81
N/A
Intermediate Image
0.91
0,80
N/A
FiLM
0.87
0.79
N/A
0.90
0.79
N/A
DMC, we see 4% performance increase when removing chunking. Hence, action chunking benefits manipulation tasks while mildly hindering locomotion tasks from experiments.
Effect of Observation History: We study the effect of using an observation history on multi-task performance. As shown in Table 3, naively using observation history where the action prediction loss is only computed for the last time step significantly degrades performance. However, since BAKU uses transformer observation encoder, allows predicting actions for all observations in the history and computing the prediction loss over all time steps. Empirically, we found this multi-step prediction loss provides richer supervision and improves the single-step loss performance by an average of 47% across all benchmarks. However, incorporating an observation history with multi-step action prediction did not noticeably improve overall policy performance compared to using no history. Therefore, our final architecture only uses the most recent observation as an input.
Effect of Goal Modality: We experiment with different goal modalities: text instruction, goal image. and intermediate goal image. The text instructions directly obtained from the task data. The goal image is obtained by randomly sampling a demonstration from the training dataset and taking the last frame. For an intermediate goal image, consider this randomly sampled task demonstration, and for every time step, treat the frame steps the future as the goal image [72]. Table 3 contains the results LIBERO-99 90 and Meta- worad, goal images do not apply to the state-based DMC tasks. We set ie to 50 steps for LIBERO-90 and 30 steps for Meta- World. Since LIBERO-90 90 has two camera views, that all three goal modalities show a similar performance with slight variations. Overall, our approach supports different goal representations with only minor variations in performance.
the hirr--person view obtain goal images. We observe
Effect of FiLM Conditioning: We examine the impact of using FiLM-conditioned vision encoder for waaguuge-guidddd multi-task policies. As shown in Table 3, on the image- based LIBERO-99 and Meta-Won World benchmarks, FiLM FiLM-conditioned vision encoder performs equally well or better than an unconditional encoder. FiLM conditioning allows modulating the vision encoder's parameters conditioned on the language input. This provides an effective way to fuse visual and linguistic information for solving tasks. Therefore, BAKU employs FiL.M-coniitioned vision encoder for our image-based experiments.

---

5 Related Work
Imitation Learning (IL.) IL [23] refers to the setting where agents learn from demonstrations without access to environment rewards. IL can be broadly categorized into Behavior Cloning (BC) [48, 68] and Inverse Reinforcement Learning (IRL) [41, 1]. BC solely learns from offline demonstrations but suffers on out-of-distributions samples [56] whereas IRL focuses on learning a robust reward function through online interactions but suffers from sample inefficiency [17, 18]. In this work, we focus on using BC for learning multi-task policies. In recent years, there have been significant advances in single-task behavior cloning with the development of multimodal action generation models using GMMs [36, 39J, EBMs [13], Be'T' [60, 11 1, 3..3. and diffusion [45, 10, 55, 7]. There has also been notable progress in solving long-how long-horizon tasks through imitation learning with some works relying solely on robot data [38. 21, 9, 77, 33, 65] while others attempt to bootstrap learning from human data [72]. Further, these advances in policy learning combined with significant strides in self-supervised representation learning [8. 42, 20J have enabled deploying these policies in messy and unpredictable environments such as our homes [61] as well as zero-shot deployment in-the- wild [62, 64]. However, despite the large body of work advancing single-task robotic policy learning, there still exists gap between single-task and multi-task performance for policy learning [75, 44, 57].
Multi-task L Learning Robotics has a long history of multi-task learning. There is a significant body of work focusing on learning policies for robotic grasping with the aim of generalizing to new tasks [32, 47, 14, 71, 12], robotic language understanding [40, 35, 66, 3]. and framing multi- task learning as a goal reaching problem [51, 28, 22]. Additionally, several works have collected varied multi-task robotics datasets [38, 34, 5, 43, 30J. Recently, there has been an increased use of transformer-based architectures for multi-task robot learning, spanning across robot navigation [62, 64]. locomotion [27, 15, 49, 50]. and manipulation [6, 78, 11. 5]. While most of these works use text conditioning for task specification, some go beyond text to use goal images [11, 16] and videos [26, 25] as well. Another emerging trend is co-training these robot policies with additional tasks, such as visual question answering and image captioning [52, 78). to develop more generalizable policies. Overall, multi-task learning has been widely applied in robotics and, more recently, using high-capacity transformer models to learn robot control policies has become common practice in the field. Despite their effectiveness, the architectures for these policies often become complicated, with the necessary components sometimes being unclear. Our proposed model, BAKU, combines key ideas from prior work into a single architecture to produce a model that is both simple and outperforms state-of-the-art methods in multi-task policy learning.
6 Conclusion and Limitations
In this work, we presented BAKU, simple transformer architecture that demonstrates improved multi-task policy learning performance on a variety of simulated and real-world domains compared to prior state-of-the-art methods. We recognize a few limitations in this work: (a) In our real-world experiments, while BAKU achieved good performance on most tasks, it struggled on some precise manipulation tasks, such as opening an oven door or placing tea bottle in the fridge. This suggests that data sharing across tasks of varying difficulty may hinder performance on more precise skills. Developing techniques learn single policy for different task complexity levels could help address this. (b) Currently, we focus on performing single skill at time. Developing algorithms capable of chaining multiple such skills can enable effective long- oong-horizon robot manipulation. work, we primarily studied the policy architecture and did not analyze the generalization benefits of multi-task learning as the number of tasks increases. A study of the emergence of such generalization with greater task diversity would be another interesting direction. Overall, we hope that BAKU serves as an important step towards developing multi-task policies capable of performing precise robotic manipulation.
(c) In this
Acknowledgments and Disclosure of Funding
We thank Nur Muhammad Shafiullah, Ulyana Piterbarg, Ademi Adeniji, Ben Evans, Gaoyue Zhou, and Irmak Güzey for valuable feedback and discussions. This work was supported by grants from Honda, Google, NSF award 2339096 and ONR awards N00014-21-1-27558 and N00014-22-1-2773. LP is supported by the Packard Fellowship.
10

---

References
[1] P. Abbeel and A. Y. Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty- first international conference Machine learning, page 1, 2004. 10
[2] J. Achiam, Adler, S. Agarwal, L. Ahmad, Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08574, 2023. 1
[3] M. Ahn, A. Brohan, N. Brown, Y. Chehotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakr- ishnan, K. Hausman, et al. Do as can, not as say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01699 2022. 10
[4] Betker, G. Goh, Jing. T. Brooks, J. Wang, L. Li, L. Ouyang, J. Zhuang, J. Lee, Y. Guo, et al. Improving image generation with better captions. Computer Science. https://edn. openai. com/papers/dall-e- pdf. 2(3):8, 2023.
[5] H. Bharadhwaj. J. Vakil, M. Sharma, A. Gupta, S. Tulsiani, and V. Kumar. Roboagent: Generalization and efficiency robot manipulation via semantic augmentations and action chunking. arXiv preprint arXiv:2309.01918 2023. 1, 3, 4, 5, 6, 10, 16, 17
[6] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, J. Dabis, C. Finn, K. Gopalakrishnan, K. Haus- man, A. Herzog, J. Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint 12.i:22212.06817.. 2022. 1, 3, 6, 10, 17
[7] L. Chen, S. Bahl, and D. Pathak. Playfusion: Skill acquisition via diffusion from language- annotated play. In Conference Robot Learning, pages 2012-2029. PMLR, 2023. 10
[8] X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9640-9649, 2021. 10
[9] Y. Chen, C. Wang. L. Fei-Fei, and C. K. Liu. Sequential dexterity: Chaining dexterous policies for oog-horizoon manipulation. arXiv preprint 20Xi:22309000987, 2023. 10
[10] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and
Systems (RS5), 2023, 2, 3, 4, 8, 10, 16
[11] Z. J. Cui, Y. Wang. N. M. M. Shafiullah, and L. Pinto. From play policy: Conditional behavior generation from uncurated robot data. arXiv preprint arXiv:22 10. 10047, 2022. 3, 10
[12] H.-S. Fang, C. Wang, H. Fang, M. Gou, J. Liu, H. Yan, W. Liu, Y. Xie, and C. Lu. Anygrasp: Robust and efficient grasp perception in spatial and temporal domains. IEEE Transactions on Robotics, 2023. 10
[13] Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, Lee, I. Mordatch, and J. Tompson. Implicit behavioral cloning, In Conference on Robot Learning. pages 158-168. PMLR, 2022. 10
[14] A. Gupta, A. Murali, D. P. Gandhi, and L. Pinto. Robot learning in homes: Improving generalization and reducing dataset bias. Advances in neural information processing systems, 31, 2018. 10
[15] A. Gupta, I.. Fan, S. Ganguli, and I.. Fei-Fei. Metamorph: Learning universal controllers with transformers. arXiv preprint arXiv:2203. /991, 2022. 10
[16] S. Haldar and L. Pinto. Polytask: Learning unified policies through behavior distillation. arXiv preprint 211002330.. 2023. 10
[17] S. Haldar, V. Mathur, D. Yarats, and L. Pinto. Watch and match: Supercharging imitation with regularized optimal transport. Conference on Robot Learning, pages 32-43. PMIR, 2023. 5, 10
11

---

[18] S. Haldar, J. Pari, A. Rai, and L. Pinto. Teach a robot to fish: Versatile imitation from one minute of demonstrations. arXiv preprint arXiv:2303.01497, 2023. 5, 10
[19] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778. 2016. 3
[20] K. He, Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000-16009, 2022. 10
[21] M. Heo, Y. Lee, D. Lee, and J. J. Lim. Furniturebench: Reproducible real-world benchmark for long-horizon complex manipulation. arXiv preprint arXiv:2305. 2023. 10
[22] D.-A. Huang, Y.-W. Chao, C. Paxton, X. Deng, L. Fei-Fei, J. C. Niebles, A. Garg, and D. Fox. Motion reasoning for goal-based imitation learning. 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 4878-4884. IEEE, 2020. 10
[23] A. Hussein, M. M. Gaber, E. Elyan, and C. Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
[24] A. Iyer, Z. Peng, Y. Dai, I. Guzey, S. Haldar, S. Chintala, and Pinto. Open teach: versatile teleoperation system for robotic manipulation. arXiv preprint arXiv:2403.0787) 2024. 3,5
[25] V. Jain, M. Attarian, N. J. Joshi, A. Wahid, D. Driess, Q. Vuong. P. R. Sanketi, P. Sermanet, S. Welker, C. Chan, I. Gilitschenski, Y. Bisk, and D. Dwibedi. Vid2rohot: End-to-end video- conditioned policy learning with cross-attention transformers, 2024. 10
[26] E. Jang, A. Irpan, M. Khansari, D. Kappler, F. Ebert, C. Lynch, S. Levine, and C. Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning, pages 991-1002. PMLR, 2022. 10
[27] M. Janner, Q. Li, and S. Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 14:1273-1286, 2021. 10
[28] T. Jurgesson, O. Avner, E. Groshev, and A. Tamar. Sub-goal trees framework for goal-based reinforcement learning. In International conference on machine learning, pages 5020-5030.
PMLR, 2020. 10
[29] A. Karpathy. mingpt: minimal pytorch re-implementation of the openai gpt. https: 2021. 4, 18
/ gi github..c.. woo//aarpthy/miaannnaaaa
[30] A. Khazatsky, K. Pertsch, S. Nair, A. Balakrishna, S. Dasari, S. Karamcheti, S. Nasiriany, M. K. Srirama, L. Y. Chen, K. Ellis, et al. Droid: A large-scale in-the-wild robot manipulation dataset.
arXiv preprini arXiv:2403. 122003.12444 12945, 2024. 10
[31] S. Lee, Y. Wang, H. Etukuru, H. Kim, N. M. M. Shafiullah, and L. Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.0318/ 2024. 2, 3, 4, 6, 8, 10, 16
[32] I. Lenz, H. Lee, and A. Saxena. Deep learning for detecting robotic grasps. The International Journal of Robotics Research, 34(4-5):705--24, 2015. 10
[33] T. Lin, Y. Zhang, Li H. Qi. B. Yi, S. Levine, and J. Malik. Learning visuotactile skills with two multifingered hands. arXiv:2404. 16823, 2024. 10
[34] B. Liu, Y. Zhu, C. Gao, Y. Feng, Q. Liu, Y. Zhu, and Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. Advances in Neural Information Processing Systems, 36, 2024. 2, 4, 5, 10, 17
[35] C. Lynch and P. Sermanet. Language conditioned imitation learning over unstructured data. arXiv preprint arXiv:2005.076488 2020. 10
[36] C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and P. Sermanet. Learning latent plans from play. Conference on robot learning, pages 1113-1132. PMIR, 2020. 4, 10, 16
12

---

[37] A. Mandlekar, Y. Zhu, A. Garg, J. Booher, M. Spero, A. Tung, J. Gao, J. Emmons, A. Gupta, E. Orbay, et al. Roboturk: A crowdsourcing platform for robotic skill learning through imitation. In Conference on Robot Learning, pages 879-893. PMLR, 2018. 1
[38] A. Mandlekar, D. Xu, R. Martín-Martín, S. Savarese, and L. Fei-Fei. Learning to generalize across long-horizon tasks from human demonstrations. arXiv preprint urXiv:2003.06085, 2020. 10
[39] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Zhu, and R. Martín-Martín. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:21088 200.0299...0 2021. 2, 10
[40] H. Mei, M. Bansal, and M. Walter. Listen, attend, and walk: Neural mapping of navigational in- structions to action sequences. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 30, 2016. 10
[41] A. Y. Ng. S. J. Russell, et al. Algorithms for inverse reinforcement learning. In Ieml, volume 1, page 2, 2000. 10
[42] M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza, Massa, A. El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023. 10
[43] A. Padalkar, A. Pooley, A. Jain, A. Bewley, A. Herzog, A. Irpan, A. Khazatsky, A. Rai, A. Singh, A. Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models, arXiv preprint ar drXiv:2310.08864, 2023. 10
[44] E. Parisotto, J. L. Ba, and R. Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. arXiv preprint arXiv:1511.06444 2015. 1, 10
[45] T. Pearce, T. Rashid, A. Kanervisto, D. Bignell, M. Sun, R. Georgeseu, S. V. Macua, S. Z. Tan, I. Momennejad, K. Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301. 10677, 2023. 4, 10, 16
[46] E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018. 1, 4, 8, 16
[47] L. Pinto and A. Gupta. Supersizing self-supervision: Learning grasp from 50k tries and 700 robot hours. In 2016 IEEE international conference on robotics and automation (ICRA), pages 3406-3415 IEEE, 2016. 10
[48] D. Pomerleau. An autonomous land vehicle in a neural network. Advances in Neural Information Processing Systems, 1, 1998. 10
[49] I. Radosavovic, T. Xiao, B. Zhang, T. Darrell, J. Malik, and K. Sreenath. Learning humanoid locomotion with transformers. arXiv e-prints, pages arXiv-2303, 2023. 10
[50] I. Radosavovic, B. Zhang, B. Shi, J. Rajasegaran, S. Kamat, T. Darrell, K. Sreenath, and J. Malik. Humanoid locomotion next token prediction. arXiv preprint 1944v 2024. 10
[51] A. Raffin, A. Hill, R. Traoré, T. Lesort, N. Díaz-Rodriguez, and D. Filliat. Decoupling feature extraction from policy learning: assessing benefits of state representation learning in goal based robotics. arXiv preprint 1901v
2019. 10
[52] S. Reed, K. Zolna, E. Parisotto, G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, et al. A generalist agent. arXiv preprint 27X1::2005.06115 2022. 10
[53] M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint 245.:2400305534 2024. I
13

---

[54] N. Reimers and I. Gurevych. Sentence-bert, Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019. URL https: http:////aaaaaoa//////////////// 10084.4
[55] M. Reuss, M. Li, X. Jia, and R. Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532 2023. 10, 16
[56] Ross, G. Gordon, and D. Bagnell. reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference artifi- cial intelligence and statistics, pages 627-637 JMLR Workshop and Conference Proceedings, 2011. 10
[57] A. A. Rusu, G. Colmenarejo, C. Gulcehre, G. Desjardins. Kirkpatrick, R. Pascanu, V. Mnih, K. Kavukcuoglu, and R. Hadsell. Policy distillation. arXiv preprint arXiv: 2015. 1, 10
[58] M. Ryoo, A. Piergiovanni. A. Arnab, M. Dehghani, and A. Angelova. Tokenlearner: Adaptive space-time tokenization for videos. Advances in neural information processing systems, 34: 12786-12797, 2021. 6
[59] C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gon- tijo Lopes. B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35: 36479-36494, 2022. 1
[60] N. M. Shafiullah, Z. Cui, A. A. Altanzaya, and L. Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing systems, 35:22955-22968, 2022. 2, 4, 8, 10, 16
[61] N. M. M. Shafiullah, A. Rai, H. Etukuru, Y. Liu, I. Misra, Chintala, and L. Pinto. On bringing robots home. arXiv preprint 1rXiv:23/1.16000 2023. 1, 10
[62] D. Shah, A. Sridhar, N. Dashora, K. Stachowicz, K. Black, N. Hirose, and S. Levine. Vint: A foundation model for visual navigation. arXiv preprint arXiv:2306. 14846, 2023. 10
[63] M. Shridhar, L. Manuelli, and D. Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning, pages 785-799, PMLR, 2023. 3
[64] A. Sridhar, D. Shah, C. Glossop, and S. Levine. Nomad: Goal masked diffusion policies for
navigation and exploration. arXiv preprint arXiv:2310.07896 2023. 10
[65] K. Sridhar, S. Dutta, D. Jayaraman, Weimer, and I. Lee. Memory neural networks for imitation learning. arXiv preprint 2rX0v:23100.6/71,, 2023. 10
[66] S. Stepputtis, Campbell, M. Phielipp. S. Lee, C. Baral, and H. Ben Amor. Language- conditioned imitation learning for robot manipulation tasks. Advances in Neural Information Processing Systems, 33:13139-131 2020. 10
[67] Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. d. L. Casas, D. Budden, A. Abdolmalcki, J. Merel, A. Lefraneq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690, 2018. 2, 5, 17
[68] F. Torabi, G. Warnell, and P. Stone. Recent advances in imitation learning from observation. arXiv preprint arXiv: 1e0.. 2019. 2, 10
[69] H. Touvron, Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv: 13971, 2023. 1
[70] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t.. Kaiser, and Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017. 4
14

---

[71] U. Viereck, A. Pas, K. Saenko, and R. Platt. Learning visuomotor controller for real world robotic grasping using simulated depth images. In Conference on robot learning, pages 291-300. PMLR, 2017. 10
[72] C. Wang, L. Fan, J. Sun, R. Zhang, L. Fei-Fei, D. Xu, Y. Zhu, and A. Anandkumar. Mimicplay: Long-horizon imitation learning by watching human play. In 7th Anmual Conference on Robot Learning, 2023. 3, 9, 10
[73] W. Wang. F. Wei, L. Dong, H. Bao, N. Yang, and M. Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers, 2020. 4
[74] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645. 2021.
[75] Yu, S. Kumar, Gupta, Levine, K. Hausman, and Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33:5824 2020. 1, 10
[76] T. Yu, D. Quillen, Z. He. R. Julian, K. Hausman, C. Finn, and S. Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100 100. PMLR, 2020. 2, 17
[77] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304. 13705, 2023. 2, 3, 4, 5, 10, 16, 17
[78] B. Zitkovich, T. Yu, S. Xu, P. Xu, T. Xiao, F. Xia, J. Wu, P. Wohlhart, S. Welker, A. Wahid, et al. Rt-2: iarion-langugg--cttie action models transfer web knowledge robotic control. In Conference on Robot Learning. pages 3165-2188 PMLR, 2023. 10
15

---

A Algorithmic Details
A.1 FiLM Conditioning
Feature-wise Linear Modulation (FiLM) [46] is a technique used for conditioning neural networks that allows the network to modulate its behavior based on an external conditioning signal, such as text instructions or observations. In the context of text conditioning for policy learning, the text instructions are first encoded into a conditioning vector. This conditioning vector is then used to modulate the activations of the neural network through FiLM layers. FiLM applies feature-wise affine transformation (scaling and shifting) to the activations of the network, conditioned on the text embedding. In other word, assuming is FiLM layer's 's input, is conditioning input, and y and are z-dependent scaling and shifting vectors.
(x) = 7(2) の x + ((z)
This allows the network adapt its computation and output based on the given text instructions, enabling tasks like instruction following or conditioning the policy on language descriptions.
A.2 Action Heads
Having separate action prediction module allows BAKU to leverage state-of-the techniques for action generation. In this work, we evaluate five different action head variants. Below we briefly describe each variant. For more details on these methods, please refer to the original publications.
Multilayer Perceptron (MLP) This is simple neural network comprising multiple dense layers. We use a two-layer MLP for our experiments.
Gaussian Mixture Model (GMM) [36] A Gaussian mixture model (GMM) action head models the policy as mixture of Gaussians, enabling multi-modal action sampling for continuous control problems. The GMM parameters are part of the learned policy network. For our experiments, employ a two-layer GMM action head with five action modes and a Softplus activation function.
Behavior Transformer (BeT) [60] The Behavior Transformer (BeT) models continuous action prediction as a two-part problem. Actions in the training data are first clustered into & bins using k-means clustering. A discrete action head classifies the cluster an action belongs to, while an offset action head predicts an offset value added to the corresponding cluster center. The discrete head uses a focal loss, while the offset head uses 1.2 loss. For our experiments, we use BeT with 64 action clusters.
Vector- Quutti-e Behavior Transformer (VQ-BeT)) [31] The Vector- cuator-uuattzee Behavior Trans- former (VQ-BeT) extends BeT by replacing k-means clustering with residual VQVAE- tok- enization, significantly improving performance over BeT. For our experiments, we employ VQ-BeT with two residual VQ layers of codebook size and latent dimension 16 and 256, respectively.
Diffusion [45, 10, 55] A diffusion action head models action prediction as a diffusion process that generates actions over time by iteratively denoising samples from a Gaussian distribution. While highly effective for multi-modal distributions, the iterative denoising during inference slows deployment speed. In this work, we use a transformer-bast diffusion head introduced by prior work [45, 10]. We use two-layer diffusion head for our experiments.
A.3 Temporal smoothing over action chunking
naïve implementation of action chunking, where a new environment observation in incorporated every à steps can be suboptimal and can result in jerky robot motion. To improve the smoothness in robot motion, we incorporate an exponential temporal ensembling technique, following prior work 5J. Instead of querying the policy every l steps, we query it at every timestep. This results an overlap in predicted action chunks and at any given timestep, there will be more than one predicted actions. Instead of using only the current action prediction, we use temporal ensemble to
16

---

combine all the past predictions. This temporal ensemble performs a weighted average over these predictions with an exponential weighing scheme w; exp(-m i), where wp is the weight for the oldest action. The speed for incorporating new observation is governed by where smaller m means faster incorporation. It must be noted that this ensembling incurs no additional training cost, only extra inference-time computation. In our experiments, similar to prior work [77, 5], we find both action chunking and temporal ensembling to be important for producing precise and smooth motion.
A.4 Hyperparameters
The complete list of hyperparameters is provided in Table 4. For RT-1 [6], we use our implementation with RT-1 action head that discretizes the continuous action into discrete bins uniformly. For MT-ACT [5], we use the open-source implementation with the default hyperparameters. We vary the action chunk length for MT-ACT for different benchmarks, the values for which have been provided in Table 4.
Training time Below we provide details about the time required to train BAKU on single NVIDIA RTX A4000 GPU.
1. LIBERO: Training for 600k steps with batch size of 64 and camera views and robot proprio- ception as input requires around 10.5 hours.
2. Meta-World: Training for 600k steps with a batch size of 64 and 1 camera view as input requires around 8 hours.
DM Control: Training for 2M steps with batch size of 128 and robot state as input requires around 26 hours.
xArm Robot:
Training for 200k steps with a batch size of 64 and 4 camera views and robot proprioception input requires around 6 hours.
B Simulation Tasks
We evaluate BAKU three simulated benchmarks: LIBERO-90 [34], MetaWorld [76], and DM Control [67]. For LIBERO-90, we directly use the datasct provided, which includes demonstrations for all 90 tasks. For details the specific LIBERO-90 tasks, please refer to the original paper [34]. For World and DM Control, we collected demonstrations from expert agents trained with reinforcement learning (RL). We include only the tasks for which were able to obtain expert demonstration data. Table 5 lists the 30 MetaWorke World tasks and 9 DM Control tasks used our experiments.
C Robot Tasks
We evaluate BAKU on 30 tasks in our real-world multi-task kitchen environment. We provide the task description along with policy deployment rollouts with BAKU for each task in Figures 5, 6, 7, and 9. The long-horizon task rollouts have been shown in Figure 10.
We deploy our learned policies at 10Hz using a high-level controller. To facilitate smooth motion on the robot, we deploy low-level Minimum-Jerk Controller 100Hz.
Robot control
D Additional Results and Analysis
D.I Real- World Task-wise Results
Table provides the task-wise performance for all tasks in our real-world multi-task kitchen environment. We collect an average of 17 demonstrations per task, with total of 520 demonstrations across all tasks. Task-wise performance for the real-world long-horizon tasks has been included in Table 7.
17

---

Table 4: List of hyperparameters.
Method
Parameter
Value
Common
Learning rate
1e 4
Image size
128 x 128 (LIBERO-90, xArm)
84 84 (Meta-World) World)
Mini-batch size
64 (LIBERO-98) Meta-World, xArm)
128 (DM Control)
Optimizer
Adam
Number training steps
600000 (LIBERO-9) Meta-Worke (oo-W
2000000 (DM Control)
200000 (xArm)
Number of demonstrations
50 (LIBERO-90)
35 (Meta-World)
500 (DM Control)
15 (xArm)
Transformer architecture
minGPT [29] (with layers and heads)
Action chunk length
10 (LIBERO-90, Meta-World)
3 (DMC)
20 (xArm)
BAKU
Observation trunk
Transformer
Action head
MLP (base)
GMM, BeT, VQ-BeT, Diffusion (variants)
Hidden dim
256
Observation history
False
Action chunking
True
Intermediate goal steps (k)
50 (LIBERO-90)
30 (Meta-World)
RT-1
Observation trunk
Transformer
Action head
MLP (base)
Hidden dim
512
Observation history
True
History length
6
Action chunking
False
MT-ACT
Observation history
False
Action chunking
True
18

---

Table 5: List of tasks in Meta-World and DM Control.
Meta- Meta-World
DM Control
basketball-v2
cartpole swingup
bin-picking-v2
cheetah run
button-press-v2
hopper stand quadruped run quadruped walk teacher easy
buttor-pres-toopooww-22
putton-pres--topdonn-walll
button-prees-wall--22
coffee-button-v2
walker stand
coffee-pull-v2
walker walk
coffee-push-v2
walker run
dial-turn-v2
disassemble-v2
door-lock-v2
door-open-v2
door-unlock-v2
drawer-close-v2
drawer-open-v2
faucet-close-v2
faucet-open-v2
hammer-v2
handle-press-v2
handle-press-side--2
handle-pull-v2
peg-
inssett-sideee
reg-umplug-side-v2
plate-slide-22
plate-slide-back--2
plate-slide-bate bra--bbckkkiee back--sde------
sate-slide-side-v2
shelf phace-place-22
soccer-v2
stick-push-v2
sweep-v2
sweep-into-v2
window-close-v2
window-open-v2
D.2 Additional Analysis
In addition to the analysis in Section 4.4, we provide further comparisons here to better justify our design choices.
Separate vs. Shared Vision Encoders On the LIBERO-90 benchmark, environment observations include images from two camera views. Table 12 compares multi-task performance using either common encoder for both views or separate view-specific encoders. While separate encoders provide 2% boost in performance, this minor gain comes at the cost of 15% increase in parameter count per camera view added (since the visual encoders comprise 1.5M parameters our 10M parameter model). For our real-world experiments involving 4 camera views, this parameter increase would be even more significant. Therefore, in BAKU, we use shared encoder for all views to keep the model compact, assisting with faster inference speeds.
19

---

Fetch glass from rack: Fetch the glass from the lower rack.
Fetch towel from rack: Fetch the towel from the lower rack.
TOR
Fetch tea bottle from rack: Fetch the bottle of green tea from the lower rack.
Fetch water bottle from rack: Fetch the bottle of vitamin water from the lower rack.
FOR
Pick blue mug: Pick the blue mug from the kitchen counter.
Pick light blue bowl: Pick up the light blue bowl from kitchen counter.
SOR Pick orange from bowl: Pick up the orange from inside the light blue bowl kept on the kitchen counter.
Figure 5: Real-world policy rollouts showing BAKU'S capability complex manipulation tasks.
Data Efficiency Analysis We analyze the performance of BAKU with varying number of demon- strations Table and Table 9. We observe that each level of data availability, BAKU shows significantly higher success rate than MT-ACT and RT-1.
20

---

Pick coffee bag: Pick up the bag of coffee from the kitchen counter.
Pick box of corn starch: Pick the box of corn starch from the kitchen counter.
Lift blue plate from rack: Lift the blue plate kept the upper rack.
Lift white plate from rack: Lift the white plate kept on the upper rack.
sod
Lift black plate from rack: Lift the black plate kept on the upper rack.
Open oven door: Open the door of the oven.
S
Close oven door: Close the door of the oven.
Figure 6: Real-world policy rollouts showing BAKU'S capability in complex manipulation tasks.
Robustness to training seeds We provide results on BAKU, RT-1, and MT-ACT across seeds in Table 10. We observe that all three methods are robust to different seed values. Further, probabilistic approaches like GMM and diffusion might be sensitive to favorable seed values, and evaluating on single seed might make the result unreliable. Thus, Table 11 includes results across seeds on BAKU
21

---

Place glass rack: Place the glass the lower rack.
20
Wipe towel: Wipe the cutting board with a lowel.
Lift pan lid: Lift the lid of the pan kept on the kitchen counter.
Put coke can in basket: Pick up the can of coke and put it in the basket.
Put cream cheese basket: Pick up the block of cream cheese put in the basket.
SQX
orange bowl: Pick up the orange and put it bowl.
Put pear in bowl: Pick the pear and pul it in the bowl.
Figure 7: Real-world policy rollouts showing BAKU'S capability complex manipulation tasks.
with different multimodal heads. We observe that BAKU with different action heads is robust to the value of the random seed. Due to limited compute and the large number of multi-task experiments, provide these results on the LIBERO-90 and Metaworld benchmarks.
22

---

TOR
10C
Put tea bottle fridge door: Pick up the bottle of green tea and place it in the door of the fridge.
Put yoghurt bottle fridge door: Pick up the bottle of yoghurt and place in the door the fridge.
Put ketchup bottle inside fridge: Pick the bottle of tomato ketchup and put inside the fridge.
Lot
Put tomato can inside fridge: Pick up the can of tomato soup and put inside the fridge.
Fetch bottle from fridge door: Take the bottle of green from the door of the fridge.
Fetch yoghurt bottle from fridge door: Take the bottle of yoghurt out from door the fridge.
Fetch tomato from fridge: Take the can of tomato soup of the fridge.
Figure 8: Real-world policy rollouts showing BAKU'S capability in complex manipulation tasks.
Observation trunk input In our proposed architecture (see Section 3.4), the encoded observations from different modalities are passed individually as tokens into the observation trunk along with the action token to output the action feature representation. An alternative approach is to concatenate all the encoded inputs into single vector and pass it through the observation trunk. As shown in Table 12, for Meta-World and DMC, which each have only single input source, there is
23

---

Fetch water bottle from fridge: Take the bottle of vitamin water out the fridge.
Fetch knife from organizer: Fetch the knife from the organizer placed on the kitchen counter.
Figure 9: Real-world policy rollouts showing BAKU'S capability in complex manipulation tasks.
Set up table: Place the white plate and the glass the kitchen counter.
Pick broom and sweep: Pick up the broom and sweep the cutting board.
Pick towel and wipe: Pick the towel from the lower rack and wipe the cutting board.
Take bowl out the oven: Take the bowl out the oven and place it on the kitchen counter.
Put yoghurt inside and take water bottle out of fridge: Put the yoghurt in the door of the fridge and take out the bottle of water from inside the fridge.
Figure 10: Real-world policy rollouts showing BAKU'S capability long-horizon manipulation tasks.
24

---

Table 6: Real task-wise performance
Number of Demonstrations
Task
Successes (out of 5)
Baku w/
RT-1
MTACT
Baku
VQ-BeT
Fetch glass from rack
5
5
Fetch towel from rack
5
5
Fetch tea bottle from rack
16
5
5
Fetch water bottle from rack
16
0
0
5
5
Pick blue mug
16
5
5
5
Pick light blue bowl
25
5
5
5
5
Pick orange from bowl
27
0
0
3
4
Pick coffee bag Pick box of corn starch
19
5
5
5
0
5
5
Lift blue plate from the rack
Lift white plate from the rack
Lift black plate from the rack
12
5
5
Open oven door
17
0
0
0
3
Close oven door
27
0
3
4
Place glass on rack
19
5
5
5
5
Wipe towel
17
5
5
Lift pan lid
18
4
Put coke can in basket
19
3
3
Put cream cheese basket
19
5
5
Put orange into bowl
14
4
5
Put pear into bowl
17
0
3
5
Put tea bottle in fridge door Put yoghurt bottle in fridge door Put ketchup bottle inside fridge Put tomato can inside fridge Fetch tea bottle from fridge door Fetch tomato can from fridge door Fetch yoghurt bottle from fridge door
18
0
0
17
5
3
5
15
5
5
5
0
0
5
5
5
5
5
10
5
4
Fetch water bottle from fridge Fetch knife from organizer
11
5
5
20
5
5
Mean
17
1.83
2.8
4.3
4.53
Mean success rate (out of 1)
0.37
0.56
0.86
0.91
difference in performance, as expected. However, for LIBERO-90. which uses two camera views and the robot's proprioceptive as inputs, there a 3% absolute improvement in performance when using separate observation tokens as compare to a single concatenated vector.
E Broader Impacts
In this work, we present BAKU, a simple and efficient transformer architecture for multi-task policy learning. This work takes an important step toward enabling more efficient training of generalist robotic agents capable of performing diverse tasks, reducing the need for large datasets of expert demonstrations which are costly and time-consuming collect. Further, BAKU focuses on improving data efficiency by maximally leveraging available training data, which is particularly valuable in robotics where data collection is expensive.
25

---

Table 7: Real task-wise performance for long-horizon tasks
Number of Demonstrations
Task
Successes (out of 5)
MTACT
Baku
Set up table
34
3
Pick broom and sweep
5
Pick towel and wipe
Take bowl of the oven
18
5
5
Put yoghurt inside and take water bottle out of fridge
17
2
4
Mean
19
3.2
4.2
Mean success rate (out of
0.64
0.84
Table 8: Data efficiency analysis the Table 9: Data efficiency analysis the Meta- LIBERO-90 benchmark.
World benchmark.
# Demos
RT-1
MT-ACT
BAKU
Demos
RT-1
MT-ACT
BAKU
5
0
0.31
0.58
0.40
0.07
0.59
10
0,01
0.48
0.71
10
0.49
0.10
0.67
25
0.04
0.49
0.83
25
0.62
0.11
0.76
50
0.16
0.54
0.9
35
0.65
0.13
0.79
Table 10: Performance of multi-task policies learned using BAKU on LIBERO-99 and GoradW We report the mean and standard deviation for each variant across 3 seeds.
LIBERO-90
Meta-World
Method
(90 tasks)
(30 tasks)
RT-1
0.14 + 0.02
0.64 = 0.01
MTACT
0.55 + 0.01
0.12 + 0.01
BAKU (Ours)
0.89 + 0.01
0.81 + 0.02
Table 11: Performance of BAKU with different action heads on LIBERO-90 and Meta- report the mean and standard deviation for each variant across 3 seeds.
Meta-World, World. We
Action Head
LIBERO-90
Meta-World
MLP
0.89 0.01
0.81 0.02
GMM
0.83 0.02
0.64 + 0.02
BeT
0.88 0.01
0.77 + 0.01
VQ-BeT
0.9 - 0.01
0.78 = 0.005
Diffusion
0.88 + 0.01
0.64 + 0.01
Table 12: Study of design decisions for the model architecture that affects multi-task performance.
Category
Variant
LIBERO-90
DMC
Separate vs. Shared Vision Encoders
Common
0.90
Separate
0.92
Observation Trunk Input
Separate
0.90
0.79
0.70
Concatenated
0.87
0.79
0.70
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section 3 explain the architecture in detail with the performance and ablations presented Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations have been discussed in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
27

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include our architecture details in Section 3.4 and experimental details in Section 4 and Appendix A.4. We will also be making our code and data public on the project website.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
28

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
We will also be making our code and data public on the project website.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ttpp.cccppbblll/// httpcccccppcccccccceeeececc guides/ CCodeSubeiieeeee Sut
ssi ssisssoololooooooo oonooiiyyyylllooooooooooomiiimil for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We include our architecture details in Section 3.4 and all experimental details in Section 4 and Appendix A.4. We will also be making our code and data public on the project website.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Error bars not reported because of insufficient computational resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
29

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The details have been included in Appendix A.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : /wwrrrsr..p.../ cc hetppali//ttiisuuiiiitpee
Answer: [Yes]..
Justification:
The paper conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We include of broader impacts statement of our work in Appendix E.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
30

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: This work poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All datasets and environments used in this work are open-source and have been appropriately cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
31

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes].
Justification: We will be releasing documented code the project website along with real-world dataset. The submission has been appropriately anonymized.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "pLQTGQS2k1", "reviewer": "Reviewer_76DH", "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.", "questions": "1. How can reported results be considered reliable given single training seed and limited evaluation rollouts?\n2. Could you describe the initial state selection process for LIBERO and MetaWorld experiments?\n3. Why are MT-ACT MetaWorld results lower compared to those reported in other aforementioned works?\n4. What key differences between BAKU and BC-ResNet-Transformer from LIBERO contribute to BAKU's superior performance?", "limitations": "Limitations are adequately addressed.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "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. \n2. 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.\n3. 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.\n4. 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.\n5. 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.\n6. 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.\n\n[PRISE]: https://arxiv.org/abs/2402.10450 \n[QueST]: https://arxiv.org/abs/2407.15840v2 \n[Libero]: https://arxiv.org/abs/2306.03310", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n\n**“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.\n\n\n**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.\n\n\n**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.\n\n\n**“... 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.\n\n\n**“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. \n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ulmblIDhM2", "reviewer": "Reviewer_eFXw", "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.", "questions": "Details are listed above.", "limitations": "YES", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n2. 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.\n2. 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?\n3. 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.\n4. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\nWe 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.\n\n\n[1] Hoffmann, Jordan, et al. \"Training compute-optimal large language models.\" Proceedings of the 36th International Conference on Neural Information Processing Systems. 2022.\n\n[2] Hernandez, Danny, et al. \"Scaling laws for transfer.\" arXiv preprint arXiv:2102.01293 (2021).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dCjkvpVPFx", "reviewer": "Reviewer_FcXE", "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.", "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.)\n\nAbout keyframe extraction, which is an important part of the success of RVT-1, how do you define keyframes on LIBERO and real-robot data? \n\nAbout action prediction, did you use the same discretization way as presented in RVT-1 paper?\n\nIs it possible to study BAKU on RLBench, so we can have a more convincing comparison?\n\nDo 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?\n\nThis work will be more convincing if the baselines can be discussed more in detail.", "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.\n\nIts 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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "The presentation is clear and easy to understand\n\nThe authors conducted extensive experiments and solid evidence to support their proposed architecture\n\nIt 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\n\nSince 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.\n\nThis 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\n**“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. \n\n\n**“...  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.\n\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7udr5zG2Uy", "reviewer": "Reviewer_RQGT", "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.", "questions": "Is RT-1 fine-tuned on the tasks, trained from scratch, or are the numbers its zero-shot performance?\n\nAre 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.\n\nThere 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.", "limitations": "The paper discusses its limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper tests on a large number of simulated and real world tasks.\n\nThe 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.\n\nThe paper doesn't include any methodological innovations, instead combining existing approaches.  While this is ok, it makes having a strong evaluation even more critical.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n**“... 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.\n\n\n**“... 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.\n\n\n**“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.\n\n\n**“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.\n\n\n**“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.\n\n\n**“... 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:\n\n\n| Architecture | 2 Views | 4 Views |  \n|-|-|-|\n| Shared Encoder| 9.14ms | 9.5ms |\n| Separate Encoders | 10.69ms | 14.91ms |\n\n\nWe 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.\n\n\nWe 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.\n\n\n[1] Chi, Cheng, et al. \"Universal manipulation interface: In-the-wild robot teaching without in-the-wild robots.\" arXiv preprint arXiv:2402.10329 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KwNDWmEzGC", "reviewer": "Reviewer_Y1eK", "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.", "questions": "- Why have you only considered two offline imitation learning baselines?\n- Why do you not report variance in your results?", "limitations": "The paper appropriately addresses its limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Strong results relative to the chosen baselines and in both simulated and real-world collections of tasks.\n- Thorough investigation of different viable architecture designs.\n- Comprehensive ablations.\n- Excellent written communication.", "weaknesses": "- No variance or confidence intervals reported for results on simulated environments.\n- Limited number of baselines, with RT-1 being somewhat outdated.\n- It would be useful to see how performance varies across data scales for the different architectures (incl. ablations and baselines).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n**“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.\n\n\n**“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.\n\n\n**“... 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. \n\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their constructive feedback. The reviewers have requested important clarifications and additional experiments. We have provided these as a detailed response to each review along with a summary of some shared concerns in this global response.

**Concerns about novelty (Reviewers 76DH, eFXw, FcXE, RQGT):** Multiple reviewers have raised concerns regarding the novelty in BAKU. We agree that our work is focused on finding the right architectural choices for multi-task policy learning. Importantly, we observe that simple design choices lead to massive improvements in policy performance, achieving a 36% improvement on the LIBERO-90 benchmark and a 35% improvement on our real-world tasks. We view this simplicity as a strength of our work that will allow the community to build on this by making minor changes to get good performance. We believe that a deeper understanding of the novelty of the architecture choices can be found in Appendix D which compares BAKU with current state-of-the-art policy architectures such as MT-A​​CT and RT-1 and highlights the architectural choices in BAKU that lead to significant 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.

**Experiments including multiple seeds (Reviewers 76DH, Y1eK):** We provide results on BAKU, RT-1, and MT-ACT across 3 seeds in Table 1 in the attached PDF. We observe that all three methods are robust to different seed values. Reviewer 76DH suggested that probabilistic approaches like GMM and diffusion are sensitive to favorable seed values and evaluating on a single seed makes the result unreliable. Thus, Table 2 in the attached PDF includes results across 3 seeds on BAKU with different multimodal heads. We observe that BAKU with different action heads is robust to the value of the random seed. Due to limited compute and the large number of multi-task experiments, we provide these results on the LIBERO-90 and Metaworld benchmarks. More seeds and the DM Control environment will be added in the final version of the paper.

**New experiment on long-horizon tasks (Reviewer RQGT):** 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. Table 3 in the attached PDF provides the results on 10 tasks in LIBERO-10 and 5 tasks in the real kitchen environment, each composed of two shorter tasks chained sequentially. We use 50 demonstrations per task for LIBERO-10 and an average of 19 demonstrations per task for the real robot. We observe that BAKU significantly outperforms our strongest baseline, MT-ACT, on these long-horizon tasks, achieving on average 19% higher success rate. This highlights BAKU's ability to learn policies that can effectively plan and execute sequences of actions over extended time horizons. We include the real-world rollouts of these long-horizon tasks in Figure 1 of the attached PDF and have also included rollout videos on our anonymous webpage (link in the paper). We provide a description of the real-world tasks below for reference.

| Task                            | Description                                                       |
|---------------------------------|-------------------------------------------------------------------|
| Set up table                    | Place the white plate and the glass on the kitchen counter        |
| Pick broom and sweep            | Pick up the broom and sweep the cutting board                     |
| Pick towel and wipe             | Pick up the towel from the lower rack and wipe the cutting board  |
| Take bowl out of the oven       | Take the bowl out of the oven and place it on the kitchen counter |
| Yoghurt in, water out of fridge | Put yoghurt inside and take water bottle out of the fridge        |
| | |

**Data efficiency analysis (Reviewer FcXE):** We analyze the performance of BAKU with varying number of demonstrations 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.

**Code availability (Reviewer FcXE):** 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 (link in the paper). Following the updated NeurIPS instructions, we have also privately shared the code with the AC.

We have addressed other concerns specific to each reviewer in their rebuttal section. We hope that these updates to our paper inspire further confidence in our work. At the same time, we invite any further questions or feedback that you may have on our work.

Meta-review: The paper presents a multitask Behavior Cloning (BC) model for robot learning, comparing various components of recent BC models and combining their best aspects. 
Some of the concerns remains on the novelty of the algorithm presented in this paper is limited, as all components of the BAKU model are derived from other papers. However, the authors provide a very detailed set of experiments to 
Key design choices include observation trunk, model sizes, action heads, goal modalities, and action chunking. Experiments on simulated benchmarks and real-world tasks validate their design. The study highlights the effectiveness of simpler components like Transformer trunks and action chunking, while questioning the need for more complex elements. 

Furthermore, the statistical significance of the experiments is debatable due to limited rollouts and potential sensitivity to initial states. Furthermore some of the reviewers also pointed out time efficiency.  The rebuttal phase resulted in a number of clarifications, additional and subsequent score updates. 
Overall the meta reviewer finds that there is a significant  technical value in summarizing the empirical ideas into best practices, while trimming the unnecessary algorithmic components. This papers does a good job at it and provides a strong experimental evaluation.   The authors are request to release clean and reusable code along with model checkpoints for reproducibility of results and evaluations.  Furthermore, the authors should also update the papers using feedback from the rebuttal phase and additional experiments.
Decision: Accept (poster)