Policy Learning with a Language Bottleneck

Megha Srivastava; Cedric Colas; Dorsa Sadigh; Jacob Andreas

cs.LGMay 7, 2024

Policy Learning with a Language Bottleneck

Authors: Megha Srivastava, Cedric Colas, Dorsa Sadigh, Jacob Andreas

Organizations: Stanford University · Massachusetts Institute of Technology Inria

Abstract

Modern AI systems such as self-driving cars and game-playing agents can achieve superhuman performance, but often lack human-like generalization, interpretability, and inter-operability with human users. Inspired by the rich interactions between language and decision-making in humans, we introduce Policy Learning with a Language Bottleneck (PLLB), a framework enabling AI agents to generate linguistic rules that capture the high-level strategies underlying rewarding behaviors. PLLB alternates between a rule generation step guided by language models, and an update step where agents learn new policies guided by rules, even when a rule is insufficient to describe an entire complex policy. Across five diverse tasks, including a two-player signaling game, maze navigation, image reconstruction, and robot grasp planning, we show that PLLB agents are not only able to learn more interpretable and generalizable behaviors, but can also share the learned rules with human users, enabling more effective human-AI coordination. We provide source code for our experiments at https://github.com/meghabyte/bottleneck .

Figures & tables

Figure 1 : Policy Learning with a Language Bottleneck ( PLLB ) alternates between two steps: 1) gen_rule generates a linguistic rule $\mathcal{L}_{i}$ explaining the agent’s best behaviors by prompting a language model with contrastive (positive and negative) episodes; 2) update learns a new policy conditioned on $\mathcal{L}_{i}$ . PLLB strengthens human-AI coordination by constraining policies to be more interpretable and generalize better.

Figure 2 : PLLB can be applied to a diversity of domains. In all domains, it iterates between gen_rule , a function that prompts an LM (top gray boxes) to extract a linguistic rule ( $\mathcal{L}$ , blue) by contrasting high and low reward episodes from past experience (green and red boxes), and update , a function for updating an agent’s policy through interaction with the environment conditioned on $\mathcal{L}$ . PLLB can be applied to multi-step decision making tasks and visual or robotics tasks alike with minimal implementation variations (replacing the LM by a VLM in gen_rule , or replacing policy regularization by simple instruction conditioning in update . Gray boxes represent prompts gen_rule and update , and full prompt details are in Appendix Section B.7 .

Figure 3 : SaySelect . a) Bottleneck agents learn as fast as TabularQ and InstructRL and faster than agents using adversarial rules Adversarial . b) Unlike TabularQ , they learn human-interpretable policies; this without relying on an external instruction like InstructRL . c) When faced with speakers enforcing a non-human-interpretable policy, Bottleneck converges faster than InstructRL .

Figure 4 : Results in Maze . a) Bottleneck agents learn as fast as the non-linguistic Baseline agents , but faster than Adversarial and LinearQ agents. b) When faced with a new maze with similar structure, Bottleneck agents learn faster than TabularQ (which does not perceive color) and LinearQ (which does, but learns slower). c) When faced with a maze of a different structure, Bottleneck agents adapt swiftly while LinearQ cannot recover. We do not evaluate Adversarial for the Generalization or Adaptation experiments due to its poor performance in the standard setting.

Figure 5 : Participants given PLLB rules solve new mazes faster than those given either visual or no aid (control). They also self-reported linguistic feedback as more useful than visual. * marks statistical significance at level 0.05 with two-sided Mann-Whitney test after Bonferroni correction.

Figure 6 : Builder : Bottleneck listeners improve upon Baseline descriptions over time, unlike an ( Adversarial ) listener that inaccurately represents rewards. Human users reconstructed target images with b. higher accuracy, and c. less time, when provided Bottleneck instruction vs. original Baseline descriptions. Birds : PLLB rules (3 iterations) help d. automated and e. human speakers generate better descriptions for naturalistic image reconstruction for three distinct reward functions.

Figure 7 : Results in Grasp . a) Bottleneck agents learn faster than the non-linguistic RL Baseline agents , as well as Adversarial and an Oracle zero-shot prompt of a VLM without any internal policy update. b) Bottleneck agents generalize faster than both Baseline (which over-relies on image details) and Oracle to unseen bag designs and locations, (c) shows the fixed linguistic rule d) Example grasps, shaded by reward.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8 : Overview of SaySelect game, reproduced from Hu and Sadigh (2023) .

Figure 9 : Visual aid provided to participants in the human subject study for the Maze task.

Figure 10 : Interface used in the human subject study for the Maze task. Cells not visited by the participants are hidden in gray.

Figure 11 : Interface used in the human subject study for the Builder task.

Figure 12 : Task environment for Grasp , showing new location and distractors for the Generalization setting.

Figure 13 : Both the builder and birds tasks consist of speaker and listener agents. At test time, the speaker needs to provide a language description to the listener that helps them recreate the image accurately. For both tasks, PLLB helps improve listener accuracy over time.

Setting	$\mathcal{L}_{1}$	$\mathcal{L}_{2}$	$\mathcal{L}_{5}$
Standard	I should follow the strategy of choosing action 4 whenever the observation is 2, 3, 4, or 5.	I should follow the strategy of choosing actions based on the observation provided by AGENT 1. If the observation is 1, I should choose action 1. If the observation is 2, I should choose action 3. If the observation is 3, I should choose action 3. If the observation is 4, I should choose action 4. If the observation is 5, I should choose action 5.	I should follow the strategy of choosing the same action as AGENT 1 for observations 1, 2, 3, 4, and 5.
Fixed Speaker	I should choose action 1 when observation is 1 or 2 or 4 or 5. I should choose action 2 when observation is 3.	I should choose action 1 when observation is 1, 2, 3 or 5. I should choose action 2 when observation is 4.	I should choose action 1 when observation is 1. I should choose action 4 when observation is 2. I should choose action 5 when observation is 3. I should choose action 2 when observation is 4. I should choose action 3 when observation is 5.

Appendix

Table 1 : Example

\mathcal{L}

rules generated for the SaySelect environment, for the Standard setting (both Listener and Speaker agents are RL agents trained from random initialization) and a Fixed Speaker agent.

Size	$\mathcal{L}_{1}$	$\mathcal{L}_{2}$	$\mathcal{L}_{5}$	$\mathcal{L}_{8}$
7x7	I should start by moving north, then continue moving north until I reach a wall, then turn east and move until I reach a wall, then turn south and move until I reach a wall, then turn west and move until I reach the end of the maze.	I should prioritize moving North, East, and South over moving West, as the HIGH REWARD EXAMPLES show a higher frequency of these actions. 2. I should avoid moving West unnecessarily, as it is rarely seen in the HIGH REWARD EXAMPLES and often leads to LOW REWARD EXAMPLES	I should take the EAST action when I see WHITE, and then repeat the EAST action until I see RED, then take the SOUTH action, and then repeat the EAST action until I see BLUE, then take the NORTH action, then take the EAST action until I see RED again, then take the SOUTH action, and finally take the WEST action.	I should take the EAST action when I observe WHITE, and keep taking EAST actions until I observe RED, then take the SOUTH action. After taking the SOUTH action, I should take the EAST action again until I observe BLUE, then take the NORTH action. After taking the NORTH action, I should take the EAST action until I observe RED again, then take the SOUTH action.

Appendix

Table 2 : Example

\mathcal{L}

rules generated for the Maze environment for the 7x7 maze size.

Setting	$\mathcal{L}_{1}$	$\mathcal{L}_{2}$	$\mathcal{L}_{3}$
Re-Initialization	1. Be specific with location details: Agent 1 should provide detailed location information for each element in the image, such as x and y coordinates. 2. Use descriptive language for elements, such as "red dot" or "green triangle".	1. Use specific coordinates when instructing Agent 2 to draw shapes. 2. Use descriptive language to specify the color and shape of each element. For example, "a green triangle" instead of "a green thing".	1. Be specific with location coordinates: provide specific coordinates for the location of each shape, using the format x=0.XX, y=0.YY. 2. Use descriptive shape names: Instead of using generic terms like "dot" or "square," use more descriptive names that indicate the shape’s color and size, such as "green triangle" or "red square."
Continual Training	1. Be specific and detailed in your instructions. High reward examples have specific coordinates and shapes, while low reward examples have more general descriptions. 2. Use a consistent format for your instructions. High reward examples have a consistent format for listing coordinates and shapes, while low reward examples have a more free-form format.	1. Provide explicit coordinates for each element in the image, using the format (x, y). 2. Use specific colors when referring to elements in the image, such as "red", "green", or "blue". Avoid using vague terms like "colored" or "shaded".	1. Use a consistent format for describing shapes, such as always listing the x-coordinate first, followed by the y-coordinate. For example, instead of "one green square at the point x=0.53, y=0.24", use "one green square at (0.53, 0.24)". 2. Avoid using vague terms like "various shades of green". Instead, use specific colors, such as "green" or "blue". Additionally, use specific shapes, such as "square" or "triangle", rather than vague terms like "rectangle".

Appendix

Table 3 : Example

\mathcal{L}

rules generated for the Builder environment.

Reward	$\mathcal{L}_{1}$	$\mathcal{L}_{2}$	$\mathcal{L}_{3}$
color	Describe the bird’s color, species, and any distinctive markings or patterns.	Describe the bird’s coloration accurately.	Describe the bird’s coloration accurately.
background	Include details about the bird’s surroundings, such as the type of branch or post it is on, and any additional elements in the background.	Include the bird’s action (perched, flying, standing) and its location (on a branch, railing, pole, etc.)	Describe the bird’s action (flying, perching, standing) and the environment it is in (sky, tree, water).
species	Describe the bird’s color, markings, and any distinctive features.	Describe the subject’s unique features, such as coloration, beak shape, or other distinguishing characteristics.	Include specific details about the bird’s appearance, such as the color of its feathers, beak, or eyes, and any distinctive markings or patterns.

Appendix

Table 4 : Example

\mathcal{L}

rules generated for the Birds environment demonstrate reward-specificity over time.

Reward	$\mathcal{L}_{10}$	$\mathcal{L}_{30}$	$\mathcal{L}_{50}$
Shoulder Torque	Select keypoints that are away from sharp edges or corners of the object to avoid potential damage and improve stability during grasping.	Prioritize grasp keypoints that are closer to the center of the object for more stable and precise pickup.	Prioritize grasping at the middle of the object’s surface rather than the top edges or handles to ensure a stable grip.

Appendix

Table 5 : Example

\mathcal{L}

rules generated for the Grasp environment.

Episode	500	1000	6000
Reward ( Original, temp=0.5 )	$0.63\pm 0.1$	$0.8\pm 0.1$	$0.96\pm 0.03$
Reward ( Original, temp=0.9 )	$0.47\pm 0.2$	$0.68\pm 0.1$	$0.96\pm 0.05$
Reward ( Original, temp=0.1 )	$0.66\pm 0.1$	$0.70\pm 0.09$	$0.96\pm 0.02$
Reward ( No Format Instruction )	$0.51\pm 0.2$	$0.77\pm 0.05$	$0.96\pm 0.02$
Reward ( Low Context )	$0.56\pm 0.2$	$0.73\pm 0.1$	$0.96\pm 0.02$
Reward ( Rephrase )	$0.64\pm 0.1$	$0.81\pm 0.2$	$0.96\pm 0.02$

Appendix

Table 6 : Reward across ablations to probe sensitivity of rule generation to prompt variation and temperature. Results are averages across 5 trials.

Episode	500	1000	6000
Interpretability ( Original, temp=0.5 )	$0.77\pm 0.03$	$0.87\pm 0.04$	$0.97\pm 0.02$
Interpretability ( Original, temp=0.9 )	$0.72\pm 0.05$	$0.80\pm 0.05$	$0.93\pm 0.03$
Interpretability ( Original, temp=0.1 )	$0.87\pm 0.04$	$0.90\pm 0.1$	$0.95\pm 0.02$
Interpretability ( No Format Instruction )	$0.67\pm 0.06$	$0.91\pm 0.03$	$0.94\pm 0.02$
Interpretability ( Low Context )	$0.77\pm 0.04$	$0.84\pm 0.05$	$0.97\pm 0.03$
Interpretability ( Rephrase )	$0.8\pm 0.04$	$0.97\pm 0.04$	$0.97\pm 0.02$

Appendix

Table 7 : Human interpretability across ablations to probe sensitivity of rule generation to prompt variation and temperature. Results are averages across 5 trials.

Episode	500	1000	6000
Reward ( Original )	$0.63\pm 0.1$	$0.8\pm 0.1$	$0.96\pm 0.03$
Reward ( Non-contrastive )	$0.53\pm 0.02$	$0.7\pm 0.2$	$0.91\pm 0.03$
Reward ( Adversarial )	$0.21\pm 0.1$	$0.45\pm 0.03$	$0.8\pm 0.05$

Interpretability ( Original )	$0.77\pm 0.03$	$0.87\pm 0.04$	$0.97\pm 0.02$
Interpretability ( Non-contrastive )	$0.75\pm 0.02$	$0.8\pm 0.06$	$0.88\pm 0.02$

Appendix

Table 8 : Reward and human-interpretability scores for contrastive and non-contrastive versions of PLLB in SaySelect . Results are averages across 5 trials.

Table 9 : Noncontrastive ablation in comparison with main results for Birds task across 3 reward functions. Results are averaged across 5 trials.

Explore similar work

Jun 25, 2026cs.AI

Joint Learning of Experiential Rules and Policies for Large Language Model Agents

For LLM agents in multi-step interactive environments, a key challenge is to make effective use of accumulated interaction experience. Existing work has typically separated two uses of such experience: keeping it outside the model as natural-language rules for later prompting, or using trajectories and feedback to update the model parameters. The former is easy to interpret but can fall out of sync with the evolving policy; the latter improves the policy more broadly but provides only limited correction for local mistakes in sparse-reward settings. We present Joint Learning of Experiential Rules and Policies for LLM Agents (JERP), which updates a long-term experiential-rule pool and the policy from the same interaction trajectories. At decision time, JERP retrieves task-relevant rules and conditions the agent on them together with the interaction history. After each episode, it uses the collected trajectories both to optimize the policy and to revise the rule pool by comparing current rollouts with reference successful trajectories. This coupling keeps the rule pool aligned with the evolving policy while allowing stable and effective behaviors to be gradually absorbed into the model itself. Experiments on AlfWorld and WebShop show that JERP yields consistent gains in decision performance for complex interactive tasks.

Shicheng Ye, Chao Yu

Sun Yat-sen University

Continual Learning for LLM Agents LLM Agent Training Policy Learning RL for LLM Agents Memory-Augmented Language Model Agents

Ask PDF

Jun 7, 2026cs.AI

Distilling LLM Reasoning into an Interpretable Policy Tree for Human-AI Collaboration

Constructing efficient and reliable policies to assist humans is indispensable for human-AI collaboration. Existing methods mainly follow two lines of work. Most prior work relies on multi-agent reinforcement learning (MARL) to learn black-box policies, which limits interpretability and raises safety concerns. Recent methods query large language models (LLMs) at each decision step, causing slow responses and high inference costs. We propose Collaboration Policy Tree (Co-pi-tree), a closed-loop method that learns an executable policy tree consisting of a partner-behavior prediction tree and an agent-action selection tree. Co-pi-tree constructs a policy by distilling LLM reasoning into policy tree code. It then evaluates the policy through partner interaction, obtains feedback, and uses natural language to summarize the interaction feedback to improve problematic branches. Experiments in Overcooked-AI show that Co-pi-tree improves average reward by 35.4% over the baseline average, while reducing the number of LLM queries by 77.7% and test-time latency by 97.1%. Project page: https://beiwenzhang.github.io/Co-pi-tree/

Beiwen Zhang, Yongheng Liang, Guowei Zou +2

Sun Yat-sen University

Reinforcement Learning Interpretability in RL Policy Distillation Human-AI Collaboration

Ask PDF

Sep 27, 2026cs.CL

NLPG: Natural-Language Policy Gradients for Self-Evolving Language Agents

Large language model agents increasingly rely on compound programs for retrieval, tool use, reasoning, and verification, yet their failures often arise from local procedural decisions. Existing reinforcement-learning and prompt-optimization approaches typically rely on scalar rewards or repeatedly modify entire prompts, making it difficult to capture and reuse procedural improvements while preserving a frozen agent. To address this problem, We propose Natural-Language Policy Gradients (NLPG), an external policy-memory method for improving a fixed agent without changing its model parameters or program structure. NLPG diagnoses execution traces, propagates downstream feedback backward through the module graph, and converts recurring failures into route-local natural-language corrections that are aggregated into bounded policy updates for subsequent executions. Across six benchmarks covering memory, reasoning, instruction following, and evidence verification, NLPG also outperforms the strongest listed baseline for each benchmark by 8.71 percentage points on average. These results provide evidence that evaluated procedural experience can be transformed into local and interpretable policy updates, enabling continual improvement of frozen agents.

Xu Liu, WenZhang Wei, Jun Cao +4

School of Electronic Information, Wuhan University, Wuhan 430079, China · School of Remote Sensing and Information Engineering, Wuhan University, Wuhan 430079, China · State Key Laboratory of Information Engineering in Surveying, Mapping and Remote Sensing, Wuhan University, Wuhan 430079, China

Continual Learning for LLM Agents Policy Optimization Language Model Self-Improvement Policy Learning Memory-Augmented Language Model Agents

Ask PDF