STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

Summary

The paper introduces the Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture. STORM proposes to encode image inputs using a stochastic variational autoencoder, and predicts latent state using a GPT-like sequential model. It then trains dynamics and policy based on the outputs of these stochastic variational autoencoder and sequential model. The authors conducted experimental comparisons with some classic baseline methods in the Atari100k environment. The results indicate that STORM has a faster speed and achieves better performance.

Strengths

1. Although the individual components proposed by STORM have been introduced before, the authors designed a solid structure to combine these parts into an effective method. 2. The experimental results demonstrate the effectiveness of STORM, and more importantly, they improve the time efficiency of Model-Based Reinforcement Learning methods.

Weaknesses

1. In the experimental section, please add a comparison with Speedy Zero [1]. Speedy Zero is also a model-based RL method which is proposed recently and has also achieved good time efficiency and performance on Atari. 2. Beyond Atari, can STORM be applied to other tasks, such as MuJoCo, DMC, MetaWorld, etc.? Dreamerv3 can achieve very good results in a wide range of different environments. If STORM cannot, the practical significance of this paper will be greatly reduced. 3. It would be better if further analysis could be provided on the benefits of using a stochastic variational autoencoder. Reference: [1]. Mei et.al. "SpeedyZero: Mastering Atari with Limited Data and Time." In ICLR 2023 https://openreview.net/forum?id=Mg5CLXZgvLJ.

Questions

See weakness above

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The generalizability of STORM could be further studied.

Summary

This paper proposes a world model architecture (STORM) to train RL agents in imagination. The world model is composed of an autoencoder with categorical latents and a Transformer. These modules are trained jointly with a reconstruction loss, a next latent state prediction loss, as well as reward, episode termination, and representation losses. Experiments in the Atari 100k benchmark indicate that the approach is effective.

Strengths

- The method is technically sound with empirical results to back its effectiveness (outperforms other methods based on learning in imagination). - It is faster to train than similar methods (e.g. twice as fast as TWM). - The paper is well-written and easy to follow.

Weaknesses

- The method lacks novelty and the incremental improvements over previous work are not properly explored. - STORM is a variant of TWM with two minor modifications: (1) a vanilla Transformer is used instead of a TransformerXL, (2) latent state and actions tokens are fused instead of having separate tokens. - Currently, it is not clear why STORM achieves better results than TWM. Is it due to implementation details or the proposed modifications actually matter? And if they matter, why? - The results have limited significance. Although STORM outperforms TWM, it only yields marginal improvements over DreamerV3, which was not specifically optimized for Atari100k. Moreover, recent model-free methods [1, 2] achieve similar or better results in the benchmark. - Ablations and additional experiments fail to provide additional insights: - lines 237-238: IRIS also relies on an autoencoder trained with a reconstruction loss and it still obtains good results on Breakout and Pong. - l261: comparing the number of layers for STORM and IRIS/TWM is not straightforward since the hidden dimension is not the same (256 vs 512). - 5.2: it is hard to draw any conclusions as only a few environments were considered and we do not know how they were picked. Also it is not clear if including z_t yields statistically signficant improvements. - 5.3: the premise of this section is interesting but again too few environments are considered and results have limited statistical significance. The improvement on Freeway over previous methods relies on the addition of an extra demonstration while the other methods use exploration strategies that do not involve expert data. It would be interesting to know whether STORM leverages demonstrations better than other world models.

Questions

### Do you use sticky actions? --- I ran your code, on the 4 following games, with and without sticky actions: | Environment | Reported | Without sticky actions | With sticky actions | | --- | --- | --- | --- | | BankHeist | 641 | 1044.0 | 218.5 | | Breakout | 16 | 29.3 | 11.2 | | MsPacman | 2673 | 2942.0 | 1921 | | PrivateEye | 7781 | 4458.4 | 100.0 | Can you clarify whether you used the Atari environments with (`ALE/<envname>-v5`) or without (`<envname>NoFrameskip-v4`) sticky actions? It is not mentioned in your paper but the code seems to use the v5 environments by default (TWM, IRIS and DreamerV3 do *not* use sticky actions). It seems that my `v5` results are significantly below the reported results, and that the `v4` are on-par/slightly better. ### Ablations to investigate the differences with TWM --- The paper would greatly benefit from a thorough investigation of the differences with TWM. It should be made clear why STORM performs better than TWM as it seems like a variant with two minor modifications. - Is the substantial performance improvement explained by implementation differences (in particular, there is [an open issue about reproducibility](https://github.com/jrobine/twm/issues/3) on TWM’s repo), or by the two modifications? - If these modifications matter, can you run ablations to demonstrate their effectiveness? - Maybe it is due to the incorporation of tricks for policy learning from DreamerV3? I am keen to significantly increase my overall rating if this concern is properly addressed during the rebuttals. ### Other concerns --- - In my opinion, there would be a substantial gain to experiment in more complex environments, e.g. Crafter [3], Minecraft [4], or Memory Maze [5]. Such results would expand our knowledge of what Transformer-based world models can achieve. - Can you include the recent work on the benchmark [1, 2] in the related work section? --- ### References - [1] *Sample-Efficient Reinforcement Learning by Breaking the Replay Ratio Barrier*. D'Oro, Pierluca and Schwarzer, Max and Nikishin, Evgenii and Bacon, Pierre-Luc and Bellemare, Marc G and Courville, Aaron. The Eleventh International Conference on Learning Representations, 2022. - [2] *Bigger, Better, Faster: Human-level Atari with human-level efficiency*. Schwarzer, Max and Obando-Ceron, Johan and Courville, Aaron and Bellemare, Marc and Agarwal, Rishabh and Castro, Pablo Samuel. arXiv preprint arXiv:2305.19452, 2023. - [3] *Benchmarking the Spectrum of Agent Capabilities*. Hafner, Danijar. International Conference on Learning Representations, 2021. - [4] *Minerl diamond 2021 competition: Overview, results, and lessons learned*. Kanervisto, Anssi and Milani, Stephanie and Ramanauskas, Karolis and others. NeurIPS 2021 Competitions and Demonstrations Track, 2022. - [5] *Evaluating Long-Term Memory in 3D Mazes*. Pasukonis, Jurgis and Lillicrap, Timothy P and Hafner, Danijar. The Eleventh International Conference on Learning Representations, 2022.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors discussed some technical limitations in the conclusion.

Summary

The authors propose several modifications to the recently proposed Transformer based world model for Model-based reinforcement learning. Specifically, they come up with a single latent stochastic state and treat action as an explicit input to the state as opposed to a token (as in previous works) and show significant performance and speed in limited regime of Atari 100k tasks.

Strengths

1. Improved performance as well as significant time reduction in STORM's agent training which is crucial in MBRL agents alongside sample efficiency. 2. Well written paper overall and was easy for me to read in a single go (NOTE: I'm very familiar with this MBRL space so that is another reason that has contributed to this but overall the flow of the paper was very intuitive.)

Weaknesses

1. **Highlighting differences between TWM, Dreamerv3, Transdreamer, IRIS, SimPLe**: (a) It would be much better to have a table indicating different aspects such as (a) RNN (b) multiple tokens (c) Transformer variant etc. This would be much easier for the reader to read. Additionally, I did not find the motivation behind some of these modeling choices (see below for detailed descriptions). (b) On L127 the authors write "In contrast to IRIS [12] that employs multiple tokens, STORM utilizes a single stochastic latent variable to represent an image." This is correct but TWM encoding style and STORM's encoding style are the same -- so I find it a bit misleading to omit TWM as all three (TWM, IRIS and STORM) are transformer models (c) On L129, re: "STORM follows a vanilla Transformer [15] structure, while TWM [11] adopts a Transformer- XL [19] structure." -- it would be helpful why one would like to use a vanilla Transformer as opposed to Transformer-XL? What are the benefits? (d) On L131, "TWM [11] treats observation, action, and reward as three separate tokens of equal importance." -- I am not sure if this statement is true. The input to the transformer in the case of TWM is (obs, action, rew) but that doesn't imply that they are equal -- the attention weights of the transformer would be the ones deciding whether to consider these tokens or not. See Figure 6 of TWM paper showing the attention map over $(s, a, r)$ and it is clear that not all are weighted equally. (e) On L134, "Unlike Dreamer [8] and TransDreamer [16], which incorporate hidden states, STORM reconstructs the original image without utilizing this information." -- What is the additional benefit of not using the hidden state? I don't think there is a significant reduction in time for the reconstruction. 2. **Performance**: (a) What specific component according to the authors attributes to the superior performace of STORM. L235-237 mention the self-attention mechanism -- which is a valid argument against the RNN based methods. However, it is unclear to me what is helpful in STORM when compared to IRIS or TWM. (b) L237-238: what do authors mean by the "nature of autoencoders?" For example, the encoding style in TWM and STORM is very similar. It is unclear to me why specifically to single object games does STORM perform poorly. I'd like the authors to elaborate on this. 3. **"Decoder at rear" experiment**: I am not sure what exactly the purpose of this experiment was. From what I understand, the model can be formulated as TSSM (as mentioned in TransDreamer paper), and hence the the reconstruction would be using the posterior $z_t$ and not the prior $\hat{z}_t$. Is there something that I'm missing here? 4. **Impact of trajectory**: (a) How exactly was the trajectory used in the replay buffer? Was there any prioritized replay buffer or was the sampling of trajectories for the training of world model uniform? (b) Inclusion of trajectory is *not* specific to STORM -- so I'm not sure how this experiment validates the usefulness of STORM specifically and not apply to other world models (TWM, IRIS etc). **[Very minor comment -- not considered for rating the paper]** 5. In appendix, Table 9, it would be more clearer to denote the "Imagination batch size" by another symbol or $B \times T$ as that is effectively what happens during imagination. ---- **Rationale for my rating** I think the model details mentioned in the paper are important to share with the broader community, however I do think that motivation for several choices made in the modeling of STORM that has *not* been explicitly provided. Additional it is not clear to me that experiments such as "Decoder at rear" or the inclusion of trajectory are fundamentally because of STORM and seem more of a study on Dreamer-*like* paradigm. I would like the authors to address my concerns above and for the time being, I have leaned towards Borderline reject. I will update my rating post-rebuttal and discussion with the authors. ----- **Post-rebuttal rating** Based on the rebuttal provided by the authors, I've decided to bump up my rating to *Weak Accept* as they have addressed a majority of my concerns. The only section that I'm currently unconvinced is the inclusion of demonstrations which is not the major contribution of this work and would require its own analysis instead (a single section won't suffice in my opinion). I once again thank the authors for their comprehensive rebuttal.

Questions

See **Weaknesses** section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA

Summary

The paper presented a Transformer-based model-based RL framework. As with earlier approaches, online data is gathered into the replay buffer with the learned reactive policy, the Transformer-based world model is trained with segments sampled from the replay buffer, then the policy is optimized with imaged data generated from the world model. The improved results were reported on the Atari 100K benchmark. The authors also conduct a thorough ablation study to justify heir architectural design choices.

Strengths

This paper is well-organized and easy to follow.

Weaknesses

- What criteria were employed for task selection in the ablation studies? They are not consistent over all the subsections, nor do they have the same trend in Table 1 compared to baseline models. The proposed method performs better in some tasks, while it doesn't in others. To provide a more robust evaluation of their proposed method, it would be better to consider extending the ablation studies to a wider range of tasks - It would also strengthen the paper if the authors could provide the converged results on the full Atari benchmark.

Questions

Please see the weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

yes.

Thanks for the rebuttal!

I thank the authors for their response. I appreciate the effort made to compare STORM with TWM and the inclusion of missing papers in the related work. Overall, the idea of the method suffers from a lack of novelty compared to TWM. However, given that the proposed method is faster and that the results are significantly better than TWM (and obtained in environments with sticky actions), I am convinced that this work would be of interest to the community. I am updating my rating from 5 to 7.

Response

Thank the authors for the clarification. I would like to keep my score as it is.

Thank you for the rebuttal. Most of my concerns have been addressed!

Thanks to the authors for the detailed rebuttal and response to all my questions. I appreciate their effort on this end. **1c** I agree with the rationale and for a benchmark on ATARI more than 512 tokens in unnecessary. However, for more long-horizon tasks where memory plays a key role, Transformer-XL can be good choice (as shown in TransDreamer). **1d** Thanks for acknowledging this. I do agree that computationally it does get expensive to send an additional *n* tokens for rewards -- especially if we are thinking of real-world experiment or might lead to negative learning of policy. **3** Thanks for the detailed explanation on this and adding this to the main paper would be incredibly helpful to the reader. A majority of my concerns (1-3) are addressed however I am not fully on-board with the rationale behind Section 5.3 **Impact of the demonstration trajectory**. The benefits of inclusion of a demonstration is trivial and has been shown in [1]. However, as most of my concerns are addressed I'd like to increase my vote to Weak Accept and vote towards acceptance of the paper. Contrary to reviewer E5z3, I do believe that Atari 100k is a challenging and sufficient benchmark to evaluate this work. Additional experiments on other environments would be nice to have but the current experiments, in my opinion, are sufficient to back up the claims. ---- References: [1] Multi-View Masked World Models for Visual Robotic Manipulation, Younggyo Seo, Junsu Kim et al, ICML 2023