Dynamics-Aligned Latent Imagination (DALI)

[OpenReview] [arXiv] [Code]

Frank Röder*, Jan Benad, Manfred Eppe, Pradeep Kr. Banerjee* (* equal contribution)
All the listed authors are members of the Institute for Data Science Foundations.

---

This is the official implementation of Dynamics-Aligned Latent Imagination (DALI) based on 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:

@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 to setup your python environment.

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

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

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

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

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
• DreamerV3: Hafner et al. 2023 - https://github.com/danijar/dreamerv3
• CARL benchmark: Benjamins et al. 2023 - https://github.com/automl/CARL