Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization

arXiv 24.3 2024 81.3 method

TLDR

Proposes cRSSM, a contextual world model for Dreamer, improving zero-shot generalization to unseen dynamics in contextual RL.

Reasoning

Strengths include a clear problem formulation and a novel modification to a well-known world model, with empirical validation on two tasks. Weaknesses are limited evaluation scope (only two tasks) and reliance on observable context, which may not hold in more complex settings.

Read-first score

Read-first score 81.3, weighted from topical fit, citation, graph, method, reproducibility, and recency signals. Original total remains 50.

Methodology quality 25%

100

Screens visible abstract and analysis fields for experiment, dataset, baseline, metric, and limitation evidence. markers=analysis,baseline,benchmark,evaluation,experiment

Reproducibility 25%

Screens links and visible text for paper, code, dataset, artifact, and repository signals. pdf=True; code=True; dataset=False; markers=code,github

Recency 8%

75.1

Uses a gentle age decay so recent papers surface without erasing older foundations. 2024

Topical relevance 42%

71.4

Uses existing LLM keyword relevance scores normalized to 0-100. world model,world simulator,generative world model,interactive world model,video world model,world dynamics prediction,model-based reinforcement learning world model

Field roles

Methodology anchorReproducibility anchor

Rank sensitivity

Stability: volatile; rank range: 13.

Keyword Scores

world model

model-based reinforcement learning world model

generative world model

world dynamics prediction

world simulator

interactive world model

video world model

Deep Analysis

Innovations

Proposes the contextual recurrent state-space model (cRSSM) that modifies Dreamer v3 to incorporate context for inferring latent Markovian states and modeling latent dynamics.
Demonstrates that systematic incorporation of context improves zero-shot generalization (ZSG) of policies trained on world model dreams.
Shows qualitatively that the approach disentangles latent state from context, enabling extrapolation of dreams to unseen contexts.

Methodology

The paper introduces cRSSM, a modification to the Dreamer v3 world model that explicitly conditions the latent state inference and dynamics on observable context values (e.g., mass or dimensions). The model is trained and evaluated on two tasks from the CARL benchmark suite, which is designed for contextual reinforcement learning. Policies are trained on imagined trajectories ('dreams') from the world model, and zero-shot generalization is measured by testing on unseen context values.

Key Results

Experiments show that incorporating context into the world model significantly improves zero-shot generalization of the learned policies compared to the baseline Dreamer v3. Qualitative analysis indicates that cRSSM learns to disentangle the latent state from context, allowing the model to generate accurate dreams for unseen contexts.

Limitations

The approach assumes observability of context values, which may not be available in all real-world settings.
Evaluation is limited to only two tasks from the CARL benchmark suite, leaving broader generalization unverified.
The paper addresses the simpler setting of contextual RL, not the full challenge of zero-shot generalization to completely unknown dynamics.

Resources

Paper Code