Improving Transformer World Models for Data-Efficient RL
Authors: Antoine Dedieu
Joseph Ortiz
Xinghua Lou
Carter Wendelken
J Swaroop Guntupalli
Wolfgang Lehrach
Miguel Lazaro-Gredilla
Kevin Patrick Murphy
Venue: ICML 2025 poster (2025)
Keywords: Model Based Reinforcement Learning
Background Planning
Transformer World Model
Primary area: reinforcement_learning->online
TL;DR: We present a ladder of simple improvements to vision-based model-based RL agents which, when taken together, achieve a significantly higher reward and score on the challenging Craftax benchmark.

Abstract: We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open-world 2D survival game that requires agents to exhibit a wide range of general abilities---such as strong generalization, deep exploration, and long-term reasoning. With a series of careful design choices aimed at improving sample efficiency, our MBRL algorithm achieves a reward of 69.66% after only 1M environment steps, significantly outperforming DreamerV3, which achieves $53.2\%$, and, for the first time, exceeds human performance of 65.0%. Our method starts by constructing a SOTA model-free baseline, using a novel policy architecture that combines CNNs and RNNs.
We then add three improvements to the standard MBRL setup: (a) "Dyna with warmup", which trains the policy on real and imaginary data, (b) "nearest neighbor tokenizer" on image patches, which improves the scheme to create the transformer world model (TWM) inputs, and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep.

Improving Transformer World Models for Data-Efficient RL
Antoine Dedieu"
Joseph Ortiz"
Xinghua Lou't Miguel Lázaro-Gredilla'
Carter Wendelken
J. Swaroop Guntupalli Kevin Murphy
Wolfgang Lehrach
Abstract
Humn
We present an approach to model-based RI. that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open- world survival game that requires agents to exhibit a wide range of general abilities-uuc as strong generalization, deep exploration, and long-term reasoning. With series of careful design choices aimed at improving sample effi- ciency, our MBRL algorithm achieves a reward of 69.66% after only 1M environment steps, signifi- cantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human per- formance of 65.0%. Our method starts by con- structing SOTA model-free baseline, using novel policy architecture that combines CNNs and RNNs. We then add three improvements to the standard MBRL setup: (a) "Dyna with warmup", which trains the policy on real and imaginary data, (b) "nearest neighbor tokenizer" on image patches, which improves the scheme create the transformer world model (TWM) inputs, and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep.
A-Mi
Our DreamerV3
Previous bast MFRL
Our bast MBRL (clow)
Reward
30
0.0
0.2
0.4
0.6
0.8
1.0
Environment steps
Figure 1: Reward Craftax-classic. Our best MBRL and MFRL agents outperform all the previously published MFRL and MBRL results, and for the first time, surpass the reward achieved by human expert. We display published methods which report the reward at I M steps with horizon- tal line from 900k to 1M steps.
Model-based RL (MBRL) aims to reduce the amount of data needed to train the policy by also learning a world model (WM), and using this WM to plan "in imagination".
To evaluate sample-efficient RL algorithms, it is common use the Atari-100k benchmark (Kaiser et al., 2019). How- ever, although the benchmark encompasses variety of skills (memory, plamning, etc), each individual game typi- cally only emphasizes one or two such skills. To promote the development agents with broader capabilities, we focus on the Crafter domain (Hafner, 2021), a 2D version of Minecraft that challenges single agent master diverse skill set. Specifically, ment (Matthews al., 2024), a fast, near-replica of Crafter, implemented in JAX (Bradbury et al., 2018), Key features of Craftax-classic include: (a) procedurally generated stochas- tic environments (at each episode the agent encounters a new environment sampled from common distribution); (b) partial observability, as the agent only sees 63 x 63 pixel image representing local view of the agent's environment, plus visualization of its inventory (see Figure ([left)): and (c) an achievement hierarchy that defines a sparse reward signal, requiring deep and broad exploration.
1. Introduction
Reinforcement learning (RL) (Sutton & Barto, 2018) pro- vides framework for training agents to act in environments so as to maximize their rewards. Online RL algorithms interleave taking actions in the environment-colleciie servations and rewards-and updating the policy using the collected experience. Online RL algorithms often employ model-free approach (MFRL), where the agent learns direct mapping from observations actions, but this require lot of data to be collected from the environment.
the Craftax-classic environ-
ob-
Equal contribution 'Google DeepMind. Correspondence to: Antoine Dedieu <adedicu @ google.com. com>, Joseph Ortiz <jocor- @iz@ (w@@goolo.com..
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
In this paper, we study improvements to MBRL methods, hased on transformer world models (TWM), in the context

---

Improving Transformer World Models for Data-Efficient RL
ture map, whose elements are then quantized into discrete tokens, using codebook. The sequence of observation tokens across timesteps is used, along with the actions and rewards, to train the WM. We propose two improvements to the tokenizer. First, instead of jointly quantizing the image, we split the image into patches and independently tokenize each patch. Second, we replace the VQ- with simpler nearest-neighbor tokenizer (NNT) for patches. Unlike VQ- VAE, NNT ensures that the "meaning" of each code in the codebook is constant through training, which simplifies the task of learning reliable WM.
9.2 7 11
4 9
2
1
Figure 2: [Left] The Craftax- observation is a 63 63 pixel image, composed of 9 x patches of 7 x 7 pixels. The observation shows the map around the agent and the agent's s health and inventory. Here we have rendered the image at 144x x 144 pixels for visibility. [Right] 64 different patches.
Our third contribution addresses the way the world model is trained. TWMs are trained by maximizing the log likeli- hood of the sequence of tokens, which is typically generated autoregressively both over time and within timeslice. We propose an alternative, which we call block teacher forc- ing (BTF), that allows TWM to reason jointly about the possible future states of all tokens within timestep, before sampling them in parallel and independently given the his- tory. With BTF, imagined rollouts for training the policy are both faster to sample and more accurate.
of the Craftax-classic environment. We make contributions across the following three axes: (a) how the TWM is used (Section 3.4); (b) the tokenization scheme used to create TWM inputs (Section 3.5); (c) and how the TWM is trained (Section 3.6). Collectively, our improvements result in an agent that, with only 1M environment steps, achieves a Crastax-classic reward of 69.66% and score of 31.77%, significantly improving over the previous state of the art (SOTA) reward of 53.20% (Hafner et al., 2023) and the previous SOTA score of 19.4% (Kauvar et al., 2023)'.
Our final contributions are some minor architectural changes to the MFRL baseline upon which our MBRL approach is based. These changes are still significant, resulting in a simple MFRL method that is much faster than Dreamer V3 and yet obtains a much better average reward and score.
Our first contribution relates to the way the world model is used: in contrast to recent MBRL methods like IRIS (Micheli et al., 2022) and DreamerV3 (Hafner et al., 2023), which train the policy solely on imagined trajectories (gen- erated by the world model), we train our policy using both imagined rollouts from the world model and real experi- ences collected in the environment. This is similar to the original Dyna method (Sutton, 1 1990), although this tech- nique has been abandoned in recent work. In this hybrid regime, we can view the WM as a form of generative data augmentation (Van Hasselt et al., 2019).
Our improvements are complementary to each other, and can be combined into a "ladder of improvements" similarr to the "Rainbow" paper's (Hessel et al., 2018) series of improvements on top of model-free DQN agents.
2. Related Work
In this section, we discuss related work in MBRL - see e.g. Moerland et al. (2023); Murphy (2024); Open ab for more comprehensive reviews. We can broadly divide MBRL along two axes. The first axis is whether the world model (WM) is used for background planning (where it helps train the policy by generating imagined trajectories), or decision- time planning (where it is used for lookahead search at inference time). The second axis is whether the WM is generative model of the observation space (potentially via latent bottleneck) or whether is latent-only model trained using self-prediction loss (which is not sufficient to generate full observations).
Our second contribution addresses the tokenizer which con- verts between images and tokens that the TWM ingests and outputs. Most prior work uses vector quantized variational autoencoder -VQ-AA, Van Den Oord et al. 2017), e.g. IRIS (Micheli et al., 2022), DART (Agarwal et al., 2024). These methods train a CNN to process images into a fea-
'The score S is given by the geometric mean of the success rate for each of the N = 22 achievements; this puts more weight on occasionally solving many achievements than on consistently solving a subset. More precisely, the score is given by S
exp (++ EM In(1 + a.)) -1, where s; E [0, 100) is the success percentage for achievement (i.e., fraction of episodes in which the achievement was obtained at least once). By contrast, the rewards are just the expected sum of rewards, or in percentage, the arithmetic mean R- = N >>> St (ignoring minor contributions to the reward based on the health of the agent). The score and reward
Regarding the first axis, prominent examples of decision- time planning methods that leverage a WM include MuZero (Schrittwieser et al., 2020) and EfficientZero (Ye et al., 2021), which use Monte-Carlo tree search over a discrete action space, as well as TD-MPC2 (Hansen et al., 2024), which uses the cross-entropy method over continuous ac- tion space. Although some studies have shown that decision-
correlated, but not the same. Unlike some prior work, report both metrics to make comparisons easier.
2

---

Improving Transformer World Models for Data-Efficien RL
time planning can sometimes be better than background planning (Alver& & Precup, 2024), it is much slower, es- pecially with large WMs such as transformers, since it re- quires rolling out future hypothetical trajectories at each decision-making step. Therefore in this paper, we focus on background planning (BP), Background planning originates from Dyna (Sutton, 1990), which focused on tabular Q- learning. Since then, many papers have combined the idea with deep RL methods: World Models (Ha & Schmidhu- ber, 2018b), Dreamer agents (Hafner et al., 2020a;b; 2023), SimPLe (Kaiser et al., 2019), IRIS (Micheli et al., 2022), A- IRIS (Micheli et al., 2024), Diamond (Alonso et al., 2024), DART (Agarwal et al., 2024), etc.
penalty to ensure sufficient exploration. The CNN used is a modification of the Impala ResNet (Espeholt al., 201 2018a).
3.2. MFRL Improvements
We improve on this MFRL baseline by both increasing the model size and adding RNN (specifically GRU) to give the policy memory. Interestingly, we find that naively in- creasing the model size harms performance, while combin- ing a larger model with a carefully designed RNN helps (see Section 4.3). When varying the ratio of the RNN state dimension to the CNN encoder dimension, we observe that performance is higher when the hidden state is low- dimensional. Our intuition is that the memory is forced to focus on the relevant bits of the past that cannot be extracted from the current image.
Regarding the second axis, many methods fit generative WMs of the observations (images) using a model with low- dimensional latent variables, either continuous (as in VAE) or discrete (as in a VQ-VAE). This includes our method and most background planning methods above 2. In contrast, other methods fit non-generative WMs, which are trained using self predictetit loss see Ni et al. (2024) for a detailed discussion. Non- -onn-generative WMs are more lightweight and therefore well-suited to decision-time planning with its large number of WM calls every decision-makinn step. However, generative WMs are generally preferred for background planning, since it is easy to combine real and imaginary data for policy learning, as we show below.
We concatenate the GRU output to the image embedding, and then pass this to the actor and critic networks, rather than directly passing the GRU output. Algorithm 2, Appendix 1. presents a pseudocode for our MFRL agent.
With these architectural changes, we increase the reward to 55.49% and the score to 16.77% This result is notable since our MFRL agent beats the considerably more com- plex (and much slower) Dreamer V3 agent, which obtains a reward of 53.20% and a score of 14.5. also beats other MBRL methods, such as IRIS (Micheli et al., 2022) (reward of 25.0%) and A-IRIS (Micheli et al., 2024) (reward of 35.0%). In addition, our MFRL agent only takes 15 minutes to train for 1M environment steps on one A100 GPU.
In terms of the architecture of the WM, many state-of-the- art models use transformers, e.g. IRIS (Micheli et al., 2022),
A-IRIS (Micheli al., 2024). DART (Agarwal et al., 2024). Notable exceptions are DreamerV2/5 (Hafner et al., 2020b; 2023), which use recurrent state space models, although improved transformer variants have been proposed (Robine et al., 2023; Zhang et al., 2024; Chen et al., 2022).
3.3. MBRL baseline
We now describe our MBRL baseline, which combines our MFRL baseline above with transformer world model (TWM)-as as in IRIS (Micheli et al., 2022). Following IRIS, our MBRL baseline uses a VQ-VAE, which quantizes the 8 x feature map Z, of a CNN to create a set of latent codes, (q. of = enc (O,). where I. = 64, of E {1....., K} is discrete code, and K 512 the size of the codebook. These codes are then passed to TWM, which is trained using teacher forcing-see Equation (2) below. Our MBRL baseline achieves a reward of 31.93%, and improves over the reported results of IRIS, which reaches 25.0%.
3. Methods
Here, we describe the components of our system, each of which improves performance, as we show in Section 4.
3.1. MFRL Baseline
Our starting point is the previous SOTA MFRL approach which was proposed as baseline in Moon al. (2024)3. This method achieves a reward of 46.91% and a score of 55.60% after 1M environment steps. This approach trains a stateless CNN policy without frame stacking using the PPO method (Schulman et al., 2017), and adds an entropy
Although these MBRL baselines leverage recent advances in generative world modeling, they are largely outperformed by our best MFRL agent. This motivates us to enhance our MBRL agent, which we explore in the following sections. AA notable exception is Diamond (Alonso et al., 2024), which fits diffusion world model directly in pixel space, rather than learning latent WM.
"This is consistent with results on Atari- 100k- which show that well-tuned model-free methods, such as BBF (Schwarzer et al., 2023), can beat more sophisticated model-based methods.
"The authors' main method uses external knowledge about the achievement hierarchy of Crafter, so cannot be compared with other general methods. We use their baseline instead.
3

---

Improving Transformer World Models for Data-Efficient RL
3.4. MBRL using Dyna with warmup
Algorithm 1 MBRL agent. See Appendix A.3 for details.
As discussed in Section we propose to train our MBRL agent on mix of real trajectories (from the environment) and imaginary trajectories (from the TWM), similar to Dyna (Sutton, 1990). Algorithm 1 presents the pseudocode for our MBRL approach. Specifically, unlike many other re- cent MBRL methods (Ha & Schmidhuber, 8a; Micheli et al., 2022; 2024; Hafner et al., 2020b; 2023) which train their policies exclusively using world model rollouts (Step 4), we include Step 2 which updates the policy with real trajectories. Note that, if we remove Steps 3 and 4 in Al- gorithm 1, the approach reduces to MFRL. The function rollout(O1, .T. TM) returns a trajectory of length T gen- erated by rolling out the policy To from the initial state O, in either the true environment Mony or the world model Me. trajectory contains collected observations, actions and rewards during the rollout т (O1:T+1, 01:7; "- Algorithm 4 in Appendix A.3 details the rollout procedure. We discuss other design choices below.
Input: number of environments Newv. environment dynamics Meny.
rollout horizon for environment Temv and for TWM TWM. background planning starting step TaP, total number of environment steps Tioal. number of TWM updates Niters and policy updates Niters
Initialize: observations Of Meny for n 1 Nenv. data buffer D 0, TWM model .M and parameters θ AC model π and parameters 4 number of environment steps t = 0.
repeat
// I. Collect data from environment
Teav rollout(O", πT. Tenr, Mem); n 1 Non
D DU TeiN ; O... T [-1] ; t+ Tenv
// 2. Update policy on environment data
4 = PPO-update-policyy (2 -aiv>)
PPO. Since PPO (Schulman et al., 2017) is an on-policy algorithm, trajectories should be used for policy updates immediately after they are collected or generated. For this reason, policy updates with real trajectories take place in Step immediately after the data is collected. An alternative approach is to use off- off-policy algorithm and mix real and imaginary data into the policy updates in Step 4, hence removing Step 2. We leave this direction as future work.
// 3. Update world model for it = 1 to Niters do WM
T 7 replaa - sample-trajectory(2 TWM), n 1: Nem
θ update- Trelaay
end for
// 4. Update policy on imagined data 2 TAP then
for it 1 to Nrtters do AC
Rollout horizon. We set Twm < Temv. to avoid the problem of compounding errors due to model imperfections (Lambert et al., 2022). However, we find beneficial to use TwM 1, consistent with Holland et al. (2018); Van Hasselt et al. (2019), who observed that the Dyna approach with TwM 1 is no better than MFRL with experience replay.
Ô = sample-obs(D), n = 1: Nenv
TWM = vollou((f To, TWM: Me), n = 1 New
n
T = IPPO-udate-policyye TWM
1::oo)
end for end if
until f Tinal
Multiple updates. Following IRIS, we update TWM Niters times and the policy on imagined trajectories Niters times.
scribe our alternative which leverages a property of Craftax- classic: each observation is composed of 9 x 9 patches of size 7 7 each (see Figure 2[left)). Hence we propose to (a) factorize the tokenizer by patches and (b) use a simpler nearest-neighhor style approach to tokenize the patches.
Warmup. When mixing imaginary trajectories with real ones, need to ensure the WM is sufficiently accurate so that it does not harm policy learning. Consequently, we only begin training the policy on imaginary trajectories after the agent has interacted with the environment for T&P steps, which ensures it has seen enough data to learn a reliable WM. We call this technique "Dyna with warmup"" In Section 4.3, we show that removing this warmup, and using TBP = 0, drops the reward dramatically, from 67.42% to 33.54%. We additionally show that removing the Dyna method (and only training the policy in imagination) drops the reward to 55.02%.
Patch factorization. Unlike prior methods which process the full image O into tokens (q'.....g) .²²) ene(O), we first divide O into L non-overlapping patches (p² (p'....,, p²) which are independently encoded into L tokens:
(q'.....g' 9 =
To convert the discrete tokens back to pixel space, we just decode each token independently into patches, and rearrange to form a full image:
3.5. Patch nearest-neighhor tokenizer
Many MBRL methods based on TWMs use a VQ-VAE to map between images and tokens. In this section, we de-
(6 (......,p" .,.,, =
4

---

Improving Transformer World Models for Data-Efficient RL
Factorizing the VQ-VAE on the L 81 patches of each observation boosts performance from 43.36% 58.92%.
Supervision
Attention mask
Teacher forzing with causal attention
Nearest-neighbor tokenizer. On top of patch factorization, we propose a simpler nearest-neighhor tokenizer (NNT) to replace the VQ-VAE. The encoding operation for each is similar to a nearest neighbor
TWM
qi
qi
patch p [0, classifier w.r.t the codebook. The difference is that, if the nearest neighbor is too far away, we add new code equal to p to the codebook. More precisely, let us denote CNN (e1,..., ,..) the current codebook, consisting of K codes e E [0, 1|xx0x3, and т threshold on the Euclidean distance. The NNT encoder is defined as:
2
43
teacher forcing with black causal attention
TWM
Qi
91
42
Figure 3: Approaches for TWM training with L = 2, T 2. of denotes token & of timestep t. Tokens in the same timestep have the same color. We exclude action tokens for simplicity. [TTop] Usual autoregressive model training with teacher forcing. [Bottom] Block teacher forcing predicts token o4+1 from input token of with block causal attention.
argmin ||p e:112 if min c=ll 1SISK
S
4 = enc(p)
K +
otherwise.
(1)
The codebook can be thought of as greedy approximation to the coreset of the patches seen so far (Mirzasoleiman et al., 2020). To decode patches, we simply return the code associated with the codebook index, i.e. dec(g') Egi.
previous timesteps, BTF independently predicts all the latent tokens the next timestep, removing the conditioning on previously generated tokens from the current step:
L E1, El = .
A key benefit of NNT is that once codebook entries are added, they are never updated. A static yet growing code- book makes the target distribution for the TWM stationary, greatly simplifying online learning for the TWM. In contrast, the VQ- VAE codebook is continually updated, meaning the TWM must learn from a non-stationn non-stationary distribution, which results in worse WM. Indeed, we show in Section 4.1 that with patch factorization, and when h = 7 meaning that the patches are aligned with the diservation-replaatii the VQ-VAE with NNT boosts the agent's 'ss reward from 58.92% to 64.96%. Figure 2[right] shows an example of the first 64 code patches extracted by our NNT.
LBTE = log IIII 1=
(3)
Importantly BTF uses a block causal attention pattern (see Figure 3), in which tokens within the same timeslice are decoded in-parallel in single forward pass. This attention structure allows the model to reason jointly about the pos- sible future states of all tokens within timestep, before sampling the tokens with independent readouts. This prop- erty mitigates autoregressive drift. As a result, BTF returns more accurate TWMs than fully AR approaches. Overall, adding BTF increases the reward from 64.96% to 67.42%. In addition, we find that BTF is twice as fast, even though in theory, with key-value caching, BTF and AR both have complexity (T.?T) for generating all the I. tokens at one timestep. and 2 for generating the entire rollout. Finally, BTF shares a similarity with Retentive Environment Models (REMs) (Cohen et al., 2024) in their joint predic- tion of next-frame tokens. However, while REMs employ retentive network (Sun et al., 2023), BTF offers broader applicability across any transformer architecture.
The main disadvantages of our approach are that (a) patch tokenization can be sensitive to the patch size (see Fig- ure 5[left]), and (b) NNT may create large codebook if there is a lot of appearance variation within patches. In Craftax-classic these problems are not very severe due to the grid structure the game and limited sprite vocabulary (although continuous variations exist due to lighting and texture randomness).
3.6. Block teacher forcing
4. Results
Transformer WMs are typically trained by teacher forcing which maximizes the log likelihood of the token sequence generated autoregressively over time and within a timeslice:
In this section, we report our experimental results on the Craftax-classic benchmark. Each experiment is run on 8 H100 GPUs All methods are compared after interacting with the environment for Tomal 1M steps. All the methods collect trajectories of length Teny 96 in New = 48 envi- ronment (in parallel). For MBRL methods, the imaginary rollouts are of length TWM = 20, and we start generating these (for policy training) after Tpp = 200k environment steps. We update the TWM Ni = 500 times and the policy Niterr 150 times. For all metries, we report the
T
L
CTT = log IIII Ci e = p(gh aillllg Ii Ser..e41 an))
(2)
We propose more effective alternative, which we call block teacher forcing (BTF). BTF modifies both the supervision and the attention of the TWM. Given the tokens from the
5

---

Improving Transformer World Models for Data- RL
Table 1: Results on Craftax-classic after IM environment interactions. denotes results on Crafter, which may not exactly match Craftax classic. - means unknown. +denotes the reported timings on a single A100 GPU. Our DreamerV3 results are based on the code from the author, but differ slightly from the reported number, perhaps due to hyperparameter discrepancies.
IRIS and A A-IRIS do not report standard errors for the score.
Method
Parameters
Reward (%)
Score (%)
Time (min)
Human Expert
NA
+65.0 +: 10.5
+50.5 + 6.8
NA
MI: Baseline
60.0M
31.93 + 2.22
4.98 1 0.50
560
M2: M1 Dyna
60.0M
43.36 1: 1.84
8.85 + 0.63
563
M3: M2 patches
56.6M
58.92 + 1.03
19.36 + 1.42
746
M4: M3 NNT
58,5M
64.96 + 1.13
25.55 0.86
1328
M5: M4 + BTF Our best MBRL (fast)
58,5M
67.42 0.55
27.91 + 0.63
759
M5: M4 + BTF, Our best MBRL (slow)
58.5M
69.66 + 1.20
31.77 + 1.43
2749
Previous best MFRL (Moon et al., 2024)
4.0M
*6.91 + 2,41
#15.60 + 1.66
Previous best MFRL (our implementation) Our best MFRI,
4.0M
47.40 + 0.58
10.71 + 0.29
26
55.6M
55.49 + 1.33
16.77 + 1.11
15
33eamerV3 (Hafner et al., 2023)
201M
+53.2
14.5 + 1.6
Our D3eamer
201M
47.18 3.88
2100
IRIS (Micheli al., 2022)
48M
225.0 + 3.2
*6.66
+8330
A A-IRIS (Micheli et al., 2024)
25M
+35.0 + 3.2
*9.30
+833
Curious Replay (Kauvar et al., 2023)
#19.4 + 1.6
mean and standard error over 10 seeds as .x+y).
codebook size of 512. For M4 and M5, which use NNT, we found it critical to use a larger codebook size of K 4096 and a threshold of T = 0.75. Interestingly, when training in imagination begins (at step TвP = 200k), there is temporary drop performance as the TWM rollouts do not initially match the true environment dynamics, resulting in distribution shift for the policy.
4.1. Climbing up the MBRL ladder
First, we report the normalized reward (the reward divided by the maximum reward of 22) for a series of agents that progressively climb our "MBRL ladder" of improvements in Section 3. Figure 4 shows the reward vs. the number of environment steps for the following methods, which we detail in Appendix A.2:
70
Iris
M1: Baseline. Our baseline MBRL agent, described in
M1: Baseline
M2: MI Dyna
Section 3.3, reaches a reward of 31.93%, and improves over IRIS, which gets 25.0%.
M3: patches
M4: M3 + NNT
Reward (%)
M3:
MBRL (fast)
M2: M1 + Dyna. Training the policy on both (real) environment and (imagined) TWM trajectories, as described in Section 3.4, increases the reward to 43.36%.
M5: M4 GTF Our besi MBRL Islow)
30
M3: M2 + patches. Factorizing the VQ-VAE over the 81 observation patches, as presented in Section 3.5, increases the reward to 58.92%.
10
M4: M3+ + NNT. With patch factorization, replacing the VQ- VAE with NNT, as presented in Section 3.5, further boosts the reward to 64.96%.
0.0
0.2
0.4
0.6
0.8
1.0
Environment steps
Figure 4: The ladder of improvements presented in Sec- tion 3 progressively transforms our baseline MBRL agent into sfate-of-the-arr method on Craftax-classic, reaching a reward of 69.66 (averaged over 10 seeds) after 1M envi- ronment steps. Training in imagination starts at step 200k, indicated by the dotted vertical line.
M5: M4 + BTF, Our best MBRL (fast): Incorporating BTF, as described in Section 3.6, leads to our best agent. achieves a reward of 67. 67.42% while BTF reduces the training time by factor of two.
M5: M4 + BTF. Our best MBRL (slow): By increasing the number of TWM training steps to WM = 4k, we obtain our best agent, which reaches a reward of 69.66%. However, due to substantial training times (~ 2 days), we do not include this agent our ablation studies (Section 4.3) and comparative studies (Section 4.4).
4.2. Comparison to existing methods
Figure I compares the performance of our best MBRL and MFRL agents against various previous methods. See also Figure 9 in Appendix for a plot of the score, and Table I for a detailed numerical comparison of the final performance. First, we observe that our best MFRL. agent outperforms
As in IRIS (Micheli et al., 2022), methods M1-3 use
6

---

Improving Transformer World Models for Data-Efficient RL
60
60
60
JoM AMY
Reward (%)
Reward (%)
Reward (%)
40
40
40
Start BP at OK
Dur best MBRL Host, quartired)
(fast)
Start BP st 100e
20
20
Remove Dyna
Our best Mann, ffust,
5c5 quartived
Berove NKT
Start SP 200k)
Redi tood continuous
Remove NKT patches
Start BP at. 400k
Remove
Start 600k
01
0.0
0.2
0.4
0.6
0.8
10
0.0
0.2
0.4 0.6
0.8
1.0
0.2
0.8
1.0
Environment steps
Environment steps
Environment steps
Figure 5: [Left] MBRL. performance decreases when NNT uses patches of smaller or larger size than the ground truth, but it remains competitive. However, performance collapses if the patches are not quantized. [Middle] Removing any rung of the ladder of improvements leads to a drop in performance. [Right] Warming up the world model before using train the policy on imaginary rollouts is required for good performance. BP denotes background planning. For each method, training in imagination starts at the color-coded vertical line, and leads an initial drop in performance.
almost all of the previously published MFRL and MBRL results, reaching reward of 55.49% and a score of 6.7% Second, our best MBRL agent achieves new SOTA reward of 69.66% and score of 31.77% This marks the first agent to surpass human-level reward, derived from 100 episodes played by 5 human expert players (Hafner, 2021). Note that although we achieve superhuman reward, our score is significantly below that of a human expert.
When to start training in imagination? Training the policy on imaginary TWM rollouts requires a reasonably accurate world model. This is why background planning (Step 4 in Algorithm 1) only begins after TвP environment steps. Figure S[right] explores the effect of varying TEP. Initiating imagination training too early (ТвP = 0) leads to perfor- mance collapse due to the inaccurate TWM dynamics.
MFRL ablation. The final 3 rows in Table 2 show that either removing the RNN or using a smaller model as in Moon al. (2024) leads drop in performance.
4.3. Ablation studies
We conduct ablation studies to assess the importance of several components of our proposed MBRL agent. Results are presented in Figure and Table 2. All the TWMs are trained for Niters 500 steps.
Table 2: Ablations results.
Method
Reward (%)
Score (%)
Our best MBRL (fast)
67.42 0.55
27.91 + 0.63
Impact of patch size. We investigate the sensitivity of our approach to the patch size used by NNT. While our best results are achieved when the tokenizer uses the oracle- provided ground truth patch size of 7 x 7, Figure S[1eft] shows that performance remains competitive when using smaller (5 5) or larger (9 9) patches.
5 5 quantized
57.28 + 1.14
18.26 + 1.18
9 9 quantized
45.55 + 0.88
10.12 0.40
7 x 7 continuous
21.20 + 0.55
2.43 0.09
Remove Dyna Remove NNT
55.02 + 3.34
18.79 + 2.14
60.66 + 1.38
21.79 + 1.33
Remove NNT & patches
45.86 + 1.42
10.36 + 0.69
Remove BTF
61.96 + 1.13
25.55 + 0.86
The necessity of quantizing. Figure 5[left] shows that, when the 7 x 7 patches are not quantized, but instead the TWM is trained to reconstruct the continuous 7 x 7 patches, MBRL performance collapses. This is consistent with find- ings in DreamerV2 (Hafner, 2021), which highlight that quantization is critical for learning an effective world model.
Use T&p 0
33.54 + 10.09
12.86 + 4.05
Best MFRL
55.49 + 1.33
16.77 1 1.11
Remove RNN
41.82 + 0.97
8.33 1 0.44
Smaller model
51.35 1 0.80
12.93 1 0.56
Annealing the number of policy updates: We linearly increase the number of policy updates on imaginary rollouts in Step 4 of Algorithm from Nilerrs 0 (when Tioal 0) to AC = 300 (when Tioal = 1M). This annealing technique achieves a reward of 65.71% while removing the drop in performance observed when we start training in imagination. See Figure 10 Appendix C.
Each rung matters. To isolate the impact of each individual improvement, we remove each individual "rung" of our ladder from our best MBRL agent. As shown in Figure 5mmiddee], each removal leads to a performance drop. This underscores the importance of combining all our proposed enhancements to achieve SOTA performance.
'The only exception is Curious Replay (Kauvar et al., 2023), which builds on DreamerV3 with prioritized experience replay (PER) to train the WM. However, PER is only better on few achievements; this improves the score but not the reward.
4.4. Comparing TWM rollouts
In this section, we compare the TWM rollouts learned by three world models in our ladder, namely M1, M3 and our
7

---

Improving Transformer World Models for Data-Efficien RL
PLADE
DOWN
NOOP
M1: Sapcline
True rollout
Symbol accuracy (1)
0.95
M3:M1 Dyna - patches
45:M3 NNT BTR Our best MBRL (tart)
0.90
M1
0.85
24
hare
M3
15
20
TWM rallaut horizon
t.ve
Map inconsistency
Rollout annotations
Feasible hallucination
MS: our best
Infessible hallucination
Figure 6: Rollout comparison for world models MI, MB and M5 (fast). [Left] Symbol accuracies decrease with the TWM rollout step. The stationary NNT codebook used by M5 makes it easier to learn a reliable TWM. [Right] Best viewed zoomed in. Map. All three models accurately capture the agent's 's motion. All models can struggle to use the history to generate consistent map when revisiting locations, however only MI makes simple map errors in successive timesteps. Feasible hallucinations. M3 and M5 generate realistic hallucinations that respect the game dynamics, such as spawning mobs and losing health. Infeasible hallucinations. MI often does not respect game dynamics; MI incorrectly adds wood inventory, and incorrectly places a plant at the wrong timestep without the required sapling inventory. M3 exhibits some infeasible hallucinations in which the monster suddenly disappears or the spawned cow has an incorrect appearance. M5 rarely exhibits infeasible hallucinations. Figure 12 in Appendix D.4 shows more rollouts with similar behavior.
best model M5 (fast). To do so, we first create an evaluation dataset of Neval 160 trajectories, each of length Teval TwM= = 20, collected during the training of our best MFRL 1:NNes 1:NNou
still be useful for learning in imagination-as long as they respect the game dynamics. Visual inspection of rollouts in Figure 6[right] reveals (a) map inconsistencies, (b) feasible hallucinations that respect the game dynamics and (c) infea- sible hallucinations. MI can make simple mistakes in both the map and the game dynamics. M3 and M5 both generate feasible hallucinations of mobs, however M3 more often hallucinates infeasible rollouts.
1:Tea-1: "1:Tes >i:Toa J. We evaluate the quality of imagined trajectories generated by each TWM. Given TWM checkpoint at IM steps and the wihh trajectory in Deval, execute the sequence of actions TT starting from Op, to obtain a rollout trajectory ÓTWM. +1
agent: Deval
{
Quantitative evaluations. For evaluation, we leverage an appealing property of -rastx-classii: each observation O comes with array of ground truth symbols S, (S!:R), with R- 145. Given 100k pairs Se). we train CNN fu. to predict the symbols from the observation; fy achieves a 99% validation accuracy. Next, we use fu to predict the symbols from the generated rollouts. Figure 6[left] displays the average symbol accuracy at each timestep t:
4.5. Craftax Full
Table 3: Results on Craftax after 1M environment interac- tions. The previous SOTA scores are unknown.
Method
Reward (%)
Score (%)
Prev. SOTA MFRL
2.3 (symbolic)
Our best MFRL
4.63 + 0.20
1.22 + 0.07
Prev. SOTA MBRI.
6.59
Our best MBRL (slow)
7.20 1 0.09
2.31 0.04
Neal
A Noval
(O "), Vt.
Table 3 compares the performance of various agents on the full version of Craftax (Matthews et al., 2024), a signifi- cantly harder extension of Crastax-classic, with more levels and achievements. While the previous SOTA agent reached 2.3% reward (on symbolic inputs), our MFRL agent reaches 4.63% reward. Similarly, while the recent SOTA MBRL (Cohen et al., 2025) reaches 6.59% reward our MBRL agent reaches a new SOTA reward of 7.20%. See Appendix E for implementation details.
where ((x; y) iff. y (and 0 S. denotes the ground truth rth symbol in the array ST associated with OF, and OO..M. ") its prediction for the rollout observation. As expected, symbol accuracies decrease with t as mistakes compound over the rollouts. Our best method, which uses NNT, achieves the highest accuracies for all timesteps, as best captures the game dynamics. This highlights that a stationary codebook makes TWM learning simpler.
We include two additional quantitative evaluations in Ap- pendix D, showing that M5 achieves the lowest tokenizer reconstruction errors and rollout reconstruction errors.
4.6. Additional experiments on MinAtar
To further validate the robustness of our approach, we con- duct additional experiments on MinAtar (Young & Tian, 2019), another grid world environment. MinAtar imple- ments four simplified Atari 2600 games. Each game has
Qualitative evaluations. Due to environment stochasticity, TWM rollouts can differ from the environment rollout but
8

---

Improving Transformer World Models for Data-Efficient RL
Asterix
Breakout
Freeway
Spacelnvaders
50
80
60
150-
30
60
40
100
20
40
10
20
20-
50
0
0.00
0.25
0.50
0.75
1.00 106
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00 is
Environment steps
Environment steps
106
Environment steps
106
Environment steps
Baseline
Dyna
patches
NNT
MBRL
Figure Our best MBRL agent leads to significant gains over our tuned MFRL agent on each MinAtar game.
symbolic binary observations of size 10 x 10 x K (K is the number of objects of the game) and binary rewards.
the significant performance improvements achieved by our proposed MBRL agent.
We first tune our model-free RL agent on the MinAtar games, keeping the same architecture as described our paper, with minor adjustments to the PPO hyperparameters, detailed in Appendix F. Second, we develop our model- based RL agent as in Craftax-classie, by integrating our three proposed improvements. We retain the majority of the MBRL hyperparameters from Craftax-classic, with minor modifications, which we detail in Appendix F.
Table 4: Best MFRI and best MBRI rewards after 1 MM steps each MinAtar game.
Game
MFRL
MBRL.
Asterix
7.47 L.02
44.81 3.54
Breakout
77.8 2.28
93.92 1.44
Freeway
65.3 1.16
71.12 186.16 1.25
SpaceInvaders
+.3.9
Figure 7 displays the evaluation performance of our pro- posed methods M1-5 (defined in Section 4.1) each game after 1 million environment steps, averaged over 10 seeds. Every 50k training steps, we evaluate each agent on 32 environments and 2k steps per environments. Figure summarizes these results by first (a) normalizing each game such that the MFRL agent achieves a reward of 1.0, before (b) averaging the performance of all agents across the games. Notably, our MBRL agents' performance increase as we climb the ladder on MinAtar, highlighting the gen- erality of our three proposed improvements. Furthermore, our best MBRL agent significantly outperforms our best MFRL agent, achieving an average normalized reward of 2.41 across the four MinAtar games.
5. Conclusion and future work
In this paper, we present three improvements vision- based MBRL agents which use transformer world models for background planning: Dyna with warmup, patch nearest- neighbor tokenization and block teacher forcing. We also present improvements to the MFRL. haseline, which may be of independent interest. Collectively, these improve- ments result MBRL agent that achieves a significantly higher reward and score than previous SOTA agents on the challenging Craftax-classie benchmark, and surpasses ex- pert human reward for the first time. Our improvements also transfer to MinAtar environments. In the future, we plan to examine how well our techniques generalize beyond grid-world environments. However, we believe our current results will already be of interest to the community.
MinAtar summary
2.5
RL
M1: Baseline
M2: M1 Dyna
We see several paths build upon method. Prioritized experience replay is promising approach to accelerate TWM training, and an off-policy RI. algorithm could im- prove policy updates by mixing imagined and real data. In the longer term, we would like to generalize our tokenizer to extract patches and tokens from large pre-trained models, such as SAM (Ravi al., 2024) and Dino-Va V2 (Oquab et al., 2024). This inherits the stable codebook of our approach, but reduces sensitivity to patch size and "superficial" ap- pearance variations. To explore this direction, and other non-reconstructive world models which cannol generate fu- ture pixels, we plan to modify the policy to directly accept latent tokens generated by the TWM.
Normalized Mormalizeed reward
2.0
M3: M2
patches
M4: M3 NNT
M5:
MBAL
1.5
1.0
0.5
0.0
0.0
0.2
0.6 Environment steps
0.4
0.8
1.0
les
Figure 8: Averaged normalized reward on MinAttr.
Finally, Table 4 compares the performance of our best MBRL and MFRL, agents at 1M steps, further emphasizing

---

Improving Transformer World Models for Data-Efficient RL
Acknowledgments
References
We thank Pablo Samuel Castro for useful discussions during the preparation of this manuscript.
Agarwal, P., Andrews, S., and Kahou, S. E. Learning to play atari in a world of tokens. ICML, 2024.
Alonso, E., Jelley, A., Micheli, V., Kanervisto, A., Storkey, A., Pearce, T., and Fleuret, F. Diffusion for world model- ing: Visual details matter in atari. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL htttps://openreveew
Impact Statement
This paper proposes new techniques for creating Al agents that can rapidly learn to play well in open-world environ- ments. As our best agent surpasses human-level perfor- mance in challenging benchmark game with limited in- teraction, this work paves the way for developing more sample-efficient and general-purpose Al agents capable of tackling real-world problems that demand strong general- ization, exploration, and long-term reasoning.
openrevieeeeeeeeeeeeeeeeee / nnt//
forum? Maadiiid--aWWDSSS
Alver, S. and Precup, D. A look at value-based decision- time background planning methods across different settings. In Seventeenth European Workshop Re- inforcement Learning, October 2024. URL https:
/ //openreview net. wtt..:/pdddiieeeeEEet......
Ba, J. I.., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint /o0.::::0070.445500 2016.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman- S.,, and Zhang, Q. JAX: composable transformations of Python+NumPl programs, 2018. URL http: /ttt::///ittttt www.cooo////.//.. jtap//aaalllmmml mll jax.
Chen, C., Wu, Y.-F,, Yoon, J., and Ahn, S. Transdreamer: Reinforcement learning with transformer world models. URL hutp:/arxiv. org/abs/2202, 9481, 2022.
Cohen, L., Wang, K., Kang. B., and Mannor, S. Improv- ing token-based world models with parallel observation prediction. arXiv preprint arXiv:2402. 05643,, 2024.
Cohen, L., Wang, K., Kang, B., Gadot, U. and Mannor, S. M3: A modular world model over streams of tokens. arXiv preprint arXiv:2502. 11337, 2025.
Espeholt, L., Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dun- ning, I., Legg. S., and Kavukcuoglu, K. IMPALA: Scalable distributed deep-RL RL with importance weighted actor-learner architectures. In ICML, pp. 1407-1416. PMLR, July 2018a. URL https: http:////occeeinn...... / httpeeepppppttteeeeeeeeeee http:reee//v///eeppoll1111/1/1
mlr. w.r.ppress/ee mlr.press/ v8C
html.
Espeholt, L., Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, I. et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. Interna- tional conference on machine learning, pp. 1407-1416. PMLR, 2018b.
Farebrother, J., Orbay, J., Vuong, Q., Taiga, A. A., Chebotar, Y., Xiao, T., Irpan, A., Levine, S., Castro, P. S., Faust, A., Kumar, A., and Agarwal, R. Stop regressing: Training value functions via classification for scalable deep RI. In
10

---

Improving Transformer World Models for Data-Efficient RL
Forty-first International Conference on Machine Learn- ing, June 2024. URL httpt:t////ccc...t. http:p/:///penreeee....oo pdf2id-dVVFFKfFFF
Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W. Recurrent experience replay in distributed reinforcement learning. In International conference learning representations, 2018.
Ha, D. and Schmidhuber, J. World models. In NIPS, 2018a. URL htttt://aaxiiii org/ abs/ 1801883.000 10122.
Kauvar, I., Doyle, C., Zhou, L., and Haber, N. Curious replay for model-based adaptation. In ICML, June 2023. URL https: wtwpiviiivvv//////a///...//aaavv wttpo/abbb........2/22
Ha, D. and Schmidhuber, J. Recurrent world models facil- itate policy evolution. Advances in neural information processing systems, 31, 2018b.
Kingma, D. P. Adam: method for stochastic optimization.
arXiv preprint 14 2014.
Hafner, D. Benchmarking the spectrum of agent capabilities. arXiv preprint 1rXiv:2109006880 2021.
Lambert, N., Pister, K., and Calandra, R. Investigat- ing compounding prediction errors in learned dynam- ics models. arXiv [cs.LG], March 2022. URL http: / //arxiv.org abgg/////////// / 2203. 09637.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. ICLR, 2020a. URL https: http:////peeeeeeeeeeeeee net./ Forum?id-S11orT
Lei Ba, J., Kiros, R., and Hinton, G. E. Layer normaliza- tion. ArXiv e-prints, pp. arXiv-1607, 2016.
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. Mas- tering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020b.
Lu, C., Kuba, J., Letcher, A., Metz, L., Schroeder de Witt, C., and Foerster, J. Discovered policy optimisation. Ad- vances in Neural Information Processing Systems, 35: 16455-16468, 2022.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. Mastering diverse domains through world models. arXiv preprint SXXiv:3301.04004 2023.
Matthews, M., Beukman, M., Ellis, B., Samvelyan, M., Jackson, M., Coward, S., and Foerster, Craftax: lightning-fast benchmark for open-ended reinforcement learning. arXiv preprint arXiv:2402. 6801, 2024.
Hansen, N., Su, H., and Wang. X. TD-MPC2: Scalable, robust world models for continuous control. 2024. URL http:// arxiv. org/ abs/ /2100.110.. 16828.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual leamn ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Micheli, V., Alonso, E., and Fleuret, F. Transform- ers sample-efficient world models. arXiv preprint 2:Xi::2209.00888 2022.
Micheli, V., Alonso, E., and Fleuret, F. Efficient world models with context-aware tokenization. arXiv preprint arXiv:2406.1993 19320, 2024.
Hessel, M. Modayil, J., van Hasselt, H., Schaul, T., Os- trovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. In AAAI, 2018. URL http:/ htap:://aaxvvv org/ abs /tss///11... 02298.
Mirzasoleiman, B., Bilmes, J., and Leskovec, J. Coresets for data-efficient training of ma- chine learning models. In ICML, 2020. URL mi http:p/ooceeiii.mll.l...momm
Holland, G. Z., Talvitie, E. J., and Bowling, M. The ef- feet of planning shape on dyna-style planning in high- dimensional state spaces, arXiv [cs.Al], June 2018. URL http:// //riviv.//riv.///// wwtgabs/////oooooooe/..gg/..eese
.df.
abs / 111.//806.000 01825.
Moerland, T. M., Broekens, J., Plaat, A., and Jonker, C. M. Model-based reinforcement learning: survey, Foundations and Trends Machine Learning, 16(1):1- 118, 2023. URL htts:t /ttpxxxxvvv.////////ee////aaa.va ort.vv./0rg/bb//00666. 16712.
Ioffe, S. and Szegedy, C. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448- 456. pmlr, 2015.
Moon, S., Yeom, J., Park, B., and Song, H. O. Discovering hierarchical achievements in reinforcement learning via contrastive learning. Advances in Neural Information Processing Systems, 36, 2024.
Kaiser, L.,, Babaeizadeh, M., Milos, P., Osinski, B., Camp- bell, R. H., Czechowski, K., Erhan, D., Finn, C., Koza- kowski, P., Levine, S., Mohiuddin, A., Sepassi, R., Tucker, G., and Michalewski, H. Model-based reinforce- ment learning for atari. arXiv (cs. March 2019. URL http://arxe www....rgrg..s....... 00374.
Murphy, K. Reinforcement learning: An overview, arXiv preprint arXiv:24/2.05266 2024.

---

Improving Transformer World Models for Data-Efficient RL
Ni, T., Eysenbach, B., Seyedsalehi, E., Ma, M., Gehring, C., Mahajan, A., and Bacon, P.-.. Bridging state and history representations: Understanding self-predictive RL. In ICLR, January 2024. URL http://arxiy org/abs/ 2401 08898.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W. and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Sun, Y., Dong, L.,, Huang. S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.0862 2023.
OpenDILab. Awesome Model-Based Reinforcement Learning. https://gitie wttphhub.ooomeeeoooommme..eeeeee oppdddiab/ awesome model-bol bosell-aaaed-Ree RL.
Sutton, R. S. Integrated architectures for learning, plan- ning, and reacting based on approximating dynamic pro- gramming. In Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
Oquab, M., Darcet, T., Moutakanni, T. Vo, H. V., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., Assran, M., Ballas, N., Galuba, W., Howes, R., Huang, P.-Y,, Li, S.-W., Misra, I,, Rab- bat, M., Sharma, V., Synnaeve, G., Xu, H., Jegou, H., Mairal, J., Labatut, P., Joulin, A., and Bojanowski. DINOv2: Learning robust visual features without super- vision. Transactions Machine Learning Research, www.:/ww.eeeeee.eeeeeeeeeeeeeeee wwt.eee/toooooe nwt./orrmttft
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT press, 2018.
Toledo, E., Midgley, L., Byrne, D., Tilbury, C. R.. Macfar- lane, M., Courtot, C., and Laterre, A. Flashbax: Stream- lining experience replay buffers for reinforcement learn- ing with jax, 2023. URL https: http:/ggttthhcccc com/ instadeepai satrepaaa/laasaaa.... flashbax/.
2024. URL
id=a68SUt 62Ft.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
Ramachandran, P., Zoph, B., and Le, Q. V. Swish: self-gated activation function. arXiv preprint arXiv: 1710.0541, (1):5,2017.
Van Hasselt, H. P., Hessel, M., and Aslanides, J. When to use parametric models in reinforcement learning? Ad- vances in Neural Information Processing Systems, 32, 2019.
Ravi, N., Gabeur, V., Hu, Y.-T., Hu, R., Ryali, C., Ma, T., Khedr, H., Rädle, R., Rolland, C., Gustafson, L., al. Sam 2: Segment anything in images and videos. arXiv preprint arXivv2408.007 00714, 2024.
Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. Mas- tering atari games with limited data. In NIPS, Novem- ber 2021. URL https: httpp::://ppeeeieew... net/ddf? id=OKeNPG3xR31
Robine, J., Höftmann, M., Uelwer, T., and Harmeling, S. Transformer-based world models are happy with 100k interactions. arXiv preprint arXi::2303.071100 2023.
Young. K. and Tian, T. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv:1903.03/70 2019.
Schritwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifree, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020.
Zhang, W. Wang, G., Sun, J., Yuan, Y., and Huang, G. Storm: Efficient stochastic transformer based world mod- els for reinforcement learning. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Schulman, J., Moritz, P., Levine, S., Jordan, M., and Abbeel, P. fligh-dimensional continuous control using generalized advantage estimation. arXiv preprint teXiv:1506.025 2015.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06377 2017.
Schwarzer, M., Obando-Ceron, J., Courville, A., Bellemare,
M., Agarwal, R., and Castro, P. S. Bigger, better, faster: Human- level atari with human-level efficiency. In ICML,
May 2023. URL http:////// / arxiv. wwtp:/s./g///.a............... 23005. 9452.
12

---

Improving Transformer World Models for Data-Efficient RL
A. Algorithmic details
A.1. Our Model-free RL agent
We first detail our new state-of-the-arr MFRL agent. As mentioned the main text, it relies on an actor-critic policy network trained with PPO.
A.I.I. MFRL ARCHITECTURE
We summarize our MFRL agent in Algorithm 2 and further detail it below.
Algorithm 2 MFRL agent
Input: Image O₁, last hidden state - parameters + Output: action dg, value t/2, hidden state ha
2, ImpalaCNN, (Or)
he, 1/1 = RNN 4. ([h--,, 2(])
ag ~ T (1ve, )+)
01 V [yt, =
Imapala CNN architecture: Each Craftax-classie image O₂ of size 63 x 63 x 3 goes through Impala CNN (Espeholt et al., 2018b). The CNN consists of three stacks with channel sizes of (64, 64, 128). Each stack is composed of (a) batch normalization (Ioffe & Szegedy, 2015), (b) convolutional layer with kernel size 3 x 3 and stride of 1, (c) max pooling layer with kemee size 3 x 3 and stride of 2, and (d) two ResNet blocks (He et al.. 2016). Each ResNet block is composed of (a) ReLU activation followed by a batch normalization, (b) a convolutional layer with kernel size 3 x 3 and stride of 1. The CNN last layer output, of size 8 x x 128 passes through a ReLU activation, then gets flattened into an embedding vector of size 8192, which we call 2t.
RNN architecture: The CNN output 2, (a) goes through layer norm operator, (b) then gets linearly mapped 256-dimensional vector, (c) then passes through ReLU activation, resulting in the new input for the RNN. The RNN then updates its hidden state, and outputs 556-dimensional vector yr. which goes through another ReLU activation.
Actor and critic architecture: Finally, the CNN output 2, and the RNN output Ye are concatenated, resulting in the 5448-dimensional embedding input shared by the actor and the critic networks. For the actor network, this shared input goes through (a) a layer normalization (Lei Ba et al., 2016), (b) fully-connected network whose 2048-dimensional output goes through ReLU, (c) two dense residual blocks whose 2048-dimensional output goes through a ReLU, (d) last layer normalization and (e) final fully-connected network which predicts the action logits.
Similarly, for the critic network, the shared input goes through (a) a layer normalization, (b) fully-connected network whose -is--dimensional output goes through a ReLU, (c) two dense residual blocks whose 2048- 1048-dimmnssioal output goes through a ReLU, (d) a last layer normalization and (c) final layer which predicts the value (which is a float).
A.1.2. PPO TRAINING
We train MFRL agent with the PPO algorithm (Schulman al., 2017). PPO is policy gradient algorithm, which briefly summarize below.
Training objective: We assume we are given trajectory, = (O1.T+1, Q1:7, T1:7, donei.r, herr) collected in the environment, where done; is a binary variable indicating whether the current state is a terminal state, and h is the RNN hidden state collected while executing the policy. Algorithm discusses how we collect such a trajectory.
Given the fixed current actor-critic parameters o PPO first runs the actor-critic network on 7, starting from the hidden state ho and returns two sequences of values U1-T-1 Ve.. (0) and probabilities T (a,) It then defines the generalized advantage estimation (GAE) as in Schulman al. (2015):
A = o, + (1 done;) = 8, + (1 done;) (726'41 +... + (71) S
"We drop the ImpalaCNN the RNN for simplicity.
13

---

Improving Transformer World Models for Data- RL
where
6₁ = (1-done, Vt.
PPO also defines the TD targets qt A Fz.
PPO optimizes the parameters @ minimize the objective value:
1 T Σ min clip(r) (vV((V(O.)) (mm
Lpro ()
(4)
where clip(u) ensures that lies in the interval [1 c, + c], r2( the probability ratio ry r+))) the entropy operator.
(a O+) and H. so (az (O.)
Algorithm: Algorithm details the PPO-update-pone ppopltte-polcc,, which called in Steps and 4 in our main Algorithm to update PPO parameters on a batch of trajectories. PPO allows multiple epochs of minibatch updates on the same batch and introduces two hyperparameters: a number of minibatches Nmb (which divides the number of environments Nemy), and a number of epochs Nepuch
Algorithm 3 PPO-update- policy
Input: Actor-critic model (x, V) and parameters &
I:Nes I:New, I:NNss Trajectories y I: (O 1:T "1:I" /1T done I:Naw L:Nas)
Number epochs Nepeec and of minibatches Amb PPO objective value parameters 7, A Learning rate Ir and max tex-grrddeent-oone norm
Moving average mean Marget, standard deviation daarget and discount factor
Output: Updated actor-critic parameters
Initialize: Define Pold = +
1.NN Compute the values 11:4+1
VoL. ( 1:1:No
Compute PPO GAEs and targets A 22:NE0 11:T
1:Non GAB(CINN 1:Nor 1:Nor
1:T 11:41:77
Standardize PPO GAES 1:NNar
A A..i meaa(( Nom
A
for ep = to Nepooh do for = 1 to Amb do sstart (A: 1) (New /nmb) + 1, Nemd & (Noms/ /bm
// Standardize PPO targer
Update Huargtt OipHrrtet + a Jmean fiTT N
Update Uurget target
(1 )td 91.
Standardize 9i.T
N
(gi.r
N
Mrarget) / Inrget
// Run the actor-critic network
Define F om ym No yonr
for to T do 2f IOpalaCNN,(OFF)
R 2f1)
for S gend
Compute V ([oP: )(]) and n (lvi,, 2x()
for n Nstart Nend
end for
// Take a gradient step Compute Cpro ( using Equation (4)
for Nsstart Nend
Noa
Define the minibatch loss Lrro ( No6 N (pro( Update $ Adam (Ø, clip-gradient V Cpro ( max-norm), Ir)
end for end for
14

---

Improving Transformer World Models for Data- RL
We make few comments below:
We use gradient clipping each minibatch to control the maximum gradient norm, and update the actor-critic parameters using Adam (Kingma, 2014) with learning rate of 0.00045.
During each epoch and minibatch update, we initialize the hidden state ho from its value ho stored while collecting the trajectory т.
introduce two changes to the standard PPO objective, described in Equation (4), First, we standardize the GAEs (ensure they are zero mean and unit variance) across the batches. Second, similar Moon et al. (2024), we maintain moving average with discount factor for the mean and standard deviation of the target q and we update the value network to predict the standardized targets.
In Algorithm
Implementation:
Note that for implementing PPO, we start from the code available in the jaxr] library (Lu et al., hppppppppppp.
2022) at
A.1.3. HYPERPARAMETERS
Table 5 displays the PPO hyperparameters used for training our SOTA MFRL, agent.
Table 5: MFRL hyperpameters
Module
Hyperparameter
Value
Environment
Number of environments New
48
Rollout horizon in environment Ten
96
Sizes
Image size
63
63
CNN output size
8 x 8 128
RNN hidden layer size
256
AC input size
8448
AC layer size
2048
PPO
?
0.925
À
0.625
clipping
0.2
TD-loss coefficient ATD
1.0
Entropy loss coefficient Xent
0.01
PPO target discount factor a
0.95
Learning
Optimizer
Adam (Kingma, 2014)
Learning rate
0.00045
Max. gradient norm
0.5
Learning rate annealing (MFRL)
True (linear schedule)
Number of minibatches (MFRL)
8
Number of epochs (MFRL)
4
MBRL experiments. We make additional changes PPO in the MBRL setting, and keep all the other hyperparameters fixed. First, we do not learning rate annealing for MBRL, while MFRL uses learning rate annealing (with linear schedule). Second, as discuss in Section A.3.3, the differences between the PPO updates real and imaginary trajectories lead to varying the number of minibatches and epochs.
Craftax experiments. We also keep all but two of our PPO hyperparameters fixed for Craftax (full), which we discuss in Appendix E.
15

---

Improving Transformer World Models for Data-Efficient RL
A.2. Model-based modules
In this section, we detail the two key modules for model-based RL: the tokenizer and the transformer world model.
1. TOKENIZER
Training objectíve: Given Craftax-classic image O, and codebook C - {el,, [el,,..,,x]] ,x,) encoder & returns a feature map Z, = (Zi...., Z.)) Each feature Z gets quantized, resulting into L tokens Q: = (g....., ........'" (d')-whic serves as Finally, decoder 'D decodes Zz back to the image space:
input to the TWM- then projected back to Zi (eg) ... ................................
Ozf.). of
Ó = =D(2). Following Van Den Oord et al. (2017); Micheli et al. (2022), we define the VQ-VAE loss as:
L Ai11O, Ód|l + A2|1O, 0.112 + A (s, 2x 112 + >a sg sg(22)l
(5)
where sg is the stop-gradient operator. The first two terms are the reconstruction loss, the third term is the codebook loss and the last term is commitment loss.
We now discuss the different VQ and VQ- VAE architectures used by the models M1-5 in the ladder described in Section 1.
Default VQ-VAE: Our baseline model MI, and our next model M2 build on IRIS VQ- VAE (Micheli et al., 2022) and fol- low the authors' code: https: / hubb com/c e looial onso/iri Siriiiii /bl ob/ abbmmaii
ma i / src/ ssre/mode mode a a/tokeni i zer/ nets. py. The encoder uses a convolutional layer (with kernel size x 3 and stride 1), then five residual blocks with two convolutional layers each (with kernel size x 3, stride and ReLU activation). The channel sizes of the residual blocks are (64, 64, 128, 128, 256). A downsampling is applied on the first, third and fourth blocks. Finally, last convolutional layer with 128 channels returns an output of size 8 8 x 128. The decoder follows the reverse architecture. Each of the L= 64 latent embeddings gets quantized individually, using codebook of size K = 512, to minimize Equation (5). We use codebook normalization, meaning that each code in the codebook C has unit L.2 norm. Similarly, each latent embedding I gets normalized before being quantized. As in IRIS, we use a1 1, A2 = 0, A3 = 1, A1 = 0.25. We train with Adam and a learning rate of 0.001.
VQ- V. For the next model M3, the encoder is two-layers MLP that maps each flattened 7 x 7 x 3 patch to a 22S-dimensional vector, using a ReLU activation. Similarly, the decoder learns a linear mapping from the embedding vector back to the flattened patches. Each embedding gets quantized individually, using a codebook of size K 512, and codebook normalization, to minimize Equation (5). Following Micheli et al. (2024), we use A 0.1, A2 1, As = A = 0.02.
Nearest neighbor tokenizer: NNT does not use Equation (5) and directly adds image patches codebook of size K- 4096, using Euclidean threshold 7 0.75.
A.2.2. TRANSFORMER WORLD MODEL
Training objective: We train the TWM on real trajectories (from the environment) of TwM = 20 timesteps sampled from the replay buffer (see Algorithm 1). We set TwM- = 20 as it is the largest value that will fit into memory on 8 H100 GPUs.
Given a trajectory 5 = (01.T41, dI., TIr,, donel.T). ..), the input to the transformer is the sequence of tokens:
.g...... of ai.... gt..... (.aa))
where enc(O,) = (gl .q.."...q an observation embedding table and an action embedding table. After several self-attention layers (using the standard causal mask or the block causal mask introduced Section 3.6), the TWM returns a sequence of output embeddings:
...) and qt E {1.... K} where K is the codebook size. These tokens are then embedded using
(E(g))...., B()b),,,,,((())...
E(ar)).
The TWM then output embeddings are then used to decode the following predictions:
(1) Following (Micheli et al., 2022), E(a;) passes through a reward head and predicts the logits of the reward ri-
(2) F(a.)) also passes through a termination head and predicts the logits of the termination state done,-
16

---

Improving Transformer World Models for Data-Efficient RL
(3) Without block teacher forcing, Elg)) passes through observation head and predicts the logits of the codebook index the same timestep E((itt) when i I. 1. Similarly E((+) passes through observation head and predicts the logits of the first codebook index at the next timestep E(()+1).
(3") With block teacher forcing, B(d)) passes through an observation head and predicts the logits of the same codebook index at the next timestep 1(
TWM is then trained by summing three losses:
(1) The first loss is the cross-entropy for the reward prediction. Note that Craftax- provides (sparse) reward of 1 for the first time each achievement is"unlocked" in each episode. In addition, it gives a smaller (in magnitude) but denser reward, penalizing the agent by 0.1 for every point of damage taken, and rewarding it by 0.1 for every point recovered. However, we found that we got better results by ignoring the points damaged and recovered, and using a binary reward target, which we implemented by setting the target reward to when the reward collected is higher than 0.5, and to 0 otherwise. This is similar to the recommendations in Farebrother al. (2024), where the authors show that value-based RI. methods work better when replacing MSE loss for value functions with cross-entropy on a quantized version of the return.
(2) The second loss is the cross-entropy for the termination predictions.
(3 The third loss is the cross-entropy for the codebook predictions, where the predicted codes vary between 1 and the codebook size K.
Architecture: We use the standard GPT2 architecture (Radford et al., 2019). We use sequence length Twm = 20 due to memory constraints. We implement key-value caching to generate rollouts fast. Table details the different hyperparameters.
Table 6: Hyperparameters for the transformer world model
Module
Hyperparameter
Value
Environment
Sequence length TwM
20
Architecture
Embedding dimension
128
Number of layers Number of heads
3
8
Mask
Causal or Block causal
Inference with key-value caching
True
Positional embedding
RoPE (Su et al., 2024)
Learning
Embedding dropout
0.1
Attention dropout
0.1
Residual dropout
0.1
Optimizer
Adam (Kingma, 2014)
Learning rate
0.001
Max. gradient norm
0.5
17

---

Improving Transformer World Models for Data- RL
A.3. Our Model-based RL agent
In this section, we detail how combine the different modules above build our SOTA MBRL agent, which is described in Algorithm 1 in the main text.
A.3.I. COLLECTING ENVIRONMENT ROLLOUT OR TWM ROLLOUT
Algorithm 4 presents the rollout method, which we call in Steps I and 4 of Algorithm 1. It requires a transition function which can either be the environment or the TWM.
Algorithm 4 Environment rollout or TWM rollout
Input: Initial observation O1,
Previous M observations Opast (0 (O-M+I,...,0, available else Opeer = 0, AC model and parameters 4
Rollout horizon T.
An environment transition Meny TWM M with parameters θ
Output: A trajectory T = (01.4+1; Q1,7, T1.7, done horr)
Initialize: hidden state hg = if O pust 9 else set h-M MM = 0
if O past # 0 then // Burn-in the hidden state for m
M do
2-M+m ImpalaCNN,(O-MHH4
h- RNN (1h --+mm 2- M+
end for end if
Initialize: == (ho) for L to T do
// Run the actor network ImpalaCNN, (O,) 2.
h = RNN& ([h zel)
(g ~ Ta (lb,, (l)
// Collect reward and next observation if environment rollout then
O rt, done, Mew (O at) else if TWM rollout then
Q. (ul..... of) enc(O.)
Q++1 ~ pe (Qu41 a1:))
0++1 = dec(Qu+1)
Fi ~ (ri Q1:e, d1:t)
done; ~ pe (done, @ lst, a1a)
end if
T+ (O,, 04, done; h4)
end for
r+ (Or41)
Below we discuss various components of Algorithm 4.
Note that in Algorithm 1, we call Algorithm in parallel starting from New observations 0 leNa. (for environment rollouts) or Ō 1:Ne (for TWM rollouts).
Parallelism.
Burn-in. The first time we collect data in the environment, we initialize the hidden state to zeros. The next time, we use burn-in to refresh the hidden state before rolling out the policy (Kapturowski et al., 2018). We do so by passing the M
18

---

Improving Transformer World Models for Data-Efficient RL
observations prior O₂ the policy, which updates the hidden state of the policy using the latest parameters. (To burn-in TWM rollout, we sample trajectory of length M 1 in Step 4 of Algorithm 1.) To enable burn-in, when collecting data, in Step 1 of Algorithm 1, we must also store the last M environment observations (0 (O-M+h..... ,,) prior to O₁,
TWM sampling. As explained in the main text, sampling from the distribution Q1+1 ~ pe pe(Qu+1 QQ, @1.)) is different when using (or not) block teacher forcing. For the former, the tokens of the next timestep (9i41..... .... are sampled in parallel, while for the latter, they are sampled autoregressively.
Maximum buffer size. To avoid running out of memory, we use maximum buffer size and restrict the data buffer 'D in Algorithm to contain at most the last 128k observations. When the buffer is at capacity, we remove the oldest observations before adding the new ones. We use flashbax (Toledo et al., 2023) implement our replay buffer in JAX.
A.3.2. WORLD MODEL UPDATE
In practice, decompose the world model updates into two steps. First, update the tokenizer update the TWM TWM times. For both updates, we use N WM implemented in Algorithm 5.
tok times. Second, we training minibatches. That is, Step 3 of Algorithm
Algorithm 5 Step 3 of Algorithm
for Viters do training do
lok
for k- = to WM
Nstart (As 1) /Neen/Nii training
+ 1, Nend = (Nem// training
+
Toplay sample-trajectory (D, Twma),
n2
Newv
e =
(
_Nii .Nm reepaay
with Equation (5)
end for end for
for it to N TWM do
for =
N mb training do
Asstart (k: 1) Nem
mmhh training WM
+ 1, Need (Nan// NWM
traininn
training
+
-
T Twm),
n 1: Newv
replay
θ TWM(t)
-N following Appendix A.2.2 replay
end for
end for
We always set Niters = 500 to perform large number gradient updates, For M1-3, we set Viters 500 as well, but for M5 we reduce Aiiters 25 for the sake of speed speed-since since NNT only adds new patches the codebook.
A.3.3. PPO POLICY UPDATE
Finally, the PPO-policy-updaa procedure called in Steps and 4 of Algorithm 1 follows Algorithm 3.
When using PPO for MBRL, we found it critical to use different numbers of minibatches and different numbers of epochs on the trajectories collected on the environment and with TWM.
In particular, as the trajectories collected in imagination are longer, we reduce the number of parallel environments, and Nmmb 8 and Nmb 1. This guarantees that the PPO updates are batches of comparable sizes-6 96 for real trajectories, and 48 x 20 for imaginary trajectories.
In addition, while the environment trajectories are limited, we can simply rollout our TWM to collect more imaginary trajectories. Consequently, we set N
4. and peppch 1.
Finally, we do not use learning rate annealing for MBRL training.
19

---

Improving Transformer World Models for Data-Efficient RL
A.3.4. HYPERPARAMETERS
Table 7 summarizes the main parameters used in our MBRL training pipeline.
Table 7: MBRL main parameters
Hyperparameter
Value
Number environments New
48
Rollout horizon in environment Tenv
96
Rollout horizon for TWM TwM Burn-in horizon M
20
5
Buffer size
128, 000
Number of tokenizer updates Nilles (with VQ-VAB)
500
Number of tokenizer updates lok (with NNT)
25
Number of TWM updates NiWM
500
Number of minibatches for TWM training AM training
3
Background planning starting step TBP
200k
Number of policy updates Nilers
150
Number PPO minibatches in enviromment Numh
8
Number PPO minibatches in imagination NWM
1
Number of epochs in environment Nenv
4
Number of epochs in imagination N
1
Learning rate annealing
False
20

---

Improving Transformer World Models for Data-Efficient RL
B. Comparing scores
Figure completes the two main Figures and by reporting the scores the different agents. Specifically, Figure 9[1eft] compares our best MBRL and MFRL agents to the best previously published MBRL and MFRL agents. Figure 9[right] displays the scores for the different agents on our ladder of improvements.
35
Iris A-tris
35
Iris
M1: Baseline
30
Previous best
30
Dyna
Our best MFRL
M3: M2
patches
25
(slow)
Score (%)
MS: M4
BTF. Our best MBRL (Fast)
20
MS: M4 BTF, Our best MBRL (slow)
Score
15
15
10
10
01
0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Environment steps
Environment steps
Figure [Left] In addition reaching higher rewards, our best MBRL and MFRL agents also achieve higher scores compared to the best previously published MBRL and MFRL results. [Right] MBRL agents' scores increase they climb up the ladder of improvements.
C. Annealing the number of policy updates
Figure 10 compares best MFRL agent (with fast training) agent trained by annealing the number of policy updates in imaginary rollouts.
70
50
Reward (%)
40
30
expert best MBRL (fast)
10
MBRL (increase BP steps)
0
0.0
0.2
0.6 Environment steps
0.4
0.8
1.0
Figure 10: Progressively increasing the number of policy updates from Niters 0 (when Total 0 env. steps) Niters 300
(when Tooal 1M) removes the drop in performance observed otherwise when we start training in imagination.
21

---

Improving Transformer World Models for Data-Efficienn RL
D. A Additional world model comparisons
This section complements Section 4.4 and presents two additional results to compare the different world models.
0.0030
M1: Baseline
Observation L2 recc (1)
0.0025
M3: M1 + Dyna + patches
0.04
BTF: best MBRL (fast)
0.0020
Tokenizer L2
0.0015
0.02
0.0010
M1: Baseline
0.0005
0.01
Dyna patches
(fast)
0.0000
0.0
0.2
0.4
0.6
0.8
1.0
5
10
15
20
Environment steps
TWM rollout horizon
Figure 11: TWM performance.[LLft]] Tokenizer L2 reconstruction error, averaged over rollouts. Lower better. By construction, our best MBRL agent, which uses NNT, constantly reaches the lowest error, as NNT directly adds observation patches to its codebook. [Right] TWM rollouts L2 observation reconstruction error, averaged over rollouts. Lower is better. M3 and M5, which both patch factorization, achieve the lowest errors.
D.1. Tokenizer reconstruction error
We first use the evaluation dataset Deval (introduced in Section 4.4) compare the tokenizer reconstruction error of our world models M1, M3, and M5 ---sinng the checkpoints 1M steps. To do so, independently encode and decode each observation O E Deval to obtain tokenizer reconstruction Ö " Figure 11[left] compares the average L2 reconstruction errors over the evaluation dataset:
Now Teat1 Σ
Σ
|10tth O 12
+ 1)
showing that all three models achieve low L2 reconstruction error. However our best model M5, which uses NNT, reaches a very low reconstruction error from the first iterations, since it directly adds image patches to its codebook rather than learning the codes online.
D.2. Rollout reconstruction error
Second, given sequence observations in a TWM rollout TWW., 1 and the corresponding sequence of observations in the environment Of: (which both have executed the same sequence of actions), Figure [[[rghh compares the observation L2 reconstruction errors at each timestep (averaged over the evaluation dataset):
Nevl
&
O.W.. O² VI.
Nevll n=1
As expected, the errors increase with as mistakes compound over the rollout. Our best method and M3, which both uses patch factorization, achieve the lowest reconstruction errors.
D.3. Symbol extractor architecture
Herein, we discuss the symbol extractor architecture introduced in Section 4.4. fa consists of (a) first convolution layer with kernel size 7 7, stride of 7, and channel size 128, which extracts feature for each patch, (b) ReL.U activation, (c) second convolution layer with kernel size 1 x 1, stride of and channel size 128, (d) second ReLU activation, (e) final linear layer which transforms the 3D convolutional output into a 2D array of logits of size 145 17 = 11445-wheee R = 145 the number ground truth symbols associated with each image of Craftax-classe and each symbol S E {1,..., 17). The symbol extractor trained with a cross-entropy loss between the predicted symbol logits and their ground truth values S. and achieves a 99.0% validation accuracy.
22

---

Improving Transformer World Models for Data- RL
D.4. Rollout comparison
In Figure 12, we show an additional rollout that exhibits similar properties those in Figure 6[right]. MI and M3 make more simple mistakes in the map layout. All models generate predictions that can be inconsistent with the game dynamics. However the errors by MI and M3 more severe, as M5's mistake relates to the preconditions of the make sword action.
MAKE_ FONL
MAKE_ROLLPIOE POOKAIC
MAKE WOOD, SWORD
NOOP
UP
DOWN
MAKE ONE, PICKAE
LEFT
SLEEP
LEFT
True
rollout
9-
9.6
M1
M3
M5:
our best
9.9
Rollout annotations
Map inconsistency
Feasible hallucination
Infeasible hallucination
Figure 12: Additional rollout comparison for world models M1, M3 and M5. Best viewed zoomed in. Map. All models exhibit some map inconsistencies. M1 make simple mistakes after the agent moves. Both M3 and M5 have map inconsistencies after the sleep actions, however the mistakes for M3 are far more severe. Feasible hallucinations. All models make feasible hallucinations when the agent exposes a new map region. The sleep action is stochastic, and only sometimes results in the agent sleeping after taking the action. As a result, M3 and M5 make reasonable generations in predicting that the agent does not sleep in the final frame. Infeasible hallucinations. MI generates cells that do not respect the game dynamics, such as spawning plant without taking the place plant action, and creating a block type that cannot exist in that location. M3 turns the agent face downwards without the down action. M5 makes the wood sword despite the precondition of having wood inventory not being satisfied.
23

---

Improving Transformer World Models for Data-Efficient RL
E. Comparing Craftax-classic and Craftax (full)
This section complements Section 4.5 and discusses the main differences between Craftax-classic and Craftax.
The first and second block Table compares both environments. Note that only use the first five parameters in our experiments in Section 4.5.
The third and fourth blocks report the parameters used by our best MFRL and MBRL agents. In Craftax (full), for MFRL, we use New = 64 environments and rollout length Teav = 64. Our SOTA MBRL agent uses Tow = 96, New = 48, and and Twm 20 as in Craftax-classic. We reduced the buffer size to 48k to fit in GPU.
the others PPO parameters are the same in Table 5.
Table 8: Environment Craftax-classic vs Craftax (full)
Module
Hyperparameter
Classic
Full
Environment (used)
Image size
63 63
130 110
Patch size
7 x 7
10 x 10
Grid size
9 x 9 17
13 x 13 43
Action space size
Max reward (# achievements)
22
226
Environment (not used)
Symbolic (one-hot) input size
1345
8268
Max cardinality of each symbol
17
40
Number of levels
10
MFRL parameters
Number of environments New Rollout horizon in environment Tew
48
64
96
64
MBRL parameters
Number environments New
48
Rollout horizon in environment Tews
96
Rollout horizon for TWM Twm
20
20
Rollout horizon for TWM TwM
20
20
Buffer size
48,000
128,
000
24

---

Improving Transformer World Models for Data-Efficient RL
F. Adapting Craftax-classic parameters to solve MinAtan tar
This section details the adaptations we made to our pipeline for solving the MinAtar environments, presented in Section 4.6. First, Table outlines the modifications our MFRL agent. Notably, we incorporate layer normalization (Ba al., 2016) and Swish activation function (Ramachandran al., 2017) within the ImpalaCNN architecture. Furthermore, found beneficial for the actor and critic networks share weights up to their distinct final linear layers. We also adjust some PPO hyperparameters.
Table 9: MFRL changes for MinAtar
Module
Parameter
Craftax
Minatar
Environment
Image size
63 x 63 x 3
10 x 10 x K
ImpalaCNN
Normalization
Batch normalization
Layer normalization Swish
Activation
ReLU
Shared network
False
True
PPO
0.925
0.95
X
0.625
0.75
PPO target discount factor
0.95
0.925
These modifications result in solid MFRL agent, whose performance is detailed in Section 4.6. We then develop decompose each MinAtar image
MBRL agent top by implementing the changes outlined in Table 10. Specifically, into 25 patches of size 2 x 2 x K each. In addition, we increase (a) the number of TWM updates to from 500 to 2k and (b) the number of policy updates in imagination from 150 to 2k. Critically, to address the high cost of bad actions in certain games (e.g. Breakout), we assign a weight 10 to the cross-entropy losses of the reward and of the done states. This strongly penalizes inaccurate predictions terminal states in imaginary rollouts. Additionally, we observe a potential issue during training in imagination where the agent could collapse to output the same action consistently. mitigate this "action collapse" and promote exploration, increase the entropy coefficient in the imagination phase from 0.01 to 0.05.
Table 10: MBRL changes for MinAtar
Module
Parameter
Craftax
Minatar
Tokenizer
Patch size
x 7 3
2 2 x K
Grid size
9 x 9
5 x 5
Training
Number of policy updates Viiters
150
2, 000
AC
Number of TWM updates TWM
500
2, 000
Termination and reward weight
1
10
PPO entropy coeff. in imagination
0.01
0.05
Note that all the MinAtar games use the same hyperparameters.
25


Reviews:
[{"review_id": "nuNMILbDiO", "reviewer": "Reviewer_fn71", "summary": "The authors propose a new model-based RL method that achieves SOTA at Crafter. The superiority of there method stems from a variety of novel insights:\n- adding a memory with low-dimensional hidden states and passing both the image embedding and the memory output to subsequent networks\n- training on a mix of real and imaginary trajectories\n- encoding image patches via nearest neighbor\n- use of block teacher forcing to train the model", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your favorable comments regarding our work.\n\nTo further validate the robustness of our approach, we have conducted additional experiments on another grid-world environment MinAtar (https://github.com/kenjyoung/MinAtar), a set of four simplified Atari 2600 games. MinAtar contains symbolic binary observations of size 10x10xK, where K is the number of objects in each game, and binary rewards.\n\nWe first tuned our model-free RL agent on these environments, keeping the same architecture as described in our paper, with minor adjustments to the PPO hyperparameters. Second, we developed our model-based RL agent, building upon our previously proposed techniques: Dyna warmup, nearest neighbor tokenizer, and block teacher forcing. We retained the majority of the MBRL hyperparameters from Craftax-classic, with the following key modifications: (a) we increased the number of WM training steps $N_{\\text{TWM}}^{\\text{iters}}$ to from 500 to 2k, (b) we increased the number of policy updates in imagination $N_{\\text{AC}}^{\\text{iters}}$  from 150 to 2k, (c)  we used patches of size 2x2xK (d) we added a weight of 10 to the cross-entropy losses of the reward and of the done states.\n\n\nThe table below compares the performance of our agents, averaged over 10 seeds after 1 million environment steps.\n\n| Game | MFRL | MBRL  | \n|--------------|-----------------|-----------------| \n| Asterix | $10.28 \\pm 1.38$ | $47.24 \\pm 10.12$ |\n| Breakout | $76.36 \\pm 1.82$ | $82.8 \\pm 8.9$ | \n| Freeway | $61.33 \\pm 3.67$ | $70.73 \\pm 0.33$ | \n| SpaceInvader | $135.88 \\pm 2.59$ | $180.7 \\pm 3.34$ |\n\nThese results appear quite competitive compared to existing approaches. Notably, our MBRL agent, after only 1M environment steps, seems to outperform the Artificial Dopamine method reported in Guan et al. [1] after 5M steps (as illustrated in their Figure 4). \n\nWe are finalizing these experiments and plan to include them in the revised version of our paper.\n\n\n**Reference:**\n[1] Guan, Jonas, et al. \"Temporal-Difference Learning Using Distributed Error Signals.\" Advances in Neural Information Processing Systems 37 (2024): 108710-108734.", "claims_and_evidence": "They perform an evaluation by incrementally adding there improvements on top of the baseline. There best model beats the baseline by more than 35% which is very significant.", "methods_and_evaluation_criteria": "They evaluate there algorithm on Crafter which is a good benchmark.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "They use 10 seeds which is standard in RL.", "relation_to_broader_scientific_literature": "The authors build on the recent line of research on transformer world models and equip them with classical insights such as Dyna. Furthermore they propose a simpler alternative to VQ-VAE.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: The paper explain precisely the different design choices they made and they perform an extensive ablation study.\nWeakness: Even if crafter is a good environment, I would like to see one or two more environments to assess the generality of the method. I am happy to raise my score once I see result for another environment.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NXCdzeodDI", "reviewer": "Reviewer_8HUm", "summary": "The authors propose a number of improvements to dreamer-style MBRL to achieve SOTA performance on the Craftax benchmark.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are pleased that you enjoyed reading the paper and that you think that reaching SOTA on Craftax is a notable achievement, on a well-known and difficult benchmark.\n\n**Claims and Evidence:**\n\n(1) You are correct in that Atari with sticky actions (frameskip) makes it stochastic.  We will remove our claim of Atari being a deterministic environment.\n\n(2) We acknowledge that we do not demonstrate a causal relation between the low dimension of the hidden state, and the fact that the policy captures control relevant information. However, we varied the ratio of the RNN state dimension to the CNN encoder dimension, and found that a lower ratio yielded better performance, a result that contrasts with prior work. We will clarify this point in our revised paper.\n\n**Weaknesses:**\n\nWe anticipate our proposed nearest neighbor tokenizer (NNT) to only work in grid-world environments. In environments with raw pixel inputs (Atari, Procgen) we believe that alternative tokenization methods, such as VQ-VAE and variants, are likely necessary. We are currently exploring this direction for future publication.\n\n**Questions:**\n\nRegarding annealing, we have conducted some follow-up experiments where we progressively increased the number of policy updates on imaginary rollouts ($N_{\\text{AC}}^{\\text{iters}}$ in Step 4 of Algorithm 1) from 0 to 300. This annealing technique achieves a reward of $65.7$% $(\\pm1.11)$, while removing the drop in performance observed when we start training in imagination at $200$k steps. We will include these results in the revised paper.", "claims_and_evidence": "The authors mention three main contributions:\n- They show that training on both real and imagined data is better than solely on imagined data\n- They embed image patches using nearest neighbors, while prior work embeds the entire image using a VQ-VAE\n- They propose to use a block-teacher forcing objective instead of a standard log likelihood objective\n\nThey provide ample evidence to back their major claims.\n\nThe authors also make smaller claims I have a bit of trouble with.\n\n> How-\never, the near-deterministic nature of Atari games allows\nagents to memorize action sequences without demonstrat-\ning true generalization (Machado et al., 2018)\n\nVirtually all modern work on Atari uses the `frameskip` variants where this is not true.\n\n> For the RNN, we find it crucial to ensure\nthe hidden state is low-dimensional, so that the memory is\nforced to focus on the relevant bits of the past that cannot\nbe extracted from the current image\n\nIs there any evidence for this claim? Just because a smaller hidden size works better does not imply this specific causation.", "methods_and_evaluation_criteria": "Craftax is an interesting, well-known, and difficult benchmark.", "theoretical_claims": "They do not make theoretical claims", "experimental_designs_or_analyses": "The authors meticulously ablate each and every component they introduce. I am quite happy with the depth of their experiments.", "relation_to_broader_scientific_literature": "The authors provide three improvements to dreamer-style MBRL. I believe other researchers focusing on such methods can integrate the authors' methods into their own models.", "essential_references_not_discussed": "I think they covered most important references.", "other_strengths_and_weaknesses": "The authors start from a simple existing model, and slowly build up tricks to reach a new SOTA. I enjoyed reading the paper, and I like that the authors ablate every single proposed change. Reaching SOTA on Craftax is a notable acheivement, as the task is open-ended and difficutl.\n\nOne major complaint I have is that the image codebook does not seem like a scalable approach for more complex problems. Yes, it works for small pixel observation spaces, but I imagine it would fail in more realistic tasks.", "other_comments_or_suggestions": "> we fo-\ncus on the Crafter domain (Hafner, 2021)\n\nincorrect citation", "questions_for_authors": "Not using imagined trajectories for some $k$ timesteps seems like a clunky solution to the problem. Did you consider reweighting imagined sampled by some variable that anneals from 0.0 to 1.0 as training goes on? The drop in Figure 4 suggests a \"hard\" warmup could hurt training.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Zel9Gi4lVz", "reviewer": "Reviewer_6wf5", "summary": "This paper proposes an approach for model-based reinforcement learning (MBRL) to improve sample efficiency and performance on the challenging Craftax-classic benchmark.\nThe method includes three improvements for both policy and transformer world model (TWM), which are “Dyna with warmup”, “nearest neighbor tokenizer” and “block teacher forcing”.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Essential References Not Discussed:**\n\nWe thank you for pointing out the REM reference: we will include it in our revised paper. You are right that both block teacher forcing (BTF) and REM predict all the next frame tokens jointly. However, while REM uses a retentive network, BTF is applicable to a broader range of transformer architectures. BTF achieves joint prediction through a modification of the causal mask and the supervision signal, as illustrated in our Figure 3 and detailed in Appendix A.2.2. \n\n**Weaknesses:**\n\nWe acknowledge that our model-based RL approach, like many others in this domain, involves a significant number of hyperparameters due to the interaction of its various components: the actor-critic policy, the tokenizer, and the world model.\nTo provide clarity, we have dedicated separate appendices to detail each of these components: Appendix A.1 for the actor-critic policy, Appendix A.2.1 for the tokenizer, and Appendix A.2.2 for the world model. Table 6, Appendix A.3.4 summarizes the main MBRL parameters which glue these parameters together.\n\nOur supplementary experiments on the MinAtar environments, detailed in our rebuttal to Reviewer fn71, indicate that the majority of our proposed pipeline transfers to these other grid-world environments.\n\n\n**Questions:**\n\n(1) We opted to primarily utilize PPO due to its well-established advantages in terms of stability and performance. PPO's clipped surrogate objective function mitigates the issues of large policy updates that can destabilize learning, a common problem with vanilla REINFORCE. Furthermore, PPO generally outperforms REINFORCE in complex environments.\n\nGiven these advantages, we believe PPO provides a robust foundation for achieving high levels of performance in Craftax-classic. Consequently, we anticipate a decrease in performance if we were to substitute REINFORCE.\n\n(2) You are right that in the worst case, our nearest neighbor tokenizer (NNT) can lead to memory overhead. However, we found that, on average, on Craftax-classic, only 2,304 codes (out of 4,096) are being used by NNT.", "claims_and_evidence": "Yes, the claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the proposed method makes sense.", "theoretical_claims": "The paper does not present any formal theoretical proofs.", "experimental_designs_or_analyses": "I have checked the experimental designs in Section 4 and Appendix B, C & D.", "relation_to_broader_scientific_literature": "The key contribution of this paper is the three improvements for model-based reinforcement learning with a transformer world model:\n\n1. “Dyna with warmup” follows the classic Dyna [1] setting.\n\n2. The “nearest neighbor tokenizer” method is an improvement on the VQ-VAE in IRIS [2].\n\n3. The “block teacher forcing” method modifies the attention mask of the GPT structure so that it can predict tokens in parallel, in a similar way as in REM [3].\n\n[1] Richard S. Sutton, et al. “Dyna, an integrated architecture for learning, planning, and reacting.” ACM Sigart Bulletin (1991).\n\n[2] Ajay Mandlekar, et al. “Transformers are sample-efficient world models.\" ICLR (2023).\n\n[3] Lior Cohen, et al. \"Improving Token-Based World Models with Parallel Observation Prediction.\" ICML (2024).", "essential_references_not_discussed": "One of the key contributions is the parallel prediction of observation tokens, but a similar parallel prediction method is also proposed, namely Algorithm 1 & 2 from REM [1] published in ICML 2024.\n\n[1] Lior Cohen, et al. \"Improving Token-Based World Models with Parallel Observation Prediction.\" ICML (2024).", "other_strengths_and_weaknesses": "* Strengths\n1. This paper is well-written with clear motivation.\n2. The three changes are sensible and effective, achieving significant performance gains in a Craftax-classic environment.\n3. The experiments are solid, especially the ablation of the components.\n\n* Weaknesses\n1. The method contains too many hyperparameters, and it is difficult to determine which hyperparameters are important, making it difficult to transfer to other environments.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. According to Figure 1, using only PPO outperforms other baseline MBRL methods, most of which are based on REINFORCE. What would happen if you also build upon REINFORCE?\n2. I am a bit confused about the NNT implementation, it looks like in the worst case NNT leads to unbounded memory overhead. How do you avoid this problem?\n\nIf you respond to these questions and address these concerns, I'll be willing to raise the score.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vp9Inc4aup", "reviewer": "Reviewer_EkAK", "summary": "This paper presents a technically sound and empirically robust contribution to MBRL, with clear innovations in tokenization and transformer training. While the Craftax-centric evaluation limits immediate generalizability, the methodological advancements (NNT, BTF) are likely to inspire follow-up work. The paper is suitable for acceptance provided the authors address the above concerns, particularly by:", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive evaluation of our work. We address your questions and comments below:\n\n**Weaknesses **\n\n(1) As detailed in our rebuttal to Reviewer fn71, we successfully trained an MBRL agent on the MinAtar benchmark, reusing our core MBRL components with little hyperparameter tuning. This agent significantly outperformed a tuned MFRL agent in these environments, highlighting the potential transferability of our approach to other grid-world settings.\n\n(2) We appreciate your suggestion for a theoretical comparison between the autoregressive and block teacher forcing (BTF) settings. We agree that this would be a valuable avenue for future investigation.\n\n(3) To quantify the impact of BTF in dynamics learning, we measured the average cross-entropy loss of the observation tokens over the last 500,000 training steps. Our TWM without BTF achieves an average CE of 0.478 (\\pm 0.04), while our best TWM with BTF reaches 0.432 (\\pm 0.004). This difference in learning dynamics suggests that BTF facilitates learning, likely by enabling the reuse of intermediate computations for next token predictions.\n\n(4) We agree that our proposed nearest neighbor tokenizer (NNT) is likely best-suited for grid-world environments. For environments with raw pixel inputs, such as Atari and Procgen, we believe that alternative tokenization methods like VQ-VAE and its variants will be necessary. We are actively exploring this direction. Our other two methods, Dyna with warmup and Block Teacher Forcing, should remain applicable in these settings.\n\n(5) Please note that we discuss TD-MPC2 and Diamond in our related work Section 2, including in our footnote 2.\n\n\n**Other comments or suggestions**\n\n(1) We would like to direct your attention to Appendix A, specifically Tables 4, 5, and 6, where we have tried to present all the hyperparameters used in our MBRL pipeline. \n\n(2) The human expert results utilized in this work have been extracted from the original Crafter paper. As detailed in Section 4.4 of [1], this dataset comprises 100 gameplay episodes recorded from five human experts who were given game instructions and several hours of practice prior to recording.\n\n\n**References**  [1] Hafner, D. Benchmarking the spectrum of agent capabilities.\narXiv preprint arXiv:2109.06780, 2021.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "NA", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "They would potentially be useful for real-world applications (e.g., robotics, automated exploration).", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "+ This paper presents a technically sound and empirically robust contribution to MBRL, with clear innovations in tokenization and transformer training. While the Craftax-centric evaluation limits immediate generalizability, the methodological advancements (NNT, BTF) are likely to inspire follow-up work. The paper is suitable for acceptance provided the authors address the above concerns, particularly by:\n\n+The incremental \"ladder of improvements\" (Table 1) and detailed ablation analyses (Table 2, Figure 5) clearly demonstrate the contribution of each component (Dyna with warmup, NNT, BTF). The sensitivity analysis of patch size and warmup duration provides actionable insights.\n+The nearest-neighbor tokenizer (NNT) and block teacher forcing (BTF) are novel and well-motivated. NNT’s static codebook addresses non-stationarity in VQ-VAE training, while BTF’s parallel token prediction mitigates autoregressive drift. These innovations are backed by quantitative metrics (symbol accuracy, reconstruction error) and qualitative rollout comparisons (Figure 6).\n+The MBRL agent’s training time (759 minutes on 8 H100 GPUs) is competitive compared to prior work (e.g., IRIS: 18330 minutes). The MFRL baseline’s efficiency (15 minutes) further underscores the practicality of the approach.\n\n\n-The evaluation is restricted to Craftax-classic and Craftax Full. While the environment’s complexity is a strength, the paper does not validate whether the proposed techniques (e.g., patch-based NNT) generalize to other domains (e.g., 3D environments, non-grid-based tasks). A discussion of applicability to broader settings (e.g., Minecraft, Atari, robotics) would strengthen the contribution, since the required codebook’s size for NNT would increase significantly in first-person-view settings, which may affect the stability of training TWM. While the empirical results are compelling, the paper lacks theoretical analysis of why BTF or NNT improve performance. For instance, how does BTF’s block causal attention reduce compounding errors compared to autoregressive sampling? This paper states that BTF improves performance by returning a more accurate TWM compared to AR methods. However, the observed gains may instead stem from more efficient representations learned through future state reasoning, rather than solely from improved accuracy. This distinction is not explicitly validated, as the ablation study does not include a direct comparison of dynamic losses. A deeper connection to existing theory (e.g., bias-variance trade-offs in world models) would add rigor.\n-The NNT assumes aligned patches (e.g., 7×7 grids), which aligns with Craftax’s design but may not generalize to environments with less structured observations (e.g., raw pixel inputs in Atari). In such cases, NNT may allocate excessive computational resources to decoding background information, potentially overlooking small, reward-relevant objects. The paper briefly acknowledges this limitation but does not explore workarounds (e.g., adaptive patching).\n- While the paper compares to IRIS and DreamerV3, it omits discussion of concurrent MBRL advances (e.g., Diamond's diffusion models, TD-MPC2’s continuous control). A broader literature review would better contextualize the work.", "other_comments_or_suggestions": "Key implementation details (e.g., GRU architecture, codebook initialization for NNT) are briefly described but lack specificity. Public code or an appendix with full hyperparameters would aid reproducibility.\n\nThe human expert baseline (65.0% reward) is derived from 100 episodes by 5 players, but the paper does not clarify how this was standardized (e.g., episode length, interaction constraints). More details would strengthen the comparison.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a novel MBRL transformer-based algorithm that leverages nearest-neighbor image tokenization, hidden state memory, and real/imaged mixture of trajectories into training. The method provides solid motivations and insights into what led to success, which most notably was SOTA in Craftax environment, surpassing Dreamer V3.

There were concerns about hyperparameters and the limitations wrt the domains studied, but I feel like the contributions are clear enough to be a contribution other works can build on in other domains, and there is sufficient ablation studies to cover HP concerns.

I therefore recommend the paper accepted.
Decision: Accept (poster)