🤖 machine learning

Task-Induced Representational Invariances Depend on Learning Objective in Deep RL

This paper demonstrates that deep reinforcement learning algorithms, specifically value-based (DQN) and policy-gradient (PPO) methods, learn distinct task-induced representational invariances—aligned with MDP homomorphism and action symmetries respectively—despite comparable performance, offering a principled framework for comparing models and gaining insights into neural coding.

Original authors: Manu Srinath Halvagal, Sebastian Lee, SueYeon Chung

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Manu Srinath Halvagal, Sebastian Lee, SueYeon Chung

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 two different students how to navigate a complex maze to find a hidden treasure. Both students are incredibly smart, they both use the exact same map, and they both eventually find the treasure perfectly every time. However, the paper reveals that they are thinking about the maze in completely different ways.

This research, titled "Task-Induced Representational Invariances Depend on Learning Objective in Deep RL," investigates how Artificial Intelligence (AI) agents "see" the world when they learn to solve problems. The authors found that the way an AI learns (its "learning objective") dictates how it organizes its internal knowledge, even if the final result is the same.

Here is a breakdown of their findings using simple analogies:

1. The Two Students: The "Value" Student vs. The "Action" Student

The researchers compared two famous AI learning methods: DQN and PPO.

The "Value" Student (DQN): This student is obsessed with the score. They ask, "How good is this specific spot in the maze?" They learn to group together spots that feel the same in terms of value.
- The Analogy: Imagine a city map where the student groups neighborhoods not by what you can do there, but by how "expensive" or "valuable" the real estate is. If two different streets have the same property value, this student treats them as identical, even if the buildings look different.
- The Paper's Finding: This student learns MDP Homomorphism Symmetries. In plain English, they learn to see the underlying mathematical structure of the maze. They realize, "Hey, this corner on the left is mathematically identical to that corner on the right because the rules of the maze are symmetrical." They compress the world into a smaller, simpler version of itself.
The "Action" Student (PPO): This student is obsessed with what to do next. They ask, "What move should I make here?" They learn to group together spots where the best move is the same.
- The Analogy: This student looks at the city and groups neighborhoods by the traffic lights. "If the light is green, I go; if it's red, I stop." They don't care if the buildings are different; they only care that the action required is the same.
- The Paper's Finding: This student learns Action Symmetries. They group states based on the optimal move. If the best move in State A is "Turn Left" and the best move in State B is also "Turn Left," they treat A and B as the same, regardless of the deeper mathematical structure of the maze.

2. The Proof: Same Goal, Different Brains

The researchers tested this in a navigation task (a digital maze).

The Result: Both students solved the maze perfectly. But when the researchers looked inside their "brains" (the internal data representations), they saw a clear split.
- The Value Student (DQN) had high similarity between states that were mathematically symmetrical (like mirror images of the maze).
- The Action Student (PPO) had high similarity between states that required the same move, even if the math behind them was different.

It's like two chefs making the exact same cake. One chef organizes their kitchen by ingredients (flour, sugar, eggs), while the other organizes by steps (mix, bake, frost). Both cakes taste the same, but their internal organization is totally different.

3. Why Does This Matter? (The Transfer Test)

The paper asks: Does this difference in "thinking style" matter when the student faces a new problem?

The Experiment: They took the students trained on one game (like Breakout) and asked them to play a new, similar game (like Pong).
The Outcome: The Value Student (DQN) was much better at transferring their skills to the new game. Because they learned the deep, structural rules of the environment (the "symmetries"), they could adapt quickly.
The Action Student (PPO) struggled more. Because they were so focused on specific moves, they had a harder time realizing that the new game was just a rotated or flipped version of the old one.

The Metaphor: If you teach someone to drive by memorizing specific turns ("Turn left at the red barn"), they might get lost if the barn is gone. If you teach them the principles of driving and road geometry, they can handle a new city with different landmarks. The Value Student learned the principles; the Action Student memorized the turns.

4. The "Chatbot" Surprise (Large Language Models)

The researchers also tested a Large Language Model (LLM)—a type of AI that powers chatbots—on the same maze task. They didn't train the chatbot; they just gave it a description of the maze in the chat.

The Finding: The chatbot's "thinking style" changed depending on how the maze was described.
- If the maze was described as a simple list of connections, the chatbot acted like the Action Student (focusing on moves).
- If the maze was described as a visual ASCII art tree (showing the structure clearly), the chatbot suddenly started acting like the Value Student, recognizing the deep symmetries.
The Takeaway: Unlike the fixed AI students, the chatbot is flexible. It can switch between "thinking styles" just by changing the prompt (the instructions).

Summary

The paper concludes that how an AI learns determines how it sees the world.

If you want an AI that understands the deep, structural rules of an environment (great for adapting to new situations), you should use Value-based methods (like DQN).
If you want an AI that is hyper-focused on the immediate best move, you might use Policy-based methods (like PPO).

This isn't just about math; it's about understanding that two systems can achieve the same perfect result while building completely different internal maps of reality. This helps scientists understand not just how AI works, but potentially how animal brains might organize learning differently depending on what they are trying to achieve.

Technical Summary: Task-Induced Representational Invariances Depend on Learning Objective in Deep RL

Problem Statement

While Deep Reinforcement Learning (RL) has achieved remarkable success in domains ranging from game playing to Large Language Model (LLM) alignment, the theoretical understanding of the learned representations within these models remains limited. This gap hinders direct comparisons between computational models and animal learning, as well as the systematic analysis of how task structures are encoded. Specifically, it is unclear how different RL algorithms (e.g., value-based vs. policy-gradient) encode task symmetries and whether these representational differences have functional consequences for transfer learning.

Methodology

The authors analyze deep RL representations through the lens of MDP reduction theory, specifically focusing on invariances to symmetries. The methodology involves:

Theoretical Framework:
- MDP Homomorphism: Utilizing group theory to identify symmetries in the state-action space where transition and reward functions remain invariant under specific transformations (e.g., reflection). This defines "MDP symmetry."
- Policy Symmetry: Defining symmetry based on whether a given policy selects the same action for different states (optimal policy symmetry, $\pi^*$ ).
- Representational Similarity Analysis (RSA): Employing RSA to compare the internal representations of different algorithms. Instead of comparing raw activations, the authors compute a Representation Similarity Matrix (RSM) based on pairwise cosine similarities between state representations. They specifically measure similarity for state pairs that are equivalent under MDP homomorphism versus those equivalent under policy symmetry.
Experimental Domains:
- Navigation Task: A labyrinth gridworld (based on Rosenberg et al.) with known MDP homomorphisms (reflection symmetry). Both tabular and visual (CNN-based) versions were used.
- Algorithms Compared:
  - DQN: A value-based method (learns a value function to derive a policy).
  - PPO: A policy-gradient method (directly optimizes the policy).
  - Deep SARSA(n): A control experiment trained on-policy (like PPO) but value-based (like DQN) to isolate the learning objective from data collection strategies.
- Complex Environments: CartPole (continuous) and Pong (high-dimensional pixels) to test generalization beyond exactly reducible MDPs.
- Transfer Learning: Experiments on Atari games (Breakout, SpaceInvaders, Pong) to test if learned symmetries aid generalization to new tasks.
- LLM Analysis: Prompting Qwen2.5-72B-Instruct to solve the navigation task using different input formats (ASCII tree vs. edge list) to observe emergent symmetries in in-context learning.

Key Contributions

Objective-Dependent Representational Invariances: The paper demonstrates that the learning objective systematically dictates the type of symmetry encoded in the representation, even when algorithms solve the same task with comparable performance.
- Value-based methods (DQN) learn representations invariant to MDP homomorphism symmetries (structural environment symmetries).
- Policy-gradient methods (PPO) learn representations invariant to action symmetries (states sharing the same optimal action).
Causal Isolation of Learning Objective: Through the Deep SARSA(n) experiment, the authors show that the divergence in symmetry encoding is driven by the learning objective (value-based vs. policy-gradient) rather than data collection strategies (off-policy vs. on-policy).
Transfer Learning Consequences: The paper provides empirical evidence that representations encoding MDP homomorphisms (DQN) facilitate better transfer learning between Atari tasks with shared environmental symmetries compared to those encoding action symmetries (PPO).
LLM Prompt-Dependence: The study reveals that LLMs can exhibit both types of symmetries, but the emergence of MDP homomorphism symmetry is prompt-dependent. It appears when the task structure is explicitly described (ASCII tree) but not when described via a less structured edge list, unlike the robust policy symmetry observed in both formats.

Key Results

Navigation Task: In the labyrinth task, DQN showed high representational similarity for MDP-symmetric state pairs but low similarity for states sharing the same optimal action. Conversely, PPO showed high similarity for action-equivalent states and low similarity for MDP-symmetric pairs.
Generalization: This dissociation held true in continuous (CartPole) and high-dimensional (Pong) environments, though the MDP-symmetry signal was weaker in Pong, likely due to local receptive fields in CNNs.
Transfer Learning: In Atari transfer experiments (e.g., Breakout $\to$ Pong), DQN agents (encoding MDP symmetries) showed significant performance gains and stability under observation rotations. PPO agents showed little to no transfer benefit and were more sensitive to observation transformations.
LLM Findings: LLMs robustly encoded policy symmetry regardless of prompt format. However, MDP homomorphism symmetry only emerged when the prompt explicitly highlighted the hierarchical structure (ASCII representation), suggesting LLMs possess a flexibility in representation that trained RL algorithms lack.

Significance and Claims

The authors claim this work provides a principled approach to comparing learned representations across RL algorithms using MDP reduction theory. The significance lies in:

Algorithm Selection: Offering a theoretical basis for choosing RL algorithms based on downstream goals. Practitioners seeking representations that capture environmental structure for transfer or generalization may prefer value-based methods, while those prioritizing action-level organization may prefer policy-gradient methods.
Neuroscience: Providing a framework to probe biological neural circuits. The authors suggest this pipeline can be extended to neural recordings to determine if animal brains preferentially encode MDP-like task symmetries or policy-like action mappings.
LLM Post-Training: Highlighting that RLHF and other post-training methods might alter LLM representations in ways that affect in-context learning capabilities, potentially introducing biases similar to those observed in PPO transfer failures.

The paper concludes modestly, acknowledging limitations such as the assumption of one-to-one state-observation mapping (challenging in POMDPs), the use of specific similarity metrics (cosine similarity), and the lack of a formal theory predicting which symmetries emerge from specific algorithmic properties. They view their findings as an initial empirical characterization rather than a complete theory.