Dynamics-Aligned Latent Imagination (DALI)

[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.

@article{Roder_DynamicsAlignedLatent_2025,
  title = {Dynamics-{{Aligned Latent Imagination}} in {{Contextual World Models}} for {{Zero-Shot Generalization}}},
  author = {R{\"o}der, Frank and Benad, Jan and Eppe, Manfred and Banerjee, Pradeep Kr},
  journal={arXiv preprint arXiv:2508.20294},
  year = {2025},
}

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.