Paper: Vision-Language Models Provide Promptable Representations for Reinforcement Learning (ICLR-2024)

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

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

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

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

Reviewer comment (Novelty):
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:  
1. Open-world multi-task control through goal-aware representation learning and adaptive horizon prediction.  
2. Video pretraining (vpt): Learning to act by watching unlabeled online videos.  
3. Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents.  
4. CLIP4MC: An RL-Friendly Vision-Language Model for Minecraft.  
5. GROOT: Learning to Follow Instructions by Watching Gameplay Videos.  
6. Learning from Visual Observation via Offline Pretrained State-to-Go Transformer.

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

Reviewer comment (Experiments):
Albeit the promises shown by the results on 3 tasks on Minecraft, I don't think the approach is thoroughly evaluated, especially given their claim on "leverage contextual prior information" and "visually-complex RL tasks" (see abstract). There are some other approaches that are designed to tackle similar issues, especially in Minecraft, e.g., [1,2,3]. Although 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.

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

Reviewer comment (Experiments):
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 very common and the tasks themselves do not seem to involve complex visual stimuli. A suggestion is to try some long-term and open-ended tasks like surviving, collecting items, etc. [7] offers a few of them and is worth taking a look at.

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

Reviewer comment (Novelty):
Minor: Some references on open-world representation learning and Minecraft agents should be cited: [1-6].

Author response:
Thank you for these suggested citations – we will include them in the extended literature review in Appendix G.

Reviewer comment (Presentation):
In Figure 2 and Figure 3, some curves do not have shadows — why is that?

Author response:
Thank you for pointing this out: we omitted the shaded error region for the MineCLIP encoder trial in Fig. 3 and the PR2L/image encoder curves in Fig. 5 for visual clarity. We have updated the paper to state this. (Here)[https://ibb.co/Lk7J8W2] is a visualization of all error regions in Fig. 3 (in the original paper, only the green was missing). Note that the Fig. 5 red and blue curves are the same as the red and blue curves in Fig 3., which do have error regions.

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

Author response:
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](https://ibb.co/PrkhKFB) 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.

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

Author response:
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](https://ibb.co/p1CL8Nx). 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.

Reviewer comment (Evaluation):
The improvements of PR2L over its various ablations appear only moderate. Furthermore, the best prompt format for the tasks are not consistent; for Spider, it is helpful to include contextual information of what a spider looks like in MineCraft, whereas for the other two tasks, it is more helpful to disregard such information. Therefore, applying PR2L to a new task may require substantial prompt engineering for the best performance.

Author response:
All ablations follow PR2L’s approach of using VLM representations for RL. Aside from the “no prompt” ablation, they are all given some form of task context. Thus, their improvement over the VLM image encoder baseline indicates that representations prompted with task context yield higher performance than non-prompted representations, albeit not as good as the main PR2L representations. 

The domain-specific auxiliary text is optional for combat spider. PR2L without it still significantly outperforms the pure image encoder baseline, as shown in ablations (green vs. blue curves in Figs. 5 and 7’s combat spider plots). However, it does even better if this auxiliary text is included (red curve). As discussed in Sec. 5.1, 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.

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. In particular, it seems like the auxiliary information for the cow / sheep tasks is detrimental because it causes the VLM’s outputs to degrade: it always answers “yes” to “is there a cow / sheep in this image?” when the auxiliary text is added (Table 2). This may be fixed by using a more advanced VLM, but for now, our evaluation approach is able to catch this systematic bias in response to auxiliary text. 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.

Reviewer comment (Experiments):
PR2L is only evaluated on 3 tasks; these tasks are also the simplest in the MineDojo benchmark. The paper would be strengthened if more tasks and domains are evaluated. Currently, it is not convincing that PR2L can be generally applied to other visuomotor control domains. 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.

Author response:
Thank you for the suggestion. The MineDojo paper evaluated their agent on 8 “programmatic” tasks (and 4 of which were “creative,” and thus do not have associated concrete reward functions). We thus ran additional experiments comparing PR2L and the VLM image encoder baseline on the remaining 5 programmatic tasks. 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](https://ibb.co/jggxzD9).

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!

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

Author response:
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](https://ibb.co/PrkhKFB)). 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.

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

Author response:
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](https://ibb.co/jggxzD9). 
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](https://ibb.co/PrkhKFB) and discussed in Appendix E.2.

Reviewer comment (Experiments):
Overall, I liked the main contribution of the paper. But I believe, in its current form, the evaluation in the paper is a bit weak and could be made more exhaustive (unseen tasks, more exhaustive MineCraft evaluation).

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

Reviewer comment (Experiments):
After reading the rebuttal, most of my questions are adequately answered. I still believe that directly training an action decoder (while keeping VLM frozen is a good baseline) and should be included in the paper.

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

Reviewer comment (Experiments):
I also feel that the paper will be stronger if they include more exhaustive Minecraft experiments (to check for generalisation on unseen tasks / objects) and more environments which are visually more complex like Habitat / AI2 Thor.

Author response:
Thank you for the suggested additional task environments. While we have developed infrastructure to run experiments in Habitat, we ran into technical difficulties when running experiments, but are working to resolve that now.

Reviewer comment (Experiments):
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. While this may work for simpler environments like Minecraft which doesn’t have a lot of clutter, it might not work for other environments / tasks (rearrangement tasks in indoor environments). Did the authors try an approach like Perceiver IO (Jaegle et al, 2021) which encodes N tokens to K (1≤K≤N) tokens?

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

Reviewer comment (Evaluation):
I think the first question — “Can promptable representations obtained via task-specific prompts enable more efficient learning than those of pretrained image encoders?” is not really answered. It’s unclear what efficiency mean (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.

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

Reviewer comment (Experiments):
I didn’t fully understand the various ablations done. Specifically, did the authors try just giving task context (and no auxiliary information)? Similarly, did the authors try giving just auxiliary information without giving task context.

Author response:
We apologize for the confusion – our ablations considered (1) removing the prompt altogether, (2) removing the VLM’s text generation abilities, and (3) adding auxiliary text if our main trial did not have it and vice-versa (task context is always there). Ablations (1) and (2) correspond to removing the middle and right-most groupings of red embeddings shown in Fig. 2 respectively. For the spider task (which has auxiliary text), (3) amounts to only providing task context (green curve in Fig. 5). The main prompts for the other tasks already only have task context (red curves in Figs. 3 and 5).

Reviewer comment (Experiments):
While I don’t expect this experiment to be performed for rebuttal, I really wish the authors evaluated their approach on tasks that are visually more complex (or the environments are more cluttered). For e.g., using the approach to perform rearrangement tasks in indoor environments like Habitat / AI2 Thor.

Author response:
Additionally, as you suggested and as mentioned in our initial reply, we have started running preliminary experiments in the Habitat environment, so as to show the applicability of PR2L to non-Minecraft embodied control domains. We consider the ObjectNav task of learning to find toilets in a household environment and find that PR2L is able to achieve over 300 successes by the end of training, while the VLM image encoder baseline only achieves ~50 on average. See [here](https://ibb.co/7XB5MVn). The prompt we use is "Is a toilet likely in this room?" We confirm that it tends to answer "yes" to this question when shown images of bathrooms, even when a toilet is not directly in view, thereby demonstrating the use of VLM contextual knowledge for an application other than pure entity recognition (as shown in Minecraft). Here are examples of images the VLM answers ["yes"](https://drive.google.com/drive/folders/1EWG8mV-OL2oFZES23u9xFE9Y0d0u05O1?usp=sharing) and ["no"](https://drive.google.com/drive/folders/1DoJd5aSEUzL8Uz4ddqEkpB2Os9QHnz53?usp=sharing) to, respectively.

These results are very preliminary, but we will be sure to include more extensive experiments in the final manuscript. We are running additional experiments now as well.