File size: 4,109 Bytes
2468f8f 4b844ed 2468f8f d0c8f4d 996842e 2af9506 996842e d0c8f4d 44d9a7f 3d0f9fd 44d9a7f d0c8f4d 64b874e d0c8f4d 93ce473 d0c8f4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | <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. |