Simulus: Combining Improvements in Sample-Efficient World Model Agents

Technical Summary: Simulus – Combining Improvements in Sample-Efficient World Model Agents

Problem Statement
World models (WMs) represent a frontier in sample-efficient reinforcement learning (RL), yet their potential is often unrealized due to the complexity of integrating promising improvements. While individual advancements exist, combining them requires significant expertise and engineering effort. Furthermore, existing token-based world models (TBWMs) are largely restricted to image observations and discrete actions, limiting their applicability to continuous control and multi-modal settings. Additionally, components such as intrinsic motivation (for epistemic uncertainty reduction) and prioritized replay are frequently overlooked in sample-efficient agents, partly due to the high cost of integration and the risk that intrinsic exploration might waste scarce interaction budgets on task-irrelevant experiences. This landscape mirrors the pre-Rainbow era in deep RL, where individually validated improvements had not been systematically combined.

Methodology
The authors introduce Simulus, a modular token-based world model agent designed to integrate and validate a suite of complementary improvements. The agent is built upon the REM architecture and comprises four distinct modules: a representation module ($V$), a world model ($M$), a controller ($C$), and a replay buffer.

1. Flexible Tokenization Framework ($V$): Simulus employs a modular tokenization system supporting arbitrary combinations of observation and action modalities. It natively handles images, continuous vectors, categorical variables, and "2D categoricals" (image-like grids of categorical variables). Encoders produce fixed-length token sequences, which are concatenated into a unified representation. This allows the agent to operate in diverse environments, including continuous control and symbolic multi-modal tasks.
2. World Model ($M$): The dynamics model utilizes a Retentive Network (RetNet) augmented with a Parallel Observation Prediction (POP) mechanism to model environment dynamics as sequence modeling over discrete tokens.
   • Intrinsic Motivation: To reduce epistemic uncertainty, Simulus employs an ensemble of next-observation prediction heads. The disagreement between these heads, measured via Jensen-Shannon Divergence (JSD), generates intrinsic rewards ($r_{int}$). These are combined with extrinsic rewards to guide the controller toward high-uncertainty regions during "imagination" (planning within the model).
   • Prioritized Replay: The replay buffer maintains world model loss values for stored examples. Training utilizes a mixture of uniform and prioritized sampling (controlled by a single parameter $\alpha$), where samples are selected proportional to their softmax-transformed losses.
3. Controller ($C$): The controller is an extended actor-critic architecture that operates entirely on simulated experience generated by the world model.
   • Regression-as-Classification (RaC): To handle unbounded reward and return signals robustly, Simulus adopts a classification network with exponential bins for value prediction. It utilizes the HL-Gauss loss method (cross-entropy based) rather than traditional mean squared error, improving stability in value estimation.
4. Training Cycle: The agent follows a cycle of data collection, representation learning, world model learning, and control learning in imagination. The controller is trained using $\lambda$-returns and a REINFORCE objective with a value baseline.

Key Contributions

• Integration of Overlooked Components: The paper identifies and integrates intrinsic motivation, prioritized world model replay, and a hybrid RaC approach. It demonstrates that these components are not only individually effective but yield synergistic gains when combined.
• Versatile Token-Based Agents: By adapting multi-modality tokenization to the sample-efficient regime, Simulus establishes TBWMs as versatile agents capable of handling visual, continuous, and multi-modal inputs, moving beyond the limitations of image-only discrete action models.
• Validation of Model-Induced Exploration: The work shows that world-model-induced exploration significantly improves performance even under tight interaction budgets, challenging the notion that intrinsic motivation is too costly for sample-efficient settings.
• Open-Source Foundation: The authors provide a strong starting point for TBWM research by open-sourcing code and model weights, lowering the barrier to entry for developing competitive world model agents.

Results
Simulus was evaluated across three diverse benchmarks with strict interaction budgets:

• Visual Atari 100K: Simulus achieved state-of-the-art (SOTA) sample efficiency among planning-free methods, outperforming baselines like DreamerV3, TWM, STORM, and DIAMOND. Notably, it is the first planning-free world model to reach human-level Interquantile Mean (IQM) and median scores, achieving superhuman performance on 13 of 26 games.
• Continuous Control (DMC Proprioception 500K): In continuous control tasks, Simulus consistently matched or slightly outperformed DreamerV3, demonstrating the effectiveness of token-based architectures in continuous domains with compact vocabularies.
• Symbolic Multi-Modal (Craftax-1M): In a partially observable 2D survival environment combining symbolic grids and continuous features, Simulus outperformed concurrent world model methods and all model-free baselines (including RND, PPO-RNN, and E3B), handling long token sequences efficiently.

Ablation studies confirmed that each component contributes individually, with intrinsic rewards proving particularly crucial. The combination of all components yielded the strongest performance, indicating significant synergy.

Significance
The paper argues that the unrealized potential of world model agents stems from the difficulty of combining modular improvements. By drawing inspiration from the Rainbow framework in DQN, Simulus demonstrates that a systematic combination of tokenization flexibility, intrinsic motivation, prioritized replay, and advanced value prediction techniques leads to state-of-the-art sample efficiency. The work establishes that token-based world models are a versatile and scalable approach for diverse RL domains, provided that the components are carefully integrated to avoid the interference often seen in monolithic architectures. The authors maintain that while token-based agents currently face training speed limitations compared to some baselines, their modular design offers distinct advantages for policy inference and future research scalability.