Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
This addresses the challenge of real-world reinforcement learning where environmental contexts are latent or hard to measure, offering a method for zero-shot generalization, though it builds incrementally on existing contextual MDP and Dreamer architectures.
The paper tackles the problem of adapting reinforcement learning agents to unseen environmental conditions without retraining by introducing Dynamics-Aligned Latent Imagination (DALI), which infers latent context representations from interactions, achieving significant gains over baselines and enabling zero-shot generalization to unseen variations.
Real-world reinforcement learning demands adaptation to unseen environmental conditions without costly retraining. Contextual Markov Decision Processes (cMDP) model this challenge, but existing methods often require explicit context variables (e.g., friction, gravity), limiting their use when contexts are latent or hard to measure. We introduce Dynamics-Aligned Latent Imagination (DALI), a framework integrated within the Dreamer architecture that infers latent context representations from agent-environment interactions. By training a self-supervised encoder to predict forward dynamics, DALI generates actionable representations conditioning the world model and policy, bridging perception and control. We theoretically prove this encoder is essential for efficient context inference and robust generalization. DALI's latent space enables counterfactual consistency: Perturbing a gravity-encoding dimension alters imagined rollouts in physically plausible ways. On challenging cMDP benchmarks, DALI achieves significant gains over context-unaware baselines, often surpassing context-aware baselines in extrapolation tasks, enabling zero-shot generalization to unseen contextual variations.