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: ```bash pip install -U huggingface_hub ``` ```python from huggingface_hub import snapshot_download snapshot_download( repo_id="darwin05/EAWM", local_dir="checkpoints" ) ``` Alternatively, you can use the Hugging Face CLI: ```bash 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](https://github.com/MarquisDarwin/EAWM/blob/master/EADream/README.md) for the EADream environment. - See [EASimulus/README.md](https://github.com/MarquisDarwin/EAWM/blob/master/EASimulus/README.md) for the EASimulus environment. Checkpoint layout: ```text 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: ```bash 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: ```bash 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_.pt` or `dmc__.pt` naming convention, for example `atari_qbert.pt` and `dmc_quadruped_walk.pt`. ### Evaluate EASimulus Checkpoints Atari example: ```bash 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: ```bash 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.