Foundation Reinforcement Learning: towards Embodied Generalist Agents with Foundation Prior Assistance

Thank you for your comments and advice! We hope the following addresses your concerns:

Weakness:

For the weakness, "acknowledging potential limitations would contribute to a more balanced discourse and pave the way for future research avenues.":

Thank you for your suggestion. We agree that the limitations and future works should be discussed in the Discussion section. We have revised our paper in this section.

Question:

For the question "Extension to multi-task RL: ...However, an intriguing extension of this work would be exploring the feasibility of training a singular policy across multiple tasks":

Thank you for your comment. We think this question is very interesting. Currently, our distilled policy prior model is multi-task, but the agent we learned is not. We think that the policies learned in each environment can be distilled, and each policy can be regarded as a skill component conditioned on tasks. In this way, it can achieve multi-task learning in a better way. We will try this in the future.

Finally, thanks again for your detailed suggestions! We have revised our paper and updated it on the website. And we highlight the changes and essential details that reviewers have mentioned in blue color.

For the weakness, "This manuscript looks like an empirical study of recent rl foundation models, without new theoretical contributions...I don't think the proposed method can be claimed as a brand new 'Foundation Reinforcement Learning' framework since it simply uses two existing foundation methods.", and

for the question, "It would be appreciated if the authors could lay out more about the uniqueness of the proposed foundation actor-critic method to the vanilla ones except for the utilization of existing foundation models.":

The main target of our work is to propose the form of embodied prior and how to leverage them for RL. We do not train any new foundation models. But we emphasize that our framework is not limited to the model UniPi or VIP. Our proposed FRL and the FAC are agnostic to the foundation prior models. We just give an example of how to leverage the proxy foundation models to assist RL, and ultimately verify that our framework is effective. We have clarified in the paper that we mainly discuss the concrete form in which to represent embodied foundation priors. We have revised our story to be more accurate in paper. We think such a meta-level question we study is very significant because:

Currently, there is a large amount of research on building foundation models for embodied AI, as we have mentioned in related works. However, these works are very different regarding the form of foundation models and are not even comparable. For example, R3M learns the visual backbone, and saycan learns task decomposer, while RT chooses VLM finetuning in an end-to-end way. They are studying to train the embodied foundation models from different perspectives, and the questions they try to answer are not the same.

Instead, our work focuses on the concrete form in which to represent embodied foundation priors, rather than the actual RL algorithms that take advantage of the foundation priors. This is because there is no widely agreed embodied foundation model form that is widely accessible. And we think defining the form of the embodied foundation model is the first priority, which is the main contribution of this work. For example, in the BERT era, most researchers believed that BERT was the universal model, but GPT proposed to build large language models in an autoregressive way, which is another form to represent the language foundation prior knowledge. Similarly, what our paper discusses is exactly the form of the foundation priors in embodied AI.

Finally, thanks again for your detailed suggestions! We have revised our paper and updated it on the website. And we highlight the changes and essential details that reviewers have mentioned in blue color.

Thank you for your comments and advice! We hope the following addresses your concerns:

For the weakness, "The so-called Reward Foundation Prior is still a 0-1 success signal as in the vanilla actor-critic framework":

Thank you for this comment. We agree that it is necessary to conduct experiments with success-reward prior. Considering there are few universal success-reward foundation models, we distill one proxy success-reward model with 50k data in total for all 8 tasks, which are labeled by the ground-truth success signals from replay buffers. The proxy model takes images as input and is conditioned on the task embeddings (multi-task), which has 1.7% false positive error and 9.9% false negative error on the evaluation datasets. Then, we run FAC with the three priors without any signals from the environment during training.

The results are attached in Fig. 6 in App.A.1 (Page 14), which we will place in the main text later. We find that compared to receiving the ground-truth success-reward signals, FAC under the 50k-images-distilled model has a limited performance drop in the tasks generally. And it can achieve much superior performance than the FAC w.o. reward. Consequently, the proposed FAC can work well under the noisy success-reward foundation prior

For the weakness, "The proposed method is compared to a limited number of baselines. Experiments on more SOTA baselines and investigation of more foundation models are welcome.":

We emphasize that the proposed Foundation Actor-Critic is based on a new setting, where all the signals for reinforcement learning come from the foundation models. Notably, there is no human manual-designed reward in the setting. Considering that there are few works targeting a similar setting, all the baselines are listed in Fig. 2. We also compare our method to the oracle setting with human-designed reward, and find that FAC achieves even better performance. It proves that we can use the foundation model in place of human-specified rewards for any new tasks.

We also add some baselines under the same settings, where only the success-reward foundation prior is provided. We choose the recent SOTA model-free RL algorithms on MetaWorld ALIX [1] and TACO [2], as well as the baseline DrQ-v2 with the success-reward only. Notably, ALIX and TACO are both built on DrQ-v2. The results are shown in Fig. 7 in App.A.2 (Page 14-15), where '*' means that only 0-1 success reward is given.

For your convenience, we summarize the results in the following table. Only FAC can achieve 100% success rates in all the environments. DrQ-v2*, ALIX*, TACO* can not work on hard tasks such as bin-picking and door-open. And FAC requires fewer environmental steps to reach 100% success rates, as shown in the Figure. The additional experiment results on the new baselines can also verify the significance and efficiency of utilizing the abundant prior knowledge for RL.

[1] Cetin, E., Ball, P. J., Roberts, S., & Celiktutan, O.. Stabilizing off-policy deep reinforcement learning from pixels. ICML 2022.

[2] Zheng, R., Wang, X., Sun, Y., Ma, S., Zhao, J., Xu, H., ... & Huang, F.. TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning. NeurIPS 2023.

Dear Reviewer AdQL,

We kindly remind you that the rebuttal period is ending soon. Please let us know if our response has addressed your concerns.

Following your suggestion, we have answered your concerns and improved the paper in the following aspects:

• We conduct experiments with success-reward prior to demonstrate that the proposed FAC can work well under the noisy success-reward foundation prior.
• We have added 2 new baselines (ALIX, TACO) in all the 8 tasks to demonstrate the advancement of our framework.
• We have made further clarification of our story as well as the contributions, and more explanations for the ablation experiments. For one thing, we emphasize that we target the concrete form in which to represent embodied foundation priors. For another, our framework is not limited to the model UniPi or VIP. Our proposed FRL and the FAC are agnostic to the foundation prior models.

All of these results and modifications have been included in the revised paper.

Thanks again for your valuable review. We are looking forward to your reply and are happy to answer any future questions.

Thank you for your comments and advice! We hope the following addresses your concerns:

Weakness:

For the weakness, "The story may be confusing. I think they did not really propose a novel framework of foundation RL, but propose novel ways to leverage and combine existing foundation models (UniPi and VIP). The authors can change the story to be more accurate":

Thank you for your suggestions.

As you understand, our work is to propose the form of embodied prior and how to leverage them for RL. We do not train any new foundation models. But we emphasize that our framework is not limited to the model UniPi or VIP. Our proposed FRL and the FAC are agnostic to the foundation prior models. We just give an example of how to leverage the proxy foundation models to assist RL, and ultimately verify that our framework is effective. We call it a novel framework in the original paper because it is a paradigm shift. As you suggested, we have clarified in the paper that we mainly discuss the question (1) "the concrete form in which to represent embodied foundation priors", not (2) "how to learn such a foundation model from the dataset". We have revised our story to be more accurate in paper. We think such meta-level question (1) we study is quite significant because:

Currently, there is a large amount of research on building foundation models for embodied AI, as we have mentioned in related works. However, these works are very different regarding the form of foundation models and are not even comparable. For example, R3M learns the visual backbone, and saycan learns task decomposer, while RT chooses VLM finetuning in an end-to-end way. They are studying to train the embodied foundation models from different perspectives, and the questions they try to answer are not the same.

Instead, our work focuses on "the concrete form in which to represent embodied foundation priors", rather than the actual RL algorithms that take advantage of the foundation priors. This is because there is no widely agreed embodied foundation model form that is widely accessible. And we think defining the form of the embodied foundation model is the first priority, which is the main contribution of this work. For example, in the BERT era, most researchers believed that BERT was the universal model, but GPT proposed to build large language models in an autoregressive way, which is another form to represent the language foundation prior knowledge. Similarly, what our paper discusses is exactly the form of the foundation priors in embodied AI.

Question:

For the question, "I am curious why you use R3M, rather than VIP, as the baseline for representation learning.":

Thank you for this comment. We use R3M because it is the earliest universal visual backbone for robotics control tasks. As you have mentioned, we agree that it is necessary to train a baseline that takes the VIP as a visual backbone. We have updated the results in Figure 2. The results show that the model with VIP as the backbone (brown curve) fails in 5 environments with 0 success rate, which is slightly better than R3M (gray curve). Both of them have much poorer performance than FAC (green curve).

Finally, thanks again for your detailed suggestions! We have revised our paper and updated it on the website. And we highlight the changes and essential details that reviewers have mentioned in blue color.

Dear Reviewer rMFS,

We kindly remind you that the rebuttal period is ending soon. Please let us know if our response has addressed your concerns.

Following your suggestion, we have answered your concerns and improved the paper in the following aspects:

• We have made further clarification of our story as well as the contributions, and more explanations for the ablation experiments. For one thing, we emphasize that we target the concrete form in which to represent embodied foundation priors. For another, our framework is not limited to the model UniPi or VIP. Our proposed FRL and the FAC are agnostic to the foundation prior models.
• We have added 1 new baseline (VIP) for representation prior, as you suggest. The model with VIP or R3M backbone has much poorer performance than FAC.

All of these results and modifications have been included in the revised paper.

Thanks again for your valuable review. We are looking forward to your reply and are happy to answer any future questions.

Thank you for your comments and advice! We hope the following addresses your concerns:

For weakness 1, "it is essential that authors present results with a foundation prior for success reward as well.":

Thank you for this comment. We agree that it is necessary to conduct experiments with success-reward prior. Considering there are few universal success-reward foundation models, we distill one proxy success-reward model with 50k data in total for all 8 tasks, which are labeled by the ground-truth success signals from replay buffers. The proxy model takes images as input and is conditioned on the task embeddings (multi-task), which has 1.7% false positive error and 9.9% false negative error on the evaluation datasets. Then, we run FAC with the three priors without any signals from the environment during training.

The results are attached in Fig. 6 in App.A.1 (Page 14), which we will place in the main text later. We find that compared to receiving the ground-truth success-reward signals, FAC under the 50k-images-distilled model has a limited performance drop in the tasks generally. And it can achieve much superior performance than the FAC w.o. reward. Consequently, the proposed FAC can work well under the noisy success-reward foundation prior.

For weakness 2 "The proposed foundation actor critic algorithm is a simple extension to DrQv2 which just adds a KL constraint to the prior foundation policy.":

Thank you for your comment. We acknowledge that such KL regularization term is simple and widely used in other algorithms.

However, in this work, the main objective of this paper is to study the concrete form of the foundation models for embodied AI, rather than the actual RL algorithms that take advantage of the foundation priors. This is because there is no widely agreed embodied foundation model form that is widely accessible. And we think defining the form of the embodied foundation model is the first priority, which is the main contribution of this work. For example, in the BERT era, most researchers believed that BERT was the universal model, but GPT proposed to build large language models in an autoregressive way, which can be quite simple but effective. Similarly, our work aims to explore the concrete form of foundation priors for embodied AI. And we propose a novel learning paradigm for this, which does RL assisted by the policy, value and success-reward prior knowledge.

It is an example of using KL regularization in DrQ-v2 regarding the policy prior, which is not the main contribution in our work. The FAC is a concrete and effective algorithm to demonstrate how to leverage the foundation priors.

For weakness 3 "how does using a policy prior pretrained on a small dataset collected using a DRQv2 perform in comparison with the UniPi data?", and weakness 6 "whether the policy prior benefits are coming purely from small in-domain finetuning or from base UniPi demonstrations that are used for distillation.,

and question "the authors can show benefits of foundation priors for policy without any finetuning on in-domain data for the tasks":

These comments focus on how the policy prior works and why it requires in-domain data.

As mentioned in the previous answer, our main contribution in this work is to investigate the concrete form of foundation priors for embodied AI. The answer in our paper is to use the foundation policy, value, and success reward model. In order to demonstrate that our framework is a valid paradigm of leveraging the foundation prior, we assume the policy prior can provide roughly correct actions (or even discretized directions). However, we find that the current policy foundation models can not work in the out-of-domain environments. (We include an example of the video prediction from the diffusion model without in-domain finetuning for your reference in the Supplementary Material). Therefore, that is why we use a small amount of data (10 videos for each task) to finetune the existing policy prior model to act as if this is the foundation prior for the general purpose. Our foundation RL framework is agnostic to any foundation prior models. However, different qualities of the foundation priors will yield different levels of guidance.

We acknowledge that these foundation priors are not particularly ready. But the main target of this paper is to investigate what is the concrete form of the foundation priors for embodied AI and the importance of the priors. Hopefully, as more researchers agree with the significance of the proposed foundation priors, they will focus on how to build and improve the foundation prior (policy/value/success reward).

For weakness 4 "The experiment section needs to be expanded and compared with other commonly used baselines and SOTA methods on MetaWorld":

We emphasize that the proposed Foundation Actor-Critic is based on a new setting, where all the signals for reinforcement learning come from the foundation models. Notably, there is no human manual-designed reward in the setting. Considering that there are few works targeting a similar setting, all the baselines are listed in Fig. 2. We also compare our method to the oracle setting with human-designed reward, and find that FAC achieves even better performance. It gives the potential that we can use foundation model in place of human-specified reward for any new tasks.

We also add some baselines under the same settings, where only the success-reward foundation prior is provided. We choose the recent SOTA model-free RL algorithms on MetaWorld ALIX [1] and TACO [2], as well as the baseline DrQ-v2 with the success-reward only. Notably, ALIX and TACO are both built on DrQ-v2. The results are shown in Fig. 7 in App.A.2 (Page 14-15), where '*' means that only 0-1 success reward is given.

For your convenience, we summarize the results in the following table. Only FAC can achieve 100% success rates in all the environments. DrQ-v2*, ALIX*, TACO* can not work on hard tasks such as bin-picking and door-open. And FAC requires fewer environmental steps to reach 100% success rates, as shown in the Figure. The additional experiment results on the new baselines can also verify the significance and efficiency of utilizing the abundant prior knowledge for RL.

[1] Cetin, E., Ball, P. J., Roberts, S., & Celiktutan, O.. Stabilizing off-policy deep reinforcement learning from pixels. ICML 2022.

[2] Zheng, R., Wang, X., Sun, Y., Ma, S., Zhao, J., Xu, H., ... & Huang, F.. TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning. NeurIPS 2023.

For weakness 5 "...button-press-topdown task shows that FAC vs FAC w.o. policy prior performs almost the same. ...on door-open task FAC vs FAC w.o value prior performs almost the same. It is unclear from this experimental setup whether each component of the framework is important.":

Thank you for your question. For most environments, it will be best to apply all the foundation priors. However, in some environments, either the policy prior or the value prior is accurate enough for solving the tasks, resulting in a few performance drops when removing the other prior. This depends on the quality of the foundation prior conditioned on the tasks. Generally, learning from the three embodied prior can be better than that from only two of them in the 8 environments. Here, we explain in detail for your mentioned cases.

In button-press-topdown, the value prior can provide enough signals to guide the arm to press the button because the value is larger when the arm is closer to the button. And the policy prior can also provide a general direction. Thus, the agent can solve the tasks given only one prior (value or policy can both work) and can work faster given two priors together.

In the door-open, the arm needs to approach the door first and then move away to open the door. In this process, the value prior provided by VIP becomes smaller when approaching the door and larger when opening it. Without policy prior, the arm has no signals to learn approaching the door since the value prior is lower when approaching the door. That is why FAC vs FAC w.o value prior performs almost the same, but FAC w.o. policy prior can not work.

Finally, thanks again for your detailed suggestions! We have revised our paper and updated it on the website. And we highlight the changes and essential details that reviewers have mentioned in blue color.

Dear Reviewer zJMD,

We kindly remind you that the rebuttal period is ending soon. Please let us know if our response has addressed your concerns.

Following your suggestion, we have answered your concerns and improved the paper in the following aspects:

• We conduct experiments with success-reward prior to demonstrate that the proposed FAC can work well under the noisy success-reward foundation prior.
• We have added 2 new baselines (ALIX, TACO) in all the 8 tasks to demonstrate the advancement of our framework.
• We have made further clarification of our story as well as the contributions, and more explanations for the ablation experiments. For one thing, we emphasize that we target the concrete form in which to represent embodied foundation priors. For another, our framework is not limited to the model UniPi or VIP. Our proposed FRL and the FAC are agnostic to the foundation prior models.

All of these results and modifications have been included in the revised paper.

Thanks again for your valuable review. We are looking forward to your reply and are happy to answer any future questions.

Thank you addressing some of my concerns. I have read authors response to all reviews.

I agree with authors that their main contribution is a concrete form of foundation priors for embodied AI. However, to validate the proposed frameworks the experiments need to focus on diverse benchmarks, show results of meaningful generalization and be broadly applicable. The current results seem promising but are tested in limited tasks and it is unclear even with additional experiments that this framework is general enough to be applied in diverse set of tasks.

Another concern I have is the proposed FAC framework is essentially saying that we should be used pretrained components for each of the policy, value and reward modules which is not a novel contribution. I don't think that using them together can make the proposed approach a "foundation reinforcement learning" framework.

We emphasize that the proposed Foundation Actor-Critic is based on a new setting, where all the signals for reinforcement learning come from the foundation models. Notably, there is no human manual-designed reward in the setting. Considering that there are few works targeting a similar setting, all the baselines are listed in Fig. 2. We also compare our method to the oracle setting with human-designed reward, and find that FAC achieves even better performance. It proves that we can use the foundation model in place of human-specified rewards for any new tasks.

In response to W4 authors highlight this experiment, I would like to point out that authors shouldn't make strong claims like "It proves that we can use foundation model in place of human-specified reward for any new tasks" as these claims are backed by experiments on limited benchmarks.

Based on the responses, I would like to thank authors to address weaknesses W1, W2, W3 and W5 through experiments and clarification. I recommend authors to incorporate this in the manuscript. I will update my rating accordingly.