Uncovering Untapped Potential in Sample-Efficient World Model Agents

TL;DR

This work tackles sample-efficient reinforcement learning with token-based world models (TBWMs) by addressing their limitations to visual inputs and discrete actions. It introduces Simulus, a modular TBWM that combines a multi-modality tokenization framework, intrinsic motivation, prioritized world-model replay, and regression-as-classification for rewards and returns. Across Atari 100K, DMC proprioception, and Craftax-1M benchmarks, Simulus achieves state-of-the-art planning-free performance and reveals strong synergy among its components through ablations. The authors provide extensive experimental details and open-source code/weights to foster broader adoption and future TBWM research.

Abstract

World model (WM) agents enable sample-efficient reinforcement learning by learning policies entirely from simulated experience. However, existing token-based world models (TBWMs) are limited to visual inputs and discrete actions, restricting their adoption and applicability. Moreover, although both intrinsic motivation and prioritized WM replay have shown promise in improving WM performance and generalization, they remain underexplored in this setting, particularly in combination. We introduce Simulus, a highly modular TBWM agent that integrates (1) a modular multi-modality tokenization framework, (2) intrinsic motivation, (3) prioritized WM replay, and (4) regression-as-classification for reward and return prediction. Simulus achieves state-of-the-art sample efficiency for planning-free WMs across three diverse benchmarks. Ablation studies reveal the individual contribution of each component while highlighting their synergy. Our code and model weights are publicly available at https://github.com/leor-c/Simulus.

Figures (11)

• Figure 1: Results overview. Simulus exhibits state-of-the-art sample-efficiency performance for planning-free methods across all three benchmarks. $\dagger$robine2023transformer, $\ddagger$dedieu2025TWM2.
• Figure 2: An illustration of the independent processing of modalities for an observation with two modalities.
• Figure 3: World model training and imagination. To maintain visual clarity, we omitted token embedding details, as well as optimization details of rewards and termination signals.
• Figure 4: Results on the DeepMind Control Suite 500K Proprioception (top) and Atari 100K (bottom) benchmarks.
• Figure 5: Craftax-1M training curves with mean and 95% confidence intervals.