mlsanityData to 7 Oct 2026

Learning Transformer-based World Models with Contrastive Predictive Coding

arXiv 2503.04416 ↗First posted 6 Mar 2025cs.LG cs.AI cs.CVcited by 20

On the map: World models → Latent dynamics models and model-based RL

Abstract (arXiv)

The DreamerV3 algorithm recently obtained remarkable performance across diverse environment domains by learning an accurate world model based on Recurrent Neural Networks (RNNs). Following the success of model-based reinforcement learning algorithms and the rapid adoption of the Transformer architecture for its superior training efficiency and favorable scaling properties, recent works such as STORM have proposed replacing RNN-based world models with Transformer-based world models using masked self-attention. However, despite the improved training efficiency of these methods, their impact on performance remains limited compared to the Dreamer algorithm, struggling to learn competitive Transformer-based world models. In this work, we show that the next state prediction objective adopted in previous approaches is insufficient to fully exploit the representation capabilities of Transformers. We propose to extend world model predictions to longer time horizons by introducing TWISTER (Transformer-based World model wIth contraSTivE Representations), a world model using action-conditioned Contrastive Predictive Coding to learn high-level temporal feature representations and improve the agent performance. TWISTER achieves a human-normalized mean score of 162% on the Atari 100k benchmark, setting a new record among state-of-the-art methods that do not employ look-ahead search.

Key facts. Cited by 20 papers in the map (data through 7 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 41 of its 56 references that are in the map. First posted on arXiv 6 Mar 2025.

Builds on

41 of 56 references.

Closest ancestors · 14

Same direction, newest first

  1. 2024Efficient World Models with Context-Aware Tokenization27 Jun
  2. 2024EfficientZero V2: Mastering Discrete and Continuous Control with Limited Data1 Mar
  3. 2023TD-MPC2: Scalable, Robust World Models for Continuous Control25 Oct
  4. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning14 Oct
  5. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  6. 2023Mastering Diverse Domains through World Models10 Jan
  7. 2022Transformers are Sample-Efficient World Models1 Sep
  8. 2022TransDreamer: Reinforcement Learning with Transformer World Models19 Feb
  9. 2021Mastering Atari Games with Limited Data30 Oct
  10. 2021DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations27 Oct
  11. 2020Mastering Atari with Discrete World Models5 Oct
  12. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
2 more
  1. 2019Exploring Model-based Planning with Policy Networks20 Jun
  2. 2018Recurrent World Models Facilitate Policy Evolution4 Sep

Other references · 27

  1. 2017Attention Is All You Needcited by 59,932
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2020A Simple Framework for Contrastive Learning of Visual Representationscited by 11,979
  4. 2018Representation Learning with Contrastive Predictive Codingcited by 7,662
  5. 2020Scaling Laws for Neural Language Modelscited by 6,411
  6. 2016Layer Normalizationcited by 5,579
  7. 2017Neural Discrete Representation Learningcited by 4,650
  8. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,016
  9. 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
  10. 2020Unsupervised Learning of Visual Features by Contrasting Cluster Assignmentscited by 2,845
  11. 2016Asynchronous Methods for Deep Reinforcement Learningcited by 2,218
  12. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,031
15 more
  1. 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,830
  2. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
  3. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,221
  4. 2018DeepMind Control Suitecited by 935
  5. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 781
  6. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 774
  7. 2020Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixelscited by 573
  8. 2019Model-Based Reinforcement Learning for Ataricited by 495
  9. 2022Temporal Difference Learning for Model Predictive Controlcited by 375
  10. 2022Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videoscited by 338
  11. 2021Benchmarking the Spectrum of Agent Capabilitiescited by 193
  12. 2019Benchmarking Model-Based Reinforcement Learningcited by 158
  13. 2023Bigger, Better, Faster: Human-level Atari with human-level efficiencycited by 145
  14. 2016The Predictron: End-To-End Learning and Planningcited by 92
  15. 2020Data Augmenting Contrastive Learning of Speech Representations in the Time Domaincited by 63

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 6; 2026: 14.

20 papers.

Sort:

By area:

  1. 2025A Comprehensive Survey on World Models for Embodied AIcited by 65
  2. 2026R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentationcited by 10
  3. 2025medDreamer: Model-Based Reinforcement Learning with Latent Imagination on Complex EHRs for Clinical Decision Supportcited by 8
  4. 2025Object-Centric World Models for Causality-Aware Reinforcement Learningcited by 3
  5. 2026Morphology-Conditioned World Model for Cross-Embodiment Quadrupedal Locomotioncited by 3
  6. 2026World Models: A Comprehensive Survey of Architectures, Methodologies, Reasoning Paradigms, and Applicationscited by 3
  7. 2025Learning Lifted Action Models From Traces of Incomplete Actions and Statescited by 2
  8. 2025DyMoDreamer: World Modeling with Dynamic Modulationcited by 2
  9. 2026Dreamer-CDP: Improving Reconstruction-free World Models Via Continuous Deterministic Representation Predictioncited by 2
  10. 2026JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learningcited by 2
  11. 2026Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamicscited by 1
  12. 2026OrbiSim: World Models as Differentiable Physics Engines for Embodied Intelligencecited by 1
8 more
  1. 2025Controllable Flow Matching for Online Reinforcement Learningcited by 0
  2. 2026Probing the Impact of Scale on Data-Efficient, Generalist Transformer World Models for Ataricited by 0
  3. 2026Differentiable Learning of Lifted Action Schemas for Classical Planningcited by 0
  4. 2026Learning Lifted Action Models from Traces with Minimal Information About Actions and Statescited by 0
  5. 2026Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learningcited by 0
  6. 2026TaskSense: Focusing on What Matters in World Modelscited by 0
  7. 2026Intrinsic Motivation in Reinforcement Learning: A Research Agenda for Adaptive Self-Organisationcited by 0
  8. 2026CyberWorld: World Models for Sample-Efficient Autonomous Cyber Defensecited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2023Facing Off World Model Backbones: RNNs, Transformers, and S4
  2. 2024One-shot World Models Using a Transformer Trained on a Synthetic Prior
  3. 2024Masked Generative Priors Improve World Models Sequence Modelling Capabilities
  4. 2025EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
  5. 2025Zero-shot World Models via Search in Memory
  6. 2024Learning to Play Atari in a World of Tokens
  7. 2025Transformer World Model for Sample Efficient Multi-Agent Reinforcement Learning
  8. 2026Contrastive World Models