Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Compositional Instruction Following with Language Models and Reinforcement Learning

Vanya Cohen, Geraud Nangue Tasse, Nakul Gopalan, Steven James, Matthew Gombolay, Ray Mooney, Benjamin Rosman

TL;DR

The paper addresses the high sample complexity of learning language-conditioned reinforcement learning tasks by introducing CERLLA, a two-phase framework that pretrains compositional world value functions (WVFs) and learns a semantic parser via RL and in-context learning to map natural language to Boolean WVF expressions. By composing WVFs with logical operators, CERLLA solves a large set of language-specified tasks with dramatically improved sample efficiency and generalization in the BabyAI/MiniGrid domain, achieving near-oracle performance on 162 tasks. A novel prompting strategy combines in-context examples with environment rollouts to train the semantic parser without demonstrations, and results demonstrate strong performance gains over non-compositional baselines, including held-out task generalization. The work highlights the practical potential of compositional RL integrated with language models for scalable, interpretable, and efficient language-grounded control.

Abstract

Combining reinforcement learning with language grounding is challenging as the agent needs to explore the environment while simultaneously learning multiple language-conditioned tasks. To address this, we introduce a novel method: the compositionally-enabled reinforcement learning language agent (CERLLA). Our method reduces the sample complexity of tasks specified with language by leveraging compositional policy representations and a semantic parser trained using reinforcement learning and in-context learning. We evaluate our approach in an environment requiring function approximation and demonstrate compositional generalization to novel tasks. Our method significantly outperforms the previous best non-compositional baseline in terms of sample complexity on 162 tasks designed to test compositional generalization. Our model attains a higher success rate and learns in fewer steps than the non-compositional baseline. It reaches a success rate equal to an oracle policy's upper-bound performance of 92%. With the same number of environment steps, the baseline only reaches a success rate of 80%.

Compositional Instruction Following with Language Models and Reinforcement Learning

TL;DR

The paper addresses the high sample complexity of learning language-conditioned reinforcement learning tasks by introducing CERLLA, a two-phase framework that pretrains compositional world value functions (WVFs) and learns a semantic parser via RL and in-context learning to map natural language to Boolean WVF expressions. By composing WVFs with logical operators, CERLLA solves a large set of language-specified tasks with dramatically improved sample efficiency and generalization in the BabyAI/MiniGrid domain, achieving near-oracle performance on 162 tasks. A novel prompting strategy combines in-context examples with environment rollouts to train the semantic parser without demonstrations, and results demonstrate strong performance gains over non-compositional baselines, including held-out task generalization. The work highlights the practical potential of compositional RL integrated with language models for scalable, interpretable, and efficient language-grounded control.

Abstract

Combining reinforcement learning with language grounding is challenging as the agent needs to explore the environment while simultaneously learning multiple language-conditioned tasks. To address this, we introduce a novel method: the compositionally-enabled reinforcement learning language agent (CERLLA). Our method reduces the sample complexity of tasks specified with language by leveraging compositional policy representations and a semantic parser trained using reinforcement learning and in-context learning. We evaluate our approach in an environment requiring function approximation and demonstrate compositional generalization to novel tasks. Our method significantly outperforms the previous best non-compositional baseline in terms of sample complexity on 162 tasks designed to test compositional generalization. Our model attains a higher success rate and learns in fewer steps than the non-compositional baseline. It reaches a success rate equal to an oracle policy's upper-bound performance of 92%. With the same number of environment steps, the baseline only reaches a success rate of 80%.
Paper Structure (15 sections, 1 equation, 5 figures, 6 tables)

This paper contains 15 sections, 1 equation, 5 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Methods
  4. Pretraining World Value Functions
  5. Compositionally Enabled RL Language Agent (CERLLA)
  6. In-Context Semantic Parsing with RL
  7. Baselines
  8. Results
  9. Simultaneous Learning of 162 Tasks
  10. Held-out Task Generalization
  11. Discussion
  12. Related Work
  13. Limitations and Future Work
  14. Conclusion
  15. Appendix

Figures (5)

  • Figure 1: Pipeline diagram of the learning process for the CERLLA agent. The agent takes in a BabyAI language mission command and a set of $10$ in-context examples that are selected using the BM25 search retrieval algorithm robertson2009probabilistic. The agent produces $10$ candidate Boolean expressions. These expressions specify the composition of the base compositional value functions. Each compositional value function is instantiated in the environment and the policy it defines is evaluated over $100$ rollouts. If the success rate in reaching the goal is greater than $92\%$, the expression is considered a valid parse of the language instruction and is added to the set of in-context examples.
  • Figure 2: Example of a task in the BabyAI domain babyai_iclr19. The agent (red triangle) needs to complete the mission -- "pick up the red key". Solving this task with compositional value functions requires using the conjunction of the pickup "red object" and "key" value functions.
  • Figure 3: Results for learning all $162$ tasks simultaneously. The mean episode success rate is plotted against the number of environment steps. Learning curves are presented for CERLLA and the non-compositional baseline agents. The Oracle agent is given the ground-truth Boolean expressions and upper bounds the attainable success rate in the environment, denoted by the dashed line at $92\%$. Our method is initialized at $19$ million steps to reflect the number of training steps used in pretraining the compositional value functions. Note the change in steps scale at $19$ million steps. Means and $95\%$ confidence intervals are reported over $10$ trials, $5$ trials for the LM Baseline.
  • Figure 4: The mean number of tasks solved is plotted against the number of environment steps. This quantity is equal to the total number of in-context examples present in CERLLA's in-context example set at that step. Because the Oracle Agent has access to the ground truth Boolean expressions for each task, it solves tasks immediately. The population of tasks remains constant, so the number of unsolved tasks decreases over time, leading to a logistic learning curve and an exponential decay in the rate at which new tasks are solved for the Oracle Agent. Means and $95\%$ confidence intervals are reported over $10$ trials.
  • Figure 5: Generalization performance for learning on $81$ randomly selected tasks (train) and evaluating on the held-out $81$ tasks (test). Learning curves are presented for our method and the baseline agents. The dashed and solid lines represent performance on the training and test sets respectively. The $x$-axis represents the fraction of the total training steps completed: $1$ million for the CERLLA and $21$ million for the LM Baseline agent. The dashed line denotes the success rate for considering the environment solved at $92\%$. Means and $95\%$ confidence intervals are reported over $15$ trials, $5$ trials for the LM Baseline.