| <div align="center"> |
|
|
| <h2>From Observations to Events: Event-Aware World Models for Reinforcement Learning</h2> |
|
|
| Zhao-Han Peng, Shaohui Li, Zhi Li, Shulan Ruan, Yu Liu, You He |
|
|
| Tsinghua University, Zhejiang University |
|
|
| **ICLR 2026** |
|
|
| <a href='https://arxiv.org/abs/2601.19336'><img src='https://img.shields.io/badge/ArXiv-2510.10125-red'></a> |
| <a href='https://github.com/MarquisDarwin/EAWM'><img src='https://img.shields.io/badge/Project-Page-Blue'></a> |
|
|
|
|
| </div> |
|
|
| ## 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. |
| <div align="center"> |
| <img src="results.png" height="220" /> |
| </div> |
|
|
| ## 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_<game>.pt` or `dmc_<domain>_<task>.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. |
