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

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

Shubham Subhnil

Proceedings of the Thirty-Fifth International Joint Conference on Artificial Intelligence
Doctoral Consortium. Pages 8341-8342. https://doi.org/10.24963/ijcai.2026/949

A vision-based robot can be faced with changing friction, gravity, or motor behavior, whose effects are observable only from pixels mid-episode. Visual model-based reinforcement learning struggles in such settings because most world models assume the latent factors driving these changes stay fixed during interaction, or fold them into a single context embedding. My thesis is developing a context-adaptive Mamba-based world model that infers K context vectors from short pixel-trajectory windows. Each vector steers a disjoint slice of the Mamba backbone through block-diagonal feature-wise linear modulation (FiLM), so distinct dynamics signals travel along parallel pathways rather than collapse into a single channel. The policy trains on imagined rollouts conditioned on the inferred context, and I show that its suboptimality is bounded by √K times the context-inference error along a Lipschitz chain through FiLM, the SSM, and the decoder. Evaluation spans within-episode shifts on DeepMind Control (Walker, Quadruped) and mode/difficulty switching on Atari, against an inferred-context baseline and a privileged-context oracle.
Keywords:
Machine Learning: Reinforcement learning
Machine Learning: Model-based and model learning reinforcement learning
Machine Learning: Representation learning
Machine Learning: Theory of deep learning