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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Röder, Frank, Benad, Jan, Eppe, Manfred, Banerjee, Pradeep Kr.
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915733481455616
author Röder, Frank
Benad, Jan
Eppe, Manfred
Banerjee, Pradeep Kr.
author_facet Röder, Frank
Benad, Jan
Eppe, Manfred
Banerjee, Pradeep Kr.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2508_20294
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
Röder, Frank
Benad, Jan
Eppe, Manfred
Banerjee, Pradeep Kr.
Machine Learning
Artificial Intelligence
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.
title Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2508.20294