Vision-Language Models Provide Promptable Representations for Reinforcement Learning

Thank you for your detailed feedback and for a positive assessment of our work. We have addressed your main concerns by adding the suggested related works, video rollouts and an explanation on why the computational cost of our method is actually much more practical than prior works that typically train entire VLMs. In the following we will answer your questions in detail.

Concerns about computational cost. Obtaining promptable representation requires running a complex VLM on every image returned from the environment, along with the prompt and answer. This cost is prohibitively high and not practical.

We agree that the computational complexity of large VLMs is a major limiting factor for their adoption into control tasks and that lowering the compute required is difficult, barring architectural changes or lowered precision. Nevertheless, several prior works have used a large (V)LM in an RL loop [ https://arxiv.org/abs/2305.13301, https://llm-rl.github.io/static/images/llarp.pdf ] or even used a full pretrained VLM to initialize a policy (where the VLM is trained as part of the policy) [ https://arxiv.org/abs/2307.15818 ]. Compared to the latter, our method actually requires less compute, as it does not require training an entire VLM (only the lightweight policy network head). We also find that PR2L is reasonable in terms of wall-clock speed, achieving ~4 Hz inference speed (sufficient for playing Minecraft in real time). Training speed is also ~4 Hz, as gradients do not flow through the VLM (only the small policy network, which adds much less overhead). Moreover, training wall clock time is further improved in offline approaches, as all offline data can be embedded by VLMs in parallel and then reused in future epochs or runs. Again, this does not lower the compute required, but drastically decreases training time. We found this to be true in our new BC experiments (see Appendix E.1), wherein we trained near-expert level policies with comparatively small training time/data and used multiple VLMs to parallelize observation embedding generation.

Concerns about the performance of the learned policy. The visually rich representation obtained at such a high computational cost should greatly improve the performance of the policy. However, the authors did not provide videos of rollouts on relevant tasks, making it difficult to judge the effectiveness of the promptable representation in a real Minecraft environment (the three tasks used in the paper are not very complex).

We include new videos visualizing 10 combat spider rollouts (here)[ https://drive.google.com/drive/folders/1GFN4C6qAEciGEK3phcqgYs3V-HxJpAOy?usp=sharing ] for both our approach and the image encoder baseline (achieving 7 and 4 successes respectively). We do note that it’s difficult to observe qualitative differences between the PR2L and baseline policies – they are more clear in reported aggregate performance metrics.

Insufficient literature review. Since the authors conducted experiments in Minecraft, they should have provided a more comprehensive discussion of articles that control and plan within Minecraft. However, the authors left out the following important literature.

Thank you for these suggested citations – we have included them in the extended literature review in Appendix G.

Please let us know if we have addressed your concerns, or if there is anything else we should consider. Otherwise, we would be extremely grateful if you could raise our score!

Thank you for your prompt response! We are glad that you find our method interesting. We would first like to clarify that our goal in this paper was to demonstrate our method PR2L can provide for an effective way to extract useful representations from pre-trained VLMs for downstream reinforcement learning. This goal realizes our long-term vision of using off-the-shelf pre-trained VLMs for RL, even though these models were never trained on data from the specific domain. As general-purpose VLMs start to become useful, we believe it is crucial to develop approaches for extracting representations from them, targeted towards downstream RL, even though this comes at a larger inference cost with currently available VLMs. While we do agree that these computational costs can be reduced specifically for the domain of Minecraft, we reiterate that the goal of this paper is not to develop better approaches for Minecraft, but to develop methodology to utilize pre-trained VLMs in RL. Developing approaches for utilizing VLMs for downstream control has been an important theme in recent works (e.g., RT-2, InstructBLIP for Robotics), despite the associated computational costs.

Specifically, for our experiments, please note that a frequency of 4 Hz is already better than prior works that combine VLMs and RL: RT-2 uses a 65B parameter model and only achieves a worse 1-3 Hz control frequency on a real robot, which they still find to be sufficient for the real world.

Likewise, we use less compute and achieve greater sample efficiency than similar works: As with our online RL experiments, LLaRP’s most performant model calls LLaMA-13B for 2e8 environment steps, thereby using a larger model (13B vs. our 7B parameters), with a longer history-dependent context length, for 1000 times more steps than PR2L. These prior approaches incur significantly larger computational costs, but are of interest to the community as they integrate foundation models with downstream RL / control algorithms.

In addition, please note that when PR2L is applied with BC, as we demonstrate in our newly added experiments, the computational cost of encoding collected offline observations with the pre-trained VLM is incurred only once and can be parallelized. Encoding observations in the offline dataset takes ~2 hours with two VLMs (and only needs to be done once), which is comparable to running BC training itself (which takes around an hour for each run).

We would appreciate it if you are willing to upgrade your rating in the light of these clarifications. We are happy to discuss further. Thank you so much!

Dear Reviewer AR8Q,

We were wondering if you have had a chance to read our reply to your feedback. As the time window for the rebuttal is closing soon, please let us know if there are any additional questions we can answer to help raise the score!

Thank you for your detailed feedback and for a positive assessment of our work. We have addressed your main concerns by adding offline imitation learning experiments and more baselines (R3M, VC-1), to which our method favorably compares. In the following we will answer your questions in detail.

There are some other approaches that are designed to tackle similar issues, especially in Minecraft, ex. [1,2,3]. Although I agree some settings could be different (RL vs. IL, etc), they all deliver some backbone design or objective functions that could facilitate better representations. Comparisons against these missing baselines would help the reader with a better understanding of the significance of the proposed method.

Thank you for the suggestion! We have run additional baselines (using VC-1 or R3M as encoders or giving the policy ground-truth entity detection) and have found PR2L outperforms all of them in all tasks. Additionally, we have also trained policies with BC on top of PR2L and image encoder representations, finding that PR2L especially excels at the “combat spider” task (doubling the baseline’s success rate and achieving near-expert performance in only one epoch).

Additionally, our primary goal was not to build a SOTA system for Minecraft tasks, but to show that the representations produced by prompting a VLM with task context are better for learning control tasks than equivalent, non-prompted representations.

In principle, our approach can be combined with other approaches for improving performance in Minecraft tasks (e.g., using a better reward function like MineCLIP or CLIP4MC). However, as we wanted to isolate performance differences to better representations elicited via prompting, we did not include such approaches.

Minecraft is indeed a challenging domain in terms of open-world and complex visual observations. However, the tasks being evaluated here (spider, cow, sheep) do not seem to be challenging enough to justify the effectiveness of the proposed method, especially on the claimed "leverage contextual prior information". These mobs are indeed very common and the tasks themselves do not seem to involve complex visual stimuli. My suggestion is to try some long-term and open-ended tasks like surviving, collecting items, etc. [7] offers a few of them and worth taking a look at.

As we just wished to investigate the effectiveness of promptable representations (rather than developing an agent that is generally better at playing Minecraft), we chose to evaluate on the same task suite that the original MineDojo paper because they were solved by simple RL algorithms (off-the-shelf PPO) with little Minecraft-specific system engineering, meaning that our approach could likewise be less Minecraft-specific. We have likewise run PR2L and the VLM image encoder baseline on all five of the remaining programmatic tasks that MineDojo evaluated on to complete the analysis. For longer-horizon or more open-ended tasks, we likely would need to adopt similar approaches to the ones presented in your suggested citations.

Some prompts shown in Table 1 require hand-crafted domain knowledge, ex. "Spiders in Minecraft are black". Is it possible to avoid this?

We find that the spider task’s domain knowledge is helpful, but not necessary. Removing it degrades performance slightly, but still outperforms using the non-promptable VLM image encoder representations (green vs. blue curves in Figs. 4 and 7’s combat spider plots). However, it does even better if this auxiliary text is included (red curve). We view the ability for VLMs to use auxiliary text as an advantage: it allows them to recognize stylized Minecraft spiders despite them being visually different from most natural images of spiders seen during pre-training. which intuitively explains why giving it this domain knowledge is helpful.

Additionally, we recognize that re-running PR2L to optimize the prompt is computationally expensive. However, we find that the simple prompt evaluation scheme in Sec 5.2 and Appendix A is an effective proxy measure for how good a prompt is for PR2L.Finally, we also acknowledge that this evaluation requires domain-specific data, but note that it uses a miniscule amount of it compared to fully fine-tuning a model – we collected all prompt evaluation data in under an hour.

[Continued in follow-up response due to character limit]

Dear Reviewer Cm8x,

We were wondering if you have had a chance to read our reply to your feedback. As the time window for the rebuttal is closing soon, please let us know if there are any additional questions we can answer to help raise the score!

Thank you for the prompt reply! We are glad we were able to address most of your concerns. We agree that there are more complex Minecraft tasks that could be tested on. Nevertheless, we find that the tasks we considered (all from the MineDojo paper) still pose an interesting and non-trivial challenge for reinforcement learning. In particular, we note that the original MineDojo paper found they needed a more complex learning algorithm than pure RL with PPO – namely, they require self-imitation learning to achieve good performance on shear sheep and milk cow (their approach fails otherwise, as shown in Figure A.8 of that paper).

We empirically confirm that RL on top of MineCLIP does poorly on these tasks in particular, despite being trained to produce representations on domain-specific data (in contrast to our VLMs, which are not trained on Minecraft data). Since these RL tasks are non-trivial even for more complex agent architectures/training algorithms that use domain-specific representations, we believe that our evaluation and contribution – showing a way to extract representations from general purpose VLMs that match or significantly exceed the performance of domain-specific ones – is still valuable.

This level of performance is not present when using the VLM image encoder baseline, which suggests that the performance increase is because of the LLM parsing the image encoder’s representations based on its prior understanding of (1) what things look like and (2) how that relates to task context, provided as a prompt. We show in our follow-up BC experiments that this results in representations that are particularly conducive to learning good actions (with our combat spider agent achieving expert-level performance after 1 epoch, doubling the success rate of the VLM image encoder baseline).

An additional point regarding evaluating on other RL benchmarks such as Atari or dm control is that, as our method does not fine-tune a VLM on domain specific data (as prior works do), we are constrained to environments that produce somewhat more realistic observations, as the VLMs are trained on natural images.

While we believe this all demonstrates how RL performance can be improved by using VLM representations that “leverage contextual prior information,” we would be happy to change the wording of that claim to something you believe to be a more accurate description of the above phenomenon – e.g., we could say that they “draw out relevant representations based on task context” (or some other suggested phrasing). Otherwise, if we have addressed your concern, we would be extremely grateful if you increase your rating to reflect this!

Thank you for your detailed feedback. We have addressed your main concerns by visualizing the representations of instructions vs. prompts and running the suggested baselines with oracle object detection and running with image encoders (R3M, VC-1) that are more suited to embodied control, to which PR2L compares favorably. In the following we will answer your questions in detail.

The only form of prompt that PR2L uses essentially amounts to object detection in the scene. This introduces a confounding factor of whether PR2L outperforms baselines because it is able to recognize objects better in the scene. Could a baseline that somehow incorporates oracle object detection information be included? This will test whether PR2L does better because it detects the object of task interest in the scene.

Thank you for the suggestion! To address the concern regarding object detection, we train baselines with oracle object detection and find that they perform worse than PR2L in all cases as seen here or Appendix E.2. Additionally, while we agree that our prompts essentially involve task-specific object detection, it's essential to highlight that this aligns with the specific requirements of our considered tasks. Detecting spiders, cows, and sheep in observations is a critical feature for accomplishing said tasks. The advantages of PR2L is that prompting (1) specifies what entities are important to detect for a task and (2) improves representations + detection of potentially OOD entities (like cartoony Minecraft spiders), obviating the need for domain-specific classifiers/representations by using general-purpose VLMs.

Relatedly, PR2L does not outperform MineCLIP on most tasks; given that MineCLIP exists and is open-sourced, PR2L's stated advantages can be better demonstrated via a new domain in which foundation pre-trained representations do not already exist. BLIP-2's vision encoder may not be the strongest baseline for pre-trained vision encoders. Several prior works such as VC-1, R3M, MVP, VIP are trained for decision-making and robotics tasks and may constitute stronger baselines in that category. Could additional pre-trained vision encoder baselines be included?

Thank you for the suggestion, to address the concern regarding more baseline comparisons, we trained additional baselines on VC-1 and R3M and found that they perform worse than PR2L in all cases (see the prior link). We note that VC-1 is the most recent of these models, and seems to be more performant than MVP and VIP, so we train on its representations rather than the latter two.

We chose to use InstructBLIP’s image encoder not because it was the best encoder for Minecraft, but to illustrate that representations that are suboptimal for control are improved by passing them through an LLM with the task-relevant prompt. The observed performance improvement of PR2L over the encoder baseline can thus entirely be attributed to the prompt extracting task-relevant information from the subpar generic encoder representations.

The paper claims that the prompts are different from instruction following; however, the prompts are still manually constructed and task-specific. It's unclear the advantage of doing so as instructions, by construction, should exist as it is a direct form of task specification.

Thanks for bringing this up. We agree that instruction-following and our prompts both specify the task and require hand-designing. However, the former often cannot be done by our chosen VLM zero-shot, resulting in lower quality representations. This is reflected in the lower performance when the VLM is given instructions in our RT-2-style baseline.

Furthermore, we find that the distribution of representations the VLM yields is very different depending on if it’s prompted with instructions or our approach, with the latter seemingly being more predictive of “good” actions to take. We observe this by plotting the first 2 principal components of the VLM features of observations from an expert policy when given instructions or our prompt, as shown here. We find that the representations from our prompt are bimodal (as the VLM usually answers “yes” or “no”) and most expert functional actions (attacking or using items) that result in rewards correspond to embeddings when the VLM answers “yes” (the left-hand clusters). This is most clear for milk cow – see the large orange points. In contrast, the PCA plots of instruction representations are much less structured; rewarding functional actions are much closer to the overall mean, and seem qualitatively harder to separate out and identify.

In summary: with our prompts, VLM representations of observations correspond strongly to actions that yield rewards for our considered tasks. This is not true of representations from instructions. More details can be found in Appendix F.

[Continued in follow-up response due to character limit]

Dear Reviewer qGuw,

We were wondering if you have had a chance to read our reply to your feedback. As the time window for the rebuttal is closing soon, please let us know if there are any additional questions we can answer to help raise the score!

Dear Reviewer qGuw,

As the rebuttal window closes tomorrow, we would again like to ask if you have had a chance to consider our reply to your feedback. Please let us know if you have any additional concerns we can address to help raise the score!

Thank you for your detailed feedback. We have addressed your main concerns by adding all remaining programmatic Minecraft tasks from the original MineDojo paper and three additional baselines to which our method favorably compares. In the following we will answer your questions in detail.

I also recommend a stronger evaluation on the Minecraft benchmark consistent with the evaluations done in MineDOJO. Currently, the paper shows results on only three tasks. For a more exhaustive evaluation, MineDOJO recommends evaluation on a collection of starter tasks (32 programmatic and 32 creative tasks).

Thank you for the suggestion! Unfortunately, to the best of our knowledge, we were unable to find the 64 core tasks mentioned in the original MineDojo paper. As the MineDojo paper evaluates their agent on only 12 tasks (four of which were “creative,” and thus do not have associated concrete reward functions), we assume that said tasks are part of the core set. Thus, we run PR2L and the VLM image encoder baseline on all 8 of the “programmatic” tasks considered in the MineDojo paper. Only “combat zombie” is currently available (as our computing cluster’s file system is down), but we find that PR2L outperforms the baseline in that case, as shown here. We have also run some additional baselines / ablations on the original three tasks, wherein the model is trained on VC-1 and R3M representations or is supplemented by ground-truth entity detection. In all cases, these perform worse than our approach, as shown here and discussed in Appendix E.2.

While I understand that the authors didn’t have the resources to train a RT-1 style baseline, would it still be possible to train an action decoder on top of the VLM to produce actions. I think having a strong RT-1 style baseline is very important to properly evaluate the question (2) mentioned in the paper — “How does PR2L compare to approaches that directly “ask” the VLM to generate the best possible actions for a task specified in the prompt?”

Great point! While we did not exactly match the architecture of RT-1/-2, we already addressed question (2) in our baseline (b), wherein the VLM was given a prompt similar to the ones used in the RT-1 and RT-2 with the task instruction and possible actions, and was then asked to produce an action. The trained Transformer-based policy then decodes the instructions and generated text into actions. Of course, unlike RT-2, this approach does not fine-tune the VLM itself on data from MineCraft, but please note that none of our methods use a fine-tuned VLM: our goal is to compare different approaches to query representations from a pre-trained model in terms of their downstream RL performance. We apologize for any confusion and have updated the paper to clarify this point.

The proposed architecture compresses the task-relevant features from the VLM into a single CLS token which can severely restrict the information available to the policy ... Did the authors try an approach like Perceiver IO (Jaegle et al, 2021) which encodes N tokens to K (1≤K≤N) tokens?

For a fair comparison against MineCLIP that produces a fixed-length representation of an observation, we chose to convert the VLM representations to a single-token bottleneck. This allows us to roughly match the size of the policy network that converts the VLM representation into an action distribution. In our experiments, we do not find this bottleneck to hinder our policy’s ability to achieve high task performance. That said, of course, PR2L can use more advanced policy architectures which do not involve such bottlenecks and we would only expect these architectures to improve performance.

It’s unclear what efficiency means (faster to train in FLOPS? faster to train measured by amount of training steps?). I think the paper can be made stronger by comparing training efficiency when using the proposed approach vs using VLM image-representations directly.

Thank you for this suggestion. We apologize for the confusion – we use the term efficiency to denote sample efficiency, following the standard textbook definitions of sample efficiency in RL. Sample efficiency is distinct from computation efficiency (measured in FLOPs), and is a standard comparison metric for RL algorithms. We have updated the paper to reflect this. The reviewer-suggested comparison is shown in the red (PR2L) and blue (VLM image representation) curves in all performance figures – when considering sample efficiency, we can simply check the returns or successes after a fixed number of training samples.

[Continued in follow-up response due to character limit]

Dear Reviewer sWjV,

We were wondering if you have had a chance to read our reply to your feedback. As the time window for the rebuttal is closing soon, please let us know if there are any additional questions we can answer to help raise the score!

Dear Reviewer sWjV,

As the rebuttal window closes tomorrow, we would again like to ask if you have had a chance to consider our reply to your feedback. Please let us know if you have any additional concerns we can address to help raise the score!