Learning to Play Atari in a World of Tokens

We thank the reviewer for their time and feedback. Please find our responses below.

From the results in Table 1, DART is worse than DreamerV3 on multiple games.

While we agree that DART may not match DreamerV3’s performance in certain games, it’s essential to note their differences. DreamerV3 relies on a strong 18 million parameters model using CNN & GRU for Atari game training, whereas DART is a transformer-based world model with 3 million parameters. DART outperforms DreamerV3 in 11 out of 26 games, achieving similar superhuman scores. It excels in median score, iqm, and optimality gap scores compared to existing models. Moreover DART has better interpretability due to its use of self-attention, as illustrated in Figure 4. In a comparison with TWM, another transformer-based world model using a Dreamer-like objective function, DART appears to be more sample-efficient.

Also, DreamerV3's results are missing from the figures.

Thanks for pointing this out. We have updated Fig. 2 and Fig. 3 with Dreamerv3 results in the revised version.

It would be better to evaluate DART's performance on multiple domains, such as robotic control, Crafter, DMLab, or even Minecraft, to show the discrete representations can generalize to different scenarios.

While running on all the environments would be challenging during the discussion period due to the limited compute budget, along with the challenges of restructuring the model for continuous action space, we trained our model on the challenging Crafter benchmark. While the training is still ongoing preliminary results after 100k iterations indicate that DART surpasses IRIS and DreamerV3 on Crafter, and is more sample efficient as updated in Section A.1 (the final results for 1M iterations would be updated soon). It is important to note that Dreamerv3 uses 200 million parameter model for training Crafter while DART is limited to 3 million parameters only.

It seems like the core contribution of the paper is from the architecture side. However, there are also many prior works that leverage the transformer architecture for world modeling. It would be good to clarify the novelties of this work and how the work differs from prior works. A system-level comparison table would be helpful.

Through this work, we propose a MBRL approach which models the world and learns the behavior using discrete tokens, resulting in improved sample efficiency. This is done by using a ViT like architecture for learning the behavior by utilising discrete tokens. These discrete tokens are predicted by the world model using a GPT-like architecture. Moreover we propose a simplified approach for modeling memory again as discrete tokens using self-attention mechanism, further helping in learning the behavior. This is an improvement over IRIS which relies on CNN + LSTM for learning the behavior using image reconstruction. We have included model configuration in Table 5 discussing the differences between the model design of recent MBRL approaches.

What's the reason for using an image-based VQ-VAE instead of a video-based VQ-VAE?

Our primary aim is to simplify the existing approaches and model a policy using more abstract representations. The image-based VQ-VAE has shown improved image synthesis, by modeling precise low-level discrete elements. As a result, we chose to focus on the Image-based VQ-VAE and rely on the transformer for encoding temporal representation. While Video VQ-VAE could be a promising direction for future exploration, it is challenging to implement it within the short timeframe of rebuttal given our limited compute resources.

We thank the reviewer again for constructive feedback. We hope we could address most of the reviewer’s concerns. We would appreciate it if they could kindly consider adjusting the score. We are happy to provide further clarification.

We thank the reviewer for their time and feedback. Please find our responses below.

I need more clarification about the motivation that long-range dependencies impede Dreamers learning. Could the authors provide experiments comparing the world model accuracy among the proposed method and prior non-transformer-based methods on some long-range tasks?

Dreamer use GRU which models temporal representation as compressed states leading to degradation for long-horizon task. Some env in Atari test for long-horizon, we compared the performance of transformer-based models with a CNN+GRU-based model on crafter. Since world model accuracy directly determines the final policy performance [1], Table 3 of the revised version show that transformer-based models (~ 3 million parameters) exhibit better performance compared to DreamerV3 (~200 million parameters for Crafter). We will report the final performance of all the models with 1M iterations in the final version of the paper as experiments are currently running.

I think the world model accuracy should be measured more in detail (i.e. future states predicting accuracy. reward predicting accuracy, etc.) to fully support the author's arguments on the RNN-based world model and Transformer-based world model.

Fig. 5 (revised version) illustrates the predicted future trajectories of the transformer-based world model. Here, we show evolution of future within the learned world model given sequence of actions. We also visualize the next state prediction accuracy, the next reward prediction accuracy for our world model. Fig. 6 compares the predicted trajectories of DART with DreamerV3 on crafter and reports higher PSNR for DART.

Could the authors also compare the model capacity among DART and other baselines?

Table 5 includes information about the model capacity as well as detailed model configuration for recent MBRL approaches. DART has almost identical model parameters as IRIS, while maintaining high performance. Dreamerv3 uses an 18M parameter model for Atari, resulting in enhanced capacity compared to other models.

Could the authors explain how the long-term dependencies are required in Atari tasks?

Not all Atari games in general require long-term dependencies, however, the use of a transformer is still beneficial as, unlike RNN which provides the compressed version of the past states, transformer can access full state information, resulting in the modeling of complex relationships for better accuracy.

Could the authors please provide the value set for some critical hyperparams?

Thanks for pointing this out. A detailed list of all hyperparameters used for each module has been added in the revised version in Section A.4.

Could the authors explain why MEM helps?

The MEM token captures temporal representation through the self-attention mechanism. Temporal representation is modeled simultaneously with spatial representation, eliminating need for extra modalities, in contrast to prior approaches that independently process spatial and temporal aspects using CNNs for spatial and RNNs for temporal aspects. This simple approach ensures that the model selectively propagates task-relevant features, as illustrated in Fig. 4. Here the agent dynamically adjusts its attention to the memory in response to the task’s demands, simplifying the training process.

Could the authors also include the STORM paper to compare and explain the advantages of DART over STORM?

STORM was accepted at Neurips 2023 and became accessible on arxiv on 14th October which is after the ICLR submission deadline. It leverages Transformer for world modeling and uses a categorical variational autoencoder that adds stochastic noise to the latent space. It demonstrates a mean human performance of 126.7% with the same superhuman score of 9. In contrast, when considering the median human normalized score, DART demonstrates better performance. However, STORM stands out for its remarkable reduction in training time, achieved through the use of a single stochastic latent variable. In contrast, DART employs discrete tokens, allowing the model to concentrate on task-relevant tokens. In the updated version, Table 9 has been extended to include the scores of STORM.

Could the authors add one ablation study on DART with continuous representation?

Evaluating DART with continuous representation will involve modifying the methodology for representation learning, world modeling, and behavior learning. This is challenging to implement and validate during the short discussions phase, due to computational resources and long training time but it is an interesting follow-up work.

[1] "High-accuracy model-based reinforcement learning, a survey." Artificial Intelligence Review.

We thank the reviewer again for their insightful comments. We hope we have addressed the concerns. Please let us know if we can provide more clarifications. Please consider increasing the score if you find the responses satisfactory.

Thank the authors for the response.

R1: long-range dependencies impede Dreamers learning

DreamerV3 performs quite well on Crafter. If compressed states lead to degradation for the long-horizon task, could the authors provide a comparison between DreamerV3 on Crafter with converged results, not limited to 100K interaction? Also, the memory token used in the proposed method is also a compressed vector; I think the argument and experiment results are not convincing to me.

R2: model accuracy comparing with baselines

Is Figure 5 & 6 the reconstruction or generation? The label in the figures indicates it is reconstruction. It's a little confusing to me. I thought the authors should provide a quantitative comparison result (as requested) with baseline models to demonstrate the proposed method learns a more accurate world model.

R3: MEM is also compressed token, could the authors explain why MEM helps?

ConvRNN can also capture both spatial and temporal relationships; why is MEM better?

R4: ablation on continuous representation.

I think this experiment is quit important to verify the effectiveness of the proposed discrete tokens.

We thank the reviewer for their insightful comments and valuable time.

The approach is simple (which is good) and integrates components used by the model already exist in literature. Learning a world models on discrete tokens has been previously introduced in IRIS and using a ViT policy (which the authors claim to be their main novelty) head has been studied by Yoon et al 2023 (https://arxiv.org/abs/2302.04419). Usage of a memory token to feed past context has also been studied in Bulatov et al 2022 (https://arxiv.org/abs/2207.06881), Didolkar et al 2022 (https://arxiv.org/abs/2205.14794), Moudgil et al 2021 (https://arxiv.org/abs/2110.14143). I would suggest to include these works in the introduction and recontextualize the work based on the above works.

In the updated version, we have revised the introduction section to incorporate the suggested works. We want to highlight that Yoon et al’s work (https://arxiv.org/abs/2302.04419) involves pretraining a transformer encoder to model object relationships. This generates representation used for downstream RL, potentially encoding task-irrelevant information. However, our approach employs a transformer encoder for policy modeling without any pretraining, allowing the agent to learn the policy by focusing on task-relevant information. Bulatov et al 2022 (https://arxiv.org/abs/2207.06881) require storing all memory tokens from the past segment for prediction, however our method, employs a self-attention mechanism to propagate information as required by the task from the current time step to the next, eliminating the need to store any past tokens (refer to Fig. 4). Our model simplifies memory representation, in contrast to Didolkar et al’s (https://arxiv.org/abs/2205.14794) two-stream network which uses a slow stream (long-term memory) network modeled using a recurrent network and a fast stream (short-term memory) modeled using a transformer which is conditioned on the slow stream.

The results dont seem strong enough. The model only differs from IRIS in the policy and in many games IRIS still performs better.Secondly, according to the IRIS paper, they achieve superhuman performance in 10 games while in Table 1 it says 9 games and DART also achieves superhuman performance in 9 games. Therefore, IRIS actually can outperform humans in more games than DART hence the usefulness of of the ViT policy is not very apparent.

DART has better performance than IRIS in 18 out of 26 games (underlined in the revised version in Table 3). From the Table 1 of original IRIS paper (https://openreview.net/pdf?id=vhFu1Acb0xb) and from Table A.1 of Schwarzer et. al (https://arxiv.org/pdf/2305.19452.pdf) IRIS outperforms humans in 9 environments only. We achieve the same superhuman score of 9, we show a significant improvement in other metrics including median, IQM, and optimality gap scores when compared to IRIS, along with better interpretability of the policy.

It would be nice to see a section where the authors compare DART to only IRIS and try to study in more detail the importance of the ViT policy.

In the revised version, in Table 1 we try to highlight DART’s superior performance compared to IRIS. We also provide a system-level comparison of IRIS and DART in Section A.3. Our modeling of states and memory in the form of tokens allows ViT to enhance interpretability as shown in Figure 4.

Figure 3b compares DART to each approach individually but I am not sure how these probabilities are calculated. Can the authors clarify this?

Introduced in [1] the metric used in Figure 3b is used to overcome the limitations of point estimates, as point estimates can vary across independent runs. We use this by collecting scores for 5 different seeds across 26 games for DART along with the model it needs to be compared with. Subsequently, the Mann-Whitney U-statistic is applied, averaging across all scores.

While the current paper and the baselines study the performance in a setting where the model is limited to 100k interactions, I think it would still be useful to compare how these approaches scale with more interactions and whether the current trends still hold with more interactions.

While it is computationally expensive to run the model for millions of interactions, we are currently training our model and have initial results for DART on a few environments with 150k interactions (Section A.2). We are currently waiting for the corresponding results with IRIS and DreamerV3 and hope to be able to share them before the end of the discussion period.

Given our limited compute resources, we hope we could address most of your concerns except the results for the longer IRIS experiment which we hope to provide soon. Kindly consider increasing the score if you find our responses satisfactory. We would be happy to answer any further questions.

[1]. "Deep reinforcement learning at the edge of the statistical precipice." Advances in neural information processing systems 34 (2021).