When do world models need context encoders for generalization?

Predictive sufficiency measures what context adds on-policy

Model-based reinforcement learning methods for generalization usually assume the agent cannot recover the latent context that governs environment dynamics from its own experience, so they supply that context externally. This paper tests the assumption instead of accepting it. Its central quantity, predictive sufficiency, asks what access to the context adds to next-step prediction under the visitation distribution an agent induces. The reference is the Bayes one-step advantage of context for memoryless and history-aware predictors, not a comparison under some fixed external policy.

The quantity is deliberately decomposed. One part is history-recoverable, meaning the agent's own trajectory already reveals the context. A second part is a residual that genuinely requires the true context. A third is the deficit contributed by a finite model. Keeping these apart matters because only the residual justifies a dedicated contextual mechanism; the history-recoverable part is information the model could already reach, and the finite-model deficit is an optimization or capacity problem rather than missing information.

The paper also classifies context-aware algorithms by the predictive risk their conditioning set can target. That classification is what makes the diagnostic practical: it matches a mechanism to the information actually available to it, and it is estimated from the ordinary trained agent without a reference policy.

Averaging over contexts produces a kernel no episode obeys

In contextual MDPs and hidden-parameter MDPs, latent parameters may be drawn i.i.d. from a shared prior, evolve through a background process, or vary without structure. A predictor that sees only the current state and action must average over the contexts still plausible at that point, producing a model mixture and reduced prediction accuracy.

The paper's MiniGrid illustration makes the failure concrete. A hidden wind variable probabilistically displaces the agent after each move, and walls block the displacement. 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. The same policy can nonetheless separate trajectories mid-episode, so the state alone becomes consistent with only one wind and identifies it implicitly, rendering contextual mechanisms partially redundant.

Visitation matters as much as the kernel. Under one and the same context, an agent crossing the open grid bends downwind and betrays the context, while an agent moving along the downwind border shows no drift because the obstacle blocks it. What remains for a context mechanism to gain therefore depends on what the agent's own trajectories already reveal and on what they cover.

Probes estimate headroom from the trained agent alone

Headroom, history recovery, and the sufficiency gap are estimated by comparing held-out next-state log losses from MLP and recurrent probes with Gaussian likelihood heads. Scoring is cross-fitted at episode level, so every transition is scored by a probe that never trained on its episode. A permutation null re-pairs episodes with contexts drawn from other episodes, and a decodability readout uses the feedforward architecture with a squared-error objective to predict context from current state and action, reporting cross-fitted R-squared.

The benchmark panel spans increasing identification difficulty. The i.i.d. settings use CARL DMC Walker Walk. The identifiable setting samples a latent horizontal wind uniformly per episode with nonzero fluid density so the wind exerts force on the body; the hard setting samples a latent actuator-strength scale uniformly around nominal with gravity fixed at nominal. An 84-context benchmark panel is held out for final evaluation. The drift setting follows the registered RWRL protocol, with one Walker parameter performing a random walk across episodes at three step sizes on frozen realizations shared by every method, chosen so consecutive-episode autocorrelation stays above 0.5 at the highest rate.

All agents share DreamerV3 hyperparameters and differ only in context configuration: a context-blind baseline, DALI-S inferring context from a recent window, cRSSM-D receiving true context through deep conditioning, a carried-state variant preserving recurrent state across episode boundaries, and a chain-latent agent maintaining a cross-episode latent. Because DreamerV3 already carries a recurrent within-episode state, comparing it with DALI isolates the value of a dedicated recent-history pathway while holding that backbone fixed. Runs use one 24GB NVIDIA L4; a 10-to-the-6-step run takes about 9 hours with 6-7GB peak allocated memory.

In the identifiable setting headroom collapses while context stays decodable

In the identifiable i.i.d. setting, estimated headroom collapsed under every profiled visitation while the context remained decodable. That combination is the paper's sharpest diagnostic result. If the context is still recoverable from the agent's own state and action, then a mechanism that supplies it is not adding missing information, and any return differences between methods reflect optimization rather than an information gap.

The paper states this implication directly: when Bayes predictive headroom is zero, any predictive gain under the same visitation reflects a reduction in finite-model error. This is a boundary on how return curves should be read. A context-aware method that outperforms a context-blind one in such a setting has not demonstrated that context was needed; it has demonstrated that its training or architecture handled the same problem better.

The contrast with the hard setting is what makes the point comparative rather than a single negative result. There, estimated headroom stayed open across the evaluated methods despite comparable returns, showing that successful control does not have to resolve predictive context. Return parity and information parity are different claims, and the probes separate them.

The same task under different priors yields different headroom

The paper's central comparative claim is that 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. In the second case the agent's behavior implicitly identifies the context, and any benefit of a contextual mechanism cannot be attributed to missing information.

This is why the class label alone is not a difficulty measure. Two settings can share a formal description while differing in whether the agent's induced visitation exposes the latent variable. The paper's recommendation for benchmark construction follows: report estimated headroom alongside an MDP's class and its benchmark definition, because the class of a setting alone does not determine its difficulty.

The non-i.i.d. results connect predictive benefits to both the drift regime and mechanism design. The contrast between chain-latent inference and carried state demonstrates the distinction between predicting episode-level context and preserving recurrent state. On the drift setting the primary endpoint is the start transient over the first 25 steps, computed post hoc from final checkpoints, with drift entries reported as means of seed means with seed ranges.

Learned state exposes persistent headroom only partially

Where headroom persists, the learned state exposes it only partially, and adding the true context still lowers the risk. This is a claim about the sufficiency gap rather than about returns: even when a contextual mechanism has something real to gain, the representation the agent learns does not surface all of it, so an oracle context retains value.

The calibration environment gives the probes a reference point. It is a two-dimensional stabilization point mass with fully observed state and a hidden regime following a slow-switching Markov chain across episodes with stay probability 0.9. The exact population score gap is 2.518, evaluated by quadrature, while the sampled evaluation gives 2.51 plus or minus 0.05. History recovery reached 0.89 plus or minus 0.03 against the exact 0.99 reference computed from the disturbance-history filter, giving a recovered fraction for the nested disturbance-only information sets rather than a certified bound on the full on-policy history optimum.

The calibration also shows what the probes are not. The probe targets are ideal diagonal-Gaussian score gaps rather than mutual information, and the corresponding information quantities are much smaller: the binary regime carries at most ln 2, about 0.693 nats, and the stationary disturbance-mixture information is 0.507. The shuffled null stays near zero. The estimated sufficiency gap is near zero for the observed-context agent, concentrates at episode start for every inference method, and is largest for carried state, whose recurrent state continues the previous episode through the reset. A short-episode variant yields an exact per-transition residual of 0.24, reproduced as a sufficiency gap that persists for the context-blind model and nearly vanishes for the observed-context one.

Removing state information restores headroom

The paper runs parallel checks that manipulate state-mediated identification rather than swapping algorithms. State projection removes current-state information in MiniGrid and Point-Mass-Attractor. State resets weaken context-dependent visitation in MiniGrid. In Point-Mass-Attractor, headroom rises when velocity is reset and remains near baseline when velocity is preserved.

These manipulations test the mechanism the paper proposes: implicit identification is produced by the agent's own behavior, so weakening what the state reveals should reopen the gap that a context mechanism can fill. The reported error bars are 95% bootstrap confidence intervals, and where shown, blue lines are full-state structural references rather than bounds after the information set or query distribution changes.

The paper also names a difficulty class it does not cover. To its knowledge, no standard benchmark makes the context identifiable only through deliberately informative actions, so every passive class is predicted equivalent there. Constructing such settings is listed among the future directions the authors consider most fruitful, alongside paired counterfactual resets where simulators permit them.

Match the mechanism to the information the agent already has

The practical criterion this paper offers is a matching rule, not a default. Before adding a context encoder, estimate predictive headroom from the ordinary trained agent and check whether the context is already decodable from state and action under the visitation that agent induces. If headroom is zero while the context stays decodable, a predictive gain under the same visitation reflects reduced finite-model error, and the mechanism's benefit cannot be attributed to missing information. If headroom stays open, the residual is real, but the learned state may expose it only partially, so supplying the true context can still lower risk.

The comparison that should guide a decision is therefore between the identifiable and hard i.i.d. settings, not between context-aware and context-blind returns. Comparable returns in the hard setting coexisted with open headroom, and collapsed headroom in the identifiable setting coexisted with decodable context. The paper's own recommendation for benchmarks follows the same logic: report estimated headroom alongside the MDP class and benchmark definition, because the same task under different priors produced headroom well above the null in one setting and not in another.

One scope condition belongs with these estimates. All estimated quantities are properties of a closed loop: each method's visitation shapes the information available to it, and that information shapes the method in turn. The authors deliberately do not compare under a shared reference policy, since no policy is neutral; random behavior collapses the visited region and competent behavior encodes the context it has already identified. The policy dependence is the mechanism under study, and it is also the conditioning that matters in practice, because a deployed mechanism operates under the visitation its training produces. The attenuation ratio is characterized on the calibration environment rather than assumed to transfer elsewhere, so the probes are best read as comparative diagnostics at the operating episode length.

Sources used in this answer