Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining
1. Paper overview
In one sentence: Sharing a transformer between world prediction and conservative Q-learning improves aggregate offline Atari performance, while short imagined searches add return at an inference-time cost. e01e02e03e04e07e09e12
| At a glance | What to know |
|---|---|
| Research problem | Source description Can a single image-based offline RL agent learn transferable dynamics and decision-making from heterogeneous replay, while avoiding the instability of scaling temporal-difference learning? The experiment spans fifteen training games and five held-out Atari games; adaptation uses offline data from each new game rather than zero-shot deployment. e01e06e08 |
| Core mechanism | |
| A key reported result | Multi-game offline Atari performance: JOWA-150M: 0.789 (78.9%); median HNS 0.456; six superhuman games. IQM human-normalized score (HNS). Fifteen pretrained games; 10% subsample from two replay runs, 10M transitions/game; 1.75M updates and batch 512. Evaluation averages sixteen episodes/game with no sticky actions and epsilon 0.001. EDT-200M 0.502; MGDT-200M 0.498; FICC-85M 0.433; SQL-80M 0.420. The 28.7 percentage-point IQM gap over EDT is a reader calculation. Aggregate advantage coexists with losses on individual games. JOWA-70M is inconsistent: Table 2 gives 0.485, while Section 5.3 and Table 4 give 0.476. e06e07e09e17 |
| Reading caution |
Core contributions
- Source description
A parallelizable beam search ranks imagined outcomes using learned rewards and terminal Q-values, connecting generative prediction to executed game actions. e04
Figure 1. One history representation supplies both world predictions and action values. Original paper, p. 4 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Read upward from an observation through the encoder and its discrete tokens, then through the shared transformer. The decoder branch reconstructs the observation; it is not an executed action. Above the backbone, p_o predicts visual tokens, while p_r and p_d predict reward and termination. The Q-head h_Q sits over the final observation-token embedding, before the current action token; the reward and termination branches use the action-position output. This placement agrees with the conditioning in Equations (4)–(6). The caption adds vocabulary, position and task embeddings, although their summation is not drawn. The colors distinguish network components, transitions, internal outputs and final outputs. e02e03e04
What it supports. The integration is architectural: both prediction and value estimation depend on the same transformer, and both losses update it. The action part supplies values for candidate actions rather than generating an action token as its training target. Actual control requires Q-based selection or the separate planning procedure.
Where the evidence stops. Figure 1 depicts forward information flow, not gradient routing or search. Joint back-propagation is established by Section 4.1.3 and Equation (9). The figure alone cannot show that this sharing improves stability or return.
2. Motivation
2.1 The problem and the proposed response
Can a single image-based offline RL agent learn transferable dynamics and decision-making from heterogeneous replay, while avoiding the instability of scaling temporal-difference learning? The experiment spans fifteen training games and five held-out Atari games; adaptation uses offline data from each new game rather than zero-shot deployment. e01e06e08
2.2 What this reading follows
JOWA asks whether learning to predict a visual environment can make a large offline value learner more useful. It turns Atari observations into discrete tokens, learns dynamics and action values through the same transformer, and uses predicted futures to refine action choice. The key distinction is between training that shapes the shared representation and planning that spends computation at decision time. The six visuals below trace that distinction through the architecture, pretrained performance, target-game adaptation and diagnostic experiments. This reading covers the supplied January 2026 arXiv v4, including its appendices, and preserves both a conflicting reported score and an inconsistency in the planning bound. e01e02e03e04e07e09e12
3. Research context
We place the paper in the collection through its world–action interface. The catalog labels and the reading’s assessment are shown separately.
| Catalog dimension | Recorded classification |
|---|---|
| Major category | Not assigned |
| Architecture | Not assigned |
| Prediction paradigm | Not assigned |
| Quadrant | Not assigned |
This table preserves the labels recorded at reading time. The current major category is Foundational work. View the current classification.
3.1 Evidence-based assessment
Insufficient evidence to decide
The recorded catalog is unassigned. Architecture evidence supports one shared world/value backbone with separate heads, not two independently trained models. Its action mechanism is Q-value-based control with forward-model planning, not inverse dynamics or direct joint generation of future observations and actions. This is an assessment for classification, not a catalog edit. e02e03e04
This is the collection’s architectural analysis, not a new related-work survey. Benchmark comparisons and their protocols appear in Section 6.
4. Problem formulation
4.1 Inputs and outputs
| Inputs | Outputs |
|---|---|
|
|
4.2 Equations and their role
5. Method in detail
5.1 Let world prediction and value learning shape the same history
Start from an image observation, which the VQ-VAE encodes as 36 visual tokens. Interleave those tokens with actions, add task and positional information, and pass the history through a transformer. The world heads learn visual-token, signed-reward and termination distributions; the action head learns a return distribution for each action. The shared parameters receive both losses. This explains why merely attaching a critic to a pretrained world representation would not reproduce JOWA's central design. Stage one trains the tokenizer and world part; stage two freezes the tokenizer and lets predictive and conservative distributional TD objectives jointly update the backbone. Table 6 then asks whether those gradients matter: stopping action-loss gradients leaves world learning intact but substantially lowers return in the six-game test. The experiment supports this coupling under its particular data and model budget. e02e03e05e10
5.2 Use a short imagined future to revise a fallible Q-ranking
A Q-value compresses long-term consequences into a single expected return. JOWA's planner checks that estimate against a short imagined trajectory: add discounted predicted rewards, then bootstrap from the best terminal Q-value. Candidate actions are limited by Q-ranking, and beam search retains only high-scoring continuations. The first action of the winning continuation is selected for the real environment. Reader analysis: this creates two possible failure points, inaccurate imagined outcomes and early pruning of useful actions. Appendix B explicitly assumes the optimal policy lies within the top-K action set. Appendix C also requires bounded errors and Lipschitz functions, and its printed final equality changes Lipschitz coefficients into error terms. Those qualifications prevent interpreting the return gains as an unconditional optimality guarantee. Table 5 additionally shows that improved decisions consume substantial inference time. e04e09e12e18
5.3 Interpret transfer as adaptation under a specific data and checkpoint policy
The five target games are excluded from pretraining, but each supplies 5k transitions for subsequent offline learning. Expert-level transitions are sampled from the final portion of DQN-Replay. Fine-tuning runs separately for each game, and the reported result comes from the highest-scoring snapshot rather than a fixed final checkpoint. JOWA's appendix describes a second fine-tuning stage that extends real segments with two planned synthetic steps, even though default pretraining avoids synthetic replay. This distinction is essential for a faithful reproduction. Table 3 shows a strong improvement over training the same nominal capacity from scratch, despite the scratch model's longer update budget. Table 21 then shows that the transfer result changes with replay quality. Thus the evidence supports efficient adaptation within this protocol, while leaving generalization across new data-collection policies and fixed deployment checkpoints open. e08e11e13e16e20
5.4 Training and inference
During training
Stage one trains the tokenizer and world part for 250k updates. Stage two freezes the tokenizer and jointly trains world and action parts for 1.5M updates. Pretraining samples 10M transitions per game from two DQN-Replay runs, totaling approximately six billion tokens. Default pretraining uses real replay, without planning or synthetic transitions. e05e06e11
Transfer separately fine-tunes each held-out game for 50k updates using 5k expert-level transitions, batch 32 and learning rate 0.00005. Appendix E.2 describes a second stage appending two planned synthetic steps to six-step real segments, but leaves its stage allocation unclear. The reported snapshot is selected by highest fine-tuning score. e08e16
During inference
At each decision, restrict candidate actions to the top Q-ranked set, expand imagined transitions, and retain beams with the largest discounted reward-plus-terminal-value total. Execute the first action of the selected path and repeat with real observations. Evaluation uses horizon two and game-specific beam widths two to four; the paper reuses K for both token count and beam width. e02e04e17
Planning is an approximate search over learned dynamics. Appendix B assumes the optimum survives top-K restriction; a finite beam can prune alternatives. The theoretical discussion therefore does not certify globally optimal executed actions. e04e12
5.5 Implementation flow
- Tokenize observations
A VQ-VAE converts each image into 36 discrete tokens from a 2048-entry vocabulary. Visual and action tokens are interleaved; vocabulary, position and task embeddings form the transformer input. Eight timesteps occupy 296 tokens. e02e05e14
- Share history, separate predictions
Observation prediction is autoregressive over earlier visual tokens and actions. Reward and termination use the current action. Figure 1 places the Q-head at the final observation token, so it can score candidate actions before consuming the chosen current action. e02
- Learn a conservative critic
The Q-head represents a return distribution whose mean is Q. Its objective combines distributional TD with CQL, which discourages unsupported actions. World heads use classification losses, including rewards mapped by sign to −1, 0 or 1. e03
6. Experiments & results
JOWA learns an Atari world model and distributional Q-function through one shared transformer, then searches short imagined futures to choose actions. Its strongest evidence is improved aggregate game return and data-efficient offline adaptation; neither universal game-wise scaling nor unconditional planning optimality is established (e02, e03, e04, e07, e08, e11, e12).
6.1 Read the original evidence
Table 2. JOWA-150M leads the aggregate score, while individual games retain clear exceptions. Original paper, p. 8 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Begin with the bottom IQM HNS row, which aggregates scores after subtracting each game's random score and dividing by its human-minus-random reference gap. Then inspect individual game rows: their entries are raw returns, not normalized percentages. The JOWA columns increase in capacity from 40M to 150M, while SQL abbreviates Scaled-QL. The comparison uses the same replay subsample, batch size and total gradient-step budget, but architectures and augmentation differ. In particular, a matched number of updates does not imply matched computation. Check median HNS and the superhuman-game count alongside IQM before interpreting the aggregate as broad dominance. e06e07e09e17
What it supports. JOWA-150M reaches IQM HNS 0.789 versus EDT-200M's 0.502, a calculated 28.7 percentage-point gap. Yet SQL scores 7023.5 on BeamRider versus JOWA-150M's 3498. The evidence supports aggregate improvement across these fifteen games, with task-specific weaknesses still visible. The median and per-game rows qualify the IQM headline.
Where the evidence stops. The 70M IQM cell is 0.485 here; Section 5.3 and Table 4 instead report 0.476. The discrepancy is unresolved. These point estimates have no confidence intervals or independent training-seed variability attached.
Table 3. Pretraining helps offline adaptation, under a best-checkpoint fine-tuning protocol. Original paper, p. 8 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Read the final IQM row to compare adaptation across five unseen games. DNS uses DQN as the reference rather than the human reference used in Table 2, so the two tables are not a common numerical scale. Each game receives 5k expert-level replay transitions and is fine-tuned separately. The far-right scratch model has the same nominal 150M capacity but receives 500k updates, versus 50k for pretrained models. Appendix E.2 selects the highest-scoring fine-tuning snapshot and describes a second stage using planned synthetic transitions. Thus these are adaptation results after learning from each target game, not zero-shot scores. e08e13e16e19
What it supports. JOWA-150M achieves IQM DNS 0.647, compared with FICC's 0.575 and the scratch model's 0.181. Increasing JOWA size raises the IQM across these variants, but the 70M model has the higher median, 0.715 versus 0.615. Different aggregates reveal different aspects of transfer.
Where the evidence stops. The stage-length split and new-task embedding initialization are unspecified. Table 18 supplies raw scores for only part of this table, omitting FICC and the smaller JOWA variants. Best-checkpoint selection also limits comparison with fixed-checkpoint deployment.
6.2 Results and evaluation conditions
| Task & protocol | Reported result | Comparison & interpretation |
|---|---|---|
| Multi-game offline Atari performance Fifteen pretrained games; 10% subsample from two replay runs, 10M transitions/game; 1.75M updates and batch 512. Evaluation averages sixteen episodes/game with no sticky actions and epsilon 0.001. | JOWA-150M: 0.789 (78.9%); median HNS 0.456; six superhuman games. IQM human-normalized score (HNS) | EDT-200M 0.502; MGDT-200M 0.498; FICC-85M 0.433; SQL-80M 0.420. The 28.7 percentage-point IQM gap over EDT is a reader calculation. Aggregate advantage coexists with losses on individual games. JOWA-70M is inconsistent: Table 2 gives 0.485, while Section 5.3 and Table 4 give 0.476. e06e07e09e17 |
| Offline adaptation to five unseen Atari games Gravitar, MsPacman, Pong, Robotank and YarsRevenge; 5k expert-level transitions/game; separate fine-tuning and best-score checkpoint selection. | JOWA-150M: 0.647. IQM DQN-normalized score (DNS) | FICC-85M 0.575; EDT-200M 0.380; JOWA-150M from scratch 0.181. Pretrained/scratch budgets are 50k/500k updates. Supports transfer after adaptation. DNS and HNS have different reference agents. Table 18 omits raw FICC and smaller-JOWA scores despite their aggregates in Table 3. e08e13e16e19 |
| Planning benefit and runtime Table 4 uses fifteen games; Table 5 uses seven games where the bare 150M agent is weak. | Fifteen-game 150M IQM: 50.8% without planning, 78.9% with planning, a 28.1 percentage-point increase. IQM HNS and reported frames per second | Seven-game beam/MCTS/no-planning IQM: 23.7%/13.4%/3%; FPS: 1.26/0.12/10.8. Beam search trades speed for return versus greedy inference. The Python MCTS comparison uses a Q-derived policy and selected games, not a general ranking of planners. e09e18 |
| Joint-training loss ablation 40M model, six named training games, 1M updates, 10% replay subsample. | Original 1.123; no world loss 0.659; stopped action-loss gradient 0.307. IQM HNS | No CQL 0.637; MSE TD 0.126; synthetic-data pretraining 0.464. Supports both shared-gradient training and distributional conservative learning in this reduced setup. Synthetic replay also changes the objective to COMBO, so its effect is confounded. e10e11 |
| Fine-tuning data-quality sensitivity JOWA-150M, five held-out games, 5k transitions from final 20%, complete, or initial 20% replay pools. | Expert/suboptimal/highly-suboptimal mean: 0.647/0.516/0.422; IQM: 0.647/0.511/0.383. Mean and IQM DNS | Same nominal transition count across quality conditions. Aggregate transfer depends on replay quality; Gravitar does not follow the aggregate ordering. e20 |
6.3 Ablations and diagnostic examples
Read component removals and qualitative examples within their stated evaluation conditions.
Tables 4 and 5. Planning improves return but adds inference cost; the two tables use different game sets. Original paper, p. 9 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Treat the left and right panels as separate experiments. On the left, Table 4 compares each model with and without planning on all fifteen pretrained games. Its improvement row uses percent signs, but the entries are differences in percentage points. On the right, Table 5 compares the 150M model on seven games selected because its bare policy performs poorly. FPS is higher-is-better throughput, while the other rows summarize return. Beam search ranks imagined reward plus terminal Q-value and runs branches in parallel. The MCTS baseline derives both value and action probabilities from Q-values, with a tuned temperature and search depth. e04e07e09e17e18
What it supports. For 150M on fifteen games, planning raises IQM HNS from 50.8% to 78.9%, or 28.1 percentage points. On the separate seven-game comparison, beam search obtains 23.7% IQM at 1.26 FPS, versus MCTS's 13.4% at 0.12 FPS. Greedy inference remains faster at 10.8 FPS.
Where the evidence stops. The seven-game result does not establish beam search superiority over MCTS across all games or real-time deployment. Both planners use Python, and satisfactory MCTS settings were not found for the other eight games. The 70M score also conflicts with Table 2.
Table 6. Stopping value-learning gradients sharply weakens the shared model in the reduced experiment. Original paper, p. 9 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Use Origin as the reference column and compare horizontally within a metric row. These are six-game, 40M experiments trained for one million updates; they are not the main fifteen-game 150M benchmark. No CQL removes the conservative penalty. No world loss removes predictive training. The sg(action-loss) column stops action-loss gradients from reaching the backbone while retaining world-model training, directly testing whether value learning should shape the shared representation. MSE replaces distributional TD with scalar squared-error training. The synthetic-data column uses half real and half imagined transitions and changes the action objective to COMBO. Values are normalized ratios, so 1.123 corresponds to 112.3%. e03e10e11
What it supports. The original IQM HNS is 1.123, falling to 0.659 without world loss and 0.307 when action-loss gradients are stopped at the backbone. This supports the usefulness of both objectives shaping a common representation in the tested setup. Distributional TD and conservatism also matter for the reported outcome.
Where the evidence stops. The ablation establishes performance differences, not a complete causal account of training stability. Synthetic replay changes both data and loss, so its 0.464 IQM cannot isolate model error from conservatism or other training effects.
Table 21. A fixed transition count does not remove dependence on replay quality. Original paper, p. 23 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. The three rightmost columns vary the source of 5k fine-tuning transitions for JOWA-150M. Expert data come from the final 20% of DQN-Replay; suboptimal data are sampled from the whole dataset; highly-suboptimal data come from its initial 20%. The game rows are raw returns, with random and DQN references beside them. Only the final mean, median and IQM rows summarize normalized scores. Compare these aggregate rows first, then look for exceptions in individual games. Gravitar illustrates why the authors' overall quality trend should not be translated into a strict ordering for every target task. e16e20
What it supports. Mean DNS drops from 0.647 with expert data to 0.516 and 0.422 in the two lower-quality conditions; IQM drops from 0.647 to 0.511 and 0.383. The headline adaptation result therefore depends partly on the selected replay pool, even when the nominal number of transitions is fixed.
Where the evidence stops. Replay-pool selection changes state coverage and behavior quality together. This experiment does not isolate which causes the aggregate difference, and its single reported table provides no uncertainty estimates for the ordering across data draws.
7. Analysis & limitations
7.1 What the evidence leaves open
Full Atari and complete-data scaling were not tested. Larger models sometimes perform worse on individual games. The authors attribute apparent Zaxxon emergence to nonlinear rewards; aggregate scaling alone cannot establish newly acquired capabilities. e11e20
The error-bound presentation is internally inconsistent: Eq. (28) changes coefficients L_r and L_Q into epsilon_r and epsilon_Q in its last equality. Eqs. (12), (24) and (29) repeat that printed form. Together with the top-K assumption, this prevents treating the stated theorem as an independently verified guarantee. e12
The main tables provide no confidence intervals or independent training-seed variability. Sixteen rollout episodes do not establish training-run robustness. Fine-tuning selects the highest-scoring snapshot, whereas pretraining uses the final snapshot; this distinction matters when reproducing transfer claims. e07e08e15e16e17
7.2 Questions for discussion
- Would planning still help when Q-ranking excludes the best action from the beam?
- How much transfer advantage remains under fixed checkpoints, repeated training seeds, and matched fine-tuning-data quality?
8. Reproducibility audit
8.1 Requirements and known gaps
Recover the specified game split and replay subsample; match the tokenizer, context length, game-conditioned embeddings and Q-head configuration. Table 13 specifies Adam/AdamW, transformer rates 0.0001 then 0.00005, discount 0.99, 51 atoms over [-10,30], and target-Q updates every 1000 steps. e13e14e15
Resolve how new task embeddings are initialized, how fine-tuning stages divide the update budget, and how checkpoints are selected before claiming a faithful transfer replication. A100 hardware is named, but GPU count and software versions are not specified in the supplied implementation details. e05e14e15e16
8.2 Proposed reproduction checks
The following checks are proposals motivated by the paper. They have not been run as part of this reading.
Check 1: Separate shared-gradient learning from the benefit of search
Reader-proposed check, not performed: use the paper's six-game 40M setup and one-million-update budget. Start matched runs from the same tokenizer and initialization, then compare the original joint objective, stopped action-loss gradients, and zero world-loss weight. Keep the replay draws, optimizer and Q-head capacity fixed, and repeat across independent training seeds. Evaluate every checkpoint with planning disabled and with horizon two using the same episode seeds. Record IQM HNS, per-game returns and TD-loss trajectories. If the original objective improves only planning-enabled return, the evidence for improved bare-critic learning is weaker; if it also improves greedy return and training consistency, that more directly supports the regularization interpretation. e03e05e10e14e15e17
Check 2: Measure when top-K pruning blocks a useful action
Reader-proposed check, not performed: freeze a trained JOWA checkpoint and save emulator states during evaluation on the fifteen training games. From identical states, compare the published top-K beam search with an exhaustive horizon-two action search, holding the terminal Q-function fixed. Repeat both searches with learned transitions/rewards and with emulator transitions/rewards, using the same signed-reward convention and discount, then execute their selected first actions in matched evaluation rollouts. Report action disagreement, realized return, model calls and wall-clock time. If exhaustive search improves return even with emulator transitions, candidate pruning is a material limitation; if replacing learned dynamics produces the larger gain, model error is the stronger bottleneck. This tests assumptions motivating the bound without treating its inconsistent printed algebra as a verified theorem. e03e04e09e12e15e17e18
8.3 Reading coverage
Visual audit: The title/byline/version on p. 1, architecture and conditioning on p. 4, objectives on p. 5, planning/training on pp. 6–7, Figure 2 and quantitative Tables 2–6 on pp. 8–9, limitations on p. 10, derivations on pp. 15–18, and implementation/evaluation/raw-score Tables 7–21 on pp. 18–23 were rendered and visually inspected. All six final original crops were separately inspected at their returned dimensions. Figure 1's arrow positions were cross-checked against Eqs. (4)–(9); its forward-flow diagram does not depict gradients or search. The 70M score discrepancy was verified on pp. 7–9; the bound inconsistency was verified on pp. 6 and 16–18. Complete source text, including references, was read across all seven chunks. Pages 2–3 and reference-only pp. 11–14 were read as text, outside this visual pass. External supplements, code and checkpoints were not inspected.
PDF pages inspected for this edition: 1, 4, 5, 6, 7, 8, 9, 10, 15, 16, 17, 18, 19, 20, 21, 22, 23. Appendix coverage: reviewed.
Original figures and tables remain the work of the source’s authors. Extractions preserve their scientific content; any HTML wrapper layout is disclosed with each figure. The surrounding reading notes are our own.
Text reading scope & known omissions
- Title, abstract and Sections 1–3: introduction, related work, distributional RL, CQL and problem setup
- Sections 4.1–4.3: architecture, objectives, planning and training pipeline
- Sections 5.1–5.6: data, baselines, main results, scaling, transfer and ablations
- Section 6: conclusion and limitations; acknowledgments and references
- Appendices A–C: Bellman derivation, constrained search and error-bound proof
- Appendices D–F: game split, implementation, fine-tuning, evaluation, raw scores and additional experiments
Outside the original text pass
- Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout.
- The complete supplied 23-page PDF text was read. The original 2024 submission was not supplied: this report reviews arXiv v4 dated 29 January 2026, labeled ICLR 2025. The observed title and all eight authors match the catalog; earlier versions were not compared (e01).
- Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout. This acquisition caveat was addressed by inspecting original PDF pages and all six final crops; exact visual scope is recorded in the edition.
- Separate supplemental material availability has not been fully verified.
- Code, linked resources and checkpoints were not inspected; no experiments were reproduced.
- Table 1 was read as text; both figures and Tables 2–21 were visually inspected. Reference-only pages were read as text.
The visual audit above records the subsequent illustrated pass.
8.4 Traceable evidence
e01PDF p. 1, title, byline, affiliations and arXiv margin
Title and eight authors match the catalog; the artifact is arXiv:2410.00564v4, 29 January 2026, labeled published at ICLR 2025.
Go to primary source ↓e02PDF p. 4, Figure 1 and Sections 4.1.1–4.1.2, Eqs. (4)–(6)
VQ-VAE tokens and actions enter a shared transformer with task embeddings. Observation, reward and termination predictors accompany a Q-value head. Q-values use the final observation-token representation before the current action.
Go to primary source ↓e03PDF p. 5, Sections 4.1.2–4.1.3, Eqs. (7)–(9)
World prediction uses cross-entropy and signed ternary rewards. Distributional TD and CQL train the action part; both losses update the transformer, with alpha and beta each 0.1.
Go to primary source ↓e04PDF pp. 5–6, Section 4.2, Eqs. (10)–(13) and beam-search description
Search combines discounted imagined rewards with a terminal maximum Q-value, restricts actions to top-K, and retains K beams. The first action of the selected path is chosen for execution. The claimed error improvement is conditional.
Go to primary source ↓e05PDF pp. 6–7, Section 4.3 and Section 5.1
Pretraining has 250k world-model/tokenizer steps and 1.5M joint steps with a frozen tokenizer. Data comprise 10M transitions per training game, 36 visual tokens per frame, length-eight sequences, approximately six billion tokens, and A100 training.
Go to primary source ↓e06PDF p. 7, Sections 5.1–5.3
Fifteen games are used for pretraining and five held out. Ten percent is sampled from two DQN-Replay runs. Baselines use batch size 512 and 1.75M steps, but differ in architecture and augmentation; FICC is adapted to offline multi-game training.
Go to primary source ↓e07PDF p. 8, Table 2 and Figure 2; p. 7, Section 5.3
Table 2 reports IQM HNS 0.789 for JOWA-150M, 0.502 EDT, 0.498 MGDT, 0.420 SQL and 0.433 FICC. JOWA-70M is 0.485 in Table 2 but 0.476 in Section 5.3. Scaling is aggregate, with exceptions in individual games.
Go to primary source ↓e08PDF p. 8, Table 3 and Section 5.5; p. 9, Section 5.5
Held-out adaptation uses 5k expert-level transitions per game. IQM DNS is 0.647 for JOWA-150M, 0.575 FICC, 0.380 EDT and 0.181 JOWA-150M from scratch. Pretrained models receive 50k steps; scratch receives 500k.
Go to primary source ↓e09PDF p. 9, Tables 4–5 and planning ablation text
On fifteen games, planning changes JOWA-150M IQM HNS from 50.8% to 78.9%; JOWA-70M is 45.1% to 47.6%. On seven selected games, beam search, MCTS and no planning achieve IQM HNS 23.7%, 13.4% and 3%, at 1.26, 0.12 and 10.8 FPS.
Go to primary source ↓e10PDF pp. 9–10, Table 6 and Section 5.6
Six-game 40M ablations use 1M updates. IQM HNS is 1.123 original, 0.637 no CQL, 0.659 no world loss, 0.307 stopped action-loss gradient, 0.126 MSE TD, and 0.464 synthetic-data pretraining. The synthetic condition also changes the action loss to COMBO and uses half imagined transitions.
Go to primary source ↓e11PDF p. 10, synthetic-data discussion and Section 6 limitations
Default pretraining excludes planning and synthetic data. Synthetic pretraining is costly and worsens the tested setup; model error and conservatism are hypotheses. Full Atari/full-data experiments were not performed, and some games decline as model size increases.
Go to primary source ↓e12PDF pp. 15–18, Appendices A–C, especially p. 16 top-K assumption and p. 17 Eqs. (26)–(29); p. 6 Eqs. (12)–(13)
Appendix B assumes the optimum remains within top-K actions. Appendix C assumes bounded estimation errors and Lipschitz functions. Eq. (28) first contains L_r and L_Q, then prints epsilon_r and epsilon_Q as their coefficients in its final equality; the same error-term form appears in Eqs. (12), (24), and (29).
Go to primary source ↓e13PDF p. 18, Appendix D and Table 7
The five held-out games are Gravitar, MsPacman, Pong, Robotank and YarsRevenge. Table 7 lists the twenty games, their difficulty labels, action-space sizes and train/fine-tune assignments.
Go to primary source ↓e14PDF pp. 18–19, Appendix E.1.1 and Tables 8–12
The tokenizer uses reconstruction, embedding/commitment and perceptual terms, vocabulary 2048, and 36 tokens per frame. The transformer has a 296-token context; 40M uses four layers, width 512 and eight attention heads, whereas 150M uses twelve layers, width 768 and twelve heads. The 150M configuration has three Q-heads.
Go to primary source ↓e15PDF p. 20, Table 13 and end of Appendix E.1
Adam trains VQ-VAE; AdamW trains other modules. Learning rates are 0.0001 then 0.00005 for the transformer, batch size 512, discount 0.99, 51 return atoms over [-10,30], and target-Q updates every 1000 steps. Final pretraining checkpoints are reported.
Go to primary source ↓e16PDF p. 20, Appendix E.2
Each held-out game is fine-tuned separately with batch 32, learning rate 0.00005 and 50k updates, reporting the highest-scoring snapshot. A second stage uses horizon/beam width two, six-step real segments and two synthesized steps; its stage-length allocation and new-task embedding initialization are not specified here.
Go to primary source ↓e17PDF p. 21, Appendix E.3 and Tables 16–17
Scores average sixteen evaluation episodes; sticky actions are disabled, evaluation epsilon is 0.001, frame skip and listed frame stack are four, and the episode cap is 108K frames. Planning horizon is two, with game-specific beam widths two to four.
Go to primary source ↓e18PDF pp. 21–22, Appendix F planning comparison and Table 19
Both planners use Python. MCTS is adapted using Q-derived values and action probabilities; tuning selects maximum Q, temperature 0.9, most valuable root child, and per-game depths. Seven games are selected where the bare policy is weak; the authors did not find satisfactory MCTS settings for the other eight.
Go to primary source ↓e19PDF pp. 22–23, Tables 18 and 20 and Appendix F ensemble/task-embedding discussion
Table 18 provides raw held-out scores for only a subset of Table 3 methods, omitting FICC and the smaller JOWA variants. Table 20 reports IQM 0.824 without task embedding versus 1.123 original, and 1.049/0.931 with equal/random Q-head ensembles.
Go to primary source ↓e20PDF p. 23, Table 21 and Appendix F non-expert-data/emergence discussions
For 5k-transition adaptation, expert, suboptimal and highly-suboptimal data yield mean DNS 0.647/0.516/0.422 and IQM 0.647/0.511/0.383. Sampling pools are the final 20%, all data and initial 20% respectively. The authors attribute the apparent Zaxxon emergence to nonlinear rewards.
Go to primary source ↓8.5 Primary sources
Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining ↗
PDF · 11,523 extracted words
Source fingerprint
2de5781c1bb2f389d662d0ea82989e63498cebe34e841f53bd82147e1e1b394d