Uncovering Untapped Potential in Sample-Efficient World Model Agents

Lior Cohen, Kaixin Wang, Bingyi Kang, Uri Gadot, Shie Mannor

Keywords: world models, deep reinforcement learning, intrinsic motivation, sample efficiency

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://anonymous.4open.science/r/Simulus-FBF5.

Primary Area: Reinforcement learning (e.g., decision and control, planning, hierarchical RL, robotics)

Submission Number: 17567