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

ContextGuard: Structured Self-Auditing for Context Learning in Language Models
May 26, 2026Hongbo Jin, Chi Wang, Haoran Tang +5Contextual Reasoning In-Context Learning
Peking University · SCUT · Tsinghua University
The Need for an External Observer Formalizing the Sufficiency Gap: A Mathematical Extension of Mixture Identifiability and Contextual Grounding in Sequence Models
May 26, 2026Francesco CorielliSufficiency Sequence Modeling
Inverting the Bellman Equation: From $Q$ -Values to World Models
Jun 19, 2026Alistair Letcher, Mattie Fellows, Alexander D. Goldie +3Model-Free Reinforcement Learning Goal-Conditioned Reinforcement Learning
FLAIR, University of Oxford · Google DeepMind · Mila, University of Montreal

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