darwin05/EAWM

From Observations to Events: Event-Aware World Models for Reinforcement Learning

Zhao-Han Peng, Shaohui Li, Zhi Li, Shulan Ruan, Yu Liu, You He

Tsinghua University, Zhejiang University

ICLR 2026

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.

Download Checkpoints

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

Using Checkpoints

The provided checkpoints are mainly intended for reproducing evaluation results. Before using them, set up the corresponding subproject environment:

• See EADream/README.md for the EADream environment.
• See EASimulus/README.md for the EASimulus environment.

Checkpoint layout:

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

Evaluate EADream Checkpoints

Atari example:

cd EADream
python3 eval.py \
  --configdir configsc.yaml \
  --game pong \
  --weights /path/to/checkpoint/dir/EADream/atari_pong.pt \
  --episodes 100 \
  --device cuda:0 \
  --result-file log/result.txt

DMC example:

cd EADream
python3 eval.py \
  --configdir configsc.yaml \
  --configs dmc_vision \
  --task dmc_cheetah_run \
  --weights /path/to/checkpoint/dir/EADream/dmc_cheetah_run.pt \
  --episodes 10 \
  --device cuda:0 \
  --result-file dmcresult.txt

EADream checkpoints follow the atari_<game>.pt or dmc_<domain>_<task>.pt naming convention, for example atari_qbert.pt and dmc_quadruped_walk.pt.

Evaluate EASimulus Checkpoints

Atari example:

cd EASimulus
python scripts/eval.py \
  --benchmark atari \
  --weights-path /path/to/checkpoint/dir/EASimulus/Atari/Pong.pt \
  --num-episodes 100 \
  --num-envs 20 \
  --seed 0 \
  --wandb-mode disabled

Craftax example:

cd EASimulus
python scripts/eval.py \
  --benchmark craftax \
  --weights-path /path/to/checkpoint/dir/EASimulus/craftax.pt \
  --num-episodes 100 \
  --num-envs 20 \
  --seed 0 \
  --wandb-mode disabled

EASimulus Atari checkpoints use ALE environment stems as filenames. For example, Pong.pt is mapped by the evaluation script to PongNoFrameskip-v4.

Batch Evaluation

You can also use the batch evaluation scripts:

• https://github.com/MarquisDarwin/EAWM/EADream/scripts/eval_atari.sh
• https://github.com/MarquisDarwin/EAWM/EADream/scripts/eval_dmc.sh
• https://github.com/MarquisDarwin/EAWM/EASimulus/scripts/eval_atari.sh
• https://github.com/MarquisDarwin/EAWM/EASimulus/scripts/eval_craft.sh

Before running these scripts, replace the checkpoint placeholder paths with real paths. For EADream, run the scripts under EADream/scripts/; For EASimulus, run the scripts under EASimulus/.

Difference From Training Resume

The .pt files under checkpoints/ are pretrained agent weights intended for evaluation only; they are not complete training checkpoints for resuming training.