Title: Uncovering Untapped Potential in Sample-Efficient World Model Agents
Venue: NeurIPS 2025
Accepted: no
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.

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 WMreplay, 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 .

1 Introduction

Sample efficiency refers to the ability of a reinforcement learning (RL) algorithm to learn effective policies using as few environment interactions as possible. In many real-world domains such as robotics, autonomous driving, and healthcare, this is particularly critical, as interactions are costly, slow, or constrained. World model agents, methods that learn control entirely from simulated experience generated by a learned dynamics model, have emerged as a promising approach to improving sample efficiency [17, 20, 37, 57].

Here, we focus on token-based world models (TBWMs) [37, 38, 11], sample-efficient RL methods that learn the dynamics entirely within a learned discrete token space, where each observation comprises a sequence of tokens. Evidently, most large scale world models [1, 12, 14] operate on multi-token observations, suggesting that such representations are advantageous at scale. TBWMs offer a clear modular design, separating the optimization of its representation, dynamics, and control models. As modular systems, TBWMs are easier to scale, develop, study, and deploy, as individual modules can be treated independently without interfering and are easier to master through divide and conquer. In addition, such separation leads to simpler optimization objectives and avoids interference between them (Appendix B.1).

However, existing TBWMs are restricted to image observations and discrete actions, such as Atari 30 games, limiting both their adoption and broader applicability, as their effectiveness in diverse 31 environments and modalities remains unclear. While multi-modality tokenization approaches exist for 32 large-scale offline settings [43, 47, 33, 34], these methods rely on large vocabularies (e.g., 33K tokens) 33 which are inefficient for online, data-limited regimes. Whether substantially smaller vocabularies can 34 preserve competitive precision and performance is still an open question, leaving these approaches 35 underexplored for sample-efficient RL. 36

Simulus

REM

DIAMOND

STORM

TWM

0.00

0.25

0.50

0.75

1.00

Human Normalized Score

0.99

0.67

0.64

0.56

0.46

0.49

Simulus

0

200

400

600

800

Mean Score

796

754

DeepMind Control 500K

(Proprioception)

Simulus

TWM

PPO-RNN

RND

E3B

0

2

4

6

6.59

5.44

2.13

2.11

1.91

Craftax-1M

Furthermore, despite compelling results [46, 49, 30], intrinsic motivation and prioritized world model 37 replay remain underexplored in sample-efficient world model agents [19, 20, 60, 11, 3], particularly 38 in combination. We conjecture that intrinsic motivation is underused as it may steer the agent toward 39 task-irrelevant regions, potentially wasting limited interaction budget. Prioritized replay [30], while 40 promising, lacks robust empirical support and proved challenging to tune in our experiments. 41

To address these limitations, we propose Simulus, a modular world model that extends a recent 42 TBWM method [11] by integrating several powerful advances from the literature: (1) a modular 43 multi-modality tokenization framework for handling arbitrary combinations of observation and action 44 modalities, (2) intrinsic motivation for epistemic uncertainty reduction [50, 49], (3) prioritized world 45 model replay [30], and (4) regression-as-classification (RaC) for reward and return prediction [16, 20]. 46

47

48

49

50

51

52

To evaluate the impact of the proposed components, we conducted extensive empirical evaluations across three diverse benchmarks, ranging from the visual Atari 100K [28], to the continuous proprioception tasks of the DeepMind Control Suite [53], to Craftax [35], which combines symbolic 2D grid maps with continuous state features. There, Simulus achieves state-of-the-art sample-efficiency for planning-free world models. Ablation studies further show the contribution and effectiveness of each component.

53

54

55

56

57

58

59

60

61

62

63

Summary of contributions:

Through extensive empirical evaluations across three diverse benchmarks, we show that intrinsic motivation, prioritized replay, and RaC significantly improve performance in the sample-efficiency setting for world model agents, particularly when combined.

Our results demonstrate the effectiveness of the multi-modality tokenization approach, as Simulus achieves state-of-the-art planning-free performance across benchmarks, establishing TBWMs as widely-applicable methods.

We propose Simulus, a versatile TBWM agent that follows a highly modular design, offering a solid foundation for future developments. To support future research and adoption of TBWMs, we open-source our code and weights.

2 Method

Notations We consider the Partially Observable Markov Decision Process (POMDP) setting. How64 ever, since in practice the agent has no knowledge about the hidden state space, consider the following 65 state-agnostic formulation. Let Ω , A be the sets of observations and actions, respectively. At every 66 step t , the agent observes o t ∈ Ω and picks an action a t ∈ A . From the agent's perspective, the envi67 ronment evolves according to o t +1 , r t , d t ∼ p ( o t +1 , r r , d t | o ≤ t , a ≤ t ) , where r t , d t are the observed 68 reward and termination signals, respectively. The process repeats until a positive termination signal 69 d t ∈ { 0 , 1 } is obtained. The agent's objective is to maximize its expected return E [ ∑ ∞ t =0 γ t r t +1 ] 70 where γ ∈ [0 , 1] is a discount factor. 71

For multi-modal observations, let o t = { o ( i ) t } | κ | i =1 where κ is the set of environment modalities and 72 o ( i ) t denotes the features of modality κ i . 73

74

75

76

77

78

Overview Simulus builds on REM [11]. The agent comprises a representation module V , a world model M , a controller C , and a replay buffer. To facilitate a modular design, following REM, each module is optimized separately. The training process of the agent involves a repeated cycle of four steps: data collection, representation learning ( V ), world model learning ( M ), and control learning in imagination ( C ).

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

2.1 The Representation Module V

V is responsible for encoding and decoding raw observations and actions. It is a modular tokenization system with encoder-decoder pairs for different input modalities. Encoders produce fixed-length token sequences, creating a common interface that enables combining tokens from various sources into a unified representation. After embedding, these token sequences are concatenated into a single representation, as described in Section 2.2. Note that encoder-decoder pairs need not be learningbased methods; however, when learned, they are optimized independently. This design enables V to deal with any combination of input modalities, provided the respective encoder-decoder pairs.

Tokenization V transforms raw observations o to sets of fixed-length integer token sequences z = { z ( i ) } | κ | i =1 by applying the encoder of each modality z ( i ) = enc i ( o ( i ) ) . Actions a are tokenized using the encoder-decoder pair of the related modality to produce z a . The respective decoders reconstruct observations from their tokens: ˆ o ( i ) = dec i ( z ( i ) ) .

Simulus natively supports four modalities: images, continuous vectors, categorical variables, and image-like multi-channel grids of categorical variables, referred to as "2D categoricals". More formally, 2D categoricals are elements of ([ k 1 ] × [ k 2 ] × . . . × [ k C ]) m × n where k 1 , . . . , k C are per channel vocabulary sizes, C is the number of channels, m,n are spatial dimensions, and [ k ] = { 1 , . . . , k } .

Following REM, we use a VQ-VAE [15, 55] for image observations. For the tokenization of continuous vectors, each feature is quantized to produce a token, as in [43]. However, to improve learning efficiency, we reduce the vocabulary size from 1024 to 125, and modify the quantization levels for optimal coverage (Appendix A.2.2). Unbounded vectors are first transformed using the symlog function [20], defined as symlog( x ) = sign( x ) ln(1 + | x | ) , which compresses the magnitude of large absolute values. Lastly, while no special tokenization is required for categorical inputs, 2D categoricals are flattened along the spatial dimensions to form a sequence of categorical vectors. The embedding of each token vector is obtained by averaging the embeddings of its entries.

2.2 The World Model M

The purpose of M is to learn a model of the environment's dynamics. Concretely, given trajectory segments τ t = z 1 , z a 1 , . . . , z t , z a t in token representation, M models the distributions of the next observation and termination signal, and the expected reward:

where θ is the parameters vector of M and ˆ r θ ( τ t ) is an estimator of E r t ∼ p ( r t | τ t ) [ r t ] .

Architecture M comprises a sequence model f θ and multiple heads for the prediction of tokens of different observation modalities, rewards, termination signals, and for the estimation of model uncertainty. Concretely, f θ is a retentive network (RetNet) [51] augmented with a parallel observation prediction (POP) [11] mechanism. All heads are implemented as multilayer perceptrons (MLP) with a single hidden layer. We defer the details about these architectures to Appendix A.3.

Embedding M translates token trajectories τ into sequences of d -dimensional embeddings X using a set of embedding (look-up) tables. By design, each modality is associated with a separate

table. In cases where an embedding table is not provided by the appropriate encoder-decoder pair, M 116 and C learn dedicated tables separately and independently. As embeddings sequences are composed 117 hierarchically, we use the following hierarchical notation: 118

where K i denotes the number of embedding vectors in X ( i ) t . Similarly, K = ∑ | κ | i =1 K i . To combine 119 latents of each z t , V concatenates their token sequences along the temporal axis based on a predefined 120 modality order. We defer the full details on the embedding process to Appendix A.3. 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

149

150

151

152

153

154

155

156

157

158

159

Sequence Modeling Given a sequence of observation-action blocks X = X 1 , . . . , X t , the matching outputs Y 1 , . . . , Y t are computed auto-regressively as follows:

where S t is a recurrent state that summarizes X ≤ t and S 0 = 0 . However, the output Y u t +1 , from which ˆ z t +1 is predicted, is computed using the POP mechanism via another call as

where X u ∈ R K × d is a learned embedding sequence. Intuitively, X u acts as a learned prior, enabling the parallel generation of multiple tokens into the future.

To model p θ (ˆ z t +1 | Y u t +1 ) , the distributions p θ (ˆ z | y ) of each token ˆ z of ˆ z ( i ) t +1 of each modality κ i are modeled using modality-specific prediction heads implemented as MLPs with a single hidden layer and an output size equal to the vocabulary size of enc i (Figure 2b). For 2D categoricals, C heads are used to predict the C tokens from each y .

Sampling

Similarly, rewards and termination signals are predicted by additional prediction heads as ˆ r t = ˆ r θ ( y ) , ˆ d t ∼ p θ ( ˆ d t | y ) , slightly abusing notations, where y is the last vector of Y u t +1 . An illustration is provided in Figure 3.

Reducing Epistemic Uncertainty via Intrinsic Motivation The world model M serves as the cornerstone of the entire system. Any controller operating within a world model framework can only perform as well as the underlying world model allows, making its quality a fundamental limiting factor. In deep learning methods, the model's performance depends heavily on the quality of its training data. Accurate dynamics modeling requires comprehensive data collection that captures the full spectrum of possible environmental behaviors. This presents a particular challenge in online RL, where the controller must systematically and efficiently explore its environment. Success depends on intelligently guiding the controller toward unexplored or undersampled regions of the dynamics space. An effective approach to this challenge involves estimating the world model's epistemic uncertainty and directing the controller to gather data from regions where this uncertainty is highest [46, 50, 49].

Our approach estimates epistemic uncertainty using an ensemble of N ens = 4 next observation prediction heads { p ϕ i (ˆ z | sg( Y u )) } N ens i =1 with parameters { ϕ i } N ens i =1 [49, 31] where sg( · ) is the stop gradient operator. To quantify disagreement between the ensemble's distributions, we employ the Jensen-Shannon divergence (JSD) [50]. For probability distributions P 1 , . . . , P n , the JSD is defined as:

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

World model training

Imagination

Input / Output

Optimization

Sampling

where H ( · ) denotes the Shannon entropy. Since observations comprise multiple tokens, we average the per-token JSD values to obtain a single uncertainty measure δ t . Training data is divided equally among ensemble members, with each predictor processing a distinct subset of each batch. Despite the ensemble approach, our implementation maintains computational efficiency, with negligible additional overhead in practice.

To guide C towards regions of high epistemic uncertainty, M provides C with additional intrinsic rewards r int t = δ t during imagination. Here, the reward provided by M at each step t is given by

where w int , w ext ∈ R are hyperparameters that control the scale of each reward type. Optimizing the controller in imagination allows it to reach areas of high model uncertainty without additional real-environment interaction.

Prioritized Replay Recent work has demonstrated that prioritizing replay buffer sampling during world model training could lead to significant performance gains in intrinsically motivated agents [30]. While their approach showed promise, it required extensive hyperparameter tuning in practice. We propose a simpler, more robust prioritization scheme for world model training.

Here, the replay buffer maintains a world model loss value for each stored example, with newly added examples assigned a high initial loss value ν 0 . During M 's training, we sample each batch using a mixture of uniform and prioritized sampling, controlled by a single parameter α ∈ [0 , 1] that determines the fraction of prioritized samples. For the prioritized portion, we sample examples proportional to their softmax-transformed losses p i = softmax( L ) i . The loss values are updated after each world model optimization step using the examples' current batch losses.

Training We use the cross-entropy loss for the optimization of all components of M . Specifically, for each t , the loss of p θ (ˆ z t | Y u t ) is obtained by averaging the cross-entropy losses of its individual tokens. The same loss is used for each ensemble member p ϕ i (ˆ z t | sg( Y u t )) . The optimization and design of the reward predictor is similar to that of the critic, as described in Section 2.3. A formal description of the optimization objective can be found in Appendix A.3.1.

2.3 The Controller C

C is an extended version of the actor-critic of REM [11] that supports additional observation and action spaces and implements regression-as-classification for return predictions.

Architecture At the core of C 's architecture, parameterized by ψ , is an LSTM [26] sequence model. 188 At each step t , upon observing z t , a set of modality-specific encoders map each modality tokens z ( i ) t to 189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

( i )

a latent vector x , where we abuse our notation x as the context of the discussion is limited to C . The latents are then fused by a fully-connected network to obtain a single vector x = g ψ ( x (1) , . . . , x ( | κ | ) ) . x t ∈ R d C is processed by C 's sequence model to produce h t , c t = LSTM( x t , h t -1 , c t -1 ; ψ ) where h t , c t are the LSTM's hidden and cell states, respectively. Lastly, two linear output layers produce the logits from which the actor and critic outputs π ( a t | h t ) , ˆ V π ( h t ) are derived. For continuous action spaces, the actor uses a categorical distribution over a uniformly spaced discrete subset of [ -1 , 1] . We defer the full details about the encoding process to Appendix A.4.

Regression as Classification for Reward and Return Prediction Robustly handling unbounded reward signals has long been challenging as they can vary dramatically in both magnitude and frequency. Hafner et al. [20] addressed this challenge by using a classification network that predicts the weights of exponentially spaced bins and employed a two-hot loss for the network's optimization. Farebrother et al. [16] studied the use of cross-entropy loss in place of the traditional mean squared error loss for value-based deep RL methods. In their work, the HL-Gauss method was shown to significantly outperform the two-hot loss method. Building on these developments, we adopt the classification network with exponential bins from [20], and apply the HL-Gauss method for its optimization. Concretely, the critic's value estimates are predicted using a linear output layer parameterized by W ∈ R m × d C with m = 128 outputs corresponding to m uniform bins defined by m +1 endpoints b = ( b 0 , . . . , b m ) . The predicted value is given by

where symexp( x ) = sign( x )(exp( | x | ) -1) is the inverse of the symlog function and ˆ b = ( b 1 + b 0 2 , . . . , b m + b m -1 2 ) are the bin centers. Given the true target y ∈ R , the HL-Gauss loss is given by

where ˜ y i = Φ( b i -symlog( y ) σ ) -Φ( b i -1 -symlog( y ) σ ) , Φ is the cumulative density function of the standard normal distribution and σ is a standard deviation hyperparameter that controls the amount of label smoothing.

Training in Imagination C is trained entirely from simulated experience generated through interaction with M . Specifically, M and C are initialized with a short trajectory segment sampled uniformly from the replay buffer and interact for H steps. An illustration of this process is given in Figure 3 (orange path). λ -returns are computed for each generated trajectory segment and are used as targets for critic learning. For policy learning, a REINFORCE [52] objective is used, with a ˆ V π baseline for variance reduction. See Appendix A.4.2 for further details.

3 Experiments

To evaluate sample efficiency, we used benchmarks that measure performance within a fixed, limited environment interaction budget. These benchmarks were also selected to address key research questions: (1) whether the proposed multi-modality approach is effective-both in continuous control settings and in handling multi-modal observations; and (2) whether the integrated components are effective across diverse environments (Section 3.3).

3.1 Experimental Setup

Benchmarks: We evaluate Simulus on three sample-efficiency benchmarks of different observation and action modalities: Atari 100K [28], DeepMind Control Suite (DMC) Proprioception 500K [53], and Craftax-1M [35].

Atari 100K has become the gold standard in the literature for evaluating sample-efficient deep RL agents. The benchmark comprises a subset of 26 games. Within each game, agents must learn from visual image signal under a tightly restricted budget of 100K interactions, corresponding to roughly two hours of human gameplay.

The DeepMind Control Suite (DMC) is a set of continuous control tasks involving multiple agent embodiments ranging from simple single-joint models to complex humanoids. Here, we follow the

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

subset of proprioception tasks used for the evaluation of DreamerV3 [20], where observations and actions are continuous vectors. At each task, the agent's interaction budget is limited to 500K steps.

Craftax is a 2D open-world survival game benchmark inspired by Minecraft, designed to evaluate RL agents' capabilities in planning, memory, and exploration. The partially-observable environment features procedurally generated worlds where agents must gather and craft resources while surviving against hostile creatures. Observations consist of a 9×11 tile egocentric map, where each tile consists of 4 symbols, and 48 state features corresponding to state information such as inventory and health. Here, we consider the sample-efficiency oriented Craftax-1M variant which only allows an interaction budget of one million steps.

Baselines On Atari-100K, we compare Simulus against DreamerV3 [20] and several methods restricted to image observations: TWM [44], STORM [60], DIAMOND [3], and REM [11]. On DMC, we compare exclusively with DreamerV3, currently the only planning-free world model method with published results on the 500K proprioception benchmark. On Craftax-1M, we compare against TWM [13], a concurrent work that proposes a Transformer based world model with a focus on the Craftax benchmark, and the baselines reported in the Craftax paper: Random Network Distillation (RND) [10], PPO [48] with a recurrent neural network (PPO-RNN), and Exploration via Elliptical Episodic Bonuses (E3B) [23]. As Craftax is a recent benchmark, there are no other published results in existing world models literature. Following the standard practice in the literature, we exclude planning-based methods [21, 57], as planning is an orthogonal component that operates on any given model, typically incurring significant computational overhead.

Metrics and Evaluation For Atari, we report human-normalized scores (HNS), calculated as agent_score -random_score human_score -random_score [39]. Following the protocol of Agarwal et al. [2] and using their toolkit, we report the mean, median, interquantile mean (IQM), and optimality gap metrics with 95% stratified bootstrap confidence intervals. For DMC and Craftax, we report the raw agent returns. We use 5 random seeds per environment. In each experiment, final performance is evaluated using 100 test episodes at the end of training and the mean score is reported.

3.2 Results

Simulus achieves state-of-the-art performance across all three benchmarks (Figure 1). On Atari 100k, it outperforms all baselines across key metrics (Figure 4). Notably, Simulus is the first planning-free world model to reach human-level IQM and median scores, achieving superhuman performance

Ball In Cup Catch

Finger Turn Hard

Pendulum Swingup

Reacher Easy

Reacher Hard

Walker Walk

Mean

0

200

400

600

800

1000

Average Return

125

960

993

1000

862

798

668

875

882

769

271

905

832

959

844

629

980

970

796

134

962

990

990

852

491

614

931

793

889

113

576

788

954

938

649

964

936

754

Simulus

1.00

1.25

1.50

1.75

TWM

STORM

DIAMOND

REM

Simulus (ours)

Mean (

)

0.25

0.50

0.75

1.00

Median (

)

0.6

0.8

1.0

Interquartile Mean (

)

0.40

0.44

0.48

0.52

Optimality Gap (

)

Human Normalized Score

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

on 13 out of 26 games (Table 9, Appendix B). Building on REM, these gains are attributed to the integration of the proposed components, demonstrating their combined effectiveness.

Effectiveness in continuous environments Figure 4 provides compelling evidence that token-based architectures can perform well in continuous domains-even with compact vocabularies: Simulus consistently matches DreamerV3 across most tasks and slightly outperforms it on average.

Effectiveness in environments with multi-modal observations

We evaluate multi-modality performance in Craftax, as it combines an image-like 2D grid map with a vector of features, involving multiple tokenizers ( V ). Simulus maintains sample-efficiency in this multi-modal environment, outperforming both concurrent world model methods (Figure 1) and all model-free baselines (Figure 5), including exploration-focused algorithms. With 444 tokens per observation arranged into sequences of 147 embeddings, even short trajectories in Craftax contain thousands of tokens, demonstrating Simulus's efficient handling of long sequences. These findings indicate that the world model ( M ) and controller ( C ) maintain strong performance under multi-modal inputs when processed by the proposed modular tokenizer.

3.3 Ablation Studies

We ablate the intrinsic rewards, prioritized replay, and regression-as-classification to demonstrate their individual contributions to Simulus's performance. In each experiment, Simulus is modified by disabling a single component. Due to limited computational resources, we consider a subset of 8 tasks for each of the Atari and DMC benchmarks, and exclude Craftax from this analysis. For Atari 100K, we used games in which significant improvements were observed. For DMC, we chose a subset that includes different embodiments. We defer the specific environment names to Appendix C.

The results are presented in Figure 6. Although all components contributed to Simulus's final performance, intrinsic rewards were especially crucial for achieving competitive performance in both benchmarks. Interestingly, the Atari 100K results indicate that combining all three components yields a significantly stronger algorithm. These findings also suggest that both prioritized world model replay and regression-as-classification enhance the effectiveness of intrinsic rewards.

More broadly, the results in Figure 6 demonstrate that encouraging the controller to explore regions of high epistemic uncertainty through intrinsic rewards significantly improves its performance in world model agents, even in reward-rich environments. This observation is non-trivial in a sample-efficient setting, where the limited interaction budget makes model-driven exploration particularly costly, as it consumes resources that could otherwise be used for task-related exploration during data collection. The latter type of exploration aims to collect new information about the true reward signal, which defines the task and its success metric. On the other hand, model-driven exploration may guide the controller towards environment regions that are irrelevant to the task at hand.

0.5

0.7

0.9

P(Simulus

>

Y

)

Simulus

REM

No Intrinsic Rewards

No Prioritized Replay

MSE Regression

Ablation

Y

1.2

1.6

2.0

2.4

IQM Human

Normalized Score (↑)

Atari 100K

0.5

0.7

0.9

P(Simulus

>

Y

)

500

600

700

800

DMC Proprioception 500K

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps

1e6

0

2

4

6

8

10

12

14

16

Average Score

Simulus

PPO-RNN

RND

E3B

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

4 Related Work

Offline Multi-Modal Methods Large-scale token sequence models for multi-modal agent trajectories have been proposed in [33, 34, 43, 47]. Gato [43] and TDM [47] tokenize inputs via predefined transformations, while Unified IO [33, 34] leverages pretrained models. These methods do not learn control through RL but rely on expert data. They also use massive models-with billions of parameters, large vocabularies, and significantly more data and compute than sample-efficient world models. Consequently, it remains unclear whether their design choices would be effective in online, sample-efficient settings with non-stationary and limited data.

World Model Agents Model-based agents that learn policies solely from simulated data generated by a learned world model were introduced by Ha and Schmidhuber [17], followed by the influential Dreamer family [18-20]. Dreamer jointly optimizes its representation and recurrent world models via a KL divergence between learned prior and posterior distributions, leading to interfering objectives (Appendix B.1) and a complex, monolithic architecture that complicates development and scaling [59]. With the rise of Transformer architectures in language modeling [56, 9], Transformer-based Dreamer variants emerged [60, 44], alongside token-based world models (TBWMs) that treat trajectories as discrete token sequences [37, 11]. However, these methods are limited to visual environments with discrete actions (e.g., Atari 100K), leaving their performance in other modalities uncertain. Recently, DIAMOND [3], a diffusion world model inspired by advances in generative modeling [45], was introduced. While it generates visually compelling outputs, it remains limited to visual domains.

Intrinsically Motivated World Model Agents Although intrinsic motivation (IM) has been extensively studied [41, 10, 22, 5, 4], its use in world models typically involves an exploration pretraining phase followed by limited task-specific fine-tuning [49, 36, 30]. While combining IM with prioritized replay has shown promise [30], it remains unexplored in standard sample-efficiency settings with external rewards.

Large-Scale Video World Models Building on recent advances in video generative modeling [25, 7, 8], recent works have introduced large-scale video world models [14, 1, 12, 54], trained offline on extensive pre-collected data to predict future frames. However, these models do not address control learning, particularly RL. While recent efforts aim to bridge this gap [59], they remain confined to visual environments and lack comprehensive empirical evaluation.

5 Limitations and Future Work

Here, we briefly highlight several limitations of this work. First, although the feature quantization approach for tokenizing continuous vectors showed promise, it leads to excessive sequence lengths. We believe that more efficient solutions can be found for dealing with continuous inputs. Second, due to the scarcity of rich multi-modal RL benchmarks, we could not extensively explore diverse modality combinations in our experiments. Lastly, token-based world model agents remain significantly slower to train than other baselines in sample-efficient RL. Nonetheless, their modular design enables faster policy inference as the controller is independent of the world model.

6 Conclusions

In this paper, we demonstrated the effectiveness of several underutilized techniques for improving world model agents. A modular multi-modality tokenization framework broadens their applicability across diverse domains, while intrinsic motivation, prioritized world model replay, and regressionas-classification enhance sample efficiency, particularly when combined. These techniques were incorporated into a token-based world model, yielding Simulus. Extensive experiments show that Simulus achieves state-of-the-art performance across diverse benchmarks, including visual, continuous, and symbolic domains. It outperforms all baselines on key metrics in both Atari-100K and the challenging Craftax benchmarks. Ablations further highlight the individual contribution of each component. We hope that the highly modular design of Simulus, along with the released code and model weights, provides a strong foundation for future work.

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

References

Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chattopadhyay, Yongxin Chen, Yin Cui, Yifan Ding, et al. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501.03575 , 2025.

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 29304-29320. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ f514cec81cb148559cf475e7426eed5e-Paper.pdf .

Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. arXiv preprint arXiv:2405.12399 , 2024.

Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the Atari human benchmark. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 507-517. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/ badia20a.html .

Adrià Puigdomènech Badia, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, Bilal Piot, Steven Kapturowski, Olivier Tieleman, Martin Arjovsky, Alexander Pritzel, Andrew Bolt, and Charles Blundell. Never give up: Learning directed exploration strategies. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id= Sye57xStvB .

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:2311.15127 , 2023.

Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL https://openai.com/research/ video-generation-models-as-world-simulators .

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ 1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf .

Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=H1lJJnR5Ym .

Lior Cohen, Kaixin Wang, Bingyi Kang, and Shie Mannor. Improving token-based world models with parallel observation prediction. In Forty-first International Conference on Machine Learning , 2024. URL https://openreview.net/forum?id=Lfp5Dk1xb6 .

Decart, Julian Quevedo, Quinn McIntyre, Spruce Campbell, Xinlei Chen, and Robert Wachen. Oasis: A universe in a transformer, 2024. URL https://oasis-model.github.io/ .

Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop 405 Guntupalli, Miguel Lazaro-Gredilla, and Kevin Patrick Murphy. Improving transformer world 406 models for data-efficient rl, 2025. URL https://arxiv.org/abs/2502.01591 . 407

408

409

410

Google DeepMind. Genie 2: A large-scale foundation world model, 2024. URL https://deepmind.google/discover/blog/genie-2-a-large-scale-foundationworld-model/ .

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12873-12883, 2021.

Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, Aviral Kumar, and Rishabh Agarwal. Stop regressing: Training value functions via classification for scalable deep RL. In Forty-first International Conference on Machine Learning , 2024. URL https://openreview. net/forum?id=dVpFKfqF3R .

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems 31 , pages 2451-2463. Curran Associates, Inc., 2018. URL https://papers.nips.cc/paper/7512-recurrent-world-modelsfacilitate-policy-evolution . https://worldmodels.github.io .

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=S1lOTC4tDS .

Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https: //openreview.net/forum?id=0oabwyZbOu .

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=Oxh5CstDJU .

Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2681-2691. PMLR, 09-15 Jun 2019. URL https://proceedings. mlr.press/v97/hazan19a.html .

Mikael Henaff, Roberta Raileanu, Minqi Jiang, and Tim Rocktäschel. Exploration via elliptical episodic bonuses. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https: //openreview.net/forum?id=Xg-yZos9qJQ .

Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units, 2017. URL https://openreview.net/forum?id=Bk0MRI5lg .

Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 , 2022.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8): 1735-1780, 1997.

Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part II 14 , pages 694-711. Springer, 2016.

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Bla˙ zej Osi´ nski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations , 2020. URL https: //openreview.net/forum?id=S1xCPJHtDB .

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5156-5165. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/katharopoulos20a.html .

Isaac Kauvar, Chris Doyle, Linqi Zhou, and Nick Haber. Curious replay for model-based adaptation. In Proceedings of the 40th International Conference on Machine Learning , ICML'23. JMLR.org, 2023.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/2017/file/ 9ef2ed4b7fd2c810847ffa5fa85bce38-Paper.pdf .

Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Autoencoding beyond pixels using a learned similarity metric. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 1558-1566, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://proceedings.mlr.press/ v48/larsen16.html .

Jiasen Lu, Christopher Clark, Rowan Zellers, Roozbeh Mottaghi, and Aniruddha Kembhavi. UNIFIED-IO: A unified model for vision, language, and multi-modal tasks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview. net/forum?id=E01k9048soZ .

Jiasen Lu, Christopher Clark, Sangho Lee, Zichen Zhang, Savya Khosla, Ryan Marten, Derek Hoiem, and Aniruddha Kembhavi. Unified-io 2: Scaling autoregressive multimodal models with vision language audio and action. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 26439-26455, June 2024.

Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Jackson, Samuel Coward, and Jakob Foerster. Craftax: A lightning-fast benchmark for open-ended reinforcement learning. In International Conference on Machine Learning (ICML) , 2024.

Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models. Advances in Neural Information Processing Systems , 34:24379-24391, 2021.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023 . OpenReview.net, 2023. URL https://openreview.net/ pdf?id=vhFu1Acb0xb .

Vincent Micheli, Eloi Alonso, and François Fleuret. Efficient world models with context-aware tokenization. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 . OpenReview.net, 2024. URL https://openreview.net/forum? id=BiWIERWBFX .

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems , volume 32, 2019.

Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning , volume 70 of Proceedings of Machine Learning Research , pages 2778-2787. PMLR, 06-11 Aug 2017. URL https://proceedings.mlr.press/v70/pathak17a.html .

Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2018. URL https://openreview.net/forum?id=SkBYYyZRZ .

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. URL https://openreview.net/ forum?id=1ikK0kHjvj . Featured Certification, Outstanding Certification.

Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. arXiv preprint arXiv:2303.07109 , 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10684-10695, June 2022.

Jürgen Schmidhuber. Formal theory of creativity, fun, and intrinsic motivation (1990-2010). IEEE Transactions on Autonomous Mental Development , 2, 2010. ISSN 19430604. doi: 10.1109/TAMD.2010.2056368.

Ingmar Schubert, Jingwei Zhang, Jake Bruce, Sarah Bechtle, Emilio Parisotto, Martin Riedmiller, Jost Tobias Springenberg, Arunkumar Byravan, Leonard Hasenclever, and Nicolas Heess. A generalist dynamics model for control. arXiv preprint arXiv:2305.10912 , 2023.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 8583-8592. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/sekar20a.html .

Pranav Shyam, Wojciech Ja´ skowski, and Faustino Gomez. Model-based active exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 5779-5788. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr.press/v97/ shyam19a.html .

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621 , 2023.

Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In S. Solla, T. Leen, and K. Müller, editors, Advances in Neural Information Processing Systems , volume 12. MIT Press, 1999. URL https://proceedings.neurips.cc/paper_files/paper/1999/ file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf .

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts , 6:100022, 2020. ISSN 2665-9638. doi: https:// doi.org/10.1016/j.simpa.2020.100022. URL https://www.sciencedirect.com/science/ article/pii/S2665963820300099 .

Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real-time game engines. CoRR , abs/2408.14837, 2024. URL https://doi.org/10.48550/ arXiv.2408.14837 .

Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu. Neural discrete representation learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/ 2017/file/7a98af17e63a0ac09ce2e96d03992fbc-Paper.pdf .

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Shengjie Wang, Shaohuai Liu, Weirui Ye, Jiacheng You, and Yang Gao. Efficientzero v2: Mastering discrete and continuous control with limited data. In Forty-first International Conference on Machine Learning , 2024. URL https://openreview.net/forum?id=LHGMXcr6zx .

ZiRui Wang, Yue DENG, Junfeng Long, and Yin Zhang. Parallelizing model-based reinforcement learning over the sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. URL https://openreview.net/forum?id= R6N9AGyz13 .

Jialong Wu, Shaofeng Yin, Ningya Feng, Xu He, Dong Li, Jianye HAO, and Mingsheng Long. ivideoGPT: Interactive videoGPTs are scalable world models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. URL https://openreview. net/forum?id=4TENzBftZR .

Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. Advances in Neural Information Processing Systems , 36, 2024.

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

A Models and Hyperparameters

A.1 Hyperparameters

We detail shared hyperparameters in Table 1, training hyperparameters in Table 2, world model hyperparameters in Table 3, and controller hyperparameters in Table 4. Environment hyperparameters are detailed in Table 5 (Atari-100K) and Table 6 (DMC).

For the DMC benchmark, we use a lower embedding dimension (Table 3) due to the significantly lower dimensionality of its observations compared to other benchmarks. As in prior work (e.g., DreamerV3), we adopt a smaller model for this setting. Note that the reduced number of Retention heads ensures a consistent head dimensionality (64).

Additionally, we used a limited decay range in DMC (Table 3) as observations effectively represent full MDP states, eliminating the need for long-term memory. By constraining the decay range, we explicitly encode this inductive bias into the model.

In the Craftax benchmark, we reduce the number of layers (Table 3) to lower computational cost. The interaction budget in Craftax is 1M steps, resulting in a particularly expensive training process. Here, we increase the decay range as an inductive bias to encourage long-term memory.

The weighting of intrinsic versus extrinsic rewards (Table 4) varies across benchmarks due to differences in reward structure and scale. For instance, DMC provides dense rewards with typical task scores reaching around 1000, whereas Craftax has extremely sparse rewards, with cumulative scores rarely exceeding 20-even over thousands of steps.

Tuning All remaining hyperparameters were tuned empirically or based on REM [11], with minimal impact on training cost and adjusted primarily for performance. Due to limited computational resources, we were unable to conduct extensive tuning, and we believe that further optimization could improve Simulus's performance.

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

A.2 The Representation Module V

A.2.1 Image Observations

Image observations are tokenized using a vector-quantized variational auto-encoder (VQ-VAE) [55, 15]. A VQ-VAE comprises a convolutional neural network (CNN) encoder, an embedding table E ∈ R n × d , and a CNN decoder. Here, the size of the embedding table n determines the vocabulary size.

The encoder's output h ∈ R W × H × d is a grid of W × H multi-channel vectors of dimension d that encode high-level learned features. Each such vector is mapped to a discrete token by finding the closest embedding in E :

where E ( i ) is the i -th row of E . To reconstruct the original image, the decoder first maps z to their embeddings using E . During training, the straight-through estimator [6] is used for backpropagating the learning signal from the decoder to the encoder: ˆ h = h +sg( E z -h ) . The architecture of the encoder and decoder models is presented in Table 7.

The optimization objective is given by

where L perceptual is a perceptual loss [27, 32], proposed in [37].

Crucially, the learned embedding table E is used for embedding the (image) tokens across all stages of the algorithm.

A.2.2 Continuous Vectors

The quantization of each feature uses 125 values (vocabulary size) in the range [ -6 , 6] , where 63 values are uniformly distributed in [ -ln(1 + π ) , ln(1 + π )] and the rest are uniformly distributed in the remaining intervals.

A.3 The World Model M 632

633

634

635

636

637

638

Embedding Details Each token in z ( i ) of each modality is mapped to a d -dimensional embedding vector X ( i ) using the embedding (look-up) table E ( i ) of modality κ i . The embedding vector that corresponds to token z is simply the z -th row in the embedding table. Formally, x ( i ) t,j = E ( i ) ( l ) , l = z ( i ) t,j where E ( l ) refers to the l -th row in E . In the special case of 2D categorical inputs, x ( i ) t,j = 1 C ∑ C n =1 E ( i ) n ( l n ) , l n = z ( i ) t,j,n where C is the number of channels and i is the index of the 2D categorical modality in κ .

639

640

641

642

643

644

645

646

647

648

To concatenate the embeddings, we use the following order among the modalities: images, continuous vectors, categorical variables, and 2D categoricals.

Prediction Heads Each prediction head in M is a multi-layer perceptron (MLP) with a single hidden layer of dimension 2 d where d is the embedding dimension.

Epistemic Uncertainty Estimation Working with discrete distributions enables efficient entropy computation and ensures that the ensemble disagreement term δ t is bounded by 1 | z t | ∑ z ∈ z t log(vocab _ size( z )) .

A.3.1 Optimization

For each training example in the form of a trajectory segment in token representation τ = z 1 , z a 1 , . . . , z H , z a H , the optimization objective is given by

where 649

650

651

652

653

654

655

656

657

658

659

660

is the average of the cross-entropy losses of the individual tokens, and L reward ( θ, r t , ˆ r t ) is the L HL-Gauss loss with the respective parameters of the reward head. Here, y i is the vector of Y u t that corresponds to z i , the i -th token of z t .

A.3.2 Retentive Networks

Retentive Networks (RetNet) [51] are sequence model architectures with a Transformer-like structure [56]. However, RetNet replaces the self-attention mechanism with a linear-attention [29] based Retention mechanism. At a high level, given an input sequence X ∈ R | X |× d of d -dimensional vectors, the Retention operator outputs

where Q , K , V are the queries, keys, and values, respectively, and D is a causal mask and decay matrix. Notably, the softmax operation is discarded in Retention and other linear attention methods. As a linear attention method, the computation can also be carried in a recurrent form:

where η is a decay factor, S t is a recurrent state, and S 0 = 0 . In addition, a hybrid form of recurrent 661 and parallel forward computation known as the chunkwise mode allows to balance the quadratic 662 cost of the parallel form and the sequential cost of the recurrent form by processing the input as a 663 sequence of chunks. We refer the reader to [51] for the full details about this architecture. 664

665

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

682

683

684

685

686

687

In our implementation, since inputs are complete observation-action block sequences X 1 , . . . , X t , we configure the decay factors of the multi-scale retention operator in block units:

where linspace( a, b, c ) is a sequence of c values evenly distributed between a and b , N h is the number of retention heads, and η min , η max are hyperparameters that control the memory decay in observation-action block units.

A.3.3 Parallel Observation Prediction (POP)

POP [11] is a mechanism for parallel generation of non-causal subsequences such as observations in token representation. It's purpose is to improve generation efficiency by alleviating the sequential bottleneck caused by generating observations a single token at a time (as done in language models). However, to achieve this goal, POP also includes a mechanism for maintaining training efficiency. Specifically, POP extends the chunkwise forward mode of RetNet to maintain efficient training of the sequence model.

To generate multiple tokens into the future at once, POP introduces a set of prediction tokens u = u 1 , . . . , u K and embeddings X u ∈ R K × d where K is the number of tokens in an observation. Each token in u corresponds to an observation token in z . These tokens, and their respective learned embeddings, serve as a learned prior.

Let X 1 , . . . , X T be a sequence of T observation-action (embeddings) blocks. Given S t -1 summarizing all key-value outer products of elements of X ≤ t -1 , the outputs Y u from which the next observation tokens are predicted are given by:

Importantly, the recurrent state is never updated based on the prediction tokens u (or their embeddings). The next observation tokens ˆ z t are sampled from p θ (ˆ z t | Y u t ) . Then, the next action is generated by the controller, and the next observation-action block X t can be processed to predict the next observation ˆ z t +1 .

To maintain efficient training, a two step computation is carried at each RetNet layer. First, all 688 recurrent states S t for all 1 ≤ t ≤ T are calculated in parallel. Although there is an auto-regressive 689 relationship between time steps, the linear structure of S allows to calculate the compute-intensive 690 part of each state in parallel and incorporate past information efficiently afterwards. In the second 691 step, all outputs Y u t for all 1 ≤ t ≤ T are computed in parallel, using the appropriate states S t -1 692 and X u in batch computation. Note that this computation involves delicate positional information 693

handling. We refer the reader to [11] for full details of this computation. 694

A.4 The Controller C 695

696

Critic The value prediction uses 128 bins in the range b = ( -15 , . . . , 15) .

697

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

Continuous Action Spaces The policy network outputs m = 51 logits corresponding to m quantization values uniformly distributed in [ -1 , 1] for each individual action in the action vector.

A.4.1 Input Encoding

The controller C operates in the latent token space. Token trajectories τ = z 1 , z a 1 , . . . , z H , z a H are processed sequentially by the LSTM model. At each time step t , the network gets z t as input, outputs π ( a t | τ ≤ t -1 , z t ) and ˆ V π ( a t | τ ≤ t -1 , z t ) , samples an action a t and then process the sampled action as another sequence element.

The processing of actions involve embedding the action into a latent vector which is then provided as input to the LSTM. Embedding of continuous action tokens is performed by first reconstructing the continuous action vector and then computing the embedding using a linear projection. Discrete tokens are embedded using a dedicated embedding table.

To embed observation tokens z , the tokens of each modality are processed by a modality-specific encoder. The outputs of the encoders are concatenated and further processed by a MLP g ψ that combines the information into a single vector latent representation.

The image encoder is a convolutional neural network (CNN). Its architecture is given in Table 8.

Categorical variables are embedded using a learned embedding table. For 2D categoricals, shared per-channel embedding tables map tokens to embedding vectors, which are averaged to obtain a single embedding for each multi-channel token vector. For both types of categorical inputs we use 64 dimensional embeddings. The embeddings are concatenated and processed by g ψ .

A.4.2 Optimization 716

λ -returns are computed for each generated trajectory segment ˆ τ = 717 ( z 1 , a 1 , ¯ r 1 , d 1 , ˆ z 2 , a 2 , ¯ r 2 , d 2 , . . . , ˆ z H , a H , ¯ r H , d H ) : 718

where ˆ V π t = ˆ V π (ˆ τ ≤ t ) . These λ -returns are used as targets for critic learning. For policy learning, a REINFORCE [52] objective is used, with a normalized ˆ V π baseline for variance reduction:

719

720

where c is an estimate of the effective return scale similar to DreamerV3 [20] and w ent is a hyper721 parameter that controls the entropy regularization weight. c is calculated as the difference between 722 the running average estimators of the 97 . 5 and 2 . 5 return percentiles, based on a window of return 723 estimates obtained in the last 500 batches (imagination). 724

725

726

727

B Additional Results

The average per-game scores for Atari-100K are presented in Table 9. The performance profile plot for Atari 100K is presented in Figure 7.

0

1

2

3

4

5

6

7

8

Human Normalized Score

(

)

0.00

0.25

0.50

0.75

1.00

Fraction of runs with score

>

TWM

STORM

DIAMOND

REM

Simulus (ours)

B.1 Interfering Objectives in RSSM Optimization 728

729

730

731

732

733

734

Here, we study the interplay between the objectives of a Dreamer-like world model, PWM [58], which uses a slightly modified version of the recurrent state space model (RSSM) of Dreamer. Concretely, we aim to understand whether the representation and sequence modeling objectives interfere by decoupling the optimization of the encoder-decoder models from that of the world model. We opted for this implementation due to its simplicity, fast runtime, and accessibility, as it is written in PyTorch [40].

735

Formally, the model consists of the following components:

We omit the reward and termination predictors and objectives for brevity. Note that in contrast to the 736 RSSM in Dreamer, the encoder and decoder models do not depend on the recurrent state h t , x t . The 737 optimization objective of PWM is given by 738

where β pred , β dyn, and β rep are coefficients and 739

To decouple the optimization, we modify the sequence model by introducing a stop-gradient operator 740 on the encoder's output during world model training: 741

Moreover, this modification allows to train the encoder-decoder models using large batches of single frames, rather than small highly-correlated batches of long trajectories. This further highlights the flexibility and advantage of a modular design.

We compare the original PWM algorithm to its decoupled variant, PWM-decoupled, across four Atari environments: Breakout , DemonAttack , Hero , and RoadRunner . These are games where PWM performed either particularly well ( Hero and RoadRunner ) or poorly ( Breakout and DemonAttack ). Each variant is trained online from scratch on each game. The results are presented in Figure 8. In addition, we present the achieved episodic returns in Figure 9, and the reconstruction quality of example episodes in Figure 10 (PWM) and Figure 11 (PWM-decoupled).

Although our results are based on a single random seed and are limited to only four environments, we observe a consistent trend. First, the reconstruction losses are consistently and significantly lower when decoupling the optimization, while the dynamics losses are significantly higher. This suggests that the objectives are interfering.

Second, we observe similar or better episodic returns (Figure 9) using the decoupled optimization, suggesting that the higher dynamics loss might not lead to worse world modeling performance in practice. Note that a higher dynamics loss in this case does not necessarily mean worse performance, as for example multiple discrete combinations could represent the same or similar frame. Thus, when the dynamics model fails to predict a specific combination, it leads to high loss values while the underlying representations are accurate.

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

Lastly, we report that similar trends were observed when training only the world model in an offline, supervised-learning fashion on pre-collected datasets. We explored this setting to eliminate complexities that may arise due to the online collection of the data.

While the presented preliminary results are noisy and limited, we believe that they uncover an 764 interesting observation on the design and optimization of current world models. 765

Step

10

2

10

1

10 0

Reconstruction Loss

Breakout

Step

5

10

15

Dynamics Loss

Breakout

Step

10 0

Reconstruction Loss

DemonAttack

Step

0

10

20

30

Dynamics Loss

DemonAttack

Step

10

1

10 0

Reconstruction Loss

Hero

Step

5

10

15

Dynamics Loss

Hero

0

20000

40000

60000

80000

100000

Step

10 0

3×10 1

4×10 1

6×10 1

2×10 0

Reconstruction Loss

RoadRunner

0

20000

40000

60000

80000

100000

Step

5

10

15

20

Dynamics Loss

RoadRunner

PWM-decoupled

PWM

0

20000

40000

60000

80000

100000

Step

0

10

20

return

Breakout

0

20000

40000

60000

80000

100000

Step

0

200

400

return

DemonAttack

0

20000

40000

60000

80000

100000

Step

0

5000

10000

return

Hero

0

20000

40000

60000

80000

100000

Step

0

10000

20000

return

RoadRunner

PWM-decoupled

PWM

003

UEE

UEE

0E94

0094

TEE

EEE

766

767

768

769

770

771

772

773

774

775

776

777

778

779

C Implementation Details

Ablation Studies For the Atari 100K benchmark, we conducted ablations on the following games: Assault , Breakout , ChopperCommand , CrazyClimber , JamesBond , Kangaroo , Seaquest , and UpNDown . For the DeepMind Control Suite, we used the tasks: acrobot-swingup , cartpole-swingup-sparse , cheetah-run , finger-turn-hard , hopper-stand , pendulum-swingup , reacher-hard , and walker-run .

Code We open-source our code and trained model weights. Our code is written in Pytorch [40].

Hardware All Atari and DMC experiments were performed on V100 GPUs, while for Craftax a single RTX 4090 was used.

Run Times Experiments on Atari require approximately 12 hours on an RTX 4090 GPU and around 29 hours on a V100 GPU. For DMC, the runtime is about 40 hours on a V100 GPU. Craftax runs take roughly 94 hours, equivalent to 3.9 days.

Craftax The official environment provides the categorical variables in one-hot encoding format. Our implementation translates these variables to integer values which can be interpreted as tokens.

Setup in Atari Freeway For the Freeway environment, we adopted a modified sampling strategy 780 where a temperature of 0.01 is used instead of the standard value of 1, following [37, 11]. This 781 adjustment helps directing the agent toward rewarding paths. Note that alternative approaches in the 782 literature tackle the exploration challenge through different mechanisms, including epsilon-greedy 783 exploration schedules and deterministic action selection via argmax policies [37]. 784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

799

800

801

802

803

804

805

806

807

808

809

810

811

812

813

814

815

816

817

818

819

820

821

822

823

824

825

826

827

828

829

830

831

832

833

834

835

836

837

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We provide extensive empirical evidence in Section 3, including ablation studies, which directly relate to our contributions and claims. The scope of our paper is sample-efficient, planning-free world model agents (RL), which is explicitly stated in the abstract and introduction, while it is also reflected by the choice of baselines in our experiments.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Section 5 explicitly discuss the limitations of our work. Additional limitations are discussed in Section 3 (e.g., the absence of ablations on Craftax due to computational limitations).

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: Our paper does not include theoretical results.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In Section 2 and in the appendix we discuss our method in full detail, including architectures, hyperparameters, etc.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: In our abstract and appendix we provide a link to the code and trained model weights. Our code has a detailed readme file for easy usage, and we also provide Docker support, which enables an easy environment setup and enhances reproducibility on any operation system.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We specify all experimental details in Section 3 and in Appendix A and C.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We utilize the rliable toolkit [2] to generate plots with appropriate error bars (Figure 4 bottom, Figure 6). Figure 5 also includes error bars. In addition, our open-sourced repository includes the full results of all runs.

Guidelines:

The answer NA means that the paper does not include experiments.

993

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide this information in Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our work follows the NeurIPS Code of Ethics. No human subjects or participants were involved. We found no special concerns beyond those related to the general topic of deep reinforcement learning.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

994

995

996

997

998

999

1000

1001

1002

1003

1004

1005

1006

1007

1008

1009

1010

1011

1012

1013

1014

1015

1016

1017

1018

1019

1020

1021

1022

1023

1024

1025

1026

1027

1028

1029

1030

1031

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

Justification: This paper presents a foundational work in the field of Machine Learning. As such, there are no direct positive or negative societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work does not pose any additional risks beyond those of common deep reinforcement learning methods. Hence, we do not introduce additional safeguards.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Our paper cites all relevant assets, and our open-sourced repository includes a credits section with the relevant credits. We follow the licenses of all assets used in our work.

Guidelines:

The answer NA means that the paper does not use existing assets.

1047

1048

1049

1050

1051

1052

1053

1054

1055

1056

1057

1058

1059

1060

1061

1062

1063

1064

1065

1066

1067

1068

1069

1070

1071

1072

1073

1074

1075

1076

1077

1078

1079

1080

1081

1082

1083

1084

1085

1086

1087

1088

1089

1090

1091

1092

1093

1094

1095

1096

1097

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Our open-sourced repository includes all new assets and is well documented. Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

1098

1099

1100

1101

1102

1103

1104

1105

1106

1107

1108

1109

1110

1111

1112

1113

1114

1115

1116

1117

1118

1119

1120

1121

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Simulus, a sample-efficient, token-based world model agent. Simulus builds upon the REM agent and integrates four additional main components: (1) a modular tokenization framework to handle multi-modal observations (e.g., images, continuous vectors, symbolic grids), (2) intrinsic motivation based on epistemic uncertainty to guide exploration, (3) a prioritized replay mechanism for world model training, and (4) a regression-as-classification approach for more robust reward and value prediction. The authors conduct an empirical evaluation across three diverse and challenging benchmarks: Atari 100K, DeepMind Control Suite 500K, and Craftax-1M. The results demonstrate that Simulus achieves state-of-the-art performance across all benchmarks. Furthermore, ablation studies are presented to validate the individual contribution and synergistic effects of the integrated components.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n* The paper is well-written, organized, and easy to follow. The methodology is explained in sufficient detail, and the motivations behind the design choices are well-articulated. A significant strength of this work is its adherence to modern standards for empirical evaluation in reinforcement learning. The authors report not just mean scores but also Interquartile Mean (IQM), median, and optimality gap metrics, using stratified bootstrap confidence intervals as recommended by Agarwal et al.. This provides a much more nuanced and reliable picture of performance compared to simple point estimates.\n* The effectiveness of Simulus is demonstrated on three distinct and relevant benchmarks, each highlighting a different capability: visual control (Atari), continuous control from proprioception (DMC), and multi-modal reasoning (Craftax). The agent is compared against a comprehensive set of strong and recent baselines, including DreamerV3, STORM, etc., which makes the claimed state-of-the-art results highly convincing.\n* The paper successfully integrates several known, powerful techniques (intrinsic motivation, prioritized replay, RaC) into a unified agent. The ablation studies compellingly demonstrate that these components not only contribute individually but also act synergistically to achieve the final reported performance.\n\n### Weaknesses\n* The primary weakness of the paper is its limited originality from a methodological standpoint. The proposed agent, Simulus, is framed as an extension of REM, and the core improvements stem from integrating existing, well-established techniques. Intrinsic motivation for exploration, prioritized replay, and RaC for value learning are all powerful but not novel ideas. While the successful engineering and integration are commendable, the work does not introduce a new fundamental algorithm or mechanism. Since the method does methodically nothing novel, a slightly stronger experimental section would strengthen the paper more to be a good paper (e.g. for the DMC tasks, which feature continuous action spaces, the authors employ a discretization strategy. I would be curious how the method performs in harder tasks with a continuous action spaces).\n* A key contribution is the modular multi-modality tokenization framework. However, of the three benchmarks used, only Craftax truly requires multi-modal input processing. To make the claim of general multi-modality more compelling, the paper would benefit from evaluation on additional multi-modal environments. The authors themselves note the scarcity of such benchmarks as a limitation with which I do not fully agree with. E.g. a simple robotic setup can be multi-modal where the agent gets the proprioceptive sensor data and RGB images as input (e.g. see [ManiSkill](https://github.com/haosulab/ManiSkill)).\n\nQUESTIONS:\n* The authors decided to discretize the continuous action space for the DMC tasks. Could you provide more insight into the potential limitations of this approach, especially for more complex robotics tasks that might require finer motor control than the ones tested?\n* The authors write in the Appendix that the DMC experiments for Simulus take around 40h. Is this the run time for a single task? Compared to DreamerV3 which has a similar performance but takes only around 3h this is much longer.\n\nLIMITATIONS:\nThe authors mention that the scarcity of multi-modal environments prevented them to evaluate further experiments with multi-modal environments. Many robotic tasks provide multi-modal observations; thus I do not fully agree with them.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper,  the author presents Simulus, a modular Tag-Based World Model agent designed to improve sample efficiency in reinforcement learning. Based on REM, Simulus integrates four main components: multimodal tokenization, intrinsic motivation for uncertainty-driven exploration, prioritized world model regeneration, and regression classification for reward prediction. It supports multiple observation modalities, including images, continuous vectors, and symbolic grids. Simulus achieves state-of-the-art performance in three benchmarks, Atari 100K, DeepMind Control Suite, and Craftax-1M, without relying on planning. Ablation studies confirm that each component makes a meaningful contribution,n and their combination produces a powerful synergy. Although training takes longer than some benchmarks, Simulus' modular design enables efficient inference and provides a solid foundation for future research.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n\n- The author has presented compelling results that demonstrate the state of the art. Notably, on Atari 100K, Simulus achieved a human-normalized quartile mean score of 0.99, significantly outperforming prior works. The proposed method also achieves competitive performance on continuous control and multimodal benchmarks.\n\n- The ablation study is well-conducted and provides clear insights into the contribution of each component. The results convincingly show that the final performance is the product of synergistic effects, not a single new trick. The discovery that intrinsic rewards are effective in reducing model uncertainty even in reward-rich environments is a valuable insight for the community. \n\n- This paper is well written, clearly structured, and easy to understand.  Also, the code and model weights are open-sourced, which enables the paper being able to be reproduced. \n\nWeakness:\n\n- The computational cost of the model will be a problem.  Just as the author said, \"token-based world model agents remain significantly slower to train than other baselines in sample-efficient RL.\"\n\n- While this paper presents a system that demonstrates good performance and is empirically validated, its core contribution lies in the combined application of existing techniques within the TBWM framework. These components have been considered separately in prior work, and this work does not introduce entirely new learning objectives or architectural mechanisms. Thus, the novelty of this paper is systematic and empirical, rather than algorithmic.\n\nQUESTIONS:\n- Can the authors quantitatively report the computational cost, such as GPU hours or FLOPs? The observed improvement may be because of the increased computation rather than the method itself, which could be unfair to the baselines. One way to ensure fairness would be to extend the baseline models' training steps to match the total computational budget.\n\n- Regarding prioritized replay, how sensitive is the model performance to the hyperparameter α? How is this hyperparameter determined?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper authors deals with the token based world models and by using them they aim to improve significantly on the sample efficiency. This paper has four contributions: possibility to use multimodal inputs, curiosity type idea for intrinsic reward, prioritized replay buffer and the use of regression in reward modeling.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n- The proposed WM obtains significant sample-efficiency in the as shown in Fig 5 (craftax)\n- WM is based on the REM paper and significantly improves on that is shown in the ablation experiments. In these experiments authors show that the modifications they invented were truly missing in the original REM. \n\nWeaknesses: \n- Experiments are a bit suspect. Different environments are experimented with a different set of algorithms. Especially problematic is that for craftax no Dreamer v3 is used! Noting that the original Crafter environment is _designed_ for Dreamer v2 and the experiments shown in the Crrafter github repo author shows that Dreamer v2 clearly wins over any other algorithm. \n- As far as I see authors have capped their intection budget ot 100k environment interactions and that itself is ok. It is ok to focus in sample-efficiency regime. But it would be also really good to see what happens model is let to use millions of steps. This kind of experiment(s) would give full disclosure to the reader.\n\nQUESTIONS:\n- Have I undersrtood correctly that world model algorithms are typically considered to be on-policy? So then could you explain how prioritized replay buffer is used?\n- In my lab we have ran a multiple WM type experiments with atari benchmarks and we have found as are also found in Dreamer set of papers that Atari is actually sometimes quite hard. Reason is that in WM you represent state by a latent code and some Atari games have really small details. Especially hard is Breakout where many times difference in successive states is ball that is only one pixel wide. That details is very easily lost in WN type latent modeling. What about your proposed method? I see some examples in the appendices, but images are too small to see what is happening.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Simulus, a token‑based world model (TBWM) agent aimed at improving sample efficiency in RL. The authors propose a modular pipeline—separating representation learning, the world model, and the controller—and weave in four key components: (1) multi‑modal tokenization to handle heterogeneous observation/action spaces, (2) intrinsic motivation targeted at reducing epistemic uncertainty, (3) prioritized replay within the world model, and (4) a regression‑as‑classification scheme for reward/return prediction. Empirically, Simulus attains state‑of‑the‑art results across several benchmarks (e.g., Atari 100K, DMC, Craftax-1M).\n\nSTRENGTHS AND WEAKNESSES:\nStrenghts:\n\nFor this reviewer, Simulus’s main strength is its modular setup: the representation, world model, and controller can each be built and tuned on their own. The paper also stitches together several underused ideas—multi‑modal tokenization, intrinsic rewards to cut down epistemic uncertainty, prioritized world‑model replay, and regression‑as‑classification—showing they work better together than alone for sample efficiency.\n\nAdditionally, on the results side, Simulus posts state‑of‑the‑art sample efficiency for planning‑free world models on Atari 100K, the DeepMind Control Suite, and Craftax‑1M. It hits human or even superhuman scores on multiple Atari titles and reliably matches or beats baselines in both continuous and multi‑modal settings.\n\nThe framework pushes TBWMs past the usual vision + discrete‑action niche. Thanks to the modular multi‑modal tokenization, it handles continuous control (DMC) and richer multi‑modal environments (Craftax), which nudges TBWMs closer to real‑world RL problems.\n\nAnother important and useful contribution of the paper is that the ablations are actually useful: intrinsic rewards, prioritized replay, and RaC return prediction are each tested in isolation. Intrinsic rewards, in particular, come out as a major contributor, and the combined effect of all pieces is clear. Moreover, the paper spells out methods, architectures, hyperparameters, and experimental details; code and pretrained weights are available and will be made public.\n\nSome of the things I find to be on the weaker side are: \n\n1. Training is expensive: tokenizing continuous vectors blows up sequence length, and the authors admit TBWMs (including Simulus) train much slower than standard baselines. Even if modularity helps inference, the upfront cost is a real barrier for groups without large compute.\n\n2. Multi‑modality isn’t fully stress‑tested: the lack of rich multi‑modal RL benchmarks limits how broadly they can validate the approach. Craftax is a start, but more varied tasks would make the claim stronger.\n\n3. While the appendix discusses interfering objectives in RSSM optimization and provides preliminary results on PWM vs. PWM-decoupled, this section feels somewhat tangential to the core contributions of Simulus itself. While interesting, it could be refined to more directly link how Simulus's design avoids or mitigates such interference, rather than just observing it in a related model.\n\nQUESTIONS:\nA few questions that the authors can please elaborate:\n\n1. Can you quantify the wall‑clock training cost of Simulus relative to key baselines (e.g., GPU hours per environment) and identify which modules dominate that cost?\n\n2. How sensitive is performance to the tokenization granularity for continuous vectors, and have you explored strategies to cap sequence length without hurting returns?\n\n3. Given the scarcity of multi‑modal RL benchmarks, do you have plans (or preliminary results) for additional domains beyond Craftax that stress different modality combinations?\n\n4. Where does Simulus underperform or fail to scale (e.g., very high‑dimensional observations or extremely long horizons), and what do you see as the most promising path to reduce the compute burden?\n\nLIMITATIONS:\nLimitations have been discussed briefly in Section 5, which are sufficient.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-07-23"}