Considering Context: When World Models Need Context Encoders

2026-10-05

AI Summary

This study addresses the unclear necessity and practical benefits of external context for hypotheses in model-based reinforcement learning. We propose a "predictive sufficiency" metric that quantifies the contribution of context to state prediction, effectively distinguishing historically recoverable information from genuine requirements to guide mechanism matching. Methodologically, we construct a practical criterion based on standard trained agents that evaluates context effectiveness without requiring a reference policy, integrating risk decomposition analysis with classification algorithms into a comprehensive evaluation framework. Our findings reveal that this metric remains invariant across MDP classes and demonstrate that external context yields no substantial gains in implicitly identifiable scenarios. These results provide both theoretical foundations and empirical guidance for the design of context mechanisms in reinforcement learning.

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.

Problem

• model-based reinforcement learning
• context encoder
• predictive sufficiency
• world models
• generalization

Innovation

• predictive sufficiency
• context encoder
• model-based reinforcement learning
• world models
• generalization