Research map / Reinforcement learning / Offline RL
Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining
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.
Introduces JOWA, a jointly-optimized world-action transformer trained offline with TD learning and world-modeling losses on 15 Atari games, plus a search-based planner, showing scaling and sample-efficient transfer to unseen games.
Problem
Scaling temporal-difference offline RL to a single generalist agent that learns both general-purpose representations and decision-making across heterogeneous multi-game Atari data, and transferring to unseen games with very little data.
Approach
A shared transformer backbone serves both a visual world model (supervised next-token, reward, and termination prediction) and a distributional Q-head trained with CQL-style TD loss, so the world-modeling loss regularizes TD learning at scale; at inference, a parallelizable beam-search planner over the learned world model corrects Q-value errors.
Key results
- JOWA-150M reaches 78.9% IQM human-normalized score on 15 pretrained games with 10% data, versus 49.8% for MGDT-200M and 50.2% for EDT-200M (planning improves IQM from 50.8% to 78.9%).
- Fine-tuned JOWA-150M attains 64.7% IQM DQN-normalized score on 5 unseen games using 5k transitions, outperforming baselines by 69.9% on average.
- Proposed planner achieves 37.8% mean HNS vs 22.1% for MCTS on 7 games and runs about 10x faster than MCTS (1.26 vs 0.12 FPS).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Scale the pretrained agent to JOWA-150M | JOWA · 15 Atari games (10% subsampled) | IQM HNS: 0.485 -> 0.789 | Helped |
| Fine-tune the pretrained JOWA-150M on 5k transitions per game rather than train from scratch | JOWA-150M · 5 unseen Atari games | IQM DNS: 0.181 -> 0.647 | Helped |
| Add inference-time planning (ablation) | JOWA-150M · 15 Atari games | IQM HNS: 50.8% -> 78.9% (+28.1%) | Helped |
| Add inference-time planning (ablation) | JOWA-70M · 15 Atari games | IQM HNS: 45.1% -> 47.6% (+2.5%) | Helped |
| Add inference-time planning (ablation) | JOWA-40M · 15 Atari games | IQM HNS: 27.1% -> 37.1% (+10.0%) | Helped |
| Use the proposed planner instead of MCTS (ablation) | JOWA-150M · 7 Atari games where bare JOWA-150M performs poorly | IQM HNS: 13.4% -> 23.7% | Helped |
| Add CQL regularization (ablation) | JOWA training loss · 6 Atari games | IQM HNS: 0.637 -> 1.123 | Helped |
| Add the world-model loss (ablation) | JOWA training loss · 6 Atari games | IQM HNS: 0.659 -> 1.123 | Helped |
| Jointly back-propagate action-loss gradients to the transformer backbone (ablation) | JOWA training · 6 Atari games | IQM HNS: 0.307 -> 1.123 | Helped |
| Use distributional TD-loss instead of MSE (ablation) | JOWA training · 6 Atari games | IQM HNS: 0.126 -> 1.123 | Helped |
| Add synthetic data during pretraining (ablation) | JOWA-40M · 6 Atari games | IQM HNS: 1.123 -> 0.464 | Hurt |
Surprising findings
- Using synthetic world-model rollouts in pretraining gave negative gains, which the authors attribute to compounding inference error and COMBO over-conservatism.
- Scaling model size does not monotonically improve raw scores on some individual games, even though IQM scales.
Limitations
- Experiments restricted to 15 pretraining games and 20-game subset rather than full Atari 2600 with complete datasets due to compute.
- Synthetic data in pretraining gave negative gains and is too costly to investigate further.
- Raw per-game scores sometimes decrease as model size grows, which the authors attribute to offline RL inherent issues.
Experimental setup
Datasets: DQN-Replay Atari dataset (10% subsample, 15 training games) DQN-Replay last 20% for 5 held-out games (5k transitions each)
Benchmarks: Atari 2600 (15 pretrained games) 5 held-out Atari games (Pong, Robotank, YarsRevenge, Gravitar, MsPacman)
Metrics: human-normalized score IQM HNS median HNS mean HNS DQN-normalized score FPS
Baselines: MTBC MGDT Elastic DT Scaled-QL FICC JOWA-150M (scratch)
Builds on: IRIS CQL C51 Scaled-QL FICC
Closest papers by meaning
- Video-Enhanced Offline Reinforcement Learning: A Model-Based Approach (ICML 2025 · similarity 0.92)
- Making Offline RL Online: Collaborative World Models for Offline Visual Reinforcement Learning (NeurIPS 2024 · similarity 0.91)
- The Generalization Gap in Offline Reinforcement Learning (ICLR 2024 · similarity 0.90)
- A Clean Slate for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
- Offline-Boosted Actor-Critic: Adaptively Blending Optimal Historical Behaviors in Deep Off-Policy RL (ICML 2024 · similarity 0.89)
- Beyond OOD State Actions: Supported Cross-Domain Offline Reinforcement Learning (AAAI 2024 · similarity 0.89)
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.