Diffusion for World Modeling: Visual Details Matter in Atari

Summary

The paper proposes a visual diffusion-based world model for reinforcement learning. The authors argue that current world models using discrete latent variables may lose important visual details, which are critical for reinforcement learning tasks. DIAMOND addresses this by generating images in the original pixel space. The paper evaluates DIAMOND on the Atari 100k benchmark and achieves strong results relative to other choices for the world model.

Strengths

- Clear and well-written paper - Strong empirical results beating prior MBRL baselines on Atari 100K - Impressive scaled-up experiments for visual diffusion world modeling, with consistent visual quality over long horizons even with a low number of diffusion steps - Good set of ablations over key design choices of the algorithm - Publicly available code for the experiments

Weaknesses

- It would be useful to include a comparison of training time for DIAMOND compared to prior baselines, e.g. what is the change in speed by using a diffusion world model v.s. Latent world model? - How does the algorithm compare to the current SOTA model-free algorithms on Atari 100K, e.g. [1]. This would help the author understand the relative utility of MBRL for this setting Minor: - Theoretical diffusion background is lengthy and the paper would be better served by including results on other environments that are currently in the appendix, or discussion on architecture used [1] Bigger, Better, Faster: Human-level Atari with human-level efficiency. Max Schwarzer, Johan Obando-Ceron, Aaron Courville, Marc Bellemare, Rishabh Agarwal, Pablo Samuel Castro.

Questions

See weaknesses.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Discussion in paper.

Summary

This paper introduces DIAMOND (DIffusion As a Model Of eNvironment Dreams), a novel RL agent trained within a diffusion-based world model. DIAMOND's world model is a diffusion model that, conditioned on past observations and actions, generates an observation at the next time step. This approach diverges from previous methods that rely on discrete latent variables, and offers a promising alternative that leverages the strengths of diffusion models, such as the ability to model multi-modal distributions. The training process is iterative, involving three key steps: data collection in the real environment, world model training on this collected data, and RL agent training with rollouts using the learned world model. The authors analyze design choices for the diffusion world model, and demonstrate that the EDM formulation is superior to DDPM for their use case through visualization of generated trajectories. They also conduct qualitative experiments to determine the optimal number of denoising steps for sampling from the world model. Finally, DIAMOND's evaluation on the Atari 100k benchmark reveals superior performance compared to other world model-based RL agents, achieving a new state-of-the-art mean human-normalized score of 1.46.

Strengths

This paper introduces a new new diffusion-based approach to world modeling in reinforcement learning, and show that it is a promising alternative to prior approaches based on discrete latent variable methods. The proposed method achieves state-of-the-art results on the Atari 100k benchmark and surpasses other world model-based RL agents. The paper also provides a thorough qualitative analysis of the design choices involved in creating an effective diffusion world model, including the choice of diffusion framework (EDM over DDPM) and the number of denoising steps. Overall, the paper is well-written and explains the concepts in a clear manner.

Weaknesses

One minor concern is the lack of quantitative analysis for the design choices. While qualitative experiments and the results on the Atari-100k benchmark effectively motivate the design choices, a more thorough quantitative analysis across the different tasks would provide stronger evidence for these choices. Such a quantitative analysis would be potentially important for applications in real-world environments.

Questions

- The performance on a few tasks (BankHeist, Frostbite, UpNDown) are considerably worse than the baselines. Have you investigated why this is the case? - Is the diffusion model retrained from scratch in each epoch of the training loop? Also, is the training data in a given epoch just what is collected in the current epoch or the union of data collected in all epochs so far? - Have you investigated how the diffusion world model evolves as the agent used for collecting the training data improves?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have included a dedicated section that discusses the limitations of the work.

Summary

This paper introduces a new world model for learning behaviors in imagination using reinforcement learning. In particular, a diffusion model is used to generate the next frame, conditioned on previous frames and actions. At each environment step, multiple denoising steps are performed to convert a noise image into the next frame. By training an RL agent on synthesized trajectories, the method achieves state-of-the-art performance on the Atari 100k benchmark.

Strengths

- S1: The world model achieves strong performance on the Atari 100k benchmark, compared with other world models. - S2: Implementing world models using diffusion models is a logical step, given the success of diffusion models in image generation.

Weaknesses

- W1: The analysis of the world model is mainly qualitative and not quantitative. The model seems to be really good at generating long trajectories without compounding errors (e.g., Figure 3(b)), but it would be nice to have a more objective measurement of this. One simple idea would be to generate long trajectories and compare the generated frames to the nearest neighbors in the replay buffer. This could also be compared with IRIS and DreamerV3. - W2: As the training is rather slow, it would be interesting to see a breakdown of the training times of the individual components. For instance, how much time is spent to generate frames compared with the reward/termination model (which is a CNN + LSTM)? Typo in L240: The single denoising step is shown in the "last row" (instead of "first row").

Questions

- Q1: The imagination horizon is set to the usual value of 15 steps. I am wondering whether longer horizons would lead to better scores, or whether this is not required for Atari?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors addressed all limitations adequately.

Summary

This paper proposes an approach for learning world models with diffusion based approaches, compared to the recently proposed ones using transformers or the ones dependent on discrete latent variables in general. The core idea is that given success of image generation using diffusion models, the visual details of game playing RL tasks can be improved if a diffusion based world model can be learnt; which can in turn lead to improved performance specifically in pixel based tasks like Atari.

Strengths

The idea of using generative models for learning the world model in RL is perhaps nothing new; past works have tried to do this with transformers or other discrete variable models; however success has not been achieved much. In contrast, while the idea of using diffusion for world models perhaps may sound oblivious given current context, this paper does a good job in demonstrating that this can work empirically in Atari. Experimental comparisons with past works such as IRIS, that are dependent on using transformers, shows that the learnt difufison based approach can achieve empirical gains.

Weaknesses

One primary bottleneck of this approach would be the dependency of the diffusion driven word model for long horizon tasks. Past works often model a per-next time step state prediction dynamics model, or uses models that can handle longer sequences. In contrast, the diffusion based approach may signifcaintly suffer from the long horizon. Experimental approach and algorithmic novelty is perhaps nothing significantly new; this paper basically does a simple plug in of diffusion model with careful fine-tuning in the context of model based RL. As stated by the authors as well, the practical choice of the diffusion approach would matter a lot here; and carefully analysis is required to ensure that the difufuson model cxan learn a good world model. Empirically, other than Atari benchmarks, can this work demonstrate that the idea can work in other domains and compare to more extensive experimental analysis with Dreamer based approaches? For example, Dreamer line of work often shows good results in locomotion or humanoid driven tasks - I think this approach can fail in those cases since it is harder to learn a world model that can generate visual details in those tasks? Does the authors have any comments around that?

Questions

Can we see some more experimental results other than Atari to demonstrate that the idea can work? More experimental comparisons with other model based approaches are required. MBRL is a huge literature in the field and lots of works have compared to different algorithms in different benchmarks. In contrast, this paper lacks empirical evidence other than Atari domains which seems concerning. Have the authors considered other non-diffusion based approaches or more recent works based on flow matching for example? Or can we see how the learnt world model differs based on the type of generative model we use? Example, if we can use different variants of diffusion or FM based models, how does the eprfomance differ? I’d like to understand in general, the significance of using generative models for learning world models. There used to be prior works on imagination augmented rollouts, or works that would learn a multi-step forward dynamics model for example; other than IRIS that are dependent on the use of transformers, can we see some more experimental comparisons with those prior approaches

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Lack of enough experimental evidence other than Atari benchmarks

Please could you respond to our rebuttal?

Dear Reviewers, Thank you again for your time and effort reviewing our paper. We are now more than halfway through the discussion period. We would be grateful if you could reply to our responses soon so that we have time to address anything remaining. Please could you let us know if we have addressed your concerns, or if there is anything else we can address for you to consider increasing your ratings of our paper? Kind regards, The Authors

Thank you for your response

Thank you for the detailed responses to my questions and comments. I have updated my score to reflect this.

Follow up on our rebuttal

Dear Reviewer hqzV, Thank you for your review and for the positive evaluation of our paper. We believe we've addressed your concerns and hope you appreciate our responses and additions. As the discussion period is coming to a close, please let us know if there are any further points we can address to improve your rating. Kind regards, The Authors