# Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization

## TL;DR

DALI is a novel framework that infers latent context representations from interactions, enabling zero-shot generalization in reinforcement learning by bridging perception and control through dynamics-aligned imagination.

## Contribution

Introduces DALI, a self-supervised latent inference method integrated into Dreamer, enhancing zero-shot generalization in cMDPs without explicit context variables.

## Key findings

- DALI outperforms baselines on challenging benchmarks.
- DALI enables counterfactual reasoning in latent space.
- DALI achieves zero-shot generalization to unseen contexts.

## Abstract

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.