Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposed a transformer-based world model. Utilizing the Transformer-XL, the agent is trained in imagination rollouts in a computationally efficient way. Though inspired by prior work, such as the Dreamer frameworks, several modifications were made and ablation studies showed benefits to the final performance. Experiments on the Atari 100K benchmark showed an overall improvement over a set of the prior model-free and model-based methods.
strength_and_weaknesses: ### Strength
- The paper is well-written generally, providing clear reasoning throughout the whole manuscript.
- Experiments on the Atari 100K benchmark showed an overall improvement over several model-free and model-based methods.
- Detailed ablation studies have been provided.

### Weakness

- Major:
    - Is the input for policy learning $z_t$ or the Transformer predicted $\hat{z}_t$ during policy learning? The description in the `Choice of Policy Input` section looks confusing. If it's $\hat{z}_t$, did you apply the frame stacking?
    - It seems this work has been inspired by the Dreamer framework a lot, why didn't the author choose to compare it with Dreamer-V2?
    - The results on the effect of taking rewards as the input seems too weak, with benefits only shown on one task. Could the authors please explain how are the tasks selected in the ablation studies?
    - Since Atari games do not require long-term memory in general, how does the Transformer help on this Atari 100K benchmark?
    
- Minor:
    - The objectives seem complex, could the authors analyze the sensitivity of the set of hyper-parameters used in the objectives?
    - How are the M trajectories generated given N sequences used to train the world model?
clarity,_quality,_novelty_and_reproducibility: The paper is well-written generally, providing clear reasoning throughout the whole manuscript. Results might be reproduced with reasonable effort.
summary_of_the_review: This paper proposed a transformer-based world model. Experiments on the Atari 100K benchmark showed an overall improvement over a set of the prior model-free and model-based methods. However, the motivation is not clear and some key statements need further explanations.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The authors investigate the effectiveness of transformers as world models. Their main innovation is to input not just states and actions (as was done in earlier works like Ha & Schmidhuber, 2018), but also to add rewards (predicted or real). This additional constraint on the input interacts well with the transformer architecture to yield sample efficient and performant agents on the Atari 100k benchmark.

UPDATE

raising my score due to the reviewers responses and work over the rebuttal period
strength_and_weaknesses: Strengths
- Their approach outperforms baselines across the board on Atari 100K.
- They train a lighter weight policy using the transformer world model, which makes inference faster than it would be otherwise.
- The method is well explained.
- The authors included code as their SI.

Weaknesses
- Missing a citation and discussion of Micheli et al., Transformers are Sample Efficient World Models. I do not believe this is published in a conference yet so we can consider it concurrent work but you should include it in your paper.
- There's no analysis of the world model! Tell me more about what the transformer has learned: is it a veridical representation of the games, are certain elements overemphasized over others? What's going on with the world model? How does its representations change when you add reward as an input?
- Results are only reported for one challenge (Atari 100K) and there's clearly a large amount of tweaking that went into this work in terms of loss and data engineering. That's understandable to an extent, but as it is paired with the limited insights into what is happening in the world model, I was left underwhelmed by this work.
- The title is flashy, but it also raises an important question which I don't believe was properly addressed in this paper: what is the sample efficiency of the transformer world model when you train on different amounts of data? For instance, try log-spaced intervals from 10 to the 100K reported in the main text. Understanding what failures occur at which dataset sizes would go a long way in illustrating the benefits (and persistent failure modes) of transformers as world models, and how introducing reward changes that sample efficiency curve.

Questions
- In your intro you write: "This can be illustrated by a simple example: in the game of Pong, the paddles and the ball move independently. A successfully trained world model is able to imagine trajectories with paddle and ball configurations that have never been observed before, which enables learning of improved behaviors." Is there evidence for this? Or are you positing that world models could help an agent solve this out-of-domain challenge?
- "In contrast to Hafner et al. (2021), the role of the observation model is to capture only non-temporal information about the current time step. However, we include short-time temporal information, since a single observation ot consists of four frames (aka frame stacking, see also Section 2.2)." I think the first sentence here is funky. Did you mean "in contrast to Hafner, *where* the role..."
- "Due to these design choices, our world model combines the benefits of transformers and recurrent neural networks..." I'm not sure any of the four bullets are areas where RNNs are better than transformers. My thought is the key benefit of RNNs over transformers is in the amount of memory needed to run them (i.e., RNNs operate on a single token at a time whereas transformers look at many at once). 
clarity,_quality,_novelty_and_reproducibility: The paper is very clear although I found the experiments (and lack of theory) underwhelming in their depth of analysis for what I think is a very cool model. The main novelties here are how the authors designed the transformer. They introduce reward into the input and utilize a bunch of data and loss engineering. The results are reported with error bars and the standard for model adjudication in RL, and so I believe these are reproducable.
summary_of_the_review: I was really excited about this paper for the possibility of understanding why transformers might make sample efficient world models. Instead, this paper is a standard benchmark paper describing a recipe for performant transformer-based world models. I am far less interested in the latter, but I can appreciate that the contribution required a lot of work and is nontrivial. So I'm borderline on this paper. It would make a fantastic workshop paper. I'm less sure if ICLR is the right venue. Looking forward to seeing what the other reviewers think.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper introduces a model-based RL approach that learns a transformer-based world model and additional techniques to further improve the sample-efficiency. Specifically, VAE is trained to learn visual information from raw pixels, and autoregressive dynamics model based on Transformer-XL architecture is trained upon a sequence consisting of states and rewards. The method is evaluated on a set of Atari tasks within 100k steps to measure sample-efficiency of the proposed method.
strength_and_weaknesses: ## Strengths

- Strong performance with simple and intuitive method
- Paper reads well

## Weaknesses

- Claims within the paper are not justified enough except that the proposed method performs well.
    - It is not clear whether transformer-based model is indeed better than previous architecture with GRU of DreamerV2. From the current results, we cannot know whether the gain comes from additional techniques or transformer-based world models or additional computes. How does it compare to DreamerV2 trained with the proposed techniques, similar computes (not parameters), and hyperparameters to be more sample-efficient?
    - The paper says training is more efficient as it's parallel, but does not provide any measurement or investigation to explicitly show this compared to previous architecture based on a recurrent architecture. Gviven that the method utilizes a quite large model consisting of 10-layers Transformers, this should be further thoroughly supported.
    - It is written that leveraging Transformer-XL [Dai et al., 2019] can be good for long-term inference, but in the paper it's not investigated whether this is indeed the case. Additional results that compared with vanilla Transformer architecture could be included to support this.
    - It is claimed that long-term dependencies can be captured with Transformers but final policy is trained on top of the representation from observation consisting of 4 stcaked frames. So it's difficult to know it's really good for capturing long-term dependencies in terms of either prediction for generating training samples or inference. It should be evaluated on more challenging tasks that require memories?
    - For thresholded entropy loss, I can understand that it can make normalized entropy stable, but is this necessarily better than changing entropy throughout the training? How does the performance is related to the change in the entropy here?
    - For sampling scheme, it would be nice to further describe and formulate how this cheme would affect the optimization objectives compared to using uniform sampling scheme. Moreover, ablation to see how this affects performance would be also nice.
- Investigation into the effect of conditioning the policy on $z_t$ is not complete. It's okay to say that this would incur more costs, but it seems an important difference to prior approaches, so there should be an investigation into this with additional results and also corresponding wall times.
- Comparison with prior works can be much more informative if it includes some metrics in terms of training costs -- as it seems like the proposed method would require much more resources than other approaches (of course it's my guess because related information is not included anywhere, correct if i'm wrong), and it also does some pre-training on initial exploration samples.
- Missing comparison with TransDreamer, and not clear why this is only included in Related Work? 

[Dai et al., 2019] Dai, Zihang, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. "Transformer-xl: Attentive language models beyond a fixed-length context." arXiv preprint arXiv:1901.02860 (2019).

[Chen et al., 2022] Chen, Chang, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. "Transdreamer: Reinforcement learning with transformer world models." arXiv preprint arXiv:2202.09481 (2022).

---

Updating the score to 6 after the discussion with AC and other reviewers. I would like to strongly encourage the authors to include results to compare with RNN-based world models or tone-down on claims on the benefit of transformer over RNN-based world models.

clarity,_quality,_novelty_and_reproducibility: ### Clarity and Quality

- The paper is clearly written and reads well

### Novelty

While the proposed architecture seems novel but it's not surprisingly novel as there has been an effort to incorporate Transformers into world models [Chen et al, 2022] and there has been recent works that also takes similar two-stage approaches in the context of video prediction [Rakhimov et al., 2020; Yan et al., 2021]. But I acknowledge the details can matter for world models so it could be a novelty of the approach if this can be thoroughly supported with analysis and additional results. 

### Reproducibility

Source code is included so it could be reproduced. But more analysis on wall time could be nice in this front.

[Chen et al., 2022] Chen, Chang, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. "Transdreamer: Reinforcement learning with transformer world models." arXiv preprint arXiv:2202.09481 (2022).

[Rakhimov et al., 2020] Rakhimov, Ruslan, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. "Latent video transformer." arXiv preprint arXiv:2006.10704 (2020).

[Yan et al., 2021] Yan, Wilson, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. "Videogpt: Video generation using vq-vae and transformers." arXiv preprint arXiv:2104.10157 (2021).
summary_of_the_review: While the paper proposed an intuitive approach that achieves strong sample-efficiency, the main weakness of the paper is in that it lacks analysis and experimental results to support the claims made in the paper, which makes me difficult to see that the current draft is strong enough to be accepted as a ICLR paper (I'm trying to avoid borderline score following the suggestion from the conference this year). Hence I'm recommending the paper to be rejected at the current status but willing to update the score per the discussion within the rebuttal period.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper investigates the Transformer-based world model for model-based reinforcement learning. 

Inspired by the successful applications of Transformers to supervised learning or generative learning, applying Transformers to reinforcement learning has been one of the most interesting topics among RL community members [1, 2, 3, 4, 5]. However, one of the pain points of Transformer applications is the much larger computation overheads when interacting with a real environment or learning in imagination [1, 2, 3]. The world model design in this paper reduces the overhead while it limits the agent to attend to the past at inference time.

In addition, they suggest reasonable techniques to train the agent on top of this world model to be more stable and better and show the ablation studies with/without the techniques too. Their model outperforms baselines, including the state-of-the-art model-free and model-based methods.


strength_and_weaknesses: Strength
- The paper is well-written and easy to follow.
- They also support ablation studies for suggested techniques, balanced dataset sampling, threshold entropy loss, and feeding back the predicted reward. Thus, the reader can understand why they applied those theoretically and empirically.
- Using $z_t$ as the input of the policy, not the $h_t$. This choice can interact with real environment more quickly, which is different from previous works [1, 3].

Weakness
- The motivation why the Transformer-based World Model will work well is not considered (e.g., in [1,2,3,4], they address the long-term dependency issue). Maybe better transition modeling could be one of the motivations. Then it should make this paper more concrete if there were empirical results about the world model qualities when using Transformer and recurrent modules.
- As pointed out above, their modeling cleverly uses $z_t$ as an input of policy, while it is limited for attending the history. It could be optional, like for the task requiring long-term dependency, using $h_t$ and usually using $z_t$, but when using $h_t$, the inference could be slower.


clarity,_quality,_novelty_and_reproducibility: It is well-written and clear, and the proposed techniques are densely analyzed through their ablation studies.

In the aspect of the novelty of their modeling, it is not too new, but they first analyzed the Transformer-based world model for the Atari 100k benchmark and showed better performance than the state-of-the-art baselines. Their techniques, balanced dataset sampling, and thresholded entropy loss are pretty new, and they showed those are important to increase performance. 

It looks easy to be reproduced.
summary_of_the_review: This paper does not consider some points about why Transformer is good for the world model and how their model can handle the long-term dependent knowledge. However, they showed impressive performance by applying Transformer to world model with their techniques, which is good enough to be shared in our community. Especially by densely analyzing their techniques, they made the reader easily understand why the techniques are required. Hopefully, the authors will revise this paper by addressing my concerns. 


[1] Parisotto, Emilio, et al. "Stabilizing transformers for reinforcement learning." International conference on machine learning. PMLR, 2020.

[2] Parisotto, Emilio, and Ruslan Salakhutdinov. "Efficient transformers in reinforcement learning using actor-learner distillation." arXiv preprint arXiv:2104.01655 (2021).

[3] Chen, Chang, et al. "Transdreamer: Reinforcement learning with transformer world models." arXiv preprint arXiv:2202.09481 (2022).

[4] Esslinger, Kevin, Robert Platt, and Christopher Amato. "Deep Transformer Q-Networks for Partially Observable Reinforcement Learning." arXiv preprint arXiv:2206.01078 (2022).

[5] Micheli, Vincent, Eloi Alonso, and François Fleuret. "Transformers are sample efficient world models." arXiv preprint arXiv:2209.00588 (2022).
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper trains a world model on Atari and trains policy on the model itself only, to achieve good performance on Atari 100k. The novelty of this paper (the setup itself existed before) is to use transformer for the world model: Consuming all the vector quantised codes of a given frame at the time, updating the transformer, sampling next set of codes. The policy is computed from the codes only. 
- The positive of this paper is showing that this works well beating previous methods published on this benchmark. 
- The negative is that the setup is not novel - a similar system has been trained on with LSTM instead of transformer (DreamerV2) but not evaluated in the same setting (100k interactions). One of the biggest drawbacks if the paper is that the authors did not rerun the same system with LSTM instead of transformer or compared to that paper in the same setting.
justification_for_why_not_higher_score: No comparison to LSTM/DreamerV2 on the same benchmark (Atari 100k) or longer training.
justification_for_why_not_lower_score: Good performance, nice to compute policy from codes.
Decision: Accept: poster