Research map / Architectures and efficiency / Attention and transformers
Learning Transformer-based World Models with Contrastive Predictive Coding
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TWISTER, a Transformer-based world model trained with action-conditioned Contrastive Predictive Coding (AC-CPC) on augmented image views, reaching 162% human-normalized mean on Atari 100k among model-based methods without look-ahead search.
Problem
Transformer-based world models trained with next-state prediction underperform RNN-based DreamerV3 on visual RL, as next-latent prediction is too easy and does not exploit Transformer representation capacity.
Approach
Extend world model prediction to K=10 future steps by applying action-conditioned CPC: a Transformer state predicts future augmented stochastic states conditioned on the sequence of future actions, with InfoNCE over the batch as negatives, and random crop-and-resize augmentation to make the task harder.
Key results
- TWISTER reaches a human-normalized mean of 162% and median of 77% on Atari 100k, the best among model-based methods without look-ahead search (DreamerV3 112%, STORM 127%).
- Removing AC-CPC drops TWISTER's Atari normalized mean from 162% to 112% and median from 77% to 44%.
- On 20 DeepMind Control tasks at 1M steps, TWISTER achieves mean 801.8 versus 739.6 for DreamerV3 and 720.9 for TD-MPC2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add AC-CPC to the world model (ablation) | TWISTER · Atari 100k | human-normalized mean (%): 112 -> 162 | Helped |
| add AC-CPC to the world model (ablation) | TWISTER · Atari 100k | human-normalized median (%): 44 -> 77 | Helped |
| use a Transformer TSSM instead of DreamerV3's RSSM (ablation) | TWISTER · Atari 100k | human-normalized mean (%): 121 -> 162 | Helped |
| condition CPC predictions on future actions (ablation) | TWISTER · Atari 100k | human-normalized mean (%): 111 -> 162 | Helped |
| condition CPC predictions on future actions (ablation) | TWISTER · Atari 100k | human-normalized median (%): 42 -> 77 | Helped |
| augment positive and negative contrastive samples (ablation) | TWISTER · Atari 100k | human-normalized mean (%): 120 -> 162 | Helped |
| augment positive and negative contrastive samples (ablation) | TWISTER · Atari 100k | human-normalized median (%): 68 -> 77 | Helped |
| add AC-CPC to the world model (ablation) | TWISTER · DeepMind Control Suite (20 tasks) | mean score: 728.7 -> 801.8 | Helped |
| add AC-CPC to the world model (ablation) | TWISTER · DeepMind Control Suite (20 tasks) | median score: 802.4 -> 907.6 | Helped |
| use 10 CPC prediction steps rather than 15 (ablation) | TWISTER · Atari 100k | human-normalized mean score | Helped |
Surprising findings
- Replacing the Transformer with DreamerV3's RSSM gives similar results without AC-CPC, but AC-CPC helps mainly the Transformer variant.
- Without future-action conditioning, CPC brings no notable performance improvement.
Limitations
- Compares only against model-based methods without look-ahead search, excluding EfficientZero V2 and BBF.
- Lower performance than IRIS/Delta-IRIS on small-object games such as Demon Attack and Assault (authors attribute to reconstruction quality).
- Results averaged over only 5 seeds; hardware and compute cost not reported.
Experimental setup
Datasets: Atari 100k DeepMind Control Suite
Benchmarks: Atari 100k DeepMind Control Suite
Metrics: human-normalized mean human-normalized median agent score number of superhuman games
Baselines: SimPLe TWM IRIS DreamerV3 STORM Delta-IRIS DreamerPro TD-MPC2
Builds on: DreamerV3 STORM Contrastive Predictive Coding DreamerV2 TD-MPC2
Closest papers by meaning
- Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement (NeurIPS 2024 · similarity 0.86)
- Zero-shot World Models via Search in Memory (NeurIPS 2025 · similarity 0.86)
- Accurate and Efficient World Modeling with Masked Latent Transformers (ICML 2025 · similarity 0.86)
- Improving Transformer World Models for Data-Efficient RL (ICML 2025 · similarity 0.86)
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling (NeurIPS 2025 · similarity 0.85)
- Do Transformer World Models Give Better Policy Gradients? (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.