Improving Transformer World Models for Data-Efficient RL

Thank you for your positive evaluation of our work. We address your questions and comments below:

**Weaknesses **

(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.

(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.

(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.

(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.

(5) Please note that we discuss TD-MPC2 and Diamond in our related work Section 2, including in our footnote 2.

Other comments or suggestions

(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.

(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.

References [1] Hafner, D. Benchmarking the spectrum of agent capabilities. arXiv preprint arXiv:2109.06780, 2021.

Essential References Not Discussed:

We 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.

Weaknesses:

We 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. To 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.

Our 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.

Questions:

(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.

Given 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.

(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.

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.

Claims and Evidence:

(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.

(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.

Weaknesses:

We 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.

Questions:

Regarding 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.

We appreciate your favorable comments regarding our work.

To 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.

We 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.

The table below compares the performance of our agents, averaged over 10 seeds after 1 million environment steps.

These 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).

We are finalizing these experiments and plan to include them in the revised version of our paper.

Reference: [1] Guan, Jonas, et al. "Temporal-Difference Learning Using Distributed Error Signals." Advances in Neural Information Processing Systems 37 (2024): 108710-108734.