Research map / Multi-agent systems and games / Human-AI and goal-conditioned
Improving Token-Based World Models with Parallel Observation Prediction
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 Parallel Observation Prediction (POP), a RetNet forward mode that predicts all observation tokens in parallel, and REM, a token-based world model agent that is 15.4x faster at imagination and reaches superhuman performance on 12 of 26 Atari 100K games.
Problem
Token-based world models such as IRIS generate observation tokens sequentially token-by-token during imagination, creating a severe bottleneck that causes long training times, poor GPU utilization, and caps the observation token length.
Approach
POP maintains dedicated prediction tokens and, using RetNet's chunkwise recurrent states summarizing each observation-action block, computes all next-observation tokens in one parallel pass; a POP-specific chunkwise forward enables batched training across all time steps despite conflicting input requirements.
Key results
- REM reaches IQM human-normalized score 0.673 on Atari 100K, the highest among all compared methods, and superhuman performance on 12 of 26 games.
- POP gives a 15.4x faster imagination compared with IRIS, with REM training in under 12 hours.
- Ablation on 8 games shows the No POP variant retains performance but is slower, and a 4x4 tokenizer (vs 8x8) lowers IQM to 1.026 from 2.201.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace POP-augmented RetNet with vanilla RetNet (ablation) | REM · Atari 100K, 8-game ablation subset | IQM human-normalized score: 2.201 -> 2.068 | No clear effect |
| Use tokenizer codebook embeddings instead of separate world model embeddings (ablation) | REM world model · Atari 100K, 8-game ablation subset | IQM human-normalized score: 1.794 -> 2.201 | Helped |
| Increase tokenizer resolution from 4 × 4 to REM's resolution (ablation) | REM tokenizer · Atari 100K, 8-game ablation subset | IQM human-normalized score: 1.026 -> 2.201 | Helped |
| Use REM's latent-code controller instead of the IRIS controller (ablation) | REM controller · Atari 100K, 8-game ablation subset | IQM human-normalized score: 1.234 -> 2.201 | Helped |
| Add action inputs to the controller (ablation) | REM controller · Atari 100K, 8-game ablation subset | IQM human-normalized score: 1.535 -> 2.201 | Helped |
| Use REM rather than STORM | Atari 100K · Atari 100K | IQM human-normalized score: 0.561 -> 0.673 | Helped |
| Use REM rather than IRIS | Atari 100K · Atari 100K | IQM human-normalized score: 0.501 -> 0.673 | Helped |
Surprising findings
- Removing POP (No POP) gives a higher mean and median on the 8-game ablation subset than full REM (mean 2.357 vs 1.947), though the authors state performance is retained.
- Using the controller without action inputs slightly outperforms full REM on some games such as Asterix and DemonAttack in the ablation table.
Limitations
- POP increases total computation per observation ((2K+1)H vs (K+1)H) in exchange for scalability.
- Evaluated only on Atari 100K; the ablations use a subset of 8 games with 5 seeds due to compute constraints.
- The recurrent state summarizes only a short context (two-frame context in implementation), not the entire history.
Experimental setup
Datasets: Atari 100K
Benchmarks: Atari 100K
Metrics: mean human-normalized score median human-normalized score IQM human-normalized score optimality gap number of superhuman games probability of improvement wall-clock run time
Baselines: IRIS SimPLe DreamerV3 TWM STORM
Builds on: IRIS RetNet DreamerV2 VQ-VAE VQ-GAN
Closest papers by meaning
- Efficient World Models with Context-Aware Tokenization (ICML 2024 · similarity 0.86)
- Reinforcement World Model Learning for LLM-based Agents (NeurIPS 2026 · similarity 0.86)
- Open-World Reinforcement Learning over Long Short-Term Imagination (ICLR 2025 · similarity 0.86)
- Learning to Play Atari in a World of Tokens (ICML 2024 · similarity 0.86)
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling (NeurIPS 2025 · similarity 0.85)
- RLVR-World: Training World Models with Reinforcement Learning (NeurIPS 2025 · 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.