Reinforcement Learning from LLM Feedback to Counteract Goal Misgeneralization

TL;DR

This work targets goal misgeneralization, an inner-alignment failure in RL where agents optimize a proxy goal under distribution shift. It introduces a framework that uses an LLM to analyze policies, suggests diverse training scenarios, and derives an LLM-informed reward model from preferences to guide RL via PPO. The approach applied to Procgen Maze demonstrates reduced misgeneralization, with notable gains when true and proxy goals are distinguishable, validating LLM supervision as scalable oversight for enhancing goal-directed learning. The findings suggest that LLMs can supervise and shape RL behavior even without task proficiency, offering a practical pathway to more robust, generalizable AI systems.

Abstract

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.

Figures (10)

• Figure 1: Example of cheese location variants. The left hand side shows an example of the fixed goal position during training. The right hand side is an example of the random test position.
• Figure 2: An example of a roll-out representation. The RGB image is the environment. The maze is its textual representation neglecting the wall padding. The trajectories are the steps taken by the agent, where (0, 0) is the upper left corner of the maze.
• Figure 3: LLM assessment and suggestions. Policy rollouts of the agent are sampled on the training dataset. They are shared with the LLM which identifies situations in which the current policy could fail, and suggest improvements to the dataset.
• Figure 4: An example of a GPT-4 preference labeling.
• Figure 5: LLM preference modelling and reward model. The RL agent is deployed on the LLM generated dataset and its rollouts are stored. The LLM compares pairs of rollouts and provides preferences, which are used to train a new reward model. The reward model is then integrated to the remaining training timesteps of the agent.