Learning Transformer-based World Models with Contrastive Predictive Coding
model-based reinforcement learning;transformer network;contrastive predictive coding
Yes
No ethics review needed.
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)?
- 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.
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.
- 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))
Yes
No ethics review needed.
- "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).
- 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.
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.
- There's no weakness as such. The paper validate all the claims made in the introduction via careful experimentation and proper ablations.
Yes
No ethics review needed.
- [Q1] Do the authors have any insight into why the performance drops when predicting 15 steps into the future?
- [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.
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.
- [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.
Yes
No ethics review needed.
- 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?
- 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.
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.
- We didn’t find a critical weakness from this paper.
We introduce TWISTER, a Transformer model-based reinforcement learning algorithm using action-conditioned Contrastive Predictive Coding to learn high-level feature representations and improve the agent performance.
@inproceedings{
anonymous2024learning,
title={Learning Transformer-based World Models with Contrastive Predictive Coding},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=YK9G4Htdew},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
model-based reinforcement learning transformer network contrastive predictive coding
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/409feaa48796a19031dad85b805219333aaece58.pdf
reinforcement learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Learning Transformer-based World Models with Contrastive Predictive Coding
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission