Context-Adaptive World Models for Visual Out-of-Distribution Generalization

1 Introduction andMotivation

Out-of-distribution (OOD) generalization in pixel-driven, partially observable reinforcement learning is a hard online problem. Latent features of the environment such as friction, actuator strength or reward-allocation rules shape what the agent sees and earns but cannot be observed from any single frame; their trajectory-level effect is the agent's context [Benjamins et al., 2023]. A robust agent must infer several such factors from pixels and stay calibrated when they shift at different times within an episode.

Existing world models leave this setting open. The Dreamer family [Hafner et al., 2025] and recent state-space model (SSM)-based world models [Wang et al., 2025; Dao and Gu, 2024] are context-agnostic. Pixel-space contextconditioned world models [Yang et al., 2025; R¨oder et al., 2025; Prasanna et al., 2024] either assume fixed context per episode or rely on a privileged context channel.

DOMINO [Mu et al., 2022] infers multi-factor context but on featurized state. Meta-RL methods that admit withinepisode drift [Chen et al., 2022; Liang et al., 2024] carry context as a single latent. No prior method jointly meets the three demands: (i) multi-factor context inference from pixels, (ii) within-episode asynchronous tracking, and (iii) propagation of the current context into the world model used for imagination-based policy training.

This work fills this gap. Architecturally, it splits the conditioning across K inferred vectors through a block-diagonal FiLM head, so each vector drives an independent dynamics pathway. Theoretically, it provides the first bound linking the context-inference error in this setting to policy suboptimality through the FiLM Lipschitz chain.

2 ProblemFormalization

The setting is a piecewise-stationary multi-factor contextual POMDP. The true context is ut = (u1t , . . . , uFt ) with each uft piecewise constant and switching asynchronously; the transition, reward, and observation models are parameterized by ut, and the agent observes only pixels and scalar rewards. The inferred context ct = (c1t , . . . , cKt ) is a set of K vectors that summarize the latent trajectory at distinct temporal scales rather than estimate individual factors. OOD evaluation probes factor values, factor combinations and switch cadences outside training.

3 Research Contributions

(1) Multi-slot context inference from pixels. An encoder infers K context vectors from a short causal pixel window via cross-attention; each slot has its own adaptive lookback and a fixed geometric temporal scale, so vectors attend to tokens over different horizons. Three losses shape the representation: a soft-label temporal contrastive loss anchors slots across same-phase steps, a persist-or-jump transition regularizer with a learned hazard smooths the inferred trajectory, and a Hilbert-Schmidt regularizer that discourages cross-slot dependence. The design extends the multi-factor context inference of [Mu et al., 2022] and the segment-aware inference of [Chen et al., 2022; Liang et al., 2024] to the pixel setting, unsupervised.

(2) Block-diagonal FiLM on a Mamba backbone. Context vectors steer a Mamba selective-SSM world model [Dao and Gu, 2024; Wang et al., 2025] through block-diagonal FiLM [Perez et al., 2018] at two recurrence-boundary sites: the per-head log time-step δt governing temporal selectivity, and the output gate gt governing output routing, replacing the encoder-input concatenation of [R¨oder et al., 2025; Prasanna et al., 2024]. The block-diagonal FiLM head assigns each vector to a disjoint subset of SSM heads, so vector k modulates only that subset and gradients through it update only vector k. A ReZero-initialized parallel pathway captures the state-context interaction that the FiLM head cannot express. Tanh-bounded scales and near-zero output initialization keep the recurrent state bounded under any context.

(3) Context-aware imagination with formal guarantees. A policy trains on imagined rollouts from replay-buffer anchors with context held fixed and appended to the policy input as a phase signal; asynchronous switching enters training through FiLM-conditioned δt on real trajectories. Three formal results ground the design: (i) FiLM at the chosen sites preserves the SSM eigenvalue bound for any context; (ii) the cost of a factored variational posterior over the F switch indicators is bounded by (F-1) h(¯λ) nats, small at sparse switch rates; (iii) policy suboptimality is bounded by √K times the context-inference error plus the world-model error along a Lipschitz chain through FiLM, the SSM and the decoder, rate-invariant in the switch cadence.

4 Evaluation Plan

Continuous control on DeepMind Control (Walker, Quadruped) uses within-episode asynchronous Bernoulli switching on physics factors and reward direction, with F ∈ {2, 3}. Discrete control on the Arcade Learning Environment uses episode-aligned mode/difficulty switching on Alien and BankHeist. OOD splits hold out factor values, joint compositional corners and switch cadences spanning 0.2× to 4× the training rate. Results are compared to DRAMA [Wang et al., 2025] (matched-backbone, contextagnostic), DALI [R¨oder et al., 2025] (inferred context, Dreamer family), and cRSSM [Prasanna et al., 2024] (privileged-context oracle). Ablations isolate the multi-scale role of slots, the chosen K, δt versus gt modulation, surplus-slot pruning at test time, and a capacity-matched control.

5 Current Status and Impact

I am in the final stage of my doctoral studies. The multi-slot conditioning above is realized on Mamba-2 and Mamba-3 SSM backbones: the slot context vectors match a privilegedcontext oracle on rare-switch joint shifts, extend gains to unseen switch cadences, and transfer from continuous DMC physics to discrete episode-boundary switching on Atari. Remaining work covers automated slot-count selection, the backbone-expressivity ceiling at joint compositional corners, and few-shot adaptation of the inferred context at evaluation time. The thesis argues that pixel-driven RL agents can be made robust to unseen non-stationarity without privileged context.

Acknowledgments

This publication has emanated from research conducted with the financial support of Taighde ´Eireann - Research Ireland under Frontiers for the Future grant number 21/FFP-A/8957. For the purpose of Open Access, the author has applied a CC BY public copyright licence to any Author Accepted Manuscript version arising from this submission.