[Paper Review] Language Reward Modulation for Pretraining Reinforcement Learning

Ademi Adeniji, Amber Xie+4 more|arXiv (Cornell University)|Aug 23, 2023

Reinforcement Learning in Robotics4 citations

TL;DR

LAMP proposes using frozen vision-language models (VLMs) to generate diverse, language-conditioned exploration rewards during pretraining for reinforcement learning, combining VLM rewards with novelty-seeking intrinsic rewards to learn a general-purpose, language-conditioned policy. This approach significantly reduces sample complexity in downstream robotic manipulation tasks on RLBench, demonstrating strong zero-shot transfer and robustness to prompting variations.

ABSTRACT

Using learned reward functions (LRFs) as a means to solve sparse-reward reinforcement learning (RL) tasks has yielded some steady progress in task-complexity through the years. In this work, we question whether today's LRFs are best-suited as a direct replacement for task rewards. Instead, we propose leveraging the capabilities of LRFs as a pretraining signal for RL. Concretely, we propose $ extbf{LA}$nguage Reward $ extbf{M}$odulated $ extbf{P}$retraining (LAMP) which leverages the zero-shot capabilities of Vision-Language Models (VLMs) as a $ extit{pretraining}$ utility for RL as opposed to a downstream task reward. LAMP uses a frozen, pretrained VLM to scalably generate noisy, albeit shaped exploration rewards by computing the contrastive alignment between a highly diverse collection of language instructions and the image observations of an agent in its pretraining environment. LAMP optimizes these rewards in conjunction with standard novelty-seeking exploration rewards with reinforcement learning to acquire a language-conditioned, pretrained policy. Our VLM pretraining approach, which is a departure from previous attempts to use LRFs, can warmstart sample-efficient learning on robot manipulation tasks in RLBench.

Motivation & Objective

• To address the challenge of sparse-reward reinforcement learning in complex environments by leveraging learned reward functions not as direct task rewards, but as pretraining signals.
• To explore whether vision-language models (VLMs) can serve as scalable, zero-shot reward generators for pretraining diverse, language-conditioned behaviors.
• To reduce sample complexity in downstream RL fine-tuning by pretraining agents on diverse, semantically shaped rewards derived from VLMs.
• To investigate the synergy between VLM-generated rewards and intrinsic exploration objectives in improving pretraining efficiency.
• To evaluate the robustness and transferability of the pretraining approach across different prompting strategies and VLM architectures.

Proposed method

• LAMP pretrains a policy by optimizing a combination of VLM-generated language-conditioned rewards and intrinsic novelty-seeking rewards from Plan2Explore.
• It computes contrastive alignment between image observations and diverse language instructions using a frozen, pretrained VLM, generating shaped exploration rewards.
• The reward is parameterized as the dot product of visual and language delta embeddings: $ r_i = (F_\phi(s_i) - F_\phi(s_0)) \cdot (L_\alpha(\mathbf{x}) - L_\alpha(\mathbf{x}_0)) $, where $ F_\phi $ and $ L_\alpha $ are frozen CLIP encoders.
• Pretraining uses a multitask, language-conditioned RL objective to learn a general-purpose policy that can be fine-tuned on downstream tasks.
• The method leverages zero-shot capabilities of VLMs to scale reward generation without human-designed reward functions or domain-specific tuning.
• It evaluates performance on RLBench tasks after fine-tuning, comparing different prompting styles and VLMs.

Figure 1 : LAMP Framework . Given a diverse set of tasks generated by hand or by a LLM, we extract VLM rewards for language-conditioned RL pretraining. At finetuning time, we condition the agent on the new task language embedding and train on the task reward.

Experimental results

Research questions

• RQ1 Can VLM-generated rewards serve as an effective pretraining signal for RL, improving downstream sample efficiency?
• RQ2 How does the diversity and semantic structure of language prompts affect pretraining performance and downstream transfer?
• RQ3 What is the synergy between VLM rewards and intrinsic exploration objectives like those in Plan2Explore?
• RQ4 How robust is the pretraining approach to variations in prompting style and choice of VLM?
• RQ5 Can LAMP generalize to unseen tasks without zero-shot transfer, and what is the role of VLM quality in performance?

Key findings

• LAMP pretraining with diverse language prompts significantly reduces sample complexity in downstream RL fine-tuning on RLBench, enabling faster convergence.
• Prompt Style 2, which uses semantically similar but diverse wordings of instructions, yields the best downstream finetuning performance.
• The combination of VLM rewards and Plan2Explore's intrinsic exploration leads to more efficient pretraining than either method alone.
• LAMP with ZeST-style CLIP rewards achieves strong performance, and results suggest the method benefits from more capable future VLMs.
• The approach is robust to different prompting strategies, with minimal performance drop across styles, indicating broad applicability.
• LAMP outperforms baselines using non-action-based prompts when combined with exploration objectives, highlighting the importance of reward quality and coverage.

Figure 2 : LAMP Method . We use R3M [ 24 ] for our VLM-based rewards. We query the R3M score predictor for pixel and language alignment, which is pretrained on the Ego4D dataset [ 15 ] . The reward model is frozen.