🤖 machine learning

Inverting the Bellman Equation: From $Q$ -Values to World Models

This paper challenges the traditional separation between model-based and model-free reinforcement learning by proving that value-based agents trained on diverse reward functions implicitly encode a unique world model, which can be explicitly extracted via a new method called $P$ -learning to achieve strong generalization on out-of-distribution tasks.

Original authors: Alistair Letcher, Mattie Fellows, Alexander D. Goldie, Jonathan Richens, Jakob N. Foerster, Oliver Richardson

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alistair Letcher, Mattie Fellows, Alexander D. Goldie, Jonathan Richens, Jakob N. Foerster, Oliver Richardson

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 to navigate a maze. Traditionally, there are two ways to do this:

The Map-Maker (Model-Based): You give the robot a blank piece of paper and ask it to draw a map of the maze, learning exactly where walls are and where doors lead.
The Muscle-Memory Learner (Model-Free): You don't ask for a map. Instead, you just tell the robot, "If you go left here, you get a cookie. If you go right, you get a shock." The robot learns a list of "good moves" for specific situations but doesn't necessarily understand why those moves work or how the maze is built.

For a long time, scientists thought these two approaches were completely separate. They believed the "Muscle-Memory Learner" only knew what to do to get a reward, but it didn't actually know the rules of the world it was living in.

The Big Discovery
This paper says: Actually, the Muscle-Memory Learner does know the map. It's just hiding it.

The authors found that if you train a robot on a wide variety of different goals (not just "get to the exit," but "get to the red spot," "get to the blue spot," "get to the corner," etc.), the robot's internal brain (its "Q-values") secretly encodes a perfect, accurate map of the entire world. It knows exactly where every action will lead, even if it was never explicitly asked to draw a map.

The Magic Trick: "P-Learning"
The paper introduces a new method called P-Learning (which stands for "Probability Learning" or "World Model Learning"). Think of it as a reverse-engineering tool.

Standard Learning (Q-Learning): You have a map (the world) and you try to figure out the best moves (the Q-values).
P-Learning: You have the best moves (the Q-values) and you try to figure out the map (the world).

The authors show that you can take a robot that has already learned how to solve many different tasks, look at its internal "to-do list" of values, and mathematically "invert" the equations to pull out the hidden map it was using all along. It's like looking at a chef's perfectly cooked meal and being able to deduce the exact recipe and ingredients they used, even if they never wrote the recipe down.

The "One Goal" vs. "Many Goals" Mystery
The paper also answers a tricky question: How many different goals does the robot need to see to learn the full map?

In a simple, predictable world (Deterministic): The robot only needs to see one goal to figure out the whole map. It's like learning a puzzle where every piece fits in only one spot; once you see the picture, you know where everything goes.
In a chaotic, unpredictable world (Stochastic): The robot needs to see many different goals (roughly as many as there are rooms in the maze) to be sure of the map. It's like trying to learn the rules of a game where dice are rolled; you need to see many different outcomes to understand the odds.

The Surprise: Hidden Superpowers
The most surprising part of the experiments is what happens when you use this "pulled-out map" to solve new problems the robot never saw before.

In one experiment, they trained a robot arm only on reaching specific positions (like "touch the red dot"). They then used P-Learning to extract the robot's hidden map. When they asked the robot to use this map to reach a specific speed (a goal it was never trained on), it succeeded!

This suggests that the robot had learned the underlying physics of the arm (how joints move to create speed) even though it was only told to care about position. It's as if you taught a person to drive only by telling them "stay in the lane," and then they could suddenly drive perfectly on a race track because they secretly understood how the car's engine and steering worked all along.

In Summary

Old Idea: Model-free agents (those that just learn rewards) don't know the world; they just know what to do.
New Idea: If you train them on enough different goals, they secretly build a perfect model of the world inside their heads.
The Tool: P-Learning is a way to "read" that hidden model out of the agent.
The Result: This hidden model is so accurate that it allows the agent to solve brand new, strange problems it was never trained for, proving that "model-free" and "model-based" learning are actually two sides of the same coin.

Technical Summary: Inverting the Bellman Equation: From Q-Values to World Models

Problem Statement

Reinforcement learning (RL) is traditionally divided into model-free and model-based paradigms. Model-free agents estimate value functions ( $Q$ ) tied to specific policies and reward functions, bypassing the explicit learning of environment dynamics (the transition kernel $P$ ). A fundamental theoretical obstacle in this domain is the value equivalence problem: distinct world models can yield identical value functions and optimal policies, implying that model-free agents do not uniquely encode the environment they inhabit.

While agents are increasingly trained on diverse goals via Goal-Conditioned RL (GCRL), it remains unclear whether the collection of value functions learned across these goals implicitly encodes a unique, accurate world model. This paper addresses the core question: When do model-free agents implicitly encode an accurate model of their environment?

Methodology: P-Learning

The authors introduce P-learning, an inverse analogue to Q-learning designed to extract an implicit world model from a trained agent.

Concept: While Q-learning iteratively updates a value function $Q$ to satisfy the Bellman equation $T^\star(Q) = Q$ for a fixed environment $P$ , P-learning treats the value function $Q$ (and associated policy $\pi$ and reward $r$ ) as fixed. It then updates a candidate world model $P_\phi$ to minimize the Bellman residual:
$L(\phi) = \| T^\pi_\phi(Q) - Q \|_d^2$
where $T^\pi_\phi$ is the Bellman operator parameterized by the world model $P_\phi$ .
Mechanism: P-learning samples from the agent's $Q$ -values, policies, and rewards to decode the internal model. In the tabular setting, the authors prove that P-learning converges to a solution involving the Moore-Penrose pseudo-inverse of a matrix $M$ derived from the rewards and value functions. Specifically, if $M$ has full rank, the extracted model is unique and given by $P_\infty = M^+ Q$ .
Differentiability: For continuous or deterministic settings, the method utilizes gradient-based optimization (or reparameterization tricks) to update the parameters of the world model, effectively "inverting" the Bellman equation.

Key Theoretical Contributions

The paper establishes sufficient conditions under which the transition kernel $P$ is uniquely identifiable from a set of goal-conditioned Q-values $(Q_g, \pi_g, r_g)_{g \in G}$ , thereby breaking value equivalence.

Finite MDPs:
- Deterministic: A single generic goal is sufficient to uniquely identify $P$ (Theorem 2). Uniqueness holds for almost all reward functions, provided the resulting matrix $M$ is column-injective.
- Stochastic: Goals must span the state space (i.e., the rank of the reward matrix equals the number of states) to ensure $P$ is uniquely identifiable (Theorem 3). This holds for generic reward functions or policies.
- Approximation: The authors provide tight error bounds showing that if Q-values are $\epsilon$ -approximate, the extracted model error scales linearly with $\epsilon$ .
Continuous MDPs:
- Deterministic: A finite number of goals can suffice under specific conditions (e.g., Gaussian rewards with termination). For indicator rewards, a goal set covering the state space with a specific radius is required (Theorem 4).
- Stochastic: Identifiability is driven by the coverage of the goal set rather than the count. For Gaussian rewards, a goal set with a non-empty interior suffices; for indicator rewards, the set must cover the state space expanded by the reward radius (Theorems 6 and 7).
Local MDPs: The paper extends results to $N$ -local MDPs (where transitions are restricted to a neighborhood), showing that fewer goals are needed if the support of the transition kernel is known.

Empirical Results

The authors validate their theory and method across several environments, demonstrating that P-learning can extract highly accurate world models even when theoretical guarantees (e.g., infinite goals) are not met.

FourRooms (Finite MDPs):
- Deterministic: Exact recovery of $P$ with $|G|=1$ .
- Windy (Local Stochastic): Quasi-perfect recovery with $|G|=4$ .
- Teleporting (Fully Stochastic): Quasi-optimal planning with $|G|=20$ , significantly fewer than the theoretical worst-case bound of $|S|=68$ .
Reacher (Continuous, Deterministic):
- Trained on only $|G|=4$ position-based goals, the extracted world model achieved high fidelity (NMSE $1.2 \times 10^{-4}$ ) despite imperfect Q-values.
- Zero-Shot Generalization: Policies trained exclusively inside the extracted implicit world model achieved near-optimal performance on out-of-distribution goals (e.g., reaching specific angular velocities or joint configurations) that were never seen during training. This suggests the agent encoded dynamics for variables (velocity, angles) it was never explicitly rewarded for.
MountainCar (Continuous, Deterministic):
- Agents trained on position-based goals and agents trained on velocity-based goals converged to near-identical implicit world models.
- Planning within these models allowed for successful zero-shot transfer to tasks involving the untrained variable (e.g., a position-trained agent solving a velocity-based goal).

Significance and Claims

The paper argues that Goal-Conditioned RL acts as an implicitly hybrid method, bridging the gap between model-free and model-based learning.

Implicit World Models: Model-free agents trained on a sufficiently rich set of goals implicitly encode a unique and accurate world model, even in stochastic and continuous environments.
Hidden Generalization: These implicit models contain "hidden generalization capabilities," allowing agents to solve tasks defined on variables orthogonal to their training rewards (e.g., learning velocity dynamics from position rewards).
Performance Correlation: There is a strong correlation ( $\rho = 0.98$ ) between an agent's performance and the accuracy of its internalized world model, suggesting that high-performing GCRL agents naturally develop better world representations.
P-Learning Utility: The proposed P-learning algorithm provides a practical tool to extract these internal models, enabling the repurposing of model-free agents for new tasks without additional environment interaction.

The authors conclude that while P-learning is not a substitute for explicit model-based RL, it reveals a fundamental connection: diverse reward-rich training allows model-free agents to "learn everything" about the environment's dynamics, effectively encoding a world model within their value functions.