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 the DALI approach using JAX.

@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 Anaylsis

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.