Source-linked AI summary
Reinforcement Learning from LLM Feedback to Counteract Goal Misgeneralization
Houda Nait El Barj, Theophile Sautory
TL;DR
Goal misgeneralization occurs when an RL agent pursues a proxy rather than its intended goal in new environments, and can be difficult to detect. This paper uses LLM analysis and preferences to build a reward model for further RL training. In maze navigation, the method improves goal generalization particularly when true and proxy goals are distinguishable and behavioral biases are evident.
Problem
Goal misgeneralization arises when an agent pursues an unintended goal despite accurate training feedback, and detecting it requires testing on substantially shifted distributions.
Source: §1 · p. 2Method
The method uses an LLM to assess agent policies and potential failure scenarios, then trains a reward model from LLM preferences over rollouts to further train the agent.
Source: §1 · p. 2, §3.3 · p. 4Results
Takeaways & Limitations
The maze-navigation results indicate that an LLM can supervise RL agents effectively despite not being proficient at the task.
Source: §1 · p. 2
Abstract
from arXiv · showhide
We introduce a method to address goal misgeneralization in reinforcement learning (RL), leveraging Large Language Model (LLM) feedback during training. Goal misgeneralization, a type of robustness failure in RL occurs when an agent retains its capabilities out-of-distribution yet pursues a proxy rather than the intended one. Our approach utilizes LLMs to analyze an RL agent's policies during training and identify potential failure scenarios. The RL agent is then deployed in these scenarios, and a reward model is learnt through the LLM preferences and feedback. This LLM-informed reward model is used to further train the RL agent on the original dataset. We apply our method to a maze navigation task, and show marked improvements in goal generalization, especially in cases where true and proxy goals are somewhat distinguishable and behavioral biases are pronounced. This study demonstrates how the LLM, despite its lack of task proficiency, can efficiently supervise RL agents, providing scalable oversight and valuable insights for enhancing goal-directed learning in RL through the use of LLMs.
1 Introduction
The introduction frames goal misgeneralization as an inner-alignment problem and motivates scalable oversight. It presents LLM feedback and an LLM-derived reward model as a way to address the problem in RL.
1 Introduction
- The introduction distinguishes outer alignment, involving accurate reward specification, from inner alignment, involving whether a trained policy adheres to preferences and goals.It describes reward hacking and specification gaming as outer-misalignment concerns and notes that RLHF can still sometimes lead to reward hacking.Source: §1 · pp. 1–2
- Goal misgeneralization occurs when an agent pursues an unintended goal despite accurate training feedback, and is difficult to detect because it requires testing on substantially shifted distributions.The introduction identifies scalable oversight and reduced human involvement in training and evaluation as concerns in complex environments.Source: §1 · p. 2
- The paper proposes an RL method in which an LLM assesses policies, suggests training modifications, and supplies feedback for a reward model.The authors report a maze-navigation evaluation and significant reductions in goal-misgeneralization instances across varied settings.Source: §1 · p. 2
- The stated contributions include counteracting goal misgeneralization without explicitly training on out-of-distribution data and using LLM preferences to supervise RL agents.The paper emphasizes that the LLM can provide this supervision despite not being proficient at the task.Source: §1 · p. 2
2 Preliminaries
The preliminaries define goal misgeneralization and introduce preference-based reward modeling. They then frame LLM feedback as a way to fine-tune an agent prone to pursuing a proxy goal.
2.1 Goal Misgeneralization
- Goal misgeneralization occurs when a policy achieves low test reward despite high training reward, pursuing a behavioral objective different from the intended goal after distribution shift.The definition associates this failure with insufficient training-environment diversity or a proxy that correlates with the intended objective during training but diverges at test time.Source: §2.1 · p. 2
2.2 Reward modeling
- Reward modeling uses paired policy rollouts with preference labels to train a reward model by minimizing cross-entropy between predicted probabilities and those labels.The preference label may favor either rollout or remain uniform when neither is preferred or both are equally preferred.Source: §2.2 · pp. 2–3
2.3 RLAIF agent
- The proposed RLAIF setup assumes an agent trained with an accurate reward can still fail out of distribution because a proxy goal correlates with its learned goal.The agent is then trained using a reward model derived from LLM feedback to increase its generalization capability.Source: §2.3 · p. 3
3 RL from LLM Feedback to reduce Goal Misgeneralization
The method trains an RL agent, uses LLM analysis to identify potential generalization failures, and gathers LLM preferences over rollouts in suggested environments. Those preferences train a reward model that is combined with the original reward during further training.
3.1 Initial RL training
- The agent first trains for τ · T timesteps, after which the method samples N I policy rollouts for LLM assessment.A rollout is a sequence of state-action pairs generated by the agent's policy at that point in training.Source: §3.1 · p. 3
3.2 LLM assessment and suggestions
- The LLM analyzes sampled rollouts for potential goal-generalization failures and suggests environment modifications; the agent is deployed in those environments to collect rollouts for reward-model training.This approach is intended for settings where the underlying reward is inaccessible or the training environment cannot be adjusted.Source: §3.2 · p. 3
3.3 LLM preference labeling
- The LLM compares pairs of policy rollouts and records its preference as a distribution over the two choices, using a uniform distribution when neither is preferred or both are equally preferable.The preference prompt includes the problem description and rollout samples, with a maze representation key and a fixed-format response ending.Source: §3.3 · p. 4
3.4 Reinforcement learning from LLM preferences
- The method trains a reward model from LLM preferences and resumes RL training using a composite of that model's reward and the original reward.The supplied description does not specify the mixing values.Source: §3.4 · p. 4
4 Dataset and experiment details
The experiments use Procgen Maze to study goal misgeneralization under a zero-shot shift in cheese locations, then apply LLM feedback to refine an RL agent. The evaluation defines training reward as capability and test reward in unseen settings as generalization capacity.
4.1 Dataset and 4.2 Goal Misgeneralization in Procgen Maze Game
- The Procgen Maze task rewards reaching cheese; training places it in an upper-right region, while test locations can occur anywhere in the maze.The training region ranges in size from 1 to 10, and Figure 1 contrasts a fixed training goal with a random test position.Source: §4.2 · p. 4, §4.1 · p. 4, Figure 1 · p. 4
4.3 RLAIF training on the Procgen Maze game
- After 75M total timesteps, the method pauses training to sample policy rollouts, which are represented as environment images, textual mazes, and agent trajectories.The LLM examines training-set rollouts to identify potential failures and suggest dataset changes.Source: §4.3 · p. 5, Figure 2 · p. 5, Figure 3 · p. 5
- 80% of GPT-4 queries identified the confound that cheese location and cheese were indistinguishable in the environment.GPT-4 also flagged other possible generalization flaws, which the study did not include.Source: §4.3 · p. 5
- Following GPT-4’s suggestion, the authors sampled 158 rollouts on mazes with cheese locations throughout the maze and formed 1,800 same-dimension rollout pairs for preference labeling.GPT-4 rated which of two candidate rollouts it preferred.Source: §4.3 · pp. 5–6, Figure 4 · p. 6
- GPT-4 preferred the second-position rollout in 56% of pairs and the first-position rollout in 43%, so the authors averaged preferences after reversing candidate order.The remaining 1% of cases were indifferent.Source: §4.3 · p. 6
4.3 RLAIF training on the Procgen Maze game and 4.4 Evaluation
- The averaged GPT-4 preferences train a reward model that is integrated into resumed PPO training; evaluation uses training reward for capability and unseen-test reward for generalization capacity.The reward model is assessed against LLM labels using precision, recall, and F1.Source: §4.3 · p. 6, Figure 5 · p. 6, §4.4 · p. 7
5 Experimental results
The LLM-informed reward model shows cautious but useful preference classification, while the trained agent improves goal generalization in several randomization regions. The gains depend on whether intended and proxy goals are distinguishable and whether the agent’s behavior remains flexible.
5.1 Reward model
- Low-prevalence preference classes influence macro-average metrics; the model has high precision but lower recall for the (0, 1) and (1, 0) classes.The two partial-preference classes each constitute 0.8% of labels, and the authors report that cautious preferences can help filter noise from LLM labeling.Source: §5.1 · p. 7, Table 1 · p. 7
5.2 RL Agent’s Performance: Capacity and Goal Generalization
- Across seven goal-randomization regions, the method reduces goal misgeneralization and matches the existing reward function in three others.The agent also achieves a maximum mean training reward of 10 across all regions; test reward exceeds the environment-reward baseline in regions {2, 3, 4, 5, 6, 7, 9}.Source: §5.2 · p. 7
- The largest improvement occurs in region 4, while the LLM-informed reward model does not improve generalization in region 1, where intended and proxy goals are indistinguishable.The region-4 agent remains flexible enough for the reward model to influence its behavior.Source: §5.2 · p. 7
- With λ = 0.5, generalization matches the environmental reward function in regions smaller than 4, but the disparity fades in larger regions.In larger regions the base agent independently learns a generalizable goal, and the LLM-informed reward becomes the primary behavioral driver.Source: §5.2 · p. 8
6 Related Work
Prior work frames goal misgeneralization, scalable oversight, and AI-generated feedback as related challenges and approaches. This paper’s context includes earlier work on eliciting misaligned behavior, improving training data, and using AI feedback for RL.
Instances of Goal Misgeneralization and potential mitigations
- Langosco et al. [2023] describe three RL examples of goal misgeneralization, including a maze task; their mitigation increases the training dataset.Shah et al. [2022] provide a broader definition and examples across deep learning; the paper uses Langosco et al.’s maze as its evaluation basis.Source: §6 · p. 8
- Prior mitigation work elicits and penalizes misaligned behavior with unrestricted adversarial examples [Song et al. 2018] and generates examples to reduce unwanted language output [Ziegler et al. 2022; Perez et al. 2022].Source: §6 · p. 8
Scalable oversight
- Scalable oversight proposals include debate [Irving et al. 2018], self-critique [Saunders et al. 2022; Leike et al. 2018], and market making [Hubinger 2020].The cited work notes that empirical evaluation is complicated by the lack of systems that exceed human capabilities [Bowman et al. 2022].Source: §6 · p. 8
RLHF & RLAIF
- Related AI-feedback work includes self-critique and self-supervision, while RLAIF uses AI models to learn human preferences for RL fine-tuning [Bai et al. 2022b].LLM preferences have also been used to train an RL agent in procedurally generated NetHack [Klissarov et al. 2023]; RLHF and RLAIF showed similar improvements over a supervised fine-tuned baseline [Lee et al. 2023].Source: §6 · pp. 8–9
7 Discussion
GPT-4 could assess whether agent trajectories reached goals more often than it could solve mazes itself. The discussion also describes position bias and identifies settings where goal confounding constrains the reward model’s impact.
7.1 LLM supervision capabilities
- GPT-4 correctly assessed whether a trajectory reached the goal 77% of the time, compared with solving mazes itself 10% of the time.The authors use this contrast to argue that supervision can rely on assessing an agent’s policy relative to its goal, without the LLM being capable of the task itself.Source: §7.1 · p. 9
7.2 Environment Confoundedness vs Reward Model Informativeness
- The agent learns location rather than a feature regardless of training method; when the goal is indistinguishable from its position, the behavior cannot be corrected without changing the dataset or model architecture.As randomization regions grow, the agent becomes less biased toward the upper-right corner, reducing the method’s impact on that bias.Source: §7.2 · p. 9
8 Conclusion
The paper presents LLM feedback as an approach to mitigate RL goal misgeneralization, with improvements especially in complex environments where true and proxy goals differ and behavioral biases are evident. The appendix examples show prompts for policy-failure analysis and rollout preferences.
8 Conclusion
- The conclusion reports improved goal generalization, particularly where proxy and true goals are distinct and behavioral biases are evident.The authors present LLM feedback as a potential source of scalable oversight and goal-focused learning.Source: §8 · p. 9
A Prompt & GPT Analyses Examples
- The prompt example gives GPT-4 a maze and trajectory sampled from the training distribution to examine potential policy failures.Other examples show GPT’s analysis responses and a prompt asking for preferences between two policy rollouts.Source: Figure 7 · p. 13, Figure 8 · p. 13, Figure 9 · p. 14
B.1 Reward model details
The appendix details reward-model training and preference evaluation, then describes the PPO objective and policy architecture used in the experiments.
B.1 Reward model details
- Reward-model training used 1,412 training rollout pairs from 142 mazes and 388 validation pairs involving 16 validation mazes.Trajectories were capped at 50 steps; training ran for 300 epochs with batch size 256 until binary cross-entropy loss plateaued.Source: §B.1 · p. 14
- Predicted rollout-preference probabilities were converted into classes using intervals presented in Table 3.Precision, recall, and F1 scores are reported for completion.Source: §B.1 · p. 15, Table 3 · p. 15
B.2 Reinforcement Learning with PPO
- PPO training estimates advantages with GAE and maximizes a clipped objective through stochastic-gradient-ascent updates.The objective uses the probability ratio between current and old policies and a clipping-range hyperparameter.Source: §B.2 · p. 15
B.3 Reinforcement Learning Model Details
- PPO hyperparameters were matched to Langosco et al. [2023], and the policy used an Impala architecture without recurrent components.The cited hyperparameter choices were intended to allow direct comparison.Source: §B.3 · p. 15
B.4 Base Agent Training
The base agent was trained using environment reward only, and Figure 10 reports its mean reward on train and test mazes.
B.4 Base Agent Training
- Table 5 lists hyperparameters for reinforcement learning.Source: Table 5 · p. 16
- Figure 10 reports mean reward on train and test mazes for the environment-reward-only base agent after 50M timesteps.Source: Figure 10 · p. 16