STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

1 Introduction

• Deep reinforcement learning requires many environment interactions, a particular challenge in real-world settings without adjustable-speed simulation.Value estimation and sparse rewards are cited as sources of low sample efficiency.Source: §1 · pp. 1–2
• Model-based methods learn a world model through self-supervision and use imagined experience to improve policies without real-environment sampling constraints.The paper describes reconstruction, action prediction, and contrastive learning as possible self-supervision approaches.Source: §1 · p. 2
• RNN-based world models impede parallel training, while prior Transformer approaches are described as slower and not outperforming DreamerV3.The comparison figure distinguishes RNN-based SimPLe and DreamerV3 from Transformer-based TWM, IRIS, and STORM.Source: §1 · p. 2, Figure 1 · p. 2
• 126.7% mean human normalized score on Atari 100k set a new record among methods without lookahead search.Training with 1.85 hours of real-time interaction took 4.3 hours on one NVIDIA GeForce RTX 3090.Source: Figure 1 · p. 2, §1 · p. 3
• STORM pairs a categorical VAE image encoder with a Transformer sequence model to improve robustness, modeling and generation quality, and training speed.The VAE is intended to reduce accumulated autoregressive prediction errors.Source: §1 · p. 2

2 Related work

• Model-based methods learn environment dynamics and use simulated experience to improve policies, while visual tasks such as video games pose modeling challenges.Recent advances in computing and deep neural networks have enabled these methods to learn dynamics and outperform model-free methods on such tasks.Source: §2 · p. 3
• VAE-LSTM methods established image-based world models; SimPLe applied them to Atari, and Dreamer variants use GRUs across several domains.These approaches improved sample efficiency or demonstrated domain capabilities, but RNN structures train slowly.Source: §2 · p. 3
• IRIS, TWM, and TransDreamer apply Transformers to world models, with prior approaches described as incurring token or attention costs or lacking limited-sample benchmark evidence.IRIS uses spatial-temporal attention over many tokens; TWM treats observations, actions, and rewards as separate tokens, while TransDreamer replaces Dreamer’s GRU.Source: §2 · p. 3
• MCTS-based model-based methods report promising Atari 100k performance but incur high computational demands; auxiliary self-supervised losses also appear in sample-efficient model-free methods.The paper focuses on improving the world model rather than further discussing those model-free methods.Source: §2 · p. 3
• STORM uses a GPT-like Transformer sequence model, a single stochastic latent variable per image, and a vanilla Transformer rather than Transformer-XL.Unlike Dreamer and TransDreamer, it reconstructs the original image without using hidden states.Source: §2 · p. 4

3.1 World model learning

• Training alternates real-policy data collection into replay, world-model updates from replay trajectories, and policy improvement using imagined experiences initialized from replay.The replay buffer stores consecutive trajectories in a FIFO queue.Source: §3.1 · p. 4
• For Atari images, a CNN VAE maps observations to stochastic categorical latents with 32 categories of 32 classes, using straight-through gradients for sampled latents.The encoder and decoder are implemented as convolutional neural networks.Source: §3.1 · p. 5
• A masked GPT-like Transformer consumes tokens formed from each latent sample and action, then predicts reward, continuation, and next-latent distributions.The causal mask lets each token attend to the sequence up to and including itself.Source: §3.1 · p. 5
• End-to-end self-supervised world-model training combines image reconstruction, reward prediction, and continuation prediction losses, with β1=0.5 and β2=0.1.The total loss is calculated over batches and sequence positions.Source: §3.1 · p. 5
• KL dynamics loss trains the sequence model to predict the next distribution, while representation loss weakly lets that prediction influence the encoder.This limits how challenging distributional-dynamics learning is for the encoder.Source: §3.1 · p. 5

3.2 Agent learning

• For imagined learning, the agent starts from replay-buffer context, samples latent states from the model prior, and uses a KV cache to accelerate Transformer inference.The initial posterior distribution is computed from the selected contextual trajectory.Source: §3.2 · p. 6
• The actor-critic learns from world-model-predicted rewards and continuation flags, using λ-returns and an exponential moving average of critic parameters to regularize value learning.The EMA regularization aids training stability and helps prevent overfitting.Source: §3.2 · p. 6

4.1 Benchmark and baselines

• Atari 100k covers 26 games; its 100k-sample limit corresponds to 400k actual frames and approximately 1.85 hours of gameplay.Human performance is measured after a player becomes familiar with each game under the same sample constraint.Source: §4.1 · pp. 6–7
• Comparisons target model-based DRL methods with similar training pipelines; lookahead-search methods are excluded because the focus is refining the world model.The authors state that lookahead search could be combined with STORM in future work.Source: §4.1 · p. 7

4.2 Results on Atari 100k

• STORM outperforms previous methods on games with large or multiple reward-related objects, but faces challenges with single small objects such as Pong and Breakout.The authors attribute the advantage to attention preserving moving-object histories; sampling randomness can potentially harm attention weights in small-object settings.Source: §4.2 · p. 7, Table 2 · p. 7

5.1 World model design and configuration

• Reconstruction loss should target the encoder output; using zt instead of ht minimally affects Pong but reduces performance on context-dependent Ms. Pacman.The latter difference concerns predicting rewards from a single frame versus several contextual frames.Source: §5.1 · p. 8, Figure 3 · p. 8
• Increasing Transformer depth did not improve final performance, and even 400k Pong samples yielded maximum reward with either 4 or 6 layers.The default Transformer uses 2 layers; the authors suggest limited Atari data and end-to-end representation learning may constrain scaling.Source: §5.1 · p. 8, Figure 3 · p. 8

5.2 Selection of the agent’s state

• Including ht improves performance on context-dependent Ms. Pacman, but makes little difference in Pong and Kung Fu Master.STORM uses the combined state [ht, zt].Source: §5.2 · p. 8, Figure 4 · p. 9
• Using ht alone in policy-evolving environments such as Pong may invite catastrophic-forgetting-like behavior; randomness from distributions such as Zt is beneficial.Source: §5.2 · pp. 8–9

5.3 Impact of the demonstration trajectory

• A single demonstration trajectory can improve robustness in sparse-reward Pong and performance in Freeway, but may hinder policy improvement in dense-reward Ms Pacman.Source: §5.3 · p. 9, Figure 5 · p. 9
• In Freeway, earning a reward requires about 70 consecutive up actions, making the first reward difficult under a uniform-action policy and motivating a demonstration to seed imagination.Source: §5.3 · p. 9
• 122.3% mean human-normalized score without Freeway’s trajectory; STORM still outperformed previous methods.The no-trajectory results are reported to support fair comparisons in future research.Source: §5.3 · p. 9

6 Conclusions and limitations

• STORM surpasses previous methods in both performance and training efficiency, using a Transformer’s sequence modeling and parallelizable training advantages.The authors say improved efficiency can broaden applicability while reducing computational costs.Source: §6 · p. 10
• Jointly optimizing the image encoder and sequence model makes the world model predict its own internal output, adding non-stationarity that may impede scalability.Source: §6 · p. 10
• Imagination starts are uniformly sampled from the replay buffer, without explicitly accounting for the on-policy distribution relevant to the policy-gradient formulation.Source: §6 · p. 10

• Figure 6 compares the method with DreamerV3 on the Atari 100k benchmark.Source: Figure 6 · p. 14
• The solid lines show the average result over five seeds.Source: Figure 6 · p. 14
• The filled areas show the minimum-to-maximum range across those five seeds.Source: Figure 6 · p. 14

Image encoder

• The image encoder uses Conv layers with kernel 4, stride 2, and padding 1, plus batch normalization and ReLU.Flatten and Reshape alter tensor indexing while preserving data and order; the submodule sizes are omitted and derivable from tensor shapes.Source: Table 3 · p. 15

Image decoder

• The image decoder uses transpose CNN layers with kernel 4, stride 2, and padding 1.Source: Table 4 · p. 15

Action mixer

• The action mixer concatenates latent state and action features; action dimension ranges from 3 to 18 across games.D denotes the Transformer's feature dimension, and LN denotes layer normalization.Source: Table 5 · p. 15

Positional encoding

• Positional encoding uses a learnable parameter matrix of shape T × D, where T is sequence length.Source: Table 6 · p. 16

Transformer block

• The Transformer block uses dropout, which the passage says can prevent overfitting.Source: Table 7 · p. 16

Sequence model

• The Transformer-based sequence model maps the embedded sequence e1:T to h1:T.Its positional encoding and Transformer block are specified separately.Source: Table 8 · p. 16

Pure MLP structures

• A one-layer MLP is a fully connected layer, and 255 is the bucket size for symlog two-hot loss.Source: Table 9 · p. 16

• The environment signals done when a life is lost but continues until the actual reset.This life-information configuration aligns with IRIS [13].Source: Table 10 · p. 17
• World-model training samples B1 trajectories of length T, while imagination starts use B2 trajectories of length C.Source: Table 10 · p. 17
• The sampling description distinguishes batches for model training from those used to start imagination.Source: Table 10 · p. 17

• Frame counts are multiplied by 4 to account for frame skipping.Source: Table 11 · p. 18
• The trajectories were gathered using pre-trained DQN agents [45].Source: Table 11 · p. 18

E Computational cost details and comparison

• V100 computational results marked with a star are extrapolated, while unstarred results were tested on actual devices.The extrapolation assumes P100 is twice as slow and A100 twice as fast, following the setup used in DreamerV3 [10].Source: Table 12 · p. 19
• The V100 extrapolation assumes P100 is twice as slow and A100 twice as fast, following the setup used in DreamerV3 [10].Source: Table 12 · p. 19

F Atari video predictions

• Eight observations and actions provide context for autoregressive imagination spanning 56 future frames in Atari games.Source: Figure 7 · p. 20