Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

Authors: lirui  wang
xinlei  chen
jialiang  zhao
kaiming  he
mit  csail
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 17:29:45

Abstract: One of the roadblocks for training generalist robotic models today is heterogeneity. Previous robot learning methods often collect data to train with one specific embodiment for one task, which is expensive and prone to overfitting. This work studies the problem of learning policy representations through heterogeneous pre-training on robot data across different embodiments and tasks at scale. We propose Heterogeneous Pre-trained Transformers (HPT), which pre-train a large, shareable trunk of a policy neural network to learn a task and embodiment agnostic shared representation. This general architecture aligns the specific proprioception and vision inputs from distinct embodiments to a short sequence of tokens and then processes such tokens to map to control robots for different tasks. Leveraging the recent large-scale multi-embodiment real-world robotic datasets as well as simulation, deployed robots, and human video datasets, we investigate pre-training policies across heterogeneity. We conduct experiments to investigate the scaling behaviors of training objectives, to the extent of 52 datasets. HPTs outperform several baselines and enhance the fine-tuned policy performance by over 20% on unseen tasks in multiple simulator benchmarks and real-world settings. See the project website (liruiw.github.io/hpt) for code and videos.

Main paper: 
 introduction: 
 Building robotic policies today is hard: it often requires collecting specific data for each robot, task, and environment, and the learned policies do not generalize beyond these specific settings. A historical lesson that has revolutionized machine learning is that pre-training [34, 27, 29] on large-scale, high-quality, and diverse data can bring general models that usually outperform specific models. Recent progress in open-source large-scale data collection [14, 76] has made this path possible, but the heterogeneity (such as varying robot hardware and different environments) present in large-scale robotic data has posed a significant challenge. A central question for the field now is how to leverage the heterogeneous robot data to pre-train robotic foundation models [3] .Foundation models from natural language processing [60, 58] and computer vision [37] have shown a paradigm to achieve general-purpose task-agnostic models through pre-training on massive amounts and diversity of data. In addition to the benefits from more data, training with diverse tasks also enforces the representation to be more generalized. These foundation models can achieve high task success rates for various tasks, are more robust to outliers, and are flexible for adapting to new tasks. These approaches map input signals from distinct domains and tasks into a high-dimensional representation space, and exhibit consistent scaling behaviors [34, 29] . After that, minimal fine-tuning is required to transfer the representation for downstream tasks to achieve good performance. It maps different embodiments, each with its own proprioception and vision sensors, onto a shared latent space by embodiment-specific tokenizers ("stems"). This aligns the heterogeneous data from different embodiments into a joint representation space. This allows us to train a shared Transformer trunk on the union of all heterogeneous datasets. The pre-trained Transformer can be transferred to a new embodiment, with a small, new tokenizer learned at transferring time.The heterogeneity in robotics presents a distinct challenge: different robots are physically different embodiments 1 of hardware acting in different environments. Each embodiment can have a distinct proprioception, including different degrees of freedom, end-effectors, motion controllers, and workspace configurations built for a specific application. Another common heterogeneity in robotics is vision heterogeneity. Robots are often equipped with different camera sensors mounted at different places (e.g. wrist and/or third-person) and the visual appearance of each robot varies dramatically due to environments and tasks. Both proprioception and vision information are crucial for complex, contact-rich, long-horizon behaviors in robotics. Poor learning of such information can lead to overfitting behaviors such as repeating motions for a particular scene and task or even trajectory.In this work, we propose to address this issue by aligning the proprioception and vision information from different embodiments to a shared "language" of policies through heterogenous pre-training (Figure 1 ). With such a shared representation, a new embodiment only requires minimal data and training to "translate" its specific setup to the shared "languages". In other words, we want to pre-train task-agnostic and embodiment-agnostic foundational models that can map raw sensor signals from individual embodiments into a shared latent space. Previous works have made significant progress in pre-training only the vision part of the policy on human videos [49, 54, 35, 62] and pre-training the full policy [6, 14, 56] with a unified model and dataset format (e.g. using languages [5] ). Additionally, they assume no proprioception in pre-training and add it post hoc in transfer learning.We introduce Heterogeneous Pre-trained Transformers (HPT), a family of architecture designed to scalably learn from data across heterogeneous embodiments. HPT modularizes a general policy network architecture (Figure 2 ) and pre-trains the policy representation of a latent transformer with supervised learning. Inspired by learning from multimodal data [1, 74, 20, 31] , we use embodimentspecific tokenizers, dubbed "stem", to align various sensor inputs such as camera views and proprioception inputs. The "trunk" is shared and pre-trained across datasets and is transferred when adapting to new embodiments and tasks that are unknown during the pre-training times. Moreover, we use task-specific action decoders, dubbed "head", to produce the action outputs. Crucially, after "tokenizing each embodiment", HPT operates on a shared space of a short sequence of latent tokens. This hierarchy is motivated by how humans handle feedback loops between specific motor responses and perceived stimuli at the level of the spinal cord's neural circuitry [69] .We extensively investigated the scaling behaviors and various designs of policy pre-training to the extent of more than 50 individual data sources (2 times more than [56] ) and model size of over 1 billion parameters. Analogous to the scaling laws [27, 29] , we found that to some extent, HPT scales with the dataset quantity and diversity as well as the model and training compute.In addition, heterogeneity can occur in different embodiment domains, such as real robot hardware, simulation domains, and human videos. We incorporate many available embodied datasets in different embodiments such as real robots [14, 76, 39] , simulation [82, 90, 50, 21, 86, 81] and internet human videos [15] in the pre-training process and demonstrate the generality of our framework including embodiments beyond expensive real-world on-robot teleoperations.Through transfer learning experiments across multiple simulation benchmarks [90, 50, 82] and real-world dexterous tasks, we compare with several baselines and the from-scratch counterparts.Overall, based on the pre-training objectives, HPT can scale with the model, data, compute, and the heterogeneity of the robotic datasets across real robots, simulations, and human videos. These pre- training procedures and models can simplify building reliable robotic policies for new embodiments and new tasks in terms of data requirements and generalized performance. As an attempt to scale heterogeneous pre-training, our code and weights are open-sourced, and we hope that HPT can shed some light on learning robot representations from heterogeneous embodiments and tasks. 
related works: 
 Pre-training and Transfer Learning. Pre-training [2] , through direct supervision [38] and/or self-supervision [57, 25, 12, 22, 10] , has been shown to learn representation useful for unseen downstream tasks in computer vision [7, 37] and natural language [58] , and their intersections [60] .The representation learned from ImagetNet [16] or web-scale data [38, 60, 18] shows robustness to distribution shifts and can be transferred to new tasks.The recent surge of foundation models [3] scales these representation learning methods [27, 29] by applying task-agnostic objectives to multitask data. Moreover, recent works [46, 42, 45] show that small projection layers can be used to align the pre-trained feature spaces of the foundation models. Different from other fields, robotics has less data quantity and diversity but much more heterogeneity.Alignment. Recent works such as Flamingo [1] , Perceiver [31] , and ImageBind [20] proposed ways to combine representations from multimodal data such as image, language, and audio by aligning these different modalities to the same latent space in the pursuit of representation learning. Our architecture design is also motivated by methods such as LLaVA [45] in the multimodal learning community. Very recently, GPT-4o [58] , Gemini [77] , MM1 [52] , X-VILA [89] , and Chameleon [75] demonstrated the capabilities of heterogeneous pre-training a universal transformer from and for multiple modalities. The idea of alignment, across modalities and/or embodiments, is important as we scale to use heterogeneous embodiments and reuse data from distinct embodiments.Representation Learning in Robotics. Representation learning has been explored in the robotic community. Previous works such as R3M [54] , VC-1 [49] , Voltron [35] , and SpatialVLM [11] investigate visual representations by training the policy with human videos and robotic data [70] . Mixture of Experts. Our architecture design is related to works in conditional computation and MoE [51, 44, 72] , where we create one expert for each embodiment, and the router (for the whole network) is determined by the embodiment. This technique has been used to scale language models to a substantial size [32] . 
heterogenoues pre-trained transformers (hpt): 
 In heterogeneous robot learning with cross embodiments, the data are generated from different domains such as simulation and real robots, across sensory modalities such as RGB images, language instructions, depth maps, 3D point clouds, and tactile images. Each robot is a unique hardware embodiment with varying degrees of freedom, end-effectors, sensor configurations, controller and action spaces, and application-specific physical setups.In the following sections, we discuss the HPT network architecture and the training procedure to address the heterogeneity above. We modularize the network architecture (Figure 2 ) into the embodiment-specific stem, the shared trunk, and the task-specific heads. Intuitively, the stems, shown in Figure 3 , are earlier layers of the neural network that align sensory inputs from heterogeneous embodiment and modalities into the shared representation space. The shared middle part of the network is called the trunk, which processes the sensory representation into a latent representation that can be used for multiple tasks. Finally, the last part of the network is the head, which maps that latent representation to the action space in individual tasks of interest. The training procedure, dubbed heterogeneous pre-training, assigns and aligns specific stem/head pairs based on the sampled embodiment and task data, and still enjoys the benefits of joint training in the shared trunk. This can be thought of as tokenizing each embodiment using neural networks and alleviating the need to unify embodiments into a homogeneous data form in standard training procedures.Stem. The stem θ stem (Figure 3 ) in HPT is composed of a proprioceptive tokenizer and a vision tokenizer. These tokenizers map heterogeneous inputs from different embodiments to a fixed number of tokens with fixed dimensions, which enables the trunk to treat them in the same manner despite large heterogeneity, as well as enjoy the scaling and inference benefits on fixed context length. The key idea is to leverage attention [78, 31, 9] to attend a fixed number of learnable tokens to features of the observations. Although we mainly focus on proprioception and vision, handling other kinds of sensor heterogeneity in tactile, 3D, and action inputs can be flexibly extended in stems.• Proprioception Tokenizers. In Figure 3 (left), for embodiment k, the proprioceptive tokenizer maps any sequence of robot proprioceptive information with dimension d k p (e.g. 7 for end-effector pose) into N p (e.g. N p = 16) tokens with dimension d with values ranging from 128 to 1024. To achieve this, we first use an MLP to map the proprioceptive input into a feature space with dimension d. We then apply sinusoidal position encoding and use attention across the state feature and the learnable tokens, to map into 16 tokens with dimension d. Proprioceptive information is critical in robot policy learning, but its usage is often as simple as feature concatenation with a vision encoder [41] .Teleop Figure 4 : Dataset Heterogeneity in Robotics. We show illustrations of dataset mixtures (each color is a distinct embodiment) from different domains including real robot teleop [14] , deployed robots [39] , simulations, and human videos [15] . See Appendix Section A for dataset mixture details. • Vision Tokenizers. In Figure 3 (right), the vision tokenizer can map any sequence of camera images (videos of multiple views) with dimension H × W × 3 into N v (we use N v = 16 by default) tokens with dimension d. To achieve this, we first use pre-trained frozen feature networks (e.g. 7 by 7 features from ResNet) and then flatten the features. After that, we again use attention across these features and learnable tokens, to map the vision input into 16 tokens with dimension d.After processing each modality individually in the time sequence order, we concatenate all modality tokens and add additional modality embeddings and sinusoidal positional embeddings. This is used as the input sequence to the trunk that we introduce below. To avoid overfitting, the stem only has a small number of parameters (one MLP and one attention layer).Related works such as Octo [56] and others [54, 49, 6] mostly focus on pre-training the vision backbone of the policy through masking or self-supervision. They often stack sequences of singleview images along channels [6] for a particular robot or use a large number of tokens (256 in [56] ).In contrast, HPT uses stems with pre-trained vision encoders to map arbitrary image sequences to a short sequence of tokens (16) . Moreover, rather than add in proprioception during transfer in related works, HPT jointly pre-trains the vision and proprioception parts, from heterogeneous datasets.Trunk. As the central component for pre-training, the trunk architecture follows a transformer, parametrized by θ trunk in the latent space with dimension d. The output token sequence length L is the same as the input token sequence length. The output token sequence is simply pooled as the final combined feature for the observation. The trunk is shared across different embodiments and tasks to capture the complex input-output relationships (i.e. the number of trunk parameters is fixed independent of the number of embodiments and tasks).Head. The policy head θ head takes the output of the trunk transformer and maps it to the action space A in each dataset. For each embodiment and task, the policy head can be an arbitrary architecture (e.g. MLP) that takes as input the pooled feature of the trunk and outputs a normalized action trajectory. The policy head is reinitialized for transferring to a new embodiment.Given a total of K datasets with heterogeneous embodiments sampled from different distributions D 1 , ..., D k , ..., D K , we letEQUATIONt } 1≤t≤T denoting the i-th trajectory of maximum length T of observation and action tuples. The objective is to minimize the following loss across datasetsEQUATIONL is behavior cloning loss computed as the Huber loss between the normalized action labels based on dataset statistics and the network's action predictions. θ = K k=1 {θ stem k , θ head k } ∪ θ trunk denotes the network parameters comprised of embodiment-specific stem and head θ stem k , θ head k for dataset k, and a single set of shared trunk parameters θ trunk across all embodiments. This training procedure has two axes of data scaling: the quantity M k for one dataset D k and the total number of datasets K. In the pre-training stage, only the trunk parameters are updated at every iteration, and the stems and heads for each heterogeneous embodiment and task are updated based on the training batch sampling. See implementation details in Appendix Section A.3.The policy transfer process is similar to aligning the features of the new domain (through pre-trained stem encoders) to the pre-trained embedding space of the trunk [42, 45] . Given a new dataset D K+1 from a new embodiment, the objective can be the same as pre-training or alternatives [13] . We reinitialize the head and stem parameters with embodiment-specific input and output dimensions (such as different proprioception and action dimensions), and freeze the weights of the trunk. 
experiments on pre-training: 
 In this section, we aim to answer the following question: Does HPT pre-training have a scaling behavior under heterogeneous data across domains? Default Setting. We use 27 robot teleoperation datasets, including a subset of the recently public Open-X Embodiment dataset [14] as the training corpus. By default, we use one camera view of the scene with the pre-trained frozen ResNet18 image encode to compute the vision features. We use proprioception information, such as end-effector poses and joint positions, whenever they are available and provided. We use a maximum of 1000 trajectories from each dataset and a total number of 16k trajectories, and a held-out validation dataset with a maximum 200 trajectories per data source. Furthermore, we use a model with a trunk size of 3.17 million parameters, which is denoted as HPT-Small (Table 1 ). The training uses a batch size of 256 for 80k iterations, which is around 0.65B tokens in the latent space that feeds into HPTs and around 5B tokens in the vision and proprioception token spaces (horizon-dependent). While we do not align or preprocess action space or observation space [56, 87] other than normalization, data cleanup and filtering would be very helpful.Scaled Setting. We use 200k trajectories with 52 datasets, including simulation (e.g. [50] ), deployed robots (e.g. [39] ), human videos (e.g. [15] ), from distinct embodiments in the training process. This includes many public and accessible robotic datasets. In addition to different tasks in different institutes, these heterogeneous mixtures of datasets (Fig. 4 and Fig. 13 ) come with multiple views, language inputs, and different observation inputs in different environments.We evaluate the HPT pre-training performance with the averaged validation loss (prediction errors on unseen trajectories) at the last iteration of pre-training. These validation datasets are fixed independent of the trajectory counts and models during training. Unless particularly noted, the validation datasets come from the same 27 datasets in the Default Setting. Note that it is unrealistic to evaluate the pre-trained models on many real-world robotic environments at scale and there are very few evaluation alternatives to measure large-scale pre-training if we ignore this objective. In fields such as NLP [30, 34] , training loss objective (e.g. perplexity) is often used to measure the progress of pre-training. Admittedly, there are several caveats to this metric including the closed-loop performance gap and the task success rate gap. We will address these issues in Section 5 on HPT transfer learning. See Appendix Section A and Section D for more details and discussions.Data Scaling. In Figure 5 (a), we observe stable and scaling validation losses even on increasingly heterogeneous embodiments. Moreover, we found the compute (e.g. samples seen per training run) and the data amounts needed to scale in tandem [34] to get closer to convergence in the training process. In the red line in Figure 5 (a), we observe better validation losses as we scale up the total number of trajectories, by using a larger model and doubling the batch size every order of magnitude increase in trajectory counts. Strictly increasing data while keeping others bottlenecked (HPT-S and fixed iterations) might cause an early plateau performance at around 1000 trajectories max per dataset, as shown in the blue line in Figure 5 . In Figure 5 (b), we also pre-train on an increasing number of datasets with a fixed number of epochs and evaluate on the fixed subset (first 10 datasets). We hypothesize that training with more embodiments contributes to the generalization of the trunk. These experiments can scale to the extent of 200k trajectories and 52 datasets.Model Scaling. In Figure 7 , we fix the number of datasets (27) in RT-X and use a maximum of 1000 trajectories for each dataset. We scale along model size (from 1M to 1B) and gradually increase the batch sizes from 256 to 2048 (doubles every order of model size increase) and use the larger dataset with 170k trajectories. We observe that when we scale to bigger models with larger amounts of compute (red line), the pre-training can achieve low validation losses until it is plateaued. We do not find a significant difference between scaling depth or scaling width.Epoch Scaling. In this experiment, we fix the number of datasets (27) and use a maximum of 1000 trajectories for each dataset. In Figure 6 , we observe that increasing batch sizes (Left), which effectively scales training tokens (Right), can generally improve the model performance until convergence. Another observation we have is to use distributed workers to load from as many datasets as possible to aggregate each batch. We hypothesize that the large variance of training on heterogeneous datasets can be reduced by using a large batch size. See Appendix B for more experiment details.We experiment beyond real-world robot teleop data, which is expensive to collect and scale. For the additional datasets, we consider 7 simulation datasets across many popular simulators Drake [82] , Mujoco [90, 50] , Isaac Sim [21] , and PyBullet [86, 81] , as well as Sapien [53] and Flex [67] , with image inputs and expert demonstrations. For the human datasets that lack proprioception and action information, we use poses and 2D positions as surrogates for the supervised policy learning objectives. We use in total 300 trajectories from EPIC kitchen [15] and PoCo [83] with a maximum trajectory length 1000. See Appendix Figure 13 and Table 4 for more details on the dataset compositions.In Figure 8 , we use a maximum of 1000 trajectories for each dataset and compare against the baseline of 27 datasets with evaluation on all the pre-trained datasets. We show that pre-training on additional embodiment datasets such as simulation and human video datasets can be possible, despite the large embodiment gaps with real robots. These datasets provide complimentary embodiment data to pure teleop data, and they illustrate how much heterogeneity can be handled in the HPT framework. 
experiments on transfer learning: 
 In the previous section, we evaluate pre-training using the validation losses. In this section, we answer the following question with task success rates in transfer learning: Can the pre-trained HPT model be transferred to new embodiments, tasks, and environments in simulation and the real world?Protocol. We evaluate the pre-trained representations on robot manipulation simulation benchmarks Meta-world [90] , RoboMimic [50] , and Fleet-Tools [82] . Each training dataset uses from 20-100 trajectories per task and each testing covers 50 episodes with different initial conditions. The policies use HPT-Small as the pre-trained trunk and reinitialize the stem and head for transferring.During the evaluation phase, we compare the following models: No Trunk uses only the stem and head without the trunk in the middle and trains from scratch as common practice [41] . From Scratch trains the entire policy from scratch with the trunk, Pretrained Frozen uses and freezes the pretrained trunk during transfer learning and Pretrained Finetuned loads the pre-trained HPT-Base trunk and finetunes the whole network end-to-end, and Pretrained Finetuned (HPT-XL) uses the same fine-tuning procedure with a pre-trained HPT-XL trunk with a lower pre-training validation loss. To reduce the variance, we conduct independent training runs and evaluations 5 times and average for each model. The inference time during transfer on an RTX 3070 GPU is 47Hz for HPT-base and 19Hz for HPT-XL, while a more recent GPU like A100 can be 3-4 times faster.Experiment. In Figure 10 (a) , we test the model on the downstream tasks in closed-loop simulation and observe improved task success rate using the pre-trained models ranging from HPT-B to HPT-XL, although pre-training for the simulation experiments only happens in the real-world embodiments.In Figure 10 (b), we run HPT on the recently released Simpler [43] Benchmark, which allows for comparing with Octo [56] , RT1-X, and RT2-X [14] on a high-fidelity simulation. We focus on three different tasks Close Drawer, Move Near, and Pick Coke Can in the Google EDR embodiment. For each task, we test several different initializations with a total of over 300 episodes for all tasks. Note that the pre-training corpus of HPT-S does not include [6] , and simulation tasks have a focus on language conditioning and do not expose proprioception inputs, which is not suitable for HPT. To address these issues, we finetune HPT on the supervised datasets with around 50 trajectories under the simulation protocol. We use HPT-base as the backbone for this experiment. We use the baseline results from [43] . See Section A.4 for more implementation and experiment details.Protocol. For the real-world experiments, we evaluate the HPTs on two different embodiments for tasks in pet care and assembly, which are not covered in the pre-training datasets [14] . In particular, for these two robots, we experiment with different observation spaces 1 camera v.s. 2 cameras as well as different action spaces relative pose v.s. absolute pose. For data collection, we experiment with both an Oculus Quest to collect relative pose control as action labels as well as kinesthetic teaching. The episode lengths of real-world teleoperation vary from 50 steps to 150 steps with 10 Hz control frequencies. We experiment with the tasks Sweep Leftover, Fill Water, Scoop Food and Switch Insertion, which require 5-20 seconds of interactions with granular or small objects with fine contacts, shown in Figure 11 . We collect around 100 demos for each task and evaluate them for 15 trials to measure the average success rate.Experiment. We adopt a similar transfer learning method in the previous section and evaluate the pre-trained HPT representations under real-world evaluation protocols. We train the policy with 20000 iterations with a batch size of 256 and a learning rate of 5e -6 . We defer implementation details to Appendix Section A.5. Quantitatively in Figure 12 , we observe pre-trained policies attain a better success rate over the No-Trunk and the From-Scratch baselines. In particular, the From-Scratch baselines in Fill-Water use the state-of-the-art diffusion policy architecture to illustrate the flexibility of the pre-trained representations. In Figure 11 , qualitatively, we observe better generalization and robustness to varying poses and numbers of granular objects, and varying camera configurations and lighting conditions with pre-trained HPT.On Table 3 , we perform an ablation study for the Sweep Leftover task. We also compare with R3M [54] , Voltron [35] , and VC-1 [49] . We use a finetuned model with the released backbone and weights. We note that these previous works focus on only pre-training the vision encoders of the policies with human videos. Finally, we compared with policies that train from scratch (From Scratch) and policies that do not use proprioception during pre-training (No Prop. Finetuned) and add in proprioception afterwards. All of our experiments use pre-trained encoders and the trainable parameters (stem and head) can be as few as 2% of the parameters. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We will open-source the code with detailed instructions to reproduce the main experimental results. The data can be generated on the fly as they come from online public datasets in Figure 13 .Guidelines:• The answer NA means that paper does not include experiments requiring code. 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, or 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 make it into the paper). 
conclusion: 
 There is room for improvement for many aspects including the dataset curation and pre-training objectives. Specifically, the embodiment splits in our balanced dataset mixture are rather simple. Moreover, careful data filtering to ensure the data quality is under-explored in this work. Also, this work has focused on supervised learning as the pre-training objective and the data size in tokens and training compute sizes in FLOPs only reach a moderate scale of LLM training to ensure full convergence. Although the model architecture and training procedure are modular and independent of embodiment setups, heterogeneous pre-training can converge slowly. For evaluation, both the simulation and real-world evaluation tasks are restricted to short-horizon manipulation tasks with a fixed embodiment, which might limit the benefits of using a higher-capacity model. Furthermore, the learned policies still do not offer very high reliability on the tested tasks (typically below 90%). See Appendix §C for some failure modes.Given the recent surge of scaled data, robot learning is still limited by its generality because of the heterogeneity, including different embodiments, tasks, and environments where the robots are operated. To handle the heterogeneity common in robotics, we propose HPT, a modular architecture and framework to embrace this heterogeneity through pre-training. We explore and scale HPT with heterogeneous datasets to over 50 available datasets. The learned representation can be transferred and improve performance in both simulation and the real world, and it shows correlations with pre-training performance. The code2 is open-source for future research. We hope this perspective will inspire future work in handling the heterogeneous nature of robotic data for robotic foundation models. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]Justification: The code of ethics is strongly followed in this work.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: The project focuses on pre-training from heterogeneous robotic data. There is no direct societal impacts associated with it.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. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).• Researchers should communicate the details of the dataset/code/model as part of 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 is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: Experiment Details. We conduct pre-training experiments across several orders of magnitudes in computes and data. The number of trajectories and transitions in the dataset is limited by the maximum number of trajectories in each of the constituent datasets. We use maximum episode counts per dataset ranging from 10 trajectories to 100000 trajectories, and the total trajectories range from around 300 trajectories and 6000 transitions to around 300k trajectories and 5 million data points.When training with 80k iterations, the approximate training epochs with fixed batch size 512 range from 200 epochs to 2 epochs. In terms of tokens, our experiment model ranges from 0.5 million to 1 billion parameters, the dataset tokens from all modalities range from approximately 32 million tokens to 5 billion tokens, and the tokens in a batch range from 0.03 million tokens to 2 million tokens (including sequence length). The compute FLOPs range from 0.03GFlops to 31GFlops. See Table 5 for more details of the scale and see Figure 13 for example lists of dataset mixtures. To facilitate future research, we will open-source the data processing scripts.Different from previous work [56, 87] , we use minimal amounts of processing and cleaning of the observation and actions in the raw trajectories. Specifically, the default training setup is to train 80000 iterations with a batch size 256, which is around 0.65B tokens in the latent space that feeds into HPTs and around 5B tokens in the perception token spaces of the raw perception inputs (such as image patches). Due to resource limits, for some bigger datasets such as Droid [76] , we did not process the full size.Real Robot Teleoperation Dataset. In total, we use a subset of 42 datasets in the Open-X Embodiment dataset [14] , including the recent Droid [76] dataset. These public datasets have high heterogeneity including distinct embodiment, environments to operate on, tasks to solve, etc.Simulation Dataset. For the additional 7 simulation dataset, we use the simulator benchmarks across all popular simulators Drake [82] , Mujoco [90, 50] , Isaac Sim [21] , and PyBullet [81] , as well as Sapien [53] and Flex [67] , with image inputs and expert demonstrations. These are used as additional training data from the simulation domains.Human Video Dataset. Since the human datasets do not contain proprioception and action information, we use hand poses and 2D positions in the image space as surrogates for the supervised learning objectives. In the PoCo [83] dataset, we use 3D positions of the hand and use 6-Dof poses extracted by ICP as actions, and for EPIC-Kitchen [15] , we use normalized 2D centers of the detection box as proprioceptions and the difference to the next frame as actions. We use in total of 2000 trajectories of video clips from EPIC kitchen with a maximum trajectory length of 500.Deployed Robot Dataset. To further increase the heterogeneity in the pre-training dataset, we consider FrodoBot-2k [39] dataset that involves plays of driving robots in the wild for gaming. This dataset is composed of deployed mobile robots in the wild. We use the front camera for this dataset. The action space of this robot is parametrized by linear and angular velocity actions and the proprioception space includes measurements from the IMU. We use a total of 150 trajectories and each trajectory contains more than 500 steps.See Figure 13 for visualization of some examples of heterogeneous dataset compositions. In practice when training on the mixture of these datasets, users can define customized sampling weights or apply stratified sampling methods. For generality, in this work, we use a balanced weight sampling method, commonly used in multitask learning. For these datasets, we process the visual features separately and save them on the disks.Stem. For vision inputs, we resize each image to the standard square size (224x224) before feeding into a ResNet18 [26] into a 7x7 vision modality token. These tokens are specifically the features before the final global pooling. If multiple views are available, we create individual projector MLP for each image view and then concatenate the vision tokens. For the vision encoders. In Figure 18 , We have experimented with multiple vision encoders such as MAE ViT base [24] , Dino V2 [59] , and CLIP ViT Base [60] . We choose ResNet and the default image size for their simplicity and common usage in policy learning. The investigation of more complex fusion and processing for vision features is left to future works.When language is used, we use T5 [64] to encode trajectory-level language modality tokens. Rather than using raw data and computing the tokens each time, these tokens are processed offline and saved as datasets before training.For low-dimensional inputs such as proprioceptions and actions, we first flatten these vectors and then apply normalization to each dimension, as a single token. We have also experimented with increasing the dimensions of these modalities by adding sinusoidal position embeddings. For multiple steps in the observation horizon, we concatenate the sequence for each modality separately. At inference time, these tokens are forwarded once and cached to avoid multiple calculations.We apply cross attention (with sinusoidal position encodings) and a shallow MLP projector [42, 46] within the stem to map different various sequences of tokens into a fixed number of tokens for that modality. The cross-attention layer has 8 heads and 64 hidden dimensions per head. We map the modality tokens into 16 tokens, 16 tokens, and 8 tokens respectively for image, proprioception, and language.Trunk. The trunk is parametrized by a decoder-only transformer architecture with embedding dimension h that we ablate from 64 to 2048, and with block numbers ranging from 16 to 80. Note that the number of parameters for the trunk scales quadratically with the dimension size and linearly with the number of layers. The trunk also supports loading from existing large language models. Refer to Table . 1 for more details on model sizes.The code is modularized so that the trunk training and transfer are independent of the encoder architecture or pre-trained weights of the stem, which can be ImageNet pre-trained ResNet [26] , R3M [54] , Voltron [35] , Dino v2 [59] , etc, and can be fintuned during transfer learning. It is also independent of the head, which can be diffusion policies, MLP, or transformer.Head. For the head architecture, we normalize the action spaces of each embodiment dataset to be between -1 and 1 element-wise, based on the dataset statistics, such that the output scale of the heads and the loss and gradients remain at similar scales. Since action trajectories and observation history can often help with the robotic problem, we pick a fixed action horizon (length 8) and observation horizon (length 4) for each embodiment. We then apply random masking along the time dimensions for each batch during training to be suitable for downstream tasks with different horizons. We support various types of policy heads in the network architectures such as standard MLP, transformer decoder, and diffusion policy. For the MLP head, we pooled the trunk feature (e.g. averaging) and then applied a 3-layer MLP. For the transformer decoder, we concatenate learnable tokens to the tokens before feeding into the trunk and use 1D convolution layer on these output tokens to regress actions. For diffusion policy in real world experiments, we use a diffusion head and train with DDPM [28] . Finally, the actions are unnormalized based on dataset statistics and Huber losses are applied for regression. The reason behind Huber loss is to balance between the "difficult frame" in robot trajectories and the easy but lengthy part of the trajectories.As discussed, HPT is a meta-level architecture where the stem, trunk, and head code are modular and each supports various architectural changes with pre-trained weights. The inference time, which includes all the pre-processing and encoder time, on the local computer with an old NVIDIA RTX 3070 GPU is 39Hz for HPT-base and 33Hz for HPT-xlarge. A modern GPU such as A100 will likely improve the speed by 3-4 times.We train HPT with AdamW [47] optimizer with a weight decay ratio 0.05, and a base learning rate of 0.0002 with a cosine learning rate schedule with warmups and dropouts. We apply proportional scaling of the base learning rate depending on the batch size. To support various horizons during Figure 14 : Flexible Head Architectures in HPT. We highlight that our HPT architecture is a meta-level architecture for policy learning, and it can work with various head architectures. The important takeaway is the scalable transformer architecture in the middle of the policy to absorb the diverse data and provide tokens for these policy heads to regress on the action outputs. transfer learning, we apply random masking along the time dimensions for each batch during training.Since action trajectories can have imbalanced losses along prediction horizons, we use Huber loss with δ = 0.1 (empirically found). We found the pre-training stage to be stable across various hyperparameters in learning rate schedules and optimizers, and the choice of the validation dataset.The code is open-sourced and the pre-released model can be downloaded easily from Huggingface.In practice, since training losses can vary across different datasets and our goal is to well on all embodiments and tasks, we apply a weighted sampling procedure for data loading. For every training iteration, we sample a dataset with inverse probabilities based on an exponential of its dataset size as a temperature. Specifically, we compute the squared root of each dataset size and sum these sizes to compute a normalization constant. For each batch item, we then sample from these dataset with the corresponding probability. This prevents large datasets from dominating a full training epoch, which is a common practice in multitask learning.Note that the stem and head for each embodiment are updated in different frequencies than the trunk, similar to a mixture-of-expert [72] training procedure. Especially under distributed training settings, each stem and head is trained with data from a particular embodiment and tasks, and the trunk will accumulate gradients from all batches from training workers. The compute resources for these pre-training experiments range from 8 V-100s to 128 V-100s and the training time spans from half a day to 1 month. The total dataset disk size is around 10Tb and the RAM memory requirement is below 50Gb. See Table 5 for summary details of the experiment.For simulation benchmarks, we use the released datasets as the expert demonstrations [90, 82, 13] . In summary, Metaworld [90] uses wrist camera view, and Robomimic [50] as well as Simpler [43] uses third-person view, with their own proprioception definitions by the dataset. Fleet-Tools [82] uses both views as inputs and uses the end effector poses as the proprioception inputs. We encode the image using pre-trained frozen ResNet features and normalize the proprioception inputs before passing them into the stem. We train single-task policies for all these simulation benchmarks except for Metaworld.For the Simpler [43] benchmark, we focus on the Close-Drawer, Move Near, and Pick Coke Can task and Google EDR embodiment with a visual matching setting. We test 9 different initializations with a total of 218 episodes. Note that the simulation tasks have a focus on language conditioning and do not expose proprioception inputs, which is not the most suitable testbed for HPT. To address these issues, we finetune HPT on the RTX supervised datasets with 79 trajectories as other simulation benchmarks. We use HPT-base as the backbone for this experiment.By default, we train with 20000 iterations with batch size 512 and small learning rate 1e -5 . The image and state stem are one-layer MLP with hidden dimension 128 and the head is two-layer MLP. We only use an observation window of length 1 and MLP as the policy head. Each training dataset uses from 10-100 trajectories per task and each test covers 50 episodes with different initial conditions. Each trajectory in the simulation has a slight difference in scene initialization. To reduce the variance, we conduct independent training runs 5 times, and the average for each baseline. In Figure 9 , we show illustrations of some simulation tasks we evaluated.Task Definition. We experiment with robotic tool-use tasks Sweep Leftover, Fill Water, Scoop Food and Switch Insertion across two different robot setups. While for both setups, we use Franka Panda as the robot, we note that the sensor locations as well as the action spaces are drastically different. We collect approximately 100 demos for each task and evaluate each task for 15 trials to measure the average success rate.During evaluation, a human supervises the robot at all times. An evaluation episode can be terminated due to safety concerns, robot faults, timeout, etc. An episode is considered successful if it accomplishes the task. In the Fill Water task, the success score of 1 means some water is poured into the bowl. In the Sweep Leftover tasks, a success score of 1 means all the piles are pushed into the plate, and a success score of 0.5 means some piles are pushed into the plate. In the Scoop Food task, a success score of 1 means some dog food is scooped and all is poured into the bowl and a score of 0.5 means some food is scooped. In the Switch Insertion task [91] , a success score of 1 means the switch is precisely inserted into the three pins on the PCB board. The robot moves to a pre-defined pose before it attempts the insertion. We pick these challenging tasks as they require contact-rich interactions with tools and granular objects, and they require high-precision as well as dense contacts. We make sure the initial conditions of the robots are the same. Due to the complexity of these tasks and human errors, the initial conditions of the object setups are not exactly the same.Transfer Learning. For the policy head in the real-world experiments, we have experimented with both MLPs and diffusion policies [13] . Fine-tuning only has active parameters of no more than 3Mb, compared to much bigger models (e.g. 100M) often used for a single task in other works. We use an observation history window of 2 with a small learning rate 2e -5 . We train with batch size 256 on a single NVIDIA RTX 2080Ti GPU for 20000 iterations (around 4 hours of wall time).In this section, we present some additional experiments and ablation studies.Based on the training curves of the four baselines in Figure 16 (a), we observe that leveraging a pre-trained HPT representation can often achieve a lower validation loss curve (lower) during fine-tuning. We do evaluation in a different implementation in unseen simulation benchmarks. Left) we show that an improvement in performance can be achieved with pre-trained HPT trunks and outperforms single-task state-ofthe-art architectures. We note that HPT trunks have not been pre-trained with diffusion heads and transformer decoder heads. Middle) we show the sample efficiency ablation study for HPT-Base. Right) We show model size ablation study in loss curves.MetaWorld RobomimicIn Figure 10 (a), we run HPT on the Simpler [43] Benchmark, which allows for comparing with Octo [56] , RT1-X, and RT2-X [14] on a high-fidelity simulation. We focus on three different tasks Close Drawer, Move Near, and Pick Coke Can in the Google EDR embodiment. For each task, we test several different initializations with a total of over 300 episodes for all tasks. Note that the pre-training corpus of HPT-S does not include [6] , and simulation tasks have a focus on language conditioning and do not expose proprioception inputs, which is not suitable for HPT. To address these issues, we generate training data using the RT-1 [6] as the expert, and finetune HPT on the RT-1X supervised datasets with around 50 trajectories and the simulation protocol. We use HPT-base with language tokenizers as the backbone for this experiment. We use the results from the paper [43] .Additionally, under the LeRobot [8] implementation, we compare HPT with state-of-the-art singletask policy architecture such as Diffusion Policy [13] and ACT [92] . In particular, we inherit a simple version of these complex policies as the head architectures (Figure 14 ). We run full training runs with 50 evaluation trials every 20000 training steps, and use the maximum task success rates during the 100000 training steps for comparisons. In Figure 17 , we compare HPT with DP on the PushT task with keypoint representations using the diffusion head and achieve similar success rates at 78%. We also compare with ACT on the Aloha Transfer Box task with the image representations using the transformer decoder head and achieve similar success rates at 60%. This showcases the flexibility of HPT architecture to work with state-of-the-art policy architectures for high-frequency control in fine-grained tasks. Moreover, in the experiment with 50 episodes in total and HPT-base, we observe an improvement in sample efficiency with pre-trained models. We also see improvement in transfer learning loss with pre-trained models at increasing scales.In Figure 10 (b), we ablate on the number of visual and proprioceptive tokens in the simulation transfer learning experiments. We observe that missing either information hurts the performance of the downstream policies. See Section A.4 for more implementation and experiment details.In this experiment, we fix the number of datasets (27) in Open-X datasets and use a maximum of 1000 trajectories for each dataset. We consider several ablations on the stem part of the HPTs. Specifically, in Figure 18 (a, b), we observe an increase in validation losses when not using either the proprioception information or the vision. Intuitively, not using such proprioception or vision information would make learning action predictions from heterogeneous datasets very challenging. This also implies that both information are critical for policy pre-training at scale.We also conduct an ablation experiment over vision backbones on a smaller subset of the training datasets among several popular vision encoders such as ViT-base [24] and DiNO [59] (Figure 18 ). Further ablating on input image resolution or joint finetuning of the vision backbone on the downstream task success rates will be interesting future work. Although the default implementation focuses on single-view visual information, the stem can naturally extend to multiple views and other modalities such as languages and action history.In Figure 15 , we conduct several ablations to pre-train the HPTs, such as using multiple views to provide 3D information implicitly, using languages for task guidance, and using previous action trajectories as additional ablations. We found these ablations to improve over the default HPT setting with a single view and vision and proprioceptive inputs. These ablations are more pronounced in certain datasets such as multiple views for insertion [68] , and language modality for Language Table [48] . Using previous action trajectories is helpful in providing additional context and embodiment information as well. We believe that integrating multiple modalities and investigating how to handle missing modalities is an exciting future direction. We leave the exploration of these ablations to future work.From the architecture perspective, we also ablate over the token sizes and observation history and do not find a big effect on the averaged validation loss. We hypothesize there is a trade-off between the information given to the policy and the desired generalization. In Figure 15(b, c ), we have also experimented with a discretized version [6] of the policy head and architecture that uses additionally learned position encodings as action tokens for transformer, with conv1D head for action regression. We opt for regression on continuous values for its generality. An initial investigation of the attention map of the transformer blocks shows that there is a dense attention weight attributed to the proprioception and vision tokens.In Figure 19 , we show some failure cases of the learned HPT policies in the real world. One of the common issues is overshooting or undershooting. For example, the policies tend to pour the water before it reaches the mug or pour the dog food a bit in front of the bowl. These issues could be due to the spatial precision of the policies and due to the data qualities failing to uncover the causal relationships. More targeted data recollection and better finetuning of the vision encoders might help address these issues.We first discuss the metric used for measuring pre-training performance, or intrinsic evaluation. Validation (or training) loss is a common metric to evaluate the progress of large-scale pre-training [34] . It is based on the goal of training a generalist model, where even fitting all the data can be a challenge [58] , as opposed to training a specialist model where overfitting is often appreciated. This is also considered a step towards more scientifically studying pre-training and scaling behavior in robotics [30] . Previously, people often rely on a single binary metric for evaluating task success rates in the real world, with only some amount of test trials.Admittedly, there are several caveats to this metric. From the evaluation perspective, the averaged validation loss depends on the overall multitask losses of all datasets. For example, increasing the number of trajectories in one dataset might lower validation loss on the associated dataset, but may not lead to an overall loss decrease. In practice, the exact subset for evaluation and the number of training steps in each dataset also play a role in the averaged validation loss metric. For example, when evaluating whether additional datasets to the default setting contribute to the representation learning, selecting the default datasets allows us to compare on the same metric. But these datasets are inevitably trained less, if we fix the number of total training iterations. Moreover, the validation loss during pre-training and the downstream policy learning performance have an evaluation gap due to task differences, and downstream policy learning and task execution have another evaluation gap due to closed-loop execution.Given the recent surge of scaled data, robot learning is still limited by its generality because of the heterogeneity, including different embodiments, tasks, and environments where the robots are operated. We propose HPT, a novel architecture and framework to embrace this heterogeneity through pre-training. We align proprioception and vision information of different embodiments through a modular stem, a shared scalable transformer trunk, and task-specific heads to actions. We explore and scale HPT with heterogeneous datasets to over 50 available datasets. The learned representation can be transferred and improved performance in both simulation and the real world. We hope this perspective will inspire future work in handling the heterogeneous nature of robotic data.Since fine-tuning is still required for robotics generalist models, future research directions include exploring different algorithms including network architectures that incorporate embodiment-specific structures, such as URDF, into network architecture as well as tokenization. One can also explore training objectives beyond supervised learning. As we move towards scaling robot pre-training, more high-quality diverse datasets with clear annotations would be crucial, including teleoperation data, simulation data, human videos, and deployed robot data. Understanding what kinds of mixture will contribute to better representations is interesting future work.Due to the complexity of real-world evaluation, large-scale unified simulation benchmarks with varying dexterity and generalization challenges would be very crucial to consistently compare among different models. For real-world policy performance, we believe that extending to longer-horizon fine manipulations with a bimanual or mobile setup would be interesting future works.Future works can also study and extend scaling laws in policy learning, and explore representations from heterogeneous data in other domains beyond robotic policy learning. Finally, leveraging more modalities and domains in robotics such as 3D point clouds, tactile data, simulation domains, and human data, etc, is worth investigating.• 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 be possible for other researchers to have some path to reproducing or verifying the results.Embodiment can be defined differently according to the context of robotics and AI. In this work, we consider robots equipped with a distinct set of sensors and actuators with the associated observation and action space to be a unique embodiment.https://github.com/liruiw/HPT and https://github.com/liruiw/lerobot 


Decision: This submission proposed a new method, called HPT, for quick (task, embodiment, and environment) adaptation by pretraining a policy representation from large datasets. In order to encourage shared learning across multi-modal data, the model divides the policy into three parts: a dataset-specific token encoder (Stem), a shared processing trunk, and dataset-specific action heads. The authors showed that this approach of dividing into three parts helps effective scaling and adaptability on both simulated and real-world tasks. The results are shown by finetuning with few demonstration in the Open-X dataset.

Overall, reviewers agreed that the method and the results are compelling and it should be shared with the community. There are several experiments and analysis reviewers asked for to further improve the submission. Please add them to your camera-ready copy.

Reviews: review: paper_summary: paper_summary: This paper presents an approach for training large transformer behavior models on diverse heterogenous data. This involves separate encoders for proprioception and image observations for each embodiment, as well as different action decoding heads. Experiments include validation loss comparisons.
strength_weakness: strength_weakness: After rebuttal - Concerns regarding validation loss vs success rate has been addressed with additional experiments, claims have been further strengthened with real world results

########################

1. Evaluation Metrics in the experiments

For training robot behavior models, the standard metric used in the community is the success rate of the trained model after being fine-tuned on a few demos for a new task [1,2,3]. The success rate measures the robot's ability to actually perform tasks of interest, using the large pre-trained weights. This is a more accurate metric than validation loss for how to judge if the large scale model training is making use of the diverse quantities of training data. Validation loss is not perfectly correlated with downstream robot performance. If it is difficult/cumbersome to do these in the real world, this analysis should atleast be done in simulation. The reason validation loss is a much weaker indicator in the robotics setting is because it doesn't measure how good the model is on the on-policy distribution, which can diverge wildly from the expert data distribution used for training. 

2. Design choice justification

The proposed approach has separate encoders and decoders for each observation and action modality. This seems like it could lead to less sharing of information across diverse datasets, leading to poorer transfer capabilities. Why should we not instead have a single action decoding head, as is done in [2]? What if the same image embedding backbone is used for all all observations - this would lead to the sharing of features potentially enabling better transfer. This choice of separate encoders/decoders for each modality is not sufficiently ablated and analyzed. 







[1]: RT1-Robotics Transformer for Real-World Control at Scale, Brohan et al.

[2]: Octo - Octo: An Open-Source Generalist Robot Policy, Octo Model Team

[3]: OpenVLA: An Open-Source Vision-Language-Action Model, Kim et al.
questions: questions: 1. How does the presented approach perform when deployed on real/simulated robots in terms of success rate, and how does it compare to previous approaches [1,2,3] ? Is it important to have separate encoders/decoder for each observation/action modality?








[1]: RT1-Robotics Transformer for Real-World Control at Scale, Brohan et al.

[2]: Octo - Octo: An Open-Source Generalist Robot Policy, Octo Model Team

[3]: OpenVLA: An Open-Source Vision-Language-Action Model, Kim et al.
limitations: limitations: NA
paper_summary: paper_summary: This paper aims to pretrain a policy representation across a variety of simulated and real robotics datasets such that it can be quickly adapted to new downstream problem instances (embodiment, environment, task). The key technical contribution is the design of a Transformer architecture that can consume a variety of robot data (i.e., observations and actions unique to each robot dataset) via a large sharable trunk + a number of small networks that translate dataset-specific information to this shared representation space. Experimental results indicate that the proposed architecture scales both in terms of dataset size / diversity and model capacity, and that it can be transferred to new real-world problems by finetuning on relatively few demonstrations (<=100).

**Post-rebuttal update:** I believe that my concerns have been addressed by the author rebuttal. I was already leaning towards acceptance and have now updated my score from 6 -> 7.
strength_weakness: strength_weakness: I'm overall quite pleased with the paper in its current form. However, there's a couple of thing that I'd like to point out:
- Since the authors are conducting transfer learning experiments in simulation (in addition the real experiments), I would have liked to see more ablations in this setting. Transfer performance seems like a way stronger signal for ablations, e.g., validating that more data does indeed lead to better transfer, not just better validation loss.
- It is not clear to me how the authors arrived at certain experimental setup choices. For example, the authors state on L282 that they use 20-100 demonstrations per task in experiments, but do not specify whether the model is transferred to each task *individually* or *jointly* (i.e., finetuning jointly on N*T demonstrations for T tasks and then evaluating vs. finetuning T models each on N demonstrations). It is also not clear to me which tasks require more demonstrations + whether this relates to the number of tasks considered in a given transfer experiment (e.g. 100 demos per task for Meta-World 5 vs. 20 demos per task for Meta-World 20 would be a reasonable comparison if the model is finetuned jointly on all tasks).
questions: questions: In addition to addressing my comments above, I would also like the authors to clarify the following:
- L282 states 20-100 demos are used, whereas L300 states that as few as 10 demos are necessary. Can the authors point to such an experiment, if available? Or is this merely an observation made during preparation of the paper? I did not find any such experiments in the appendices.
- Figure 10 (left) seemingly does not include the test loss of the larger model, whereas it is included in Figure 10 (right). Why is that?
limitations: limitations: I believe that limitations are adequately addressed.
paper_summary: paper_summary: This paper presents Heterogeneous Pre-trained Transformers (HPT), a method for training robotic models that addresses the challenge of heterogeneity across different robot embodiments and tasks. HPT pre-trains a shared neural network trunk to create a universal representation, which is then fine-tuned for specific tasks, showing improved performance and scalability with diverse datasets as well as demonstrating effective generalization to new tasks in both simulation and real-world settings.
strength_weakness: strength_weakness: Despite the difficulty of large-scale evaluation, there is no inevitable relationship between the loss values and the task success rates. For instance, the main challenge in some tasks lies in certain key frames, rather than the average error in lengthy trajectories.

In a sense, verifying the scaling behavior at the level of loss values is quite evident.

In terms of architectural design, this paper (stem, trunk and head) does not offer significant technical contributions.

Lacking comparisons with baselines (only some vision encoders on Sweep Leftover Task), such as RT-X and Octo.

In Figure 5, why is the loss lower without vision than without proprioceptive information? I am puzzled. Does "no vision" here mean that the policy has no visual information as input? Why can a reasonable loss value still be obtained without vision?
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: The paper introduces HPT a large-scale transformer model pretrained on multi embodiment robotic data. The main idea is to split the policy into three parts: a dataset specific token encoder for images and proprioceptive information called Stem. A shared Trunk, that processes all latent tokens and a dataset specific action head. By splitting the architecture into dataset specific encoders and dataset agnostic trunk the model is able to better process data from different robot embodiments and scale well using diverse data from Open-X-Embodiment Dataset (OXE). The proposed architecture is tested on pretraining on OXE and adapted for sim and real robot experiments.
strength_weakness: strength_weakness: - No adequate baselines in the real world: A comparison of a model pretrained on full robot trajectories with action prediction against a few pretrained vision encoders that are mostly trained on robotic-free first person human videos is not a fair or useful comparison and does not provide relevant insights for the readers. At least, I expect a comparison against Octo [1] as another generalist policy pretrained on OXE. Also the experiment is missing baselines trained from scratch such as Diffusion Policy or ACT. Further the experiments in the real world only cover a single robot embodiment. 
- For a method that is motivated by better learning from heterogeneous robot data, some additional results with other robots having different action spaces would be required to showcase the advantages of the proposed method
- No adequate baselines in simulation: The experiments on robotic datasets such as Robomimic and the real world again do not contain adequate comparisons. Specifically in the Robomimic experiment there is no comparison against state-of-the-art in-domain Diffusion Policy or ACT trained from scratch. The are no insights if the proposed model pretrained on diverse data improves upon these models trained from scratch. While comparison against the same architecture with different ablations exist, comparisons against other architecture are missing. Experiments on the recently proposed SimlerEnv [2] could also be interesting to compare the performance of the method against other large-scale pretrained methods such as Octo and RT-2X.
- Inference time: While the performance of larger variants of HPT is promising, I am not convinced by the idea that an $80$ layer transformer model is a suitable backbone for a policy used in a real robot context. 
- Several points about the image-encoder are not well described and miss relevant details.

[1]: Team, Octo Model, et al. "Octo: An open-source generalist robot policy." _arXiv preprint arXiv:2405.12213_ (2024).

[2]: Li, Xuanlin, et al. "Evaluating Real-World Robot Manipulation Policies in Simulation." _arXiv preprint arXiv:2405.05941_ (2024).
questions: questions: - Can you add adequate baselines such as Octo for the related experiments in sim and real to better justify the performance of HPT as mentioned in the weakness section. 
- I found the part about the image encoder a bit short of details. Is the ResNet pretrained on Image-Net in the final model or which pretrained image encoder are you using? Have you experimented with also fine-tunining it? Since ViT has a lower validation loss why did you decide to use ResNet as the default? Which pre-trained Vit are you using for the ablation? More details on this part would be helpful for the reader. 
- How fast is the inference time for the different models?
limitations: limitations: Yes

Comments: title: Thank you, Reviewer Vi4L!
Response to Reviewer wUp5 (Continued)
Response to Author Reply
Response to Reviewer wUp5
Response for Reviewer wUp5 (continued)
Response for Reviewer wUp5
Thank you, Reviewer AQbN!
Response to Review of Submission7169 by Reviewer wUp5 (Continued)
Response to Review of Submission7169 by Reviewer Vi4L (Continued)
Response to Follow-Up Questions (Continued)
Thanks for the answer. More follow-up questions.
Response to Follow-Up Questions
Thank you for the detailed answers - some followup questions
Response to Review of Submission7169 by Reviewer wUp5 (Continued)
Response to Review of Submission7169 by Reviewer Vi4L (Continued)
General Response to All Reviewers (Continued)
comment: Thank you so much for the response, Reviewer Vi4L!
Hi Reviewer wUp5,

We hope that we have addressed most of your concerns and questions. Since this is the last day of the rebuttal, please don't hesitate to let us know if you have any other final questions and/or comments. Thanks!
I thank the authors for their detailed response and the extra experiments conducted. I apologize for asking a comparison to OpenVLA [3] since it was released after the submission, this was an oversight. I had originally mentioned it in the weaknesses section of the review to point to methods that use success rate instead of validation loss to compare performance. This was to substantiate the point that validation loss is not typically used in the community to compare methods, and that the efficacy of the approach would need success rate on tasks instead. With the additional experiments in Fig 13.a as well as Fig 16.b, these concerns have been addressed. I also appreciate the extra real world experiments with a new robot embodiment. 

Given this, I am increasing my score (3->6)
As previously discussed, HPT pre-training is capable of handling multiple modalities, including images, proprioception, and language. For tasks that require specific conditions, such as those in the Simpler Benchmark, we include language instructions. While we did not explicitly explore using goal images as conditions and we do not think this is necessary for all robotic tasks such as single-object grasping, this could be another effective way to convey task information, potentially offering more detailed guidance than language instructions. Other than pre-training time, new modalities can also be added and feature backbones can be finetuned during transfer time thanks to the global pooling of the trunk.

Please do not hesitate to let us know if you have any other questions and/or comments.
Does this mean that HPT has no instruction/condition at all? I believe a multi-task policy should at least need a goal image as a condition. I'm sorry that I did not find these details in the paper.
> 3. As the authors mentioned, the architecture design of this paper is merely aimed at addressing the dimension alignment problems among different embodiments. However, some works such as SMART [3] also deal with different embodiments, using shared prediciton head and tokenizer during pre-training. I would like to know what findings and conclusions the authors have made in terms of model design?

Indeed SMART is a seminal and relevant work that also studies generalist policy, and we will add it to the reference. However, they have a slightly more targeted focus on RL with pure simulation experiments on deepmind control (DMC), originated from decision transformer. The main difference from our work is that we focus on heterogeneous pre-training with more challenging real world tasks, which includes 52 different datasets real world, simulation, and human videos. We also validate our findings in the real world transfer learning experiments. We also did not apply causality, masking, or dynamics learning in the pre-training process. For example, real world experiments demand much higher generality across RGB images by leveraging vision, as compared to simulation experiments in locomotion.

In terms of the model design, we highlight it’s **”not merely aimed to addressing the dimension alignment problems”**. We propose stem network that encodes various modalities into a shared latent space throught **cross-attention**. Because of the drastically different embodiments in the complex real world scenarios, we use **different stems and heads** for different embodiments, as opposed to “shared prediciton head and tokenizer”. In terms of findings, we ablate over different kinds of input modalities in this network architecture. For example, we find that adding language instructions and multiple views can sometimes help with the performance. We also have run extensive experiments in scaling data, training schedule, and model sizes, which are useful insights for future works.


> 4. I still cannot understand why the model can achieve a 34% success rate without visual information input. Could the authors elaborate on the details of the fine-tuning and evaluation setup? For example, what are the differences between fine-tuning and testing.

For that particular experiment using Fleet-tools on Figure 13, we adopt the HPT finetuning and evaluation procedure to reinitialize the head and stem, and train on 200 demonstrations generated by a motion planner, and we then evaluate on the unseen scenes. The robot starts at a fixed initialization and there are object initialization randomness between fine-tuning and testing. But the variation is not very significant, which is why there are non-zero success rates even without visual inputs. 

> 5. As for different datasets, what modalities are included in HPT's input? What types of task instructions are there?

We mainly use vision and proprioception as the input modalities for HPT. We can optionally use language instructions as the inputs. In the open-x embodiment datasets, there are various sorts of task instructions from high-level instruction for the entire task such as “route cable” in berkeley_route_cable dataset to low-level motion instructions such as “grasp the yellow block and turn it left” in taco_play dataset.



Thank you very much again for your constructive feedback. Please do not hesitate to let us know if you have any other questions and/or comments.
Thank you so much for the new feedback and suggestions. We will address your questions below. 

> 1. In transformer-based robot models, I believe that what is more challenging to scale up than training data and model capacity is the scale of evaluation. Before RT-X and Octo, there have already been lots of works on verifying the scaling laws of robotic learning. Compared with these works, I hope to see more in-depth experiment findings, such as which data are more important, rather than just finding that the model can scale up.

Thanks reviewers to bring these great works up and we will add these works to the reference. Indeed, the study on using transformer for pre-training has a line of research. In addition to real world and simulation transfer learning, we propose to use the large-scale validation losses to measure pre-training and scaling law similar to fields such as LLM, precisely because of the "challenges in the scale of evaluation" as you mentioned. What’s new after RT-X is the open-source large-scale heterogeneous data in the real world. Therefore, the main focus of this work is to leverage such real world expensive data, and intuitively we believe such data has less domain gaps to real world transfer learning and are considered more precious. At the same time, such data also presents huge heterogeneity in terms of the tasks and the embodiments. This is also one of the main difference from earlier work such as Gato and VIMA. For example, the real world training data can include mobile manipulation for dexterous tasks and large-scale language instructions for table-top settings.

In terms of data experiments in the paper, we have conducted experiments in scaling dataset number, data trajectories, and different domains of data including simulation and human videos. We have some findings such as increasing the dataset number can improve the trunk representations, despite its large distribution shifts. We discover that adding more trajectories (often >1000) for one specific embodiment such as in the Google Kuka dataset can often lead to marginal improvements, while the diversity of the pre-training datasets is still limited. 

The joint training performance improvement also applies to simulation and human videos, which are under-explored in the community. In the ablation, we have also investigated and presented the training performance of each specific dataset in the mixtures. We discover that there are several modes of learning, which might be due to the difficulties of that embodiment and task:  For certain tasks, the loss can also improve suddenly during training (grokking) and for some datasets, it can steadily improve or plateau. We will add more discussions on the data findings to the manuscript. 




> 2. In addition, I agree with review Vi4L's point that I believe it is inappropriate to use the validation set loss for a large number of experiments as a substitute, because the policy is multimodal, or the authors can ensure that the data in the validation set can well cover various situations. Changing the loss function will not solve this problem.

We have discussed extensively the rationale and motivations for using validation losses for **evaluating pre-training**. Please see our common answer and manuscript for details. Each validation dataset contains up to 100 trajectories which covers a wide range of test cases. As mentioned earlier, the scale of evaluation is rather challenging, so we are also open to other suggested metrics for evaluating pre-training performance. We have conducted additional transfer learning experiments to verify the correlations between representation and its downstream task performance. 

Thanks again for bringing this up. In our newly revised manuscript, we reorganize the content such that the paper focuses more on transfer learning (Section 5).

(to continue)
I have read the authors' responses and appreciate their efforts. However, I still have some concerns:

1. In transformer-based robot models, I believe that what is more challenging to scale up than training data and model capacity is the scale of evaluation. Before RT-X and Octo, there have already been lots of works on verifying the scaling laws of robotic learning. Compared with these works, I hope to see more in-depth experiment findings, such as which data are more important, rather than just finding that the model can scale up.

[1] A Generalist Agent. TMLR 2022

[2] Vima: General robot manipulation with multimodal prompts. ICML 2023

[3] SMART: SELF-SUPERVISED MULTI-TASK PRETRAINING WITH CONTROL TRANSFORMERS. ICLR 2023

[4] RT-1: ROBOTICS TRANSFORMER FOR REAL-WORLD CONTROL AT SCALE. RSS 2023

[5] Robocat: A self-improving foundation agent for robotic manipulation. 2023

2. In addition, I agree with review Vi4L's point that I believe it is inappropriate to use the validation set loss for a large number of experiments as a substitute, because the policy is multimodal, or the authors can ensure that the data in the validation set can well cover various situations. Changing the loss function will not solve this problem.

3. As the authors mentioned, the architecture design of this paper is merely aimed at addressing the dimension alignment problems among different embodiments. However, some works such as SMART [3] also deal with different embodiments, using shared prediciton head and tokenizer during pre-training. I would like to know what findings and conclusions the authors have made in terms of model design? 

4. I still cannot understand why the model can achieve a 34% success rate without visual information input. Could the authors elaborate on the details of the fine-tuning and evaluation setup? For example, what are the differences between fine-tuning and testing.

5. As for different datasets, what modalities are included in HPT's input? What types of task instructions are there?
We would like to thank reviewer AQbN for the response and for sharing your opinions in this review!
> For training robot behavior models, the standard metric used in the community is the success rate of the trained model after being fine-tuned on a few demos for a new task [1,2,3]. The success rate measures the robot's ability to actually perform tasks of interest, using the large pre-trained weights. This is a more accurate metric than validation loss for how to judge if the large scale model training is making use of the diverse quantities of training data. Validation loss is not perfectly correlated with downstream robot performance. If it is difficult/cumbersome to do these in the real world, this analysis should atleast be done in simulation.

> How does the presented approach perform when deployed on real/simulated robots in terms of success rate, and how does it compare to previous approaches [1,2,3] ?

The authors conduct transfer learning experiments in simulation + real and report success rates. They don't compare to OpenVLA [3] since it was not released at the time of submission. I don't think these comments are appropriate at all.
Thank you for responding to my questions in such detail. I have read the author rebuttal, responses to my questions, as well as those of my fellow reviewers. I believe that my concerns have been addressed and I have updated my score from 6 -> 7 to reflect that. This score is contingent on the authors incorporating reviewer feedback into the camera-ready version of their paper, especially the clarifications on experimental setup (finetuning T models each on N demonstrations), as well as the new experiments that were conducted during the rebuttal.
Dear Reviewer, we would like to thank you again for your efforts and time in providing thoughtful feedback and comments. We’ve conducted additional experiments according to your suggestions and replied to all the questions and concerns. Since the discussion period is ending soon, we would greatly appreciate it if you could let us know whether you have any additional comments. Thanks a lot!
Dear Reviewer, we would like to thank you again for your efforts and time in providing thoughtful feedback and comments. We’ve conducted additional experiments according to your suggestions and replied to all the questions and concerns. Since the discussion period is ending soon, we would greatly appreciate it if you could let us know whether you have any additional comments. Thanks a lot!
> Could you elaborate on the action spaces of the real world experiments you added for the rebuttal?

The action space for the switch insertion task is position control, as opposed to the velocity control in the pet-care tasks. It works as follows: The robot takes in the current end effector translation and observed image and input, and output the goal end effector translation that the robot needs to reach. The robot controller is an impedance controller that is tuned to executed these fine-grained insertion motions with compliance.


> Do you have any insights on fine-tuning vs Co-Training on your model that you could share? Which variant is better for your model? Given the modularity of your approach I am curious is Co-Training is the upper bound for your model too.

This is an excellent question. We believe that in the ideal case with enough engineering, co-training will work better than finetuning, as it circumvents the catastrophic forgetting issues in finetuning and allows for the full network to update. Indeed, as you pointed out, the modularity in our approach will create a new stem/head for a new domain, independent of finetuning or co-training. Therefore, the main difference lies in whether the trunk gets to share the training with other datasets and gets to update during adaptations. In our experiments, the joint training regime will work better, but a more focused sampling for the target dataset is necessary, because a balanced sampling regime might lead to under-training for the adapting domain.


Thank you very much again for your constructive feedback and the score raise. Please do not hesitate to let us know if you have any other questions and/or comments.
Thanks for the answers! I increased my score to 6.
 
I have some further questions: 

- Could you elaborate on the action spaces of the real world experiments you added for the rebuttal? 

- Do you have any insights on fine-tuning vs Co-Training on your model that you could share? Which variant is better for your model? Given the modularity of your approach I am curious is Co-Training is the upper bound for your model too.
> Is there a reason for not testing on other Simpler Tasks? As far as I know what HPT trained on both datasets, so it could be interesting to include the full results on the bridge tasks too for the final version of the paper. This comprehensive reporting could serve as a valuable baseline for future research in the field of foundation models.


Thanks the reviewer for bringing this up. We were running out of time when preparing this rebuttal. We will try to add the bridge tasks for the final version of the paper and to provide a baseline for future research.


> Training Code: While it's good that the pre-trained models will be open-sourced, I'm curious about the plans for releasing the training pipeline. Making this available could greatly benefit the robotics community.


Thanks the reviewer for bringing this up. Yes, following the great practice of Octo and OpenVLA, the entire codebase for the training pipeline and pre-trained models will all be open-source for the robotics community upon acceptance. We would love to see how the community will use our models as well as any feedback and potential improvements.


> Proprioceptive Input: Could you elaborate on why this approach was effective in your model, particularly in light of potential overfitting or causal confusion issues that have been observed in other works like Octo and also a lot in my personal experiences?

We first note that proprioception provides critical information for achieving human-like dexterous motions, especially in high-dimensional control such as dexterous hands and legged locomotion. Without appropriate architecture, extensive training, and appropriate data filtering, indeed there can be overfitting or causal confusion between the actions and proprioception. However, in this work, the proprioception is early-fusion with the vision information, both of which take equal amounts of tokens (16), which makes proprioception information usage different from the traditional usage: feature concatenation with the low-dimension vector after feature processing for the vision. We hypothesize that this design with large scale heterogeneous training might be the reason why proprioception contributes to our findings, as observed in Figure 5 and Figure 13.


Thank you very much again for your constructive feedback. Please do not hesitate to let us know if you have any other questions and/or comments.
I would like to thank the authors for their thorough and thoughtful responses to my concerns. The new experiments and details will significantly strengthen the paper.

However, I have a few remaining questions:

- Is there a reason for not testing on other Simpler Tasks? As far as I know what HPT trained on both datasets, so it could be interesting to include the full results on the bridge tasks too for the final version of the paper. This comprehensive reporting could serve as a valuable baseline for future research in the field of foundation models. 

- Training Code: While it's good that the pretrained models will be open-sourced, I'm curious about the plans for releasing the training pipeline. Making this available could greatly benefit the robotics community.

- Proprioceptive Input: Could you elaborate on why this approach was effective in your model, particularly in light of potential overfitting or causal confusion issues that have been observed in other works like Octo and also a lot in my personal experiences?
We also copy the general response here for completeness. 

- We highlight that the main motivation to tokenize embodiments and align them in the same latent space, is due to the practical constraints of heterogeneity in **proprioception** and action spaces. These input and output tensors with different dimensions cannot be handled with a **single network** technically (think about two weight matrices with different dimensions), and are **not considered** in previous works before. To handle these, we use a lightweight encoder-decoder approach in this project. The idea is partly motivated by the “experts” in the recent mixture of expert literature as well as the “projector” or “connector” in the multimodal literature. For the vision modalities, we actually **share** the vision encoder and only use a small adaptor to project the frozen vision encoders. 

- Therefore, this is also a meta-level choice that is brought by heterogeneous pre-training, and is independent of specific architecture choices (shared head or separate heads). While this HPT architecture is a specific architectural choice that we made, the idea is motivated to handle **more heterogeneity across domains and sensor modalities**. This is a more **significant point and is different** than previous works. There are more complex methods such as using GNN or a unified tokenizer across these embodiments to handle this issue, but we leave those for future work.

-  Moreover, this simple and general architecture choice is **not** conflicted with sharing information across embodiments because the encoders (stem) and decoders (head) do not contain many parameters, and are merely applied to match the dimensions and project tensors in different spaces to the shared latent space, similar to how multimodal foundation models are trained. Moreover, in all robotic generalist works including Octo and OpenVLA, finetuning is still necessary,  which means some information in the new embodiment is not learned during pre-training. 

> Lacking comparisons with baselines (only some vision encoders on Sweep Leftover Task), such as RT-X and Octo.

We have added additional experiments to compare against several generalist policies in the recently released Simpler Benchmark. See Figure 13(a) for the result.  Specifically, we compare the models across 3 different tasks in the Google GDR embodiment with the visual matching setup. We have found HPT to perform quite well on these benchmarks compared to Octo for instance. For the real-world experiments, unfortunately we have not been able to make Octo finetuning work well in our real-world setup despite weeks of effort, so its current finetuning performance does not match really well. It can be trained on the local machine but the trained policies are not doing anything meaningful. We hypothesize that there might be issues in the finetuning process.



> In Figure 5, why is the loss lower without vision than without proprioceptive information? I am puzzled. Does "no vision" here mean that the policy has no visual information as input? Why can a reasonable loss value still be obtained without vision?

Depending on the tasks, there could be coupling information from the proprioceptive input and the action output. Our hypothesis is that in some simpler tasks, merely memorizing proprioceptive history and doing extrapolation leads to better performance compared to vision-only systems. Therefore, even if no vision information is used, there can still be a loss value. We have added an additional experiment in Figure 13(b) to show in transfer learning experiments in simulation, vision tokens are still necessary to achieve good performance measured by task success rates.

Thank you very much again for your constructive feedback. Please do not hesitate to let us know if you have any other questions and/or comments.
We also copy the general response here for completeness. 

- We highlight that the main motivation to tokenize embodiments and align them in the same latent space, is due to the practical constraints of heterogeneity in **proprioception** and action spaces. These input and output tensors with different dimensions cannot be handled with a **single network** technically (think about two weight matrices with different dimensions), and are **not considered** in previous works before. To handle these, we use a lightweight encoder-decoder approach in this project. The idea is partly motivated by the “experts” in the recent mixture of expert literature as well as the “projector” or “connector” in the multimodal literature. For the vision modalities, we actually **share** the vision encoder and only use a small adaptor to project the frozen vision encoders. 

- Therefore, this is also a meta-level choice that is brought by heterogeneous pre-training, and is independent of specific architecture choices (shared head or separate heads). While this HPT architecture is a specific architectural choice that we made, the idea is motivated to handle **more heterogeneity across domains and sensor modalities**. This is a more **significant point and is different** than previous works. There are more complex methods such as using GNN or a unified tokenizer across these embodiments to handle this issue, but we leave those for future work.

- Moreover, this simple and general architecture choice is **not** conflicted with sharing information across embodiments because the encoders (stem) and decoders (head) do not contain many parameters, and are merely applied to match the dimensions and project tensors in different spaces to the shared latent space, similar to how multimodal foundation models are trained. Moreover, in all robotic generalist works including Octo and OpenVLA, finetuning is still necessary,  which means some information in the new embodiment is not learned during pre-training. 

> What if the same image embedding backbone is used for all observations - this would lead to the sharing of features potentially enabling better transfer. 

Thanks for bringing this up. Since we are not finetuning the pre-trained frozen vision encoders, all the image embedding backbones are actually shared across observations from different embodiments. Only the adaptation layers are different for embodiments, which are meant to project the vision features to the shared latent space, as commonly used in the multimodal learning literature. In this work, we only explore single-stage pre-training, finetuning, or joint training of vision encoders can potentially happen in the later stages of pre-training. We will leave those for future work.


> How does the presented approach perform when deployed on real/simulated robots in terms of success rate, and how does it compare to previous approaches [1,2,3]? 

We have added additional experiments to compare against several generalist policies in the recently released Simpler simulation benchmark. See **[Figure 13a]** for the result. Specifically, we compare the models across 3 different tasks in the Google GDR embodiment with the visual matching setup. We have found HPT to perform quite well on these benchmarks compared to Octo for instance. Also, OpenVLA is released after our submission and RT1 does not support finetuning. Hope this addresses your concern.

> Is it important to have separate encoders/decoder for each observation/action modality?

Please see our earlier answers. As for the benefit or motivation for creating separate encoders/decoders, we are targeting a simplistic approach to handling heterogeneous data and maximizing the information that can be shared. Critically, the separate encoders/decoders are relatively small (2% of the parameters) compared to the trunk. The majority of the information is contained in the pre-trained encoders as well as the shared trunk.


Thank you very much again for your constructive feedback. Please do not hesitate to let us know if you have any other questions and/or comments.
**2. Lack comparison with other generalist models [Vi4L,wUp5,KFGf]**: In **[Figure 13(a)]**, we discussed new experiments on the **Simpler [2] benchmark which demonstrated strong performance compared with other generalist models including RT-1, RT-2, and Octo.** This simulation experiment shows that our model is comparable to state-of-the-art models. We note that it can be challenging to apply the zero-shot version of these models to our setting and we also do not have access to their hardware. We also note that works like OpenVLA, brought by reviewer Vi4L, are released after our submission.

**3. Architecture Designs [Vi4L,wUp5]**: We want to provide some clarifications on our architecture design and hopefully shed light on our motivation. 

- We highlight that the main motivation to tokenize embodiments and align them in the same latent space, is due to the practical constraints of heterogeneity in **proprioception and action spaces**. Without any pre/post-processing, these input and output tensors with different dimensions cannot be handled with a **single network** (think about two weight matrices with different dimensions), and are **not considered** in previous works before. To handle these, we use a simple and efficient encoder-decoder approach in this project. The idea is partly motivated by the “experts” in the recent mixture of expert literature as well as the “projector” or “connector” in the multimodal literature. For the vision modalities, we actually **share** the vision encoder and only use a small adaptor to project the frozen vision encoders. 

- This is also a meta-level choice that is brought by heterogeneous pre-training, and is independent of specific architecture choices (shared head or separate heads). While this HPT architecture is a specific architectural choice that we made, the idea is motivated to handle **more heterogeneity across domains and sensor modalities**. This is a more **significant point and is different** than previous works. There are more complex methods such as embedding embodiment information, using paddings, GNN or a unified tokenizer across these embodiments to handle this issue, but we leave those for future work.

- Moreover, this simple and general architecture choice is **not** conflicted with sharing information across embodiments because the encoders (stem) and decoders (head) do not contain many parameters, and are used as simple methods applied to match the dimensions and project tensors in different spaces to the shared latent space, similar to how multimodal foundation models are trained. Moreover, similar to other generalist policy works including Octo and OpenVLA, finetuning is used at test time to adapt well to new embodiment and tasks.

**4. More real-world experiments and ablation [KFGf]**: We also added additional experiments on the insertion task with a **different robot embodiment**, shown in **[Figure 14]**. This task uses a different observation and action space and requires more precision than the pet-care setting. We have also made ablation studies on different pre-training models in the real world in Table 4. These new experiments further validate the observation on HPT pre-trained model. 
 

We thank the reviewers again for their raised points. I hope our new experiments and clarifications addressed some of the comments and concerns of the reviewers.