🤖 machine learning

Understanding Goal Generalisation in Sequential Reinforcement Learning

This paper investigates how reinforcement learning agents generalize goals across sequential training pipelines, revealing that salient features and early-learned goals shape out-of-distribution behavior, and introduces an interpretable "latent policy gradients" method to predict these behaviors based on the evolution of low-dimensional latent variables.

Original authors: Jason Ross Brown, Edward James Young

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jason Ross Brown, Edward James Young

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot dog to fetch. First, you train it to fetch a red ball. It learns quickly. Then, you decide to teach it a new trick: fetch a blue frisbee. You might expect the robot to simply forget the red ball and focus entirely on the blue frisbee.

But what if, when you put the robot in a new room with both a red ball and a blue frisbee, it starts chasing both? Or what if it ignores the blue frisbee entirely and goes back to the red ball, even though you told it to fetch the blue one?

This is the problem of Goal Generalization. It's about understanding how an AI's past training history shapes what it values in new, unfamiliar situations. This paper by Jason Ross Brown and Edward James Young tries to solve the mystery of why AI agents make these specific choices when they leave the training gym and enter the real world.

Here is a simple breakdown of their work:

1. The Experiment: The Maze Game

The researchers didn't use complex real-world robots. Instead, they used a simple video game: a maze.

The Agent: A little grey circle (the robot).
The Goal: A specific object, like a Red Cross or a Blue Diamond.
The Training: They trained the robot in two ways:
1. Single Stage: Train it only to find the Red Cross.
2. Two Stages: Train it to find the Red Cross first, then switch and train it to find the Blue Diamond.

After training, they put the robot in a maze with two objects (e.g., a Red Cross and a Blue Diamond) that it had never seen together before. They watched to see which one the robot chased.

2. What They Found: The "Habit" of AI

The robots didn't act randomly. They had very consistent "personalities" based on their training history. The researchers discovered three main rules:

Shape is King, Color is Queen: The robots cared much more about the shape of the object (cross vs. diamond) than the color (red vs. blue). If a robot was trained on a "Cross," it would chase any cross, even if it was a different color.
Old Habits Die Hard (Persistence): If a robot learned to love "Crosses" in the first stage, it kept that love even after being trained to find "Diamonds" in the second stage. The first lesson stuck around and influenced the second.
The "Double-Training" Effect: If a robot was trained on a "Red Cross" and then a "Red Diamond," it became super obsessed with the color Red. The color Red was reinforced twice. However, the "Diamond" shape (which only appeared in the second stage) became much less important. The strong habit of "Red" blocked the robot from learning to value "Diamonds" as much as it otherwise would have.

3. The Solution: "Latent Policy Gradients"

The researchers wanted to predict these behaviors without having to re-train the robot every time. They invented a method called Latent Policy Gradients.

Think of this like a GPS for an AI's brain.

Instead of looking at the millions of lines of code inside the robot's brain, the researchers created a simple, low-dimensional map (a set of hidden numbers, or "latents").
They imagined these numbers evolving like a ball rolling down a hill. Every time the robot was trained on a new task, it was like pushing the ball in a specific direction.
By simulating this "ball rolling" through the training history (Stage 1, then Stage 2), they could predict exactly what the robot would do in a new maze.

The Analogy:
Imagine the robot's preferences are a piece of clay.

Training Stage 1 is a sculptor pressing the clay into a "Cross" shape.
Training Stage 2 is a second sculptor trying to press it into a "Diamond" shape.
The final shape isn't just a Diamond; it's a weird hybrid because the first press left a deep impression.
The Latent Policy Gradient method is a mathematical formula that looks at the two sculptors' instructions and predicts exactly what the final, weird clay shape will look like, without actually having to sculpt it.

4. Why This Matters

The paper shows that AI behavior isn't just a mystery; it has a structure.

It's Predictable: Even though the AI is in a new environment, its behavior is a logical result of its training history.
It's Interpretable: The researchers' method isn't a "black box." They can look at their formula and say, "Ah, the robot is chasing this because it was trained on this specific shape first, and that feature is very 'sticky' for this type of AI."

Summary

The paper argues that to understand what an AI will do in the real world, you can't just look at its final training. You have to look at its entire journey. By treating the AI's learning process as a series of steps where old habits stick and new ones are built on top of them, the authors created a simple, mathematical way to predict how an AI will generalize its goals. They proved that while AI behavior can be surprising, it follows a hidden logic that we can map and understand.

Technical Summary: Understanding Goal Generalisation in Sequential Reinforcement Learning

Problem Statement

Reinforcement learning (RL) agents often exhibit unintended goal-directed behaviors when deployed in out-of-distribution (OOD) environments. While unpredictable OOD behavior is a known issue across machine learning, RL agents frequently display a specific failure mode known as goal misgeneralisation: agents perform competently but pursue objectives that yield high rewards in training environments yet fail to align with intended goals in novel settings.

This problem is particularly critical for modern frontier AI systems, which are typically trained via complex, multi-stage pipelines (e.g., pre-training, supervised fine-tuning, and RL from human feedback). Current literature lacks a principled understanding of how these sequential training histories shape an agent's values and subsequent OOD behavior. The core question addressed is: Given an agent's training pipeline and its in-distribution behavior, can we predict its behavior on OOD samples?

Methodology

Experimental Setup

The authors investigate this problem in a controlled, procedurally generated maze environment.

Task: Agents (CNN-based PPO) navigate 8x8 mazes to reach goal objects defined by color and shape.
Training Pipelines: The study utilizes over 100 distinct training pipelines, ranging from single-stage training to two-stage sequential training. In two-stage pipelines, agents are trained on a first goal object and then fine-tuned on a second.
Evaluation: Agents are evaluated in OOD environments containing pairs of objects (including novel colors and shapes not seen during training). The primary metric is the empirical preference function ( $\hat{\Pi}_\Omega$ ), defined as the frequency with which an agent selects one object over another (or neither) across 276 evaluation environments.
Scale: The analysis covers over 298 agents and 250+ OOD environments.

The Latent Policy Gradients (LPG) Method

To explain and predict the observed behaviors, the authors introduce Latent Policy Gradients, a method that models the evolution of low-dimensional latent variables ( $w$ ) during training.

Latent Representation: The agent's preferences are encapsulated by a low-dimensional vector $w$ (dimension 10, matching the number of object features), rather than the full network parameters (4.8 million).
Preference Parametrisation: Preferences are modeled using a Boltzmann-rational assumption, where the probability of choosing a goal is a logistic function of its value. The value is linear in both the goal's features ( $\phi(g)$ ) and the latent variables: $V(g) = \phi(g) \cdot Sw$ .
Simulation of Training: The method simulates the training pipeline by performing gradient ascent on a modified policy objective over the preference function. The objective includes an entropy regularisation term controlled by a hyperparameter $\tau$ $τ$ .
- For a single stage with goal $\phi(g)$ , the objective is $J(\Pi) = \Pi(\phi(g)|\phi(g), 0) + \tau h(\Pi(\phi(g)|\phi(g), 0))$ .
- For multi-stage pipelines, the latent variables are updated sequentially, with the final weights of one stage initializing the next.
Interpretability: The method learns a saliency matrix ( $S$ ) which captures feature learning rates and cross-feature confusion (where training on one feature affects the valuation of another).

Key Results

Empirical Phenomena

Analysis of the empirical data reveals three distinct patterns in goal generalisation:

Feature Salience: Not all features drive generalisation equally. Shape features (particularly the cross shape) are more salient than color features (with black being the least salient).
Value Persistence: Values learned in the first stage of a multi-stage pipeline persist even when the second stage introduces a new goal. Early training history significantly influences final OOD preferences.
Strengthening and Inhibition:
- Strengthening: If a feature is shared between training stages, its value is strengthened relative to if it appeared only in the second stage.
- Inhibition: The presence of a shared feature inhibits the formation of values for non-shared features in the second stage (a phenomenon consistent with gradient starvation).
- Ordering Effects: When goals share a feature, the order of training stages significantly alters the final agent values.

Model Performance

The Latent Policy Gradients model achieves strong predictive accuracy:

Predictive Fit: The model minimizes the KL-divergence between predicted and empirical preference functions, outperforming baselines such as "Memoryless" (ignoring history), "Simultaneous" (ignoring order), and "Diagonal S" (ignoring feature confusion).
Generalisation: A model fitted solely on single-stage pipelines successfully predicts the behavior of agents trained on two-stage pipelines. Similarly, a model trained without distractors generalises to pipelines containing them.
Coherence: Empirical agent preferences are found to be approximately coherent (Boltzmann-rational), with Elo scores achieving 89.73% directional accuracy on held-out pairs.

Significance and Contributions

The paper makes three primary contributions:

Empirical Analysis: It provides a large-scale empirical analysis of goal generalisation across over 100 training pipelines, demonstrating that OOD behavior is coherent and structured by the training history, specifically highlighting phenomena like value persistence and feature salience.
Problem Formulation: It reframes the goal generalisation problem as a prediction task: mapping training pipelines directly to OOD preference functions.
Latent Policy Gradients: It introduces a novel, interpretable method for solving this prediction problem. The method:
- Achieves high predictive accuracy with a small number of parameters.
- Offers a mechanistic interpretation where multi-stage training corresponds to iterated projection of latent weights onto hyperplanes defined by the training goals.
- Defines a feature similarity metric (via $SS^T$ ) that explains how training on one goal influences the valuation of others.

Conclusion

The authors conclude that while OOD RL agent behavior is path-dependent on the entire training pipeline, this dependence possesses an underlying structure that can be captured. The Latent Policy Gradients framework provides a principled, interpretable tool for understanding how multi-stage training shapes agent values. This work lays the groundwork for a developmental perspective on goal generalisation, suggesting that understanding the evolution of latent variables during training is key to predicting and ensuring the safety of AI systems deployed in novel environments. The authors note that while the current framework assumes linear decomposition of values over features, the mathematical structure is general and serves as a foundation for extending these insights to more complex systems.