Paper: Foundation Reinforcement Learning: towards Embodied Generalist Agents with Foundation Prior Assistance (ICLR-2024)

Reviewer comment (Experiments):
Limitation: While this paper is robust in its findings and methodology, acknowledging potential limitations would contribute to a more balanced discourse and pave the way for future research avenues. One such aspect is the comparison with methodologies like VLM that have a noted advantage in generating language-instructed policies.

Author response:
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.

Reviewer comment (Experiments):
Extension to multi-task RL: It seems the authors train a policy on each of individual tasks in Meta-World. However, an intriguing extension of this work would be exploring the feasibility of training a singular policy across multiple tasks.

Author response:
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.

Reviewer comment (Novelty):
It looks like what happened in this paper is that the authors simply substituted the value and policy functions in the vanilla actor-critic with the VIP (Ma et al., 2022) model and the UniPi (Du et al., 2023) model, respectively. The so-called Reward Foundation Prior is still a 0-1 success signal as in the vanilla actor-critic framework. By doing so, the authors peddle a concept of foundation reinforcement learning. This manuscript looks like an empirical study of recent RL foundation models, without new theoretical contributions.

Author response:
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

Reviewer comment (Novelty):
From my perspective, 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.

Author response:
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.

Reviewer comment (Novelty):
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).

Author response:
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.

Reviewer comment (Experiments):
I am curious why you use R3M, rather than VIP, as the baseline for representation learning.

Author response:
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).

Reviewer comment (Experiments):
Eventhough paper mention about using foundation priors for policy, value and success-reward the results use ground truth success-reward and do not show any results with foundation success-reward prior. Given the proposed contribution is a framework with foundation prior in policy, value and success-reward it is essential that authors present results with a foundation prior for success reward as well.

Author response:
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.

Reviewer comment (Experiments):
The experimental setup is quite simple and doesn’t present extensive comparison with other baselines. For example, how does using a policy prior pretrained on a small dataset collected using a DRQv2 perform in comparison with the UniPi data? It is unclear if using data collected from a finetuned foundation model like UniPi is better than any other policy trained using in-domain data.

Author response:
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).

Reviewer comment (Experiments):
Comparison is only presented with DRQv2, R3M and other zero-shot baselines like UniPi and Prior. The experiment section needs to be expanded and compared with other commonly used baselines and SOTA methods on MetaWorld.

Author response:
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. 

| Success rate (%) with 3 runs | FAC  | DrQ-v2* | ALIX* | TACO* |
| ---------------------------- | ---- | ------- | ----- | ----- |
| bin-picking-v2               | 100  | 0       | 0     | 0     |
| button-press-topdown-v2      | 100  | 0       | 0     | 0     |
| door-open-v2                 | 100  | 0       | 0     | 0     |
| door-unlock-v2               | 100  | 0       | 100   | 33.3  |
| drawer-close-v2              | 100  | 100     | 110   | 100   |
| drawer-open-v2               | 100  | 0       | 0     | 0     |
| hammer-v2                    | 100  | 66.7    | 66.7  | 66.7  |
| window-close-v2              | 100  | 100     | 33.3  | 66.7  |

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

Reviewer comment (Experiments):
Experiments are only presented in 1 simple environment where results on some tasks are not that convincing. For example, in figure 3 ablations results on button-press-topdown task shows that FAC vs FAC w.o. policy prior performs almost the same. Similarly, 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. Authors need to present experiments on more complex environments to clearly demonstrate importance of each component.

Author response:
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.

Reviewer comment (Experiments):
The core contribution of FRL is using foundation priors for RL training but in the experiments authors finetune the UniPi distilled policy with in-domain data which raises concerns that whether the policy prior benefits are coming purely from small in-domain finetuning or from base UniPi demonstrations that are used for distillation. It is essential that to disentangle these factors and present detailed ablations in the experiments.

Author response:
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).

Reviewer comment (Experiments):
It’d be nice if authors can show benefits of foundation priors for policy without any finetuning on in-domain data for the tasks used in experiments as that is the most exciting result in my opinion.

Author response:
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).