Learning Transformer-based World Models with Contrastive Predictive Coding

We sincerely thank you for your positive and encouraging feedback on our submission. We are pleased to hear that you found our paper engaging and enjoyable to read.

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?

The shaded lines represent cosine similarities obtained for individual games. The main curve is obtained by averaging similarities over all 26 games. We find that the game with the fastest declining state similarity is Up N Down, where the game state is highly dynamic. On the other hand, the game with the slowest declining state similarity is Amidar, where distant observations are highly correlated to current observations.

Did you test the agent performance on the Atari100k benchmark when doesn’t apply data augmentation? Only giving temporally near samples as the positive sample (without augmentation)

The paper experimented with data augmentation with the default setting of K=10 future contrastive steps. We demonstrates the utility of data augmentation for complexifying the AC-CPC objective. The use of random crops requires the world model to identify several key elements in the observations in order to accurately predict positives samples. We suppose that lowering K and removing data augmentation could significantly simplify the objective and degrade performance.

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.

Thank you for providing this work as an additional possible reference. Could you share with us the full reference of the paper? The review formatting seems to have removed the reference link.

For AC-CPC Predictor, how did you input the actions? Did you concatenate them?

Yes, we concatenate the sequence of actions along the feature dimension with the model states $s_{t}$. We modified section 3.1 in the paper to specify that actions are first concatenated to condition the predictor network.

Can you test the performance when not using reconstruction loss for the world model training?

Our paper demonstrated the impact of AC-CPC on learning Transformer-based world model representations. We are studying if AC-CPC can also learn representations for the encoder network without using reconstruction.

why they (IRIS and $\Delta$-IRIS) showed worse performance in Median Human Normalized Score measurement?

It is possible to achieve a higher Mean score with a lower Median if the best-performing games achieve high scores. This is why the Median metric is also important to evaluate the general performance of the methods. In the case of IRIS and $\Delta$-IRIS, both methods archive very good performance on Breakout, which improves the Mean score compared to TWM and STORM. The two metrics are important and complementary to effectively evaluate the performance of methods.

Thank you for your rebuttal in the limited time. Our concerns are well addressed in their rebuttal, I will keep my score.

Sorry for missing the reference, that is

Deng, Fei, Ingook Jang, and Sungjin Ahn. "Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations." International conference on machine learning. PMLR, 2022.

Thank you for reviewing our paper and providing many valuable suggestions. Please find below our response to the concerns and questions that you raised in the review.

[W1] (i) The authors do not clarify the necessity of the representation network, which is only briefly mentioned in Eq. (1).

We updated section 3.1 (lines 313-319) to further detail the utility of the representation network for computing similarities. Contrary to the original CPC paper, which experiments with continuous feature states, we use discrete latent states for the world model. The encoder network outputs discrete one hot encoded representations $z_{t}$ that are sampled from a categorical distribution with 32 groups, each with 32 classes. This requires learning a representation network to project the discrete encoded states $z_{t}$ to contrastive features $e_{t}^{k}$ and compute similarities.

[W1] (ii) Notation introduced in Eq. (1), specifically $e_{t}^{k}$ and $\hat{e}_{t}^{k}$, is not referenced further in the text.

[W1] (iii) Visualizations of the representation networks and AC-CPC predictors in Figure 3(a) could enhance clarity.

We also updated the section 3.1 to include notations introduced in Equation 1. We originally found it adequate to introduce the contrastive features as notations in Equation 1 along with other variables. The features were also referenced among other predictions on line 246. We agree that including AC-CPC predictions and targets in Figure 3 (a) improves clarity by referencing the notations. We updated Figure 3 (a) to include a more detailed visualization of the AC-CPC objective including predicted contrastive features $e_{t}^{k}$ introduced in equation 1. We hope this improves clarity.

[W2] The ablation studies section could benefit from improved consistency

Thank you for highlighting this. We updated the paper accordingly to ensure a consistent naming convention and ordering of the ablations. We increased the text size in figures to improve readability, especially Figures 7 and 8. We also increased the size of Figure 6 to limit unnecessary whitespace around plots.

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.

As stated in our global response, 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.

Do the authors have any insight into why the performance drops when predicting 15 steps into the future?

We find that increasing the number of AC-CPC steps has a positive impact on most games. However, a degradation of the results is noticed when predicting 15 steps into the future. One possible reason for these findings is that some of the games include uncontrollable randomness that makes distant target states very hard to predict compared to other negative samples. The world model will learn to identify the set of possible futures and attribute high similarity to probable future samples. However, if uncontrollable game randomness is such that a large amount of negative samples is part of the set of probable futures, the world model will not be able to accurately predict the correct future sample without additional conditions. This is why future works could explore the use of “learned latent actions” to condition world models on environment changes that the agent does not control.

We sincerely thank you for your positive and encouraging feedback on our submission.

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

We computed the average cosine similarities for the different ablations, including the dreamer baseline. You can find the plot of cosine similarities for ablations in the rebuttal supplementary materials. We find that different ablations do not have a significant impact on cosine similarities. This is as intended because the cosine similarities are computed from the encoder network outputs $z_{t}$, which are trained to be compressed encoded features for image observations.

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).

As mentioned in our global response, we are working on the application of our method to Crafter and DMC.

Did you use EMA encoder or any stop gradient operation when obtaining future latents ?

Similarly to the original CPC paper and similar approaches such as vq-wav2vec, we did not use an EMA encoder or the stop gradient operator. We let AC-CPC loss gradients back-propagate to encoder parameters like the other objectives optimized by the world model.

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 random crop and resize augmentation is applied independently for each image in the input batch. This prevents the model from easily identifying images from the same sequence, which could simplify the AC-CPC objective.

Thank you for your constructive feedback on our paper. Please find below our response to the concerns and questions that you raised in the review.

The paper did not discuss similar works that combine contrastive learning with Dreamer, notably DreamingV2 and DreamerPro. While targeting different tasks, DreamingV2 has a quite similar contrastive loss design.

Our paper discusses three main domains of related works in Section 2: We start by giving a general overview of model-based reinforcement learning methods. We then discuss the proposed Transformer-based approaches to world model learning. Finally, we introduce the CPC method and our motivations for applying the CPC objective to model-based Reinforcement learning.

We are familiar with the DreamerPro and Dreaming algorithms, which are related to contrastive-based world model learning. The two methods proposed reconstruction-free variants of the Dreamer algorithm to improve robustness for continuous control tasks. DreamerPro proposed to learn representations with SwAV while Dreaming proposed a multi-step contrastive objective. Knowing that DreamerPro was applied to Atari games in the DreamerPro paper appendix, we found it interesting to evaluate the method on the Atari 100k benchmark using the official code provided by the authors. Using 5 different training seeds, we obtained normalized Mean and Median scores of 79% and 28%, which is lower than DreamerV3. We provide a modified Table with DreamerPro results in the rebuttal supplementary materials as a reference.

However, since the focus of these works was on learning representations without reconstruction loss for continuous control tasks and using the standard recurrent-based Dreamer world model, we did not find it necessary to include them in the related works section. Although it employs a contrastive-based objective, the Dreaming method differs from CPC in key aspects. Some of the key differences are that the prediction layer is applied recurrently into the future with an overshooting distance K set to 3. Contrary to the Dreaming objective, AC-CPC directly makes the prediction by providing the concatenated sequence of future actions as inputs to the predictor network. This allows TWISTER to make predictions over longer distances without encountering possible vanishing/exploding gradients during back-propagation through time.

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

We find that AC-CPC has a positive impact on Gopher, helping the agent to accurately model the gopher’s position and tunnels for planning. As stated in section 4.2, we also find that AC-CPC improves the performance of games with small moving objects such as Breakout, Pong and Asterix. This is also the case for Up N Down, Kangaroo, and Road Runner, which significantly improve Mean and Median scores.

The Median metric analyzes the performance of average-performing games in each experiment. This metric is interesting to evaluate the general performance of algorithms by ignoring outlier tasks like Gopher. For instance, Delta-IRIS faces a similar situation with Breakout. We find that TWISTER archives a higher Median score compared to previous Transformer-based approaches, which indicates its positive impact not only on Gopher but also on average-performing games like Breakout, Pong Up N Down, Kangaroo or Road Runner.

Another known metric for comparing the performance of average-performing games is IQM. IQM shows the aggregated performance on the middle 50% of combined runs, ignoring outlier games. We attached an updated ablation figure for the number of contrastive steps in the supplementary materials showing the impact of AC-CPC on IQM metric scores. We can see that AC-CPC representation learning has a significant impact on IQM scores, validating its effectiveness on diverse games.

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 and R2I)

As stated in our global rebuttal response, 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 and backbones (S4) are considered and we are working on getting such results. However, due to limited resources, we cannot guarantee to provide new results.

We would like to thank all the reviewers for their valuable feedback and responses. We also appreciate the decision of Reviewer $**1PDj**$ and Reviewer $**z6Gc**$ to increase the rating of our paper: “Thank you for addressing all my concerns”, “Thank you for addressing all of my concerns in your revision”.

We performed experiments on the DeepMind Control Suite to evaluate TWISTER performance under continuous action spaces. We follow DreamerV3 that evaluates on 20 tasks using only high-dimensional image observations as inputs and a budget of 1M environment steps for training. The following Table compares the official DreamerV3 results with TWISTER (3 seeds) using a budget of 1M environment steps. Bold numbers indicate the best performing method for each task.

We find that TWISTER achieves state-of-the-art performance with stable learning across all tasks. The application of Transformer-based world models to continuous control tasks has so far been very limited. We hope our work will inspire researchers to further study the potential benefits of Transformer-based world models and representation learning for continuous control tasks.