Learning Transformer-based World Models with Contrastive Predictive Coding

Authors: Maxime Burchi, Radu Timofte
Primary area: reinforcement learning
Keywords: model-based reinforcement learning, transformer network, contrastive predictive coding

Abstract:
The DreamerV3 algorithm recently obtained remarkable performance across diverse environment domains by learning an accurate world model based on Recurrent Neural Networks (RNNs). Following the success of model-based reinforcement learning algorithms and the rapid adoption of the Transformer architecture for its superior training efficiency and favorable scaling properties, recent works such as STORM have proposed replacing RNN-based world models with Transformer-based world models using masked self-attention. However, despite the improved training efficiency of these methods, their impact on performance remains limited compared to the Dreamer algorithm, struggling to learn competitive Transformer-based world models. In this work, we show that the next state prediction objective adopted in previous approaches is insufficient to fully exploit the representation capabilities of Transformers. We propose to extend world model predictions to longer time horizons by introducing TWISTER (Transformer-based World model wIth contraSTivE Representations), a world model using action-conditioned Contrastive Predictive Coding to learn high-level temporal feature representations and improve the agent performance. TWISTER achieves a human-normalized mean score of 162% on the Atari 100k benchmark, setting a new record among state-of-the-art methods that do not employ look-ahead search. We release our code at https://github.com/burchim/TWISTER.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1729562615051

thread_type:
official_review

<thread_object>
creator: reviewer_trkk

<content>
summary:
In this paper, the authors investigated the Model-Based Reinforcement Learning (MBRL) agent with contrastive predictive coding. There have been many tries to apply Transformer as the world model backbone for MBRL. However, the authors argued the impacts from the tries are limited due to the limitations of the predictive loss (one-step next observation prediction), and proposed a new method equipped with the contrastive predictive coding loss. Through this, they reported the state-of-the-art performance on Atari100k benchmark among the MBRL agents without look-ahead search.

strengths:
- The breakthrough performance on Atari100k among the MBRL agents without look-ahead search. They showed clearly better performance than previous works, while the works have shown comparable performances for each other.
- Good motivation and validation of that. Their hypothesis, the predictive loss for next frame is not good enough to learn meaningful information, is reasonable and they validated it well through their proposed agent and experiments.
- Well written paper. We can enjoy to read this paper, because of well summarized related works, well discussed their motivation and diverse empirical ablation studies.
- The analyses for the experimental results. They properly discussed their experimental results, for instance, why their model shows better performance on Atari100k (because of capturing small but important objects which can be missed only through reconstruction loss). Ablations of their architectures are studied substantially.
- The detailed hyper-parameters are shared in the appendix.- We didn’t find a critical weakness from this paper.

weaknesses:
- We didn’t find a critical weakness from this paper.

questions:
- In Figure 2, you showed well why one-step predictive loss can be insufficient to learn the world model, but what means the shaded lines in the graph?
- Did you test the agent performance on Atari100k benchmark when doesn’t apply data augmentation? Only giving temporally near samples as the positive sample (without argumentation) is better I think (if the performance from that is comparable with the version equipped by the data argumentation).
- There is another study to apply contrastive loss to MBRL agent [1], considering it as a reference can be helpful to make this paper better.
- For AC-CPC Predictor, how did you input the actions? Did you concatenate them?
- Can you test the performance when not using reconstruction loss for the world model training? It can be more native version of MBRL with contrastive loss.
- In lines 426-432, your analyses why AC-CPC objective is helpful and what the previous works missed, make sense. As following your analyses, it makes sense that the reconstruction-based policy learning methods (IRIS and \delta-IRIS) show better performance than TWM or STORM in the Mean Human Normalized Score measurement. Then, why they (IRIS and \delta-IRIS) showed worse performance in Median Human Normalized Score measurement? Interestingly, your method outperforms for the two measurements. Can you explain which factors make this difference?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730207214342

thread_type:
official_review

<thread_object>
creator: reviewer_z6gc

<content>
summary:
This paper introduces a novel world model for deep reinforcement learning that integrates DreamerV3 with transformer-based architectures and contrastive predictive coding. Instead of predicting only the next state, this world model is designed to predict multiple time steps into the future. The method demonstrates state-of-the-art performance on the Atari 100k benchmark.

strengths:
- [S1] The paper presents a novel combination of established approaches, specifically transformer-based world models and CPC, which leads to strong performance results on the Atari 100k benchmark.
- [S2] Ablation studies are conducted thoroughly, providing insight into the components of the model. (Though see Weakness [W2] for potential areas of improvement.)
- [S3] The related work section on world models is thorough, with Table 1 providing a clear and valuable summary.

weaknesses:
- [W1] The explanation of AC-CPC is somewhat vague, which is problematic given the importance of CPC in the proposed method. Some specific concerns include:
  - (i) The authors do not clarify the necessity of the representation network, which is only briefly mentioned in Eq. (1).
  - (ii) Notation introduced in Eq. (1), specifically $e_t^k$ and $\hat{e}_t^k$, is not referenced further in the text.
  - (iii) Visualizations of the representation networks and AC-CPC predictors in Figure 3(a) could enhance clarity.
- [W2] The ablation studies section could benefit from improved consistency. Specifically, the order of the studies differs across Table 3, the main text, and Figure 6, which may lead to confusion. I recommend aligning the order across these references. Additionally, inconsistent naming conventions add to the lack of clarity. Suggested improvements:
  - In Figure 6(a), "Number of CPC steps" could be updated to "Number of contrastive steps" to match the main text.
  - The paragraph titled "Future Actions Conditioning" could be revised to "Action Conditioning," aligning with Figure 6(d).
- [W3] Many figures suffer from small text, making them challenging to read without zooming. Figure 6, in particular, has unnecessary whitespace around plots, which could be adjusted to improve readability.
- [W4] Including evaluations on other benchmarks, such as the DeepMind Control Suite, would strengthen the soundness of the findings. However, I acknowledge that this may be difficult to implement.

questions:
- [Q1] Do the authors have any insight into why the performance drops when predicting 15 steps into the future?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730895054371

thread_type:
official_review

<thread_object>
creator: reviewer_2k8g

<content>
summary:
The paper explores the use of contrastive predictive coding (modified objective function) for learning efficient world models. 

The paper is very well written.

Motivation: The work argues that the task of predicting the next frame may not require a complex model (as the successive video frames are somewhat similar or change in distribution is very sparse). The table 1 is very helpful on distinguishing how different methods effect each other. 

Proposed Method: The work builds upon transformer based world models and CPC loss to learn high level representations for improving the sample efficiency of the learning agent.

Ablations: The work conducts thorough ablations to understand the role of different components (number of steps in CPC loss, role of data augmentation, role of world model architecture i.e. RNNs v/s Transformer) and the effect of action conditioning.

strengths:
- The paper is very well written.
- The paper does a good job of ablating various different components used in proposed method.
- The paper does a good job of citing relevant literature.

weaknesses:
- There's no weakness as such. The paper validate all the claims made in the introduction via careful experimentation and proper ablations.

questions:
- "As shown in Figure 2, the cosine similarity between adjacent latent states of the world model is very high, making it relatively straightforward for the world model to predict the next state compared to more distant states."

In the introduction paper talks about cosine similarity for the proposed method, it will also be interesting to plot the cosine similarity for all the ablations the paper propose (number of contrastive steps, future action conditioning, effect of data augmentation) (as well as the dreamer baseline). 

- It will be helpful if authors can include another benchmark to make sure the proposed method shows gains on other environments too (like Crafter as mentioned in section 2.2 for related work).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731183624439

thread_type:
official_review

<thread_object>
creator: reviewer_1pdj

<content>
summary:
The paper proposes to combine Transformer-based world models with contrastive learning, and outperforms DreamerV3 and recent Transformer-based MBRL agents on the Atari 100k benchmark.

strengths:
- The paper provides a detailed overview of DreamerV3 and recent Transformer-based MBRL agents. Table 1 provides a good summary of implementation differences.
- There is adequate ablation study to justify the design of action-conditioned CPC loss and the CPC steps used. The finding that random crop and resize work better than random shift is quite intriguing.
- Sufficient implementation details are provided.
- The paper achieves strong results on the Atari 100k benchmark.

weaknesses:
- The paper did not discuss similar works that combine contrastive learning with Dreamer, notably [DreamingV2](https://arxiv.org/abs/2203.00494) and [DreamerPro](https://proceedings.mlr.press/v162/deng22a.html). While targeting different tasks, DreamingV2 has a quite similar contrastive loss design.
- From Table 2, it is quite intriguing that the proposed method TWISTER achieves significantly higher score on Gopher than previous methods. Further from Table 10, without the contrastive loss, the score on Gopher drops much more significantly than on other games. I am a bit concerned that the overall improvement brought by contrastive learning might be dominated by this single game.
- Based on the first two points, the paper could be strengthened by showing the general applicability of the method. For example, the paper can investigate more tasks (e.g., Crafter, DMC), or more backbones (e.g., S4/S5, see [S4WM](https://openreview.net/forum?id=GDYuzX0rwj&noteId=O5ZiA3xL4E) and [R2I](https://openreview.net/forum?id=1vDArHJ68h))

questions:
Some clarifying questions on implementation details
- Did you use EMA encoder or any stop gradient operation when obtaining future latents $z'_{t+k}$?
- Is the augmentation consistent across time (e.g., cropping the same region across time rather than choosing a random cropping region at each time step)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732044365785

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the Reviewers for their insightful comments and valuable feedback. We appreciate the recognition from Reviewers $\textbf{1PDj}$, $\textbf{2k8g}$ and $\textbf{TrKK}$ of *“good”* and *“excellent”* paper presentation. We also would like to thank Reviewer $\textbf{z6Gc}$ for the suggestions of improvements regarding the paper presentation. We thank Reviewers $\textbf{2k8g}$ and $\textbf{TrKK}$ for their positive and encouraging feedback on our paper: *“The paper validates all the claims made in the introduction via careful experimentation and proper ablations”*, *“We can enjoy to read this paper, because of well summarized related works, well discussed their motivation and diverse empirical ablation studies”*. We also would like to thank Reviewer $\textbf{1PDj}$ for the constructive comments on our paper. We address the main feedback from reviewers below and summarize the changes made to the submission.


1. Following the suggestions of Reviewer $\textbf{z6Gc}$, we updated Figure 3 (a) to include a more detailed visualization of the AC-CPC objective including target representations $e_{t}^{k}$ and predicted contrastive features $\hat{e}_{t}^{k}$ introduced in Equation 1. We also updated Section 3.1 (lines 313-319) to detail the utility of the representation network for computing similarities. We arranged the order of ablation figures and results in Table 3 to coincide with the order of subsections. The text size of figures has also been increased to improve readability.

2. We provide additional resources to illustrate our responses in the supplementary materials. Following the comments made by Reviewer $\textbf{2k8g}$, we plot cosine similarities of image observations for different ablations. We also provide a modified result Table including the [DreamerPro](https://proceedings.mlr.press/v162/deng22a.html) method for additional comparison. Finally, following the concerns highlighted by Reviewer $\textbf{1PDj}$ about the impact of outlier tasks, we attached an updated ablation figure for the number of contrastive steps showing IQM metric scores.


3. One common feedback pointed out by Reviewers $\textbf{1PDj}$, $\textbf{2k8g}$ and $\textbf{z6Gc}$ was that the findings of the paper would be strengthened by evaluating on other benchmarks such as the DeepMind Control Suite (DMC) and Crafter. Our paper reports results on the Atari 100k benchmark, which is commonly acknowledged as sufficiently diverse to verify the effectiveness of algorithms. This is in line with the common practice, the state-of-the-art methods we compare with are reporting on Atari 100k (SimPLe, TWM, IRIS and STORM). The suggested other benchmarks are considered and we are working on getting such results. However, due to limited resources, we cannot guarantee to provide new results.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors presents a significant advancement in model-based reinforcement learning by effectively combining Transformer-based world models with contrastive predictive coding, achieving state-of-the-art human-normalized performance on Atari 100k. The authors thoroughly addressed reviewer concerns by improving figure clarity, providing detailed technical explanations, and demonstrating the method's effectiveness through new DMC results. The comprehensive ablation studies, clear presentation, and strong empirical results across multiple benchmarks, combined with thorough responses during rebuttal, make this a valuable contribution worthy of acceptance. I recommend accepting this paper for ICLR 2025.