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

TL;DR

Event-Aware World Model (EAWM) tackles the brittleness of observation-centric world models by learning event-centric, kinetic representations from unlabeled observations. It introduces an automated multimodal event generator and a Generic Event Segmentor (GES) to identify informative event boundaries, paired with an Event Predictor to shape latent dynamics; the framework is compatible with diverse world-model architectures such as DreamerV3 and Simulus. Empirically, EAWM yields 10-45% gains over strong baselines across Atari 100K, Craftax 1M, DM Control 500K, and DMC-GB2, achieving new state-of-the-art performance and superior robustness to visual variations. Overall, the work demonstrates that focusing policy learning on salient events rather than raw observations can substantially improve generalization, with potential for broad applicability and future extensions to larger multimodal systems.

Abstract

While model-based reinforcement learning (MBRL) improves sample efficiency by learning world models from raw observations, existing methods struggle to generalize across structurally similar scenes and remain vulnerable to spurious variations such as textures or color shifts. From a cognitive science perspective, humans segment continuous sensory streams into discrete events and rely on these key events for decision-making. Motivated by this principle, we propose the Event-Aware World Model (EAWM), a general framework that learns event-aware representations to streamline policy learning without requiring handcrafted labels. EAWM employs an automated event generator to derive events from raw observations and introduces a Generic Event Segmentor (GES) to identify event boundaries, which mark the start and end time of event segments. Through event prediction, the representation space is shaped to capture meaningful spatio-temporal transitions. Beyond this, we present a unified formulation of seemingly distinct world model architectures and show the broad applicability of our methods. Experiments on Atari 100K, Craftax 1M, and DeepMind Control 500K, DMC-GB2 500K demonstrate that EAWM consistently boosts the performance of strong MBRL baselines by 10%-45%, setting new state-of-the-art results across benchmarks. Our code is released at https://github.com/MarquisDarwin/EAWM.

Figures (27)

• Figure 1: Imagined frames and event predictions under the framework of EAWM at imagination step 9. The first two rows show ground-truth and imagined images, while the third row highlights their pixel-wise differences. Each column corresponds to a distinct trajectory. Notably, although the imagined frames may deviate from the ground truth in object positions, the event predictor consistently localizes spatial boundaries accurately.
• Figure 2: Overview. EAWM surpasses existing model-free and model-based RL across 55 test tasks, encompassing both continuous and discrete control, as well as multi-modal observations. (a) Mean human-normalized scores and the $95\%$ stratified bootstrap confidence intervals agarwal2021metric on the 26 tasks of Atari 100K. (b) Percentage of scores against the maximum score in Craftax. (c) Mean returns on 10 challenging tasks from DeepMind Control Suite. (d) Mean returns over 6 tasks on 3 test environments from DMC-GB2.
• Figure 3: EAWM architecture that predicts the next observations and events. Given the length of the trajectory segment $k$, the sequence model outputs $\mathbf{y}_t$, which summarizes the embeddings $\mathbf{Z}_{t-1}=[\mathbf{z}_{t-k},...,\mathbf{z}_{t-1}]$ and $\mathbf{A}_{t-1}=[\mathbf{a}_{t-k},...,\mathbf{a}_{t-1}]$. The observation predictor predicts the next observation $\hat{\mathbf{o}}_t$ via the outputs of the sequence model $\mathbf{y}_t$ and the outputs of the dynamics predictor $\hat{\mathbf{z}}_t$. The representation model combines observation encodings with hidden states $\mathbf{h}_t$ to obtain embeddings $\mathbf{z}_t$. The generic event segmentor identifies event boundaries, the starting and ending points of meaningful segments of the observation stream.
• Figure 4: Ablation studies on key components of EAWMs with 5 random seeds over 6 Atari games: Assault, Breakout, Gopher, Krull, and Ms Pacman, Up N Down. The results show Simulus with the solid lines and EADream with the dashed lines.
• Figure 5: Ablation studies on key components of EAWM on DeepMind Control Suite.