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

Abstract

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.