Language Reward Modulation for Pretraining Reinforcement Learning

Baselines

Knowledge-Based Exploration algorithms are indeed not the only class of pretraining methods. However, we direct the reader to [7]. URLB conducts a systematic review of the various classes of unsupervised pretraining methods and finds that across every task evaluated, the Knowledge-Based methods (Plan2Explore, RND, and Disagreement) outperform the Data-Based methods (APT, ProtoRL) and Competence-Based (Skill Learning) methods (APS, DIAYN, SMM). This holds across all tasks from state but is even more pronounced for the pixel-based tasks where for the robotics tasks, evaluated the Knowledge-Based methods achieve ~4 times the normalized return of the Data-Based methods and ~15 times the normalized return of the Competence-Based methods (See Figure 3 and Figure 4 of [7]). This evidence guided our evaluation and selection of baselines.

[1] Exploration by Random Network Distillation (Burda et al 2018)

[2] Zero-Shot Reward Specification via Grounded Natural Language (Mahmoudieh et al 2022)

[3] Can Foundation Models perform Zero-Shot task Specification for Robot Manipulation (Cui et al 2022)

[4] Learning Transferrable Visual Models from Natural Language Supervision (Radford et al 2021)

[5] VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-training (Ma et al 2022)

[6] Language Models are Few-Shot Learners (Brown et al 2020)

[7] URLB: Unsupervised Reinforcement Learning Benchmark (Laskin et al 2021)

Primary Contribution

The primary contribution of this work is to propose a possible route towards pretraining a single generalist RL agent with limited human supervision that can efficiently adapt to many new tasks. Toward this goal we select a variety of downstream finetuning tasks for evaluation and do not include any of these tasks in our pretraining setup nor do we enforce any very strict requirements on their connection to the pretraining tasks. R3M has been used as a visual encoder to accelerate behavior learning by initializing with generalizable representations, however, to our knowledge has not been used as a reward function. Furthermore, instead of using a VLM to specify a task reward for a single language instruction or image goal as is done in [2][3], we demonstrate experimentally that by using a diverse set of language prompts with a frozen VLM, LAMP can be pretrained to do a wide array of tasks even more effectively than with the most widely-used exploration algorithms (Plan2Explore, RND). Our pretraining design is a divergence from existing approaches both in our inclusion of zero-shot VLM rewards from representation learning models and the way in which we generate them with diverse language prompting styles with loose relevance to the affordances in the pretraining scene. We also demonstrate a synergy between these pretrained components and fully unsupervised methods which to our knowledge has not been experimentally demonstrated. Lastly, we conduct a systematic evaluation of which language promoting techniques are best suited for the pretraining setting, which we anticipate to be of interest to the community as VLMs grow in popularity.

Connection to Goal-Conditioned RL: Our method is related to goal-conditioned RL in the sense that we learn a language-conditioned agent and language instructions can be considered semantic goals. However, the conventional goal-conditioned RL setup does not address the desiderata of scalable behavior pretraining. In general, this setting requires human supervision to select interesting goals for the agent to reach and actually achieving those goals often remains an exploration challenge. A contribution of LAMP is in identifying a compatibility of the relative ease of enumerating exploratory language goals compared to image goals and the constraints on human involvement for open-ended pretraining.

Pretraining Desiderata and Computational Cost

There have been considerable efforts to leverage VLMs such as [4] as a task reward. In fact, works such as [5] select datasets and architectures towards producing general-purpose reward functions. However, we argue in this work that due to issues of reward exploitation and noisy models that these LRF’s are ill-suited for directly learning downstream tasks. We instead argue that the pretraining phase is a more favorable context, which we refer to as “lower-stakes”, as we are reasonably less concerned with the sample-inefficiency of a “one-time cost” that requires limited human supervision than for a recurring cost we incur each time we desire to quickly learn a new skill. This is corroborated by [7]. Moreover, the differences in desiderata between pretraining and finetuning make the admittedly somewhat computationally intensive nature of LAMP more palatable. The paradigm of computationally expensive, large-scale pretraining in service of multi-purpose few-shot or task-specific finetuning characterizes much of modern machine learning [6].

Desired Results in Appendix

We were similarly interested to understand the impact of the reward weighting term. We had already provided these results in the Appendix. Please refer to Figure 11 in Section D.2 of the Supplementary Material. We find that LAMP works best with our choice of alpha but still exhibits very strong performance for an alpha value of 0 corresponding to training only with VLM rewards and no Plan2Explore bonus. We think this is quite an encouraging result and hope this addresses the reviewers concerns. We were additionally interested in how another novelty-seeking algorithm would perform. We direct the reviewer to an additional widely-used baseline, Random Network Distillation (RND) [1] in the Appendix (Figure 13, Section D.3.1). LAMP is able to outperform this baseline on all of the tasks evaluated.

For our thoughts on Question 7, see the end of the second to last paragraph of Section 6.1 on language prompting. We state "we hypothesize that both exploration coverage and exploration efficiency during pretraining play an important role. By separately increasing exploration coverage through Plan2Explore, the quality of the VLM rewards becomes more important for focusing the auxiliary exploration objective on useful exploratory behaviors. Thus, LAMP Prompt 2, which incorporates Plan2Explore, is trained on higher-quality, more relevant rewards, and can explore more efficiently during pretraining, and therefore enables more effective finetuning."

5. We simply observe the strong increase in the popularity and interest around large-scale computationally intensive pretraining for sample-efficient task-specific finetuning across machine learning (computer vision, natural language processing and generation, deep generative models, protein synthesis, robot learning, etc.) due to modern (last decade) advances in computation and datasets. We do not claim that this paradigm encompasses all of machine learning today. We situate our work within this broader trend characterized by notable breakthroughs throughout the field.

We also provide some additional results to the paper that we believe strengthen the claims made in the work.

We have added a new Reward Weighting Ablation section to the main text as well as added the ablation on all 5 of the evaluation tasks as opposed to just the Pick Up Cup task from the original manuscript in Section 6.2 in an effort to further enhance clarity. We arrive at a similar conclusion - LAMP is robust to this weighting and performs somewhat competitively even with an alpha value of 0 corresponding to only VLM rewards. However, the synergy of the Plan2Explore and VLM rewards leads to the optimal performance.

We have similarly added a new VLM Ablation Section to the main text and evaluated on all 5 of the evaluation tasks as opposed to just Pick up Cup from the original manuscript in Section 6.3.

We hope these updates have addressed your most pressing concerns and can lead to a rating increase. If not, please let us know what further clarifications or results you would require before the conclusion of the rebuttal period.

[8] Video Prediction Rewards for Reinforcement Learning [Escontrela et al 2023]

[9] Daydreamer: World Models for Physical Robot Learning [Wu et al 2022]

1. Much of the allure of general purpose pretrained models is that they can be used for many downstream tasks. Our pretraining setup is designed for this in that we do not assume knowledge of the diverse downstream tasks in pretraining the agent - they are completely unseen at finetuning time. As $n$ grows, the "new baseline" proposed will become increasingly more weak as the pretraining budget per task shrinks. However, since for our method the agent is only pretrained once, the sample complexity does not scale as poorly with the number of tasks. This is because, as we demonstrate in the paper, learning each of 5 very semantically and behaviorally diverse tasks is more sample-efficient from a pretrained LAMP agent. The agent performs a closing behavior, a valve turning behavior, fine-grained picking behaviors, and a pushing behavior all benefitting from the same pretrained model. We additionally note that for many robotics tasks, seemingly small modifications to the task can harm the success rate of a pretrained agent [9]. If instead of training an agent from scratch each time, we can more efficiently learn by warm-starting training from a pretrained agent, we further realize the benefits of task-agnostic finetuning.

2. The example you describe bears relation to what we propose in LAMP and as we previously stated, our method can be seen as similar to goal-conditioned RL. However, LAMP contains novel components and examines new questions outside of the default goal-conditioned RL setting including diverse language prompting, shaped VLM rewards, pretraining task and environment design, and synergies with unsupervised objectives that differentiate the work from prior work. We are happy to acknowledge the connection of our work with goal-conditioned RL, however, maintain that LAMP involves novel contributions of interest to the community.

3. There is such discussion in the paper. In the Introduction we mention: "While these methods...they come with the novel and relatively under-explored property of being a scalable means of generating many different rewards by simply prompting with different language instructions. This property is particularly compatible with the assumptions of RL pretraining where we desire to learn general-purpose behaviors with high coverage of environment affordances, require minimal human supervision..." Also see Section 4.1, the early discussion in Section 6.1, and Section 7.

Regarding the 4th concern in the original response - LAMP consistently outperforms the baselines across the very different downstream tasks evaluated. The finetuning curve can be high variance, however, the trend holds averaged over seeds and across different tasks and ablations. We have added many more experiments that evidence this consistent trend in Section 6.2.

Regarding the 5th concern in the original response, we have adjusted the language in the paper. We have removed the claim that "LAMP is not inherently reliant on a particular VLM" and instead clarify that LAMP can work well with CLIP rewards in addition to R3M rewards. Please see Section 6.3 of the updated paper.

Regarding Question 3, we take motivation from findings in [8] that maximizing rewards from pretrained models can be a challenging optimization that suffers from reward exploitation and convergence to local minima. The paper states "In the reinforcement learning setting, an additional issue is that solely optimizing a dense reward such as VIPER can lead to early convergence to local optima." and "In the reinforcement learning literature, the entropy bonus over input trajectories corresponds to an exploration term that encourages the agent to explore and inhabit a diverse distribution of sequences within the regions of high probability under the video prediction model... VIPER is agnostic to the choice of exploration reward and in this paper we opt for Plan2Explore and RND." We similarly found empirically in our setting that adding the intrinsic novelty bonus enables the RL agent to explore more modes of a noisy reward function with many local minima and thus leads to better performance.

Regarding Question 4, thank you for the catch regarding the expectation notation for the RL objective - see the revised paper with this fixed.

Important Clarification of Pretraining Setting and Contribution

It appears that Reviewer PM9W may have missed that the pretraining phase is fully finetuning task-agnostic and a “one-time cost.” During the pretraining phase, the agent optimizes different rewards and encounters many different but randomly-sampled tasks. The analysis we perform during pretraining examines various prompting styles and their relevance to the pretraining tasks, which we assume knowledge of for certain prompting styles, and not the finetuning tasks, which we assume no knowledge of. Prompt Style 2 is not based on the finetuning tasks but the pretraining tasks. VLM-guided RL pretraining is only done once, but the same, single pretrained agent is able to achieve the improved finetuning results on all the tasks evaluated. The 5 finetuning evaluation tasks are not included in the pretraining phase. This setup was intentional as this line of work is in pursuit of a single, general-purpose pretrained RL agent that can adapt to arbitrary new tasks with better sample-efficiency. We expect this important clarification to change your evaluation of the contribution as well address your questions. Please inform us if you have further questions.

Using LAMP directly as a task reward would not enable the RL agents to solve any of the tasks. As we evidence in Figure 3, and motivate in the introduction, rewards from current pretrained VLMs are far too noisy and susceptible to reward exploitation to be usable with online RL. In fact, this is central to our message - current VLMs are ill-suited for this purpose in that they lack the required precision to supervise challenging robot manipulation tasks, however, they do possess certain properties that render them very useful for pretraining RL agents with limited supervision. Scripted finetuning task rewards are simply a means to provide a consistent evaluation of the adaptation capabilities of the various pretrained agents. Our contribution is centered on proposing an RL agent pretraining setup.

We thank the reviewer for mentioning an additional baseline. While this baseline is relevant and worth citation, it makes additional assumptions on the environment - in this case access to an oracle language captioner which provides perfect natural langue descriptions of the current visual scene. We direct the reviewer to an additional widely-used baseline with more similar environment assumptions to our own in Random Network Distillation (RND), which is also baselined against in the mentioned work. We provide results in the Appendix (Figure 13, Section D.3.1). LAMP is able to outperform this baseline on all of the tasks evaluated.

[1] Exploration by Random Network Distillation (Burda et al 2018)

Desired Results in Appendix

We were similarly interested to understand the impact of the reward weighting term. We had already provided these results in the Appendix. Please refer to Figure 11 in Section D.2 of the Supplementary Material. We find that LAMP works best with our choice of alpha but still exhibits very strong performance for an alpha value of 0 corresponding to training only with VLM rewards and no Plan2Explore bonus. We think this is quite an encouraging result and hope this addresses the reviewers concerns.

Ego4D Texture Overlay

We used the Ego4D textures during pretraining to increase the difficulty of the control task with more complex visuals to more closely approximate the real world. These textures help to narrow the domain gap between the robot environment and the datasets used to pretrain the VLMs. We adopt the same texture randomization scheme for all pretraining methods we evaluate for fair comparison.

Desired Results in Appendix

We were similarly interested to understand the impact of the reward weighting term. We had already provided these results in the Appendix. Please refer to Figure 11 in Section D.2 of the Supplementary Material. We find that LAMP works best with our choice of alpha but still exhibits very strong performance for an alpha value of 0 corresponding to training only with VLM rewards and no Plan2Explore bonus. We think this is quite an encouraging result and hope this addresses the reviewers concerns.

Other Pretrained Models

VIPER is not a VLM in that it is not language-conditioned. Also our focus is on leveraging pretrained VLMs with strong generalization capabilities by providing diverse prompting to generate diverse rewards. VIPER is only shown to work as a reward function on mostly in-domain data. VIP, although it demonstrates some level of domain generalization, is also not a VLM, but rather relies on goal image alignment to compute a reward. Using such a method in a pretraining scheme quite different from that to which we propose would possibly lead to interesting results, however, would require possessing many interesting and diverse exploration goals a priori, which is a strong assumption of open-ended pretraining. Central to our contribution is that diverse language is a uniquely flexible means of generating diverse rewards with relatively little human supervision. In general, the message of our work is that providing diverse prompting with pretrained VLMs to generate rewards can serve as an effective pretraining scheme, not that R3M is the only particular pretrained model that can work.

LAMP as a task-reward

To answer simply, the success rate when using LAMP as a task reward on all of the tasks is 0%. As we evidence in Figure 3, and motivate in the introduction, rewards from current pretrained VLMs are far too noisy and susceptible to reward exploitation to be usable with online RL. In fact, this is central to our message - current VLMs are ill-suited for this purpose in that they lack the required precision to supervise challenging robot manipulation tasks, however, they do possess certain properties that render them very useful for pretraining RL agents with limited supervision.

Thank the authors for the response.

Other Baselines

It is obvious that the baselines mentioned differ from the proposed method (e.g., not VLM), but the similarity is that they could all provide reward signals based on observations. Though these reward functions are designed for in-domain tasks, they are also noisy compared to the oracle reward and thus could also serve as a pertaining method.

Therefore, it is still good to see more comparison with these recent works.

LAMP as a task-reward

To answer simply, the success rate when using LAMP as a task reward on all of the tasks is 0%.

It is good to see the points described by the authors are updated in the main paper.

Again, thank the authors for their response. I do not have further questions.

Important Clarification of Pretraining Setting and Contribution

It appears that Reviewer YXuM may have missed that the pretraining phase is fully finetuning task-agnostic and a “one-time cost.” During the pretraining phase, the agent optimizes different rewards and encounters many different but randomly-sampled tasks. The analysis we perform during pretraining examines various prompting styles and their relevance to the pretraining tasks, which we assume knowledge of for certain prompting styles, and not the finetuning tasks, which we assume no knowledge of. The above point addresses “pretraining saves us 50k steps…” and “I wonder if the method is competitive on a fixed-compute budget with single-task training.” VLM-guided RL pretraining is only done once, but the same, single pretrained agent is able to achieve the improved finetuning results on all the tasks evaluated. The 5 finetuning evaluation tasks are not included in the pretraining phase. This setup was intentional as this line of work is in pursuit of a general-purpose pretrained RL agent that can adapt to arbitrary new tasks with better sample-efficiency. Furthermore, to the third point, all 5 of the evaluations are on OOD tasks. The pretraining tasks are randomly sampled from a set of combinatorially many possible tasks using ShapeNet [2] objects. See Figure 8 in Appendix C.2 for examples. We expect this important clarification to change your evaluation of the contribution as well address your questions. Please inform us if we have misinterpreted your concerns or if you have further questions.

Ablations

Regarding the ablations, Reviewer YXuM mentions that we do ablate the choice of exploration reward with inclusion of RND in the Appendix. We agree with this and are glad the reviewer has acknowledged this ablation. Furthermore, they suggest ablating the choice of choice of temporal objective in order to prove that there is “benefit to the proposed scheme.” We do make strong claims that a particular formulation of temporal alignment is uniquely critical for pretraining, rather we propose in the paper that it be used as a language-conditioned pretraining reward and demonstrate the benefits when integrating it in LAMP. We simply employ the most natural alignment objective based on the way the VLM is trained or what has been found to work well in prior work. If the reviewer still disagrees we would appreciate further explanation and suggestions as to what other language-temporal alignment objectives they would like to see.

To address the last point, the hyperparameter choices for the Masked World Models agent follow the paper which introduced the method [1]. We treat the RL largely as a black box and do not tune RL specific agent hyperparameters as these are orthogonal to our method. We provide the default hyperparameters in the interest of enhancing simple reproducibiilty. The base RL algorithm is kept entirely constant across all baselines (Plan2Explore, RND). The justification for using Masked World Models (MWM) as opposed to a “more standard” RL algorithm such as SAC, PPO, DreamerV3 is that MWM can solve all of the challenging downstream tasks evaluated from scratch with no tuning. We found that this was not the case for the aforementioned more standard algorithms. Lastly, we point the reviewer to our ablation on the balance weighting between exploration reward and VLM reward in Appendix Section D.2 Figure 11 as this hyperparameter is related to our contribution.

[1] Masked World Models for Visual Control (Seo et al 2022)

[2] ShapeNet: An Information Rich 3-D Model Repository (Chang et al 2015)

Setting and Contribution

We are glad that the reviewer understands that the pretraining is a one-time cost aimed at pretraining a general purpose agent. In this work, we simply build on a wide-body of literature aimed at improving general purpose RL pretraining with limited human supervision. We mention many of these kinds of works in the Pretraining for RL section of the Related Works and situate our contribution therein. The reviewer mentions orders of magnitude gains in efficiency for language tasks by referencing Internet-scale language pretraining. They go on to reason that our much more modest gains do not warrant the additional complexity of our method. However, it is safe to say that the domains of reinforcement learning and robotics are far behind language in terms of the adaptive capabilities of general-purpose models. Recent work such as [3] makes this very apparent. A longstanding challenge in RL and robotics is to build similarly capable models, though current methods fall well short. We far from claim to rival these levels of capability, rather we propose an approach that improves on current methods and reveals potential avenues for further improvements.

Improved Clarity and New Experimental Ablations

We have clarified our experimental setting by adding "We do not include any of the downstream evaluation tasks that the agent finetunes on in the pretraining setup. The pretraining tasks are exclusively comprised of ShapeNet objects. All of the downstream tasks are unseen tasks from the original RLBench task suite." Please see the updated Section 5.1.

Regarding the use of MWM, in Section 4.2, we justify our decision with "we use Masked World Models (MWM) (Seo et al., 2022), an off-policy, model-based method with architectural inductive biases suitable for fine-grained robot manipulation."

Thank you for confirming the additional clarity provided by the reward weighting ablation. We have added a new Reward Weighting Ablation section to the main text as well as added the ablation on all 5 of the evaluation tasks as opposed to just the Pick Up Cup task from the original manuscript in Section 6.2 in an effort to further enhance clarity. We arrive at a similar conclusion - LAMP is robust to this weighting and performs somewhat competitively even with an alpha value of 0 corresponding to only VLM rewards. However, the synergy of the Plan2Explore and VLM rewards leads to the optimal performance.

We have similarly added a new VLM Ablation Section to the main text and evaluated on all 5 of the evaluation tasks as opposed to just Pick up Cup from the original manuscript in Section 6.3.

We hope these updates have addressed your most pressing concerns and can lead to a rating increase. If not, please let us know what further clarifications or results you would require before the conclusion of the rebuttal period.

[3] GPT-4 Technical Report [OpenAI 2023]