[Paper Notes] Inverting the Bellman Equation: From Q-Values to World Models
Published:
Inverting the Bellman Equation asks a deceptively simple question: if a model-free, value-based agent has learned (Q(s,a,g)) for many goals, how much of the environment dynamics is already encoded inside those values?
My read: the paper turns the usual RL direction around. Classical value learning fixes the transition kernel (P) and solves for (Q). This work fixes (Q), the policy (\pi), and the reward family (r), then solves for a compatible transition kernel (P). The resulting method, (P)-learning, is both an extraction algorithm and a conceptual bridge between model-free RL, model-based RL, and goal-conditioned RL.
Paper Info
The paper is “Inverting the Bellman Equation: From (Q)-Values to World Models” by Alistair Letcher, Mattie Fellows, Alexander D. Goldie, Jonathan Richens, Jakob N. Foerster, and Oliver Richardson from FLAIR, University of Oxford, Google DeepMind, and Mila. It appears on arXiv as arXiv:2606.21173, submitted on June 19, 2026. The paper reports code at github.com/aletcher/inverting-bellman.
The high-level claim is that goal-conditioned value functions can carry more than task-specific preference. Under suitable goal coverage, they can identify the environment transition kernel. Even when the theory’s sufficient conditions are stronger than the experiments satisfy, the extracted world models are accurate in Reacher, MountainCar, and stochastic FourRooms variants.
Why This Is Interesting
| Model-free RL is usually described as reward-specific: the agent learns values or policies useful for the training reward, while a model-based agent learns dynamics (P(s’ | s,a)). The value-equivalence problem explains why this distinction is not trivial. Many different dynamics can induce the same value function for a single reward, so a value function alone may underdetermine the world. |
Goal-conditioned RL changes the information content. A collection of values ({Q_g}_{g\in G}) gives many Bellman constraints on the same unknown transition kernel. Each goal contributes a different “probe” of the next-state distribution. If these probes are rich enough, the transition kernel becomes identifiable.
The paper’s central question is:
[ \text{When do value-based agents implicitly encode an accurate model of their environment?} ]
P-Learning: Bellman Inversion
For a fixed goal-conditioned agent, assume we know:
- a (Q)-function (Q(s,a,g)),
the policy (\pi(a’ s’,g)), usually induced by the (Q)-function, - the reward family (r(s,g)),
- the discount factor (\gamma).
Instead of learning (Q) from environment transitions, (P)-learning learns a candidate world model (P_\phi) by minimizing the Bellman residual:
[ L(\phi)=\left|T^\pi_{P_\phi}(Q)-Q\right|_d^2. ]
Here the Bellman operator depends on the candidate model:
[ T^\pi_{P_\phi}(Q)(s,a,g) = \mathbb{E}{s’\sim P\phi(s,a),\,a’\sim \pi(s’,g)} \left[ r(s’,g)+\gamma Q(s’,a’,g) \right]. ]
So (P)-learning searches for dynamics under which the already-learned (Q)-values satisfy the Bellman equation. This is the inverse analogue of (Q)-learning:
| Usual direction | Inverted direction |
|---|---|
| Fix (P), learn (Q) such that (T_P(Q)=Q) | Fix (Q), learn (P) such that (T_P(Q)=Q) |
| Environment samples provide next states | The candidate model samples next states |
| Target moves as (Q) changes | Objective is stationary for fixed (Q,\pi,r) |
For finite state spaces, the paper rewrites the Bellman equation as a linear system. Define
[ M(s’,g)=r(s’,g)+\gamma V(s’,g), ]
where (V(s’,g)=\mathbb{E}_{a’\sim\pi(s’,g)}[Q(s’,a’,g)]). For each state-action pair:
[ M P(s,a)=Q(s,a). ]
This makes the extraction problem especially clear. The rows of (M) are Bellman probes indexed by goals; the unknown is the next-state distribution (P(s,a)). If (M) has enough rank, (P(s,a)) is determined.
The tabular (P)-learning update is:
[ P_{n+1}(s,a)=P_n(s,a)-\alpha M^\top(MP_n(s,a)-Q(s,a)). ]
Theorem 1 shows that this converges to:
[ P_\infty(s,a)=M^+Q(s,a)+(I-M^+M)P_0(s,a), ]
where (M^+) is the Moore-Penrose pseudo-inverse. The second term is the remaining value-equivalence ambiguity. If (M) has full column rank, that ambiguity vanishes and the solution becomes unique.
When Are Dynamics Identifiable?
The theoretical section gives sufficient conditions for recovering (P) from (Q,\pi,r). A useful way to remember the results is:
| Setting | What goals need to do | ||||
|---|---|---|---|---|---|
| Deterministic finite MDP | A single generic goal can be enough | ||||
| Stochastic finite MDP | Goals need to span the state space; ( | G | \ge | S | ) is a clean sufficient condition |
| Local stochastic finite MDP | Fewer goals can suffice if transition support is local or known | ||||
| Deterministic continuous MDP | Coverage or structured Gaussian goals can separate successor states | ||||
| Stochastic continuous MDP | The paper proves coverage-based results mainly for unconditional policies |
The intuition is test functions. For each goal (g), the Bellman equation gives:
[ Q(s,a,g)=\mathbb{E}_{s’\sim P(s,a)}[M_g(s’)]. ]
So each (M_g) is a known test function whose expectation under the unknown next-state distribution must match (Q(s,a,g)). One goal gives one expectation constraint. Many sufficiently different goals can determine the whole distribution, much like enough moments can determine a finite distribution.
This also clarifies the deterministic case. If dynamics are deterministic, the unknown next state is a point, so the columns (M(\cdot,g)) only need to distinguish states. A single generic reward can make those columns distinct. For stochastic dynamics, the unknown is a distribution over states, so spanning or rank conditions are stronger.
Experiments
The experiments test whether this extraction works when the theoretical conditions are only partly met. Agents are trained with PQN plus Hindsight Experience Replay, then (P)-learning extracts a world model from the converged (Q)-values.
Reacher. The agent is trained on only four position goals on the unit circle. The continuous state includes joint angles, angular velocities, and fingertip position, but the training rewards depend only on fingertip position. Despite imperfect learned (Q)-values, the extracted world model is very accurate: the paper reports world-model NMSE around (1.2\times 10^{-4}). Policies trained entirely inside the extracted model are quasi-optimal on unseen goals, including velocity goals that were outside the position-only reward family.
The architecture sweep is also important. Across 42 Reacher architectures and 10 seeds each, better goal-conditioned agents tend to have more accurate extracted world models. The paper reports a Spearman correlation of (-0.98) between agent return and world-model error, and (+0.95) between agent return and the return of policies trained inside the extracted model on unseen goals.
MountainCar. The position-trained agent uses four sparse position goals. The extracted dynamics match the true transition kernel well, with NMSE around (6.7\times 10^{-3}). Planning inside the extracted world model works on out-of-distribution goals such as velocity targets and constrained hill-climbing. A second agent trained on velocity goals yields a world model extremely close to the position-trained agent’s world model, suggesting that different reward views of the same environment can lead to similar implicit dynamics.
| FourRooms. The finite-state experiments test deterministic, windy, and teleporting gridworlds. In deterministic FourRooms, one generic training goal recovers the world model exactly, matching the finite deterministic theory. In windy FourRooms, four goals bring the world-model-derived policy within about 1% of optimal return. In teleporting FourRooms, 20 goals work well even though the worst-case stochastic finite-state guarantee would suggest ( | G | = | S | =68). |
Strengths and Limitations
The strongest part of the paper is the clean inversion view. It gives a concrete procedure for asking what dynamics are already latent in a trained value function. That makes the result relevant to transfer, interpretability, auditing, and hybrid RL: a model-free agent may already contain a usable model, even if it never explicitly represented one during training.
The second strength is the theory-experiment contrast. The theorems explain when recovery is guaranteed, while the experiments show that useful recovery can happen with far fewer goals than the worst-case sufficient conditions require. This gap is productive: it suggests that practical environments and learned (Q)-functions have structure that the broad theory does not yet exploit.
The limitations are real. The continuous stochastic theory is restricted mainly to unconditional policies. The work does not cover partially observable MDPs or model-free agents that do not learn (Q)-values. The experiments are still small relative to modern large-scale RL systems, and (P)-learning is framed as an extraction and analysis tool, not as a replacement for model-based RL.
Takeaway
The reusable idea is: a goal-conditioned (Q)-function is a set of Bellman probes into the transition kernel. With enough goal diversity, those probes can identify the dynamics. (P)-learning turns that observation into an algorithm: fit a world model whose Bellman operator makes the agent’s existing (Q)-values self-consistent. The result is a sharper picture of goal-conditioned RL as an implicitly hybrid method: it can look model-free on the surface while carrying a recoverable world model inside its values.