https://github.com/frankroeder/dali
Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
https://github.com/frankroeder/dali
contextual-reinforcement-learning counterfactual-consistency counterfactuals dreamer neurips-2025 representation-learning worldmodel zero-shot-generalization
Last synced: 4 months ago
JSON representation
Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
- Host: GitHub
- URL: https://github.com/frankroeder/dali
- Owner: frankroeder
- License: mit
- Created: 2025-06-16T15:34:00.000Z (about 1 year ago)
- Default Branch: main
- Last Pushed: 2026-01-17T06:53:53.000Z (5 months ago)
- Last Synced: 2026-01-17T17:58:41.052Z (5 months ago)
- Topics: contextual-reinforcement-learning, counterfactual-consistency, counterfactuals, dreamer, neurips-2025, representation-learning, worldmodel, zero-shot-generalization
- Language: Python
- Homepage: https://openreview.net/forum?id=41bIzD5sit
- Size: 162 KB
- Stars: 2
- Watchers: 0
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Dynamics-Aligned Latent Imagination (DALI)
#### [[OpenReview]](https://openreview.net/forum?id=41bIzD5sit) [[arXiv]](https://arxiv.org/abs/2508.20294) [[Code]](https://github.com/frankroeder/DALI)
[Frank Röder*](https://frankroeder.github.io/), [Jan Benad](https://scholar.google.com/citations?user=FWy1Ov0AAAAJ), [Manfred Eppe](https://scholar.google.de/citations?user=fG0VWroAAAAJ), [Pradeep Kr. Banerjee*](https://scholar.google.com/citations?user=cnSjMBwAAAAJ) (* equal contribution)
All the listed authors are members of the [Institute for Data Science Foundations](https://www.tuhh.de/dsf/homepage).
---
This is the official implementation of **Dynamics-Aligned Latent Imagination** (**DALI**) based on [Jax](https://github.com/jax-ml/jax).
DALI introduces a Dreamer-based architecture with a self-supervised context encoder designed to infer latent environmental
parameters, such as gravity or friction. This integration facilitates robust zero-shot adaptation to diverse and previously
unseen dynamics. Furthermore, manipulating the learned context representation enables physically plausible counterfactual imagination directly within the world model.
### Citation
If this code or paper aids your research, please cite it as follows:
```bib
@inproceedings{
Roder_DynamicsAlignedLatent_2025,
title={Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization},
author={Frank R{\"o}der and Jan Benad and Manfred Eppe and Pradeep Kr. Banerjee},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
url={https://openreview.net/forum?id=41bIzD5sit}
}
```
## Setup
Use [uv](https://github.com/astral-sh/uv) to setup your python environment.
```bash
uv sync
uv pip install -e ./dreamerv3_compat
uv pip install -e ./
```
## Training
Run scripts in `./local_scripts/` to generate results for experts, random policies, and DALI variants.
## Record Data for Analysis
```bash
uv run -m contextual_mbrl.dreamer.record_context --logdir logs/carl_dmc_walker_double_box_enc_img_dec_img_ctxencoder_transformer_normalized/1337
```
### Log data for counterfactual dreams
```bash
uv run -m contextual_mbrl.dreamer.record_counterfactual_plausibility --logdir logs/carl_dmc_ball_in_cup_double_box_enc_img_dec_img_ctxencoder_transformer_grssm_normalized/1337 --jax.platform cpu
```
### Log data for imagined counterfactual obs trajectories
```bash
uv run -m contextual_mbrl.dreamer.record_counterfactual_plausibility_obs --logdir logs/carl_dmc_walker_double_box_enc_img_dec_img_ctxencoder_transformer_normalized/1337 --jax.platform cpu
```
### Record dataset for counterfactual obs analysis
```bash
uv run -m contextual_mbrl.dreamer.record_counterfactual_plausibility_obs_dataset --logdir logs/carl_dmc_walker_double_box_enc_img_dec_img_ctxencoder_transformer_normalized/1337 --jax.platform cpu
```
## Plots
To generate the plots, run the scripts in the `./analysis` directory.
## Acknowledgments
- cRSSM variants: Prasanna et al. 2024 - [https://github.com/sai-prasanna/dreaming_of_many_worlds](https://github.com/sai-prasanna/dreaming_of_many_worlds)
- DreamerV3: Hafner et al. 2023 - [https://github.com/danijar/dreamerv3](https://github.com/danijar/dreamerv3)
- CARL benchmark: Benjamins et al. 2023 - [https://github.com/automl/CARL](https://github.com/automl/CARL)