README.md

Overview

EAWM (Event-Aware World Models) is an event-aware world model framework for reinforcement learning. In addition to conventional observation prediction, EAWM explicitly models events so that the world model can learn sparser and more interpretable environment dynamics. The repository includes experiments on Atari, DeepMind Control Suite, DMC-GB2, and Craftax.

This repository provides two implementations:

• EADream: a DreamerV3-style implementation for Atari 100K, DMC Vision, and DMC-GB2.
• EASimulus: a Simulus-based implementation for Atari 100K and Craftax.

The visualization files in this repository show policy rollouts, outputs from the automated event generator, and open-loop future-observation predictions from the world model. Raw experiment results are available in results/.

Environment Setup

Please choose one of the instantiations of EAWM from EADream and EASimulus, and follow the instructions in the EADream/README.md or EASimulus/README.md to configure the environment.

Atari

Visualization of Policy

Automated Event Generator

Imagination of Future Observations

DMC-GB2

Visualization of Policy

Automated Event Generator

Open-loop prediction in train environments

EAWM	DreamerV3	DyMoDreamer

Open-loop prediction in unseen test environments

EAWM	DreamerV3	DyMoDreamer

As demonstrated above, we observe that observation prediction in current MBRL world models generalizes poorly to unseen test environments due to limited interaction diversity in the training environments. In contrast, event prediction is inherently more tractable and interpretable, owing to the sparsity and well-defined semantic structure of events.

📈 Results

The folder results contains raw scores of world models (for each game, and for each training run).

Download Checkpoints

The pretrained checkpoints are hosted on Hugging Face: darwin05/EAWM.

You can download the full checkpoint repository with the Hugging Face Hub Python API:

pip install -U huggingface_hub

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="darwin05/EAWM",
    local_dir="checkpoints"
)

Alternatively, you can use the Hugging Face CLI:

hf download darwin05/EAWM \
  --repo-type model \
  --local-dir checkpoints

After downloading, the checkpoint directory should follow this structure:

checkpoints/
|-- EADream/
|   |-- atari_pong.pt
|   |-- atari_breakout.pt
|   `-- dmc_cheetah_run.pt
`-- EASimulus/
    |-- Atari/
    |   |-- Pong.pt
    |   `-- Breakout.pt
    `-- craftax.pt

Citation

@inproceedings{
    Peng2026from,
    title={From Observations to Events: Event-Aware World Models for Reinforcement Learning},
    author={Zhao-Han Peng and Shaohui Li and Zhi Li and Shulan Ruan and Yu Liu and You He},
    booktitle={The Fourteenth International Conference on Learning Representations},
    year={2026},
    url={https://openreview.net/forum?id=OWkkFaq1IZ}
}

Acknowledgements

The code in this repository is built upon the following repositories.

• https://github.com/danijar/dreamerv3
• https://github.com/leor-c/Simulus
• https://github.com/fkodom/yet-another-retnet
• https://github.com/google-research/rliable
• https://github.com/wandb/wandb
• https://github.com/thuml/HarmonyDream
• https://github.com/NM512/dreamerv3-torch
• https://github.com/aalmuzairee/dmcgb2