When Does Dynamics Randomization Test Context Inference?

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

cs.LGOct 23, 2026

When Does Dynamics Randomization Test Context Inference?

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

Organizations: King AI Labs, Microsoft Gaming · King AI Labs, Microsoft Gaming & KTH Royal Institute of Technology

Abstract

Environments with randomized dynamics are widely used to ask whether world models need explicit context inference. However, randomization can leave little predictive information for a context encoder to add, because distinct latent parameters may produce similar transitions at visited states, or those states may already reveal the active context. We formalize the remaining one-step information as predictive headroom and compare kernels under fixed and visited state--action distributions to distinguish changes in query coverage from changes in context weights. We analyze discrete and continuous environments and find low headroom both when kernels differ little and when visitation reduces an otherwise large discrepancy. At the same time, resetting state components increases headroom, showing that context evidence accumulates in state. We propose reporting this headroom alongside the randomized parameters, so that gains from explicit context inference can be interpreted relative to the information missing from the current state and action.

Figures & tables

Figure 1: Dynamics variation and predictive headroom across perturbation strengths. (a,b) Wind and Resistance mechanisms. (c–e) The structural reference $\mathcal{I}_{\mathrm{struct}}$ and estimated visited headroom $\widehat{\mathcal{I}}_{\mathrm{emp}}$ differ substantially on Wind and Attractor but remain close on Resistance. Attractor’s structural curve uses the interior Gaussian reference. Shading shows bootstrap intervals.

Setting	Fixed	Matched	Visited
	$\mathcal{I}_{\mathrm{struct}}$	$\mathcal{I}_{\mathrm{match}}$	$\widehat{\mathcal{I}}_{\mathrm{emp}}$
Wind ( $v_{\max}=1$ )	0.871	0.620	$0.077\pm 0.002$
Resistance ( $\epsilon_{\max}=0.9$ )	0.062	0.059	$0.055\pm 0.001$
Attractor ( $k=0.3$ )	1.006	1.006	$0.109\pm 0.002$

Table 1: Comparisons at representative settings. Fixed and matched columns retain prior context weights, while visited headroom uses fitted visitation-dependent weights. Attractor’s identical prior references use its query-invariant interior kernel. Uncertainties are the larger one-sided deviations from the

95\%

bootstrap intervals.

Reset	Wind	Attractor
None	$0.078\pm 0.002$	$0.112\pm 0.002$
Full, every 10 steps	$0.411\pm 0.006$ $(+430\%)$	$0.473\pm 0.004$ $(+324\%)$
Full, every 2 steps	$0.723\pm 0.017$ $(+833\%)$	—
Position only, every 10 steps	—	$0.143\pm 0.003$ $(+28\%)$

Table 2: State resets increase predictive headroom without changing episode context. Full resets resample position and direction in Wind. In Attractor they resample position and zero velocity. Changes in parentheses use each intervention dataset’s own no-reset baseline, collected separately from the sweep datasets.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Point-Mass-Attractor force fields. The first four panels show the context-dependent force

-k(\mathbf{p}-\theta_{i})

for each kernel

T_{\theta_{i}}

. The final panel shows the prior-averaged force

-k\mathbf{p}

Contextual RL

See all 20 papers

Considering Context: When World Models Need Context Encoders
Oct 5, 2026Oleg Smirnov, Sofiane Ennadir, John Pertoft +2Reinforcement Learning World Model Learning
King AI Labs, Microsoft Gaming
Not All Answers Are Contextually Persuadable: Inference Dynamics in Large Language Models under Contextual Influence
Oct 3, 2026Zongye Hu, Weiqing Luo, Yanjie Fu +3LLM Inference Framing Effects in Language Models
Arizona State University, Tempe, AZ, United States · University of Maryland, College Park, MD, United States · Morgan Stanley, New York, NY, United States
Dynamic Context Scheduling: Learning Beyond the Static Universe
Aug 21, 2026Martin Mráz, André BiedenkappContextual RL Reinforcement Learning
University of Freiburg

Oct 5, 2026J/K move · Enter open · S save

Considering Context: When World Models Need Context Encoders

Oleg Smirnov, Sofiane Ennadir, John Pertoft, Bjartur Hjaltason+1

King AI Labs, Microsoft Gaming

Reinforcement Learning World Model Learning Generalization in RL Model-Based RL Contextual RL

Open paper AskPDF arXiv

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 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.

Abstract

Contextual RL

Explore similar work

Considering Context: When World Models Need Context Encoders

Not All Answers Are Contextually Persuadable: Inference Dynamics in Large Language Models under Contextual Influence

Dynamic Context Scheduling: Learning Beyond the Static Universe