Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

Thank you so much for the extensive feedback and suggestions. We have revised the manuscript and conducted new experiments in the attached PDF and will address your questions below.

No adequate baselines in the real world .... 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.

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.

The Fill-Water task in the real-world experiment in Figure 14,15 is actually conducted with the diffusion policy architecture. Therefore the from-scratch baseline represents state-of-the-art single-task diffusion policies. We use the relative policy action output for the diffusion policies. We observe an improvement by using the pre-trained trunk compared to the train-from-scratch baselines. We highlight again that the ACT or diffusion policies are only considered heads of specialist policy in this framework, so it’s orthogonal to our contribution and is only used to showcase the flexibility of the pre-trained representations to improve model performance. In the Appendix, we also investigate the difference between using different heads in pre-training. We will clarify these in future versions of the writing.

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.

Thanks for bringing this up. We have added additional experiments on a different embodiment conducting real-world insertion tasks. The action space of this robotic setting is absolute pose, rather than the relative pose used in the pet care setting. Please Figure 14 and Figure 15 left for more details and visualizations. We hope this answers your question.

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. How fast is the inference time for the different models?

The inference time, which includes all the pre-processing and encoder time, on the local computer with an ancient NVIDIA RTX 3070 GPU is 39Hz for HPT-base and 33Hz for HPT-xlarge. It is expected to be 3-4 times faster on a more modern A100 used in OpenVLA. Indeed, the largest model can sometimes be overkill for local deployment of robot policies. But the HPT-XL model can easily run in real-time with our limited hardware in our real-world experiments.

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.

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.

I found the part about the image encoder a bit short of details ..... More details on this part would be helpful for the reader. Several points about the image-encoder are not well described and miss relevant details.

Thanks for bringing this up. We have added additional details to the network section in the appendix. We have copied it here for completeness. For vision inputs, we resize each image to the standard square size (224x224) before feeding into an Imagenet pre-trained ResNet18 into a 7x7 vision modality token, which are the features before the final pooling layer. If multiple views are available, we create individual projector MLP for each image view and then concatenate the vision tokens. We use ResNet mostly for its simplicity and legacy reasons (as the majority of the dataset is only processed once at the beginning of the project).

Due to the computation resource limitation and other practical reasons, we did not finetune it as in other VLM works by default. We were using the ViT-base from the masked autoencoder and CLIP for vision encoder ablation. We hypothesize that the vision features for robotic control do not differ by that much across these features learned from the standard images found in vision benchmarks or the internet. We think exploring this further can be interesting future work.

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 extensive feedback and suggestions. We have revised the manuscript and conducted new experiments in the attached PDF and will address your questions below.

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.

Thank you for bringing this up. Indeed, we use the smooth L1 loss to balance between the “difficult part” and the easy but lengthy part of the trajectories. I agree that it’s not optimal to use the validation loss and therefore we conduct many transfer learning experiments and only use the validation loss to evaluate the pre-training performance. Please see the paper for more discussions on using this metric. Since the submission, we have also conducted new experiments on a new real-world embodiment and did additional experiments in simulation benchmarks.

We also copy the general response here for completeness.

We highlight that the experiments in this work are three-fold: pre-training, evaluation in the simulation benchmark (Figure 16), and evaluation in the real world (Figure 14, 15). Only in the pre-training part, where we need to extensively measure the scaling behaviors of the learned representations, do we use the validation losses. It’s not an ideal choice, but it’s a common choice adopted by foundation model training in both LLM and multimodal foundation models. We note that the standard practice of robot learning (specialist model) is to overfit on ~100 trajectories, in which cases validation loss might not be that indicative, and evaluation is not that difficult. But in our case, we are training representations on massive amounts of robot embodiments and trajectories (generalist model), in which faithful evaluation of such representation without expensive rollouts is still an open question. Compared to earlier works that ignore this metric, we believe validation loss is one step towards that goal and should be representative of the learning progress in this process than single-task training. We also investigate their correlations with transfer learning performance in later experiments, where we show that bigger and more well trained models indeed perform better.

Additionally, we have proposed a simple fix (using smooth L1 loss) to make it more balanced between the difficult and easy parts of the trajectories. This does not include the high costs of evaluating representations in robotics. Previously, validation losses were simply ignored by practitioners, and our work is one step towards that direction. How to evaluate robot representations without running a large amount of rollouts is still an open research question. We believe that validation loss can be a more continuous reflection of representation at a large scale, and we will further revisit and investigate the correlations at a large scale in future work.

Last but not least, we did many experiments in the two sections on transfer learning in simulation and the real world. The simulation experiment is more extensive than previous works with thousands of runs. These preliminary experiments have also showcased the correlation of validation loss and transfer learning performance (Table 4 and Figure 16b). We want to highlight that previous works pre-trained on Open-X have evaluated mostly using the training embodiments, which we do not have access to. Our real-world evaluation tasks are also more challenging than previous works.

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

Thank you for bringing this up. This scaling behavior analysis has become a norm in the LLM domain, but in robotics, there has not been any work investigating similar problems to the best of our knowledge. Under such heterogeneous training across diverse embodiments and tasks, it is unclear to us beforehand if the training will even converge or proceed without instability issues. We will highlight the part on model transfer learning in the revised version.

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

Thank you for bringing this up. Indeed, the architecture was designed from a simple and general mindset to tackle the problem of heterogeneous learning. It is different from a monolithic model such as Octo and OpenVLA, in order to handle the heterogeneity in embodiments, and is also aligned with the multimodal learning literature. Besides the heterogeneous pre-training and architecture contributions, our work also provides insights into the scaling behavior investigations and conduct evaluation in simulations. (continued)

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.

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)

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.

Thank you so much for the extensive feedback and suggestions. We have revised the manuscript and conducted new experiments in the attached PDF and will address your questions below.

Since the authors conduct transfer learning experiments in simulation (in addition to 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.

In addition to the real-world experiments in [Table 4] where we show a better-trained trunk can achieve better performance and compare with several ablations. In [Figure 16], we also show that a better-trained HPT-XL model can outperform an HPT-Base model across all the simulation benchmarks. We also have done additional simulation experiments [Figure 13b] on transferring with and without vision and proprioception tokens to illustrate their benefits. Hope this addresses your concerns.

We also copy the general response here for completeness.

We highlight that the experiments in this work are three-fold: pre-training, evaluation in the simulation benchmark (Figure 16), and evaluation in the real world (Figure 14, 15). Only in the pre-training part, where we need to extensively measure the scaling behaviors of the learned representations, do we use the validation losses. It’s not an ideal choice, but it’s a common choice adopted by foundation model training in both LLM and multimodal foundation models. We note that the standard practice of robot learning (specialist model) is to overfit on ~100 trajectories, in which cases validation loss might not be that indicative, and evaluation is not that difficult. But in our case, we are training representations on massive amounts of robot embodiments and trajectories (generalist model), where validation loss should be more representative of the learning progress.

Additionally, we have proposed a simple fix (using smooth L1 loss) to make it more balanced between the difficult and easy parts of the trajectories. This does not include the high costs of evaluating representations in robotics. Previously, validation losses were simply ignored by practitioners, and our work is one step towards that direction. We believe that validation loss can be a continuous metric of representation without running infinite physical experiments at a large scale, and we will further revisit and investigate the correlations at a large scale in future work.

Last but not least, we did many experiments in two sections on transfer learning in simulation and the real world in the submitted version. The simulation experiment is more extensive than previous works. These preliminary experiments have showcased the correlation of validation loss and transfer learning performance (Table 4 and Figure 16b)We wanted to highlight that previous works pre-trained on Open-X have evaluated mostly using the training embodiments, which we do not have access to. Our real-world evaluation tasks are also more challenging than previous works.

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).

Other than the meta-world setup (where single-task evaluation becomes too easy), all experiments, including the real-world task, have training conducted on each task individually, as a standard setup. So we use “finetuning T models each on N demonstrations setup.”

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).

The only joint multitask experiment is meta-world. In that case, we try to use a demonstration number such that the trained policies performance is within a reasonable range, between 50 and 90 percent success rates. indeed, when we train with more tasks, we do not have to use the same number of trajectories for each task. There is no particular rule for which tasks require demonstrations in practice. It mostly depends on the learner's performance.

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?

The 10 demo requirement is an observation that we made on the robomimic Lift task, which is arguably a simple simulation benchmark task.

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?

For scaling ablation experiments, we mostly use the pre-training experiments as well as the task performance success rates as the metric, as the success rate has been brought up to be more important for robotics. We have completed a new set of experiments with different scales of pre-trained models. Please see Figure 16(b) for more details.

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 extensive feedback and suggestions. We have revised the manuscript and conducted new experiments in the attached PDF and will address your questions below.

Evaluation Metrics in the experiments. 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 at least be done in simulation.

Thanks for bringing this up and we are aware of the limitation on validation loss. Therefore, significant experiments have been done in simulation (Figure 16 b). We have shown that the more well-trained model HPT-XL can outperform HPT-base in the simulation benchmark. Moreover, we have added additional experiments to compare against several generalist policies in the recently released Simpler Benchmark. See Figure 13(a) for the result. While evaluation and metrics are still huge issues for robotics and generalist models, we believe that our simulation experiments are more extensive than previous works such as RT, Octo, and OpenVLA. Moreover, we kindly note that OpenVLA is released after the paper is submitted.

We also copy the general response here for completeness.

• We highlight that the experiments in this work are three-fold: pre-training, evaluation in the simulation benchmark (Figure 16), and evaluation in the real world (Figure 14, 15). Only in the pre-training part, where we need to extensively measure the scaling behaviors of the learned representations, do we use the validation losses. It’s not an ideal choice, but it’s a common choice adopted by foundation model training in both LLM and multimodal foundation models. As an analogy, there are also potential issues raised on the validity of perplexity for a small dataset in language modeling (which can also have out-of-distribution issues), but it is generally acceptable that perplexity is a good metric for studying large-scale pre-training these days.

• We note that the standard practice of robot learning (specialist model) is to overfit on ~100 trajectories with one day of training, in which cases validation loss might not be that indicative, and evaluation is not that difficult. But in our case, we are training with hundreds of GPUs and weeks of time with massive amounts of robot embodiments and trajectories (generalist model), where validation loss should be more representative of the learning progress.

• Additionally, we have proposed a simple fix (using smooth L1 loss) to make it more balanced between the difficult and easy parts of the trajectories. This does not include the high costs of evaluating representations in robotics. Previously, validation losses were simply ignored by practitioners, and our work is one step towards that direction. We believe that validation loss can be a more continuous reflection of representation at a large scale, and we will further revisit and investigate the correlations at a large scale in future work.

• Last but not least, we did many experiments in section 5 on transfer learning in simulation and the real world in the submitted version. The simulation experiment is more extensive than previous works with thousands of runs. These experiments have showcased the correlation between validation loss and transfer learning performance (Table 4 and Figure 16b), where larger and more well-trained models indeed can perform better. We wanted to highlight that previous works pre-trained on Open-X have been evaluated mostly using the training embodiments, which we do not have access to and we argue that measuring representation learning performance in a novel environment is also a valid choice. Our real-world evaluation tasks are also more challenging than previous works.

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 Octo?

The key motivation is that different observation and action spaces are very heterogeneous: they have different dimensions and live in different spaces. We want to separate out the part that can be potentially shared, which is the middle part of the network that contains more abstract representations. The proposed method is one architecture choice to handle such issues.

Encoder: It is technically impractical to map different dimensions of tensors (in proprioception) to the same space with the same network, unless we have a unified tokenization method. Our proposed cross-attention (or resampler) is one method to unify different inputs. In Octo, they do not consider proprioception and therefore can use a domain-specific token to represent the heterogeneity. When they need to handle the proprioception, they still need to reinitialize a new encoder for such domain and input (see their Berkeley Insertion experiment).

Decoder: We want to highlight that each of the action spaces can be drastically different from one another in robotics (such as pose and joint angles). Moreover, a finetuning procedure is still needed for any generalist model that we discuss in robotics. Octo trains a shared decoder based on transformer architecture and needs to reinitialize the task token during finetuning. In some cases, they still have to finetune the decoder head. The compute efficiency and the generalist design are not that different from our case. Indeed, in the early part of the project, we also experimented with Octo-like architecture with task-specific decoder tokens, but we did not find a big difference. We prefer to not expose the task information to the trunk to ensure it can learn task-agnostic representations.

(continued)

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.