Considering Context: When World Models Need Context Encoders

Oleg Smirnov; Sofiane Ennadir; John Pertoft; Bjartur Hjaltason; Sara Karimi

cs.LGOct 5, 2026

Considering Context: When World Models Need Context Encoders

Authors: Oleg Smirnov, Sofiane Ennadir, John Pertoft, Bjartur Hjaltason, Sara Karimi

Organizations: King AI Labs, Microsoft Gaming

Abstract

Methods for generalization in model-based reinforcement learning typically assume that an agent cannot recover the latent context governing the environment dynamics from its own experience, and therefore supplies it externally. We formalize and test this assumption with \emph{predictive sufficiency}, which quantifies what access to the context adds to next-step prediction under the visitation distribution an agent induces, and separates that quantity into a history-recoverable part, a residual requiring the true context, and the deficit added by a finite model. We classify context-aware algorithms by the predictive risk their conditioning set can target and demonstrate across environments of increasing identification difficulty that the headroom does not follow the MDP class. The same task under different priors leaves predictive headroom in one setting and nothing distinguishable from zero in another, where the agent's behavior implicitly identifies the context and any benefit of such a mechanism cannot be attributed to missing information. Where headroom persists, the learned state exposes it only partially, and adding the true context still lowers the risk. Our contribution is a practical criterion for matching contextual mechanisms to the information available to them, estimated from the ordinary trained agent without a reference policy.

Figures & tables

Panel of Figure 1 — Figure 1: Contextual variation in MiniGrid, where a hidden wind variable probabilistically displaces the agent after each move, and walls block the displacement. (a) A model trained on pooled experiences with different winds learns the averaged kernel, which has no net displacement and inflated spread and governs no individual episode. (b) Under the same policy the trajectories separate mid-episode, so the state alone is consistent with only one wind and identifies it implicitly, rendering contextual mechanisms partially redundant. (c) Under one and the same context, the trajectory of an agent crossing the open grid (blue) bends downwind and betrays the context, while an agent moving along the downwind border (orange) shows no drift at all, since the obstacle blocks it. What remains for context mechanisms to gain depends on what the agent’s own trajectories already reveal, and on what they cover.

Figure 2: Controlled diagnostics. Wind applies a sampled drift after the native action (a), whereas Resistance either commits or voids the action (b). Across perturbations (c–e), $\mathcal{I}_{\mathrm{struct}}$ separates from predictive headroom on Wind and Point-Mass-Attractor but remains close to it on Resistance.

State	Setting	$\mathcal{I}_{\mathrm{struct}}$	$\mathcal{I}_{\mathrm{struct}}^{\pi}$	$\mathcal{I}_{\mathrm{emp}}$
Discrete	Wind ( $v_{\max}=1$ )	0.871	0.620	$0.077\pm 0.002$
Discrete	Resistance ( $\epsilon_{\max}=0.9$ )	0.062	0.059	$0.055\pm 0.001$
Continuous	Attractor ( $k=0.3$ )	1.006	1.006	$0.109\pm 0.002$

Table 1: Representative MiniGrid and Point-Mass-Attractor settings. The matched column sets

\mu(s,a)=p^{\pi}(s,a)

while retaining prior context weights. CIs are 95% episode-level bootstrap. All estimates are conservatively biased upward, so finite-sample error can inflate headroom.

Class	Conditioning set (bound)	Sufficient when	Methods
Memoryless	$(s_{t},a_{t})$ ( $\mathcal{R}_{\mathrm{M}}$ )	$\theta$ decodable from $(s_{t},a_{t})$	Base probes
Transition window	last $K$ transitions ( $\geq\mathcal{R}_{\mathrm{H}}$ )	recent evidence	DALI-S , CaDM, UP-OSI, GrBAL, PEARL, IIDA
Recurrent	within-episode summary ( $\geq\mathcal{R}_{\mathrm{H}}$ )	evidence within the episode	DreamerV3 , ReBAL, VariBAD
Cross-episode chain	cross-episode summary ( $\geq\mathcal{R}_{\mathrm{H}^{+}}$ )	episodes correlated	Carried state , LILAC , RL 2 , L2RL
Ground-truth context	realized $\theta$ ( $\mathcal{R}_{\mathrm{O}}$ )	always	cRSSM-D , cRSSM-C

Table 2: Method classes ordered by the predictive risk their conditioning set can target. Bold denotes the evaluated representatives. The memoryless class is instantiated by the measurement protocol itself.

Setting	Environment	Context process	Difficulty level	Tests
Calibration	Point mass	persistent two-state chain, i.i.d. and short-episode variants	within/across episodes	P3 (control)
Identifiable	CARL Walker	i.i.d. per episode	current state and action	P1, P4
Hard	CARL Walker	i.i.d. per episode	episode history	P2, P4
Drift	RWRL Walker	non-i.i.d. Markov chain across episodes	across episodes	P3

Table 3: Environment settings, ordered by the level at which the context becomes identifiable.

	Identifiable			Hard
Method	Return $\uparrow$	Headroom	Resid. gain	Return $\uparrow$	Headroom	Resid. gain
DreamerV3	$712\pm 10$	$-0.00\pm 0.12$	$-0.13\pm 0.11$	$946\pm 6$	$2.39\pm 0.70$	$0.38\pm 0.08$
DALI-S	$756\pm 72$	$0.03\pm 0.13$	$-0.15\pm 0.13$	$941\pm 14$	$2.10\pm 1.76$	$0.44\pm 0.19$
LILAC-style	$589\pm 88$	$-0.09\pm 0.08$	$0.00\pm 0.09$	$\mathbf{953\pm 4}$	$3.49\pm 0.68$	$0.73\pm 0.29$
cRSSM-D	$\mathbf{862\pm 17}$	$-0.14\pm 0.16$	$-0.06\pm 0.09$	$948\pm 6$	$3.64\pm 1.71$	$-0.22\pm 0.80$

Table 4: Computed quantities over i.i.d. environments. Identifiable and Hard settings report return, probe headroom, and the residual context gain of the linear context readout. Entries are three-seed IQM with 95% bootstrap confidence intervals, method differences use paired seed resampling.

	Start transient $\downarrow$			Start prior loss $\downarrow$
Method	Low	Mid	High	Low	Mid	High
DreamerV3	$\mathbf{2.64\pm 0.45}$	$2.80\pm 0.15$	$2.79\pm 0.86$	$\mathbf{10.15\pm 0.82}$	$10.77\pm 0.59$	$11.24\pm 0.83$
Carried state	$4.06\pm 0.36$	$3.92\pm 0.69$	$3.93\pm 0.51$	$13.30\pm 0.56$	$12.51\pm 0.75$	$12.75\pm 0.77$
DALI-S	$3.12\pm 0.19$	$\mathbf{2.52\pm 0.13}$	$3.05\pm 0.89$	$10.68\pm 0.05$	$\mathbf{10.45\pm 1.23}$	$11.37\pm 0.33$
LILAC-style	$2.94\pm 0.92$	$3.19\pm 0.86$	$\mathbf{2.46\pm 0.61}$	$11.14\pm 0.93$	$11.54\pm 0.90$	$\mathbf{10.22\pm 0.91}$
cRSSM-D	$2.84\pm 0.35$	$3.05\pm 0.41$	$2.55\pm 0.71$	$10.48\pm 0.31$	$10.68\pm 1.05$	$10.38\pm 0.33$

Table 5: Computed quantities over the non-i.i.d. environment: start transient and start prior loss at three drift rates. Entries are means

\pm

sample stdev across seeds.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Point-Mass-Attractor transition fields. Each episode uses one hidden attractor

\theta_{i}

, inducing a distinct field

T_{\theta_{i}}

; prior averaging produces a center-seeking field

\bar{T}_{\mathrm{prior}}

that governs no episode.

Figure 4: Point-Mass-Attractor trajectories under uniform random actions. The hidden attractor is fixed within an episode, causing visited states to become informative about context.

Figure 5: Parallel checks of state-mediated identification. State projection removes current-state information in MiniGrid (a) and Point-Mass-Attractor (c). State resets weaken context-dependent visitation in MiniGrid (b). In Point-Mass-Attractor (d), headroom rises when velocity is reset and remains near baseline when velocity is preserved. Error bars denote

95\%

bootstrap CIs. Where shown, blue lines are full-state structural references, not bounds after the information set or query distribution changes.

Agent setting	Value
Backbone	DreamerV3, block-GRU
Parameters (Walker baseline)	$9.7$ M ( $4.3$ M transition)
Deterministic state	$2048$
Stochastic state	$32\times 16$ categorical
RSSM hidden width	$256$
Encoder / decoder	$3$ layers, width $256$

Appendix

Table 6: Agent and probe hyperparameters.

Figure 6: Analytical calibration environment. The hidden regime determines the rotation of the autoregressive disturbance. The physical state resets between episodes while the regime follows a Markov chain. The prior mixture averages the two conditional disturbance fields but is not a realized regime.

(a) Probe calibration
Quantity	Exact	Est.
Context score gap	$2.51\pm 0.05$	$0.48\pm 0.03$
History recovery	$0.99$	$0.89\pm 0.03$
Shuffled null	$0$	$-0.04\pm 0.02$

Appendix

Table 7: Probe calibration on the analytical point-mass environment. The context score gap is a diagonal-Gaussian log-score difference: the exact column is a sampled evaluation of the

2.518

-nat population gap, and the estimate column reports the estimate-to-target ratio. History recovery reports a recovered fraction and the shuffled null a probe-risk difference. Entries are IQM

\pm

the larger 95% bootstrap CI half-width.

Figure 7: Drift metrics at the three rates on frozen realizations shared across methods. Lower is better in (a)–(c), higher in (d); lines are seed means and bands show sample stdev. The lowest mean startup losses shift from the context-blind baseline to the transition-window and chain-latent agents as drift increases.

Explore similar work

May 26, 2026cs.CL

ContextGuard: Structured Self-Auditing for Context Learning in Language Models

Recent benchmarks reveal that despite strong reasoning capabilities, large language models (LLMs) still struggle to faithfully apply complex contextual knowledge. These failures are often not wholesale reasoning collapses: in context-rich tasks, models may follow the central reasoning path while missing peripheral, persistent, or format-sensitive requirements.

Hongbo Jin, Chi Wang, Haoran Tang +5

Peking University · SCUT · Tsinghua University

Contextual Reasoning In-Context Learning Reasoning Paths

Ask PDF

May 26, 2026cs.CL

The Need for an External Observer Formalizing the Sufficiency Gap: A Mathematical Extension of Mixture Identifiability and Contextual Grounding in Sequence Models

We construct a binary mixed-regime process with one deterministic textual regime and one random regime governed by an unobserved latent state. Even an ideal infinite-capacity sequence predictor that exactly recovers the text-only marginal law can become overconfident when the observed prefix is compatible with the wrong latent regime. The resulting entropy difference is not an ordinary optimization error; it is a sufficiency gap caused by marginalization over an unobserved state. We then formalize retrieval, tool use, and external grounding through an auxiliary binary signal with fidelity

γ\in [1/2,1]

. The resulting Bayesian update yields a contextual dominance threshold: a corrective signal reverses the posterior odds induced by the textual history exactly when its fidelity exceeds the text-only posterior weight assigned to the misleading regime. This threshold reduces, but does not generally eliminate, the sufficiency gap; complete closure requires perfect revelation of the relevant latent state or an equivalent verification mechanism. The analysis clarifies why temperature scaling cannot restore missing context, why grounding mechanisms must be both informative and learnably usable by the model, and why autonomous sequence models require structurally decoupled observers or verifiers in high-stakes domains.

Francesco Corielli

Sufficiency Sequence Modeling Intermediate Latent States Identifiability Contextual Grounding

Ask PDF

Jun 19, 2026cs.LG

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

Model-based and model-free reinforcement learning are traditionally viewed as separate paradigms: instead of learning a model of the transition kernel

P

, model-free agents typically estimate value functions tied to a specific policy and reward. In this paper, we challenge this dichotomy by proving that value-based agents trained on a sufficiently rich set of reward functions, e.g. using goal-conditioned RL, implicitly encode a unique and accurate world model. To extract this model in practice, we introduce \textit{

P

-learning}, an inverse analogue to

Q

-learning that samples from an agent's

Q

-values, policies and rewards to decode its internal model of the environment. We then provide sufficient conditions on the type and number of goals for which agents encode the true kernel

P

, covering both stochastic and deterministic MDPs over finite or continuous state spaces. Even when our assumptions are violated, we empirically demonstrate that agents trained on a handful of reward functions encode accurate dynamics in

\texttt{Reacher}

\texttt{MountainCar}

and stochastic variants of

\texttt{FourRooms}

. Surprisingly, we find that policies trained exclusively on a \texttt{Reacher} agent's implicit world model are quasi-optimal on out-of-distribution, velocity-based goals despite position-only training -- suggesting that agents contain hidden generalisation capabilities and providing a new lens into the connection between model-based, model-free, and goal-conditioned RL.

Alistair Letcher, Mattie Fellows, Alexander D. Goldie +3

FLAIR, University of Oxford · Google DeepMind · Mila, University of Montreal

Model-Free Reinforcement Learning Goal-Conditioned Reinforcement Learning Reward Functions Bellman Equation World Models World+1

Ask PDF

Abstract

Explore similar work

ContextGuard: Structured Self-Auditing for Context Learning in Language Models

The Need for an External Observer Formalizing the Sufficiency Gap: A Mathematical Extension of Mixture Identifiability and Contextual Grounding in Sequence Models

Inverting the Bellman Equation: From QQQ-Values to World Models

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