Context-Aware World Models for Task-Agnostic Control

Busra Tugce Gurbuz, Hafez Ghaemi, Christopher C Pack, Shahab Bakhtiari, Eilif B. Muller

Track: Extended Abstract Track

Keywords: world models, reinforcement learning, self-supervised learning, neuroscience

Abstract: Building accurate and generalizable world models requires aligning an agent’s inductive biases with the structure of the environments it inhabits. Realistic environments are inherently multi-task, containing both shared global dynamics and context-specific variations that change with the task. We argue that a world model should reflect this structure and enable efficient separation of context-specific dynamics while integrating shared regularities that support generalization. As a first step towards this, we introduce Context-Aware World Models (CaWM), which aligns the agent’s world model with the multi-task structure of the environment. In contrast to existing model-based approaches that assume access to ground-truth task labels, CaWM learns to infer latent task contexts directly from its environment interactions via a self-supervised objective, and uses inferred contexts to modulate the world model representations and enable task-agnostic control. To benchmark CaWM, we present Multi-FoE, a multi-task visual foraging environment with egocentric partial observations and boundary-free task switching. Empirically, CaWM achieves higher performance and success rate compared to context-free baselines, and approaches the performance of an oracle with ground-truth task labels.

Submission Number: 130