Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning

Summary

This paper presents a method for learning world models from in-the-wild videos. By utilizing a context encoder to capture contextual information, the proposed method explicitly models both the context and dynamics to facilitate knowledge transfer across scenes. Experiments are performed on various simulation benchmarks such as Meta-world, DMC Remastered, DeepMind Control Suite, and CARLA.

Strengths

The motivation for learning better world models by disentangling context information and dynamics is clear and seems reasonable to me; The proposed method for learning context information is simple and straightforward; The paper is well-organized and easily read.

Weaknesses

The proposed framework is based on the assumption that the context information lies equally in each frame, however, it is very likely that some moving objects might be occluded at some time. There are no specific designs for handling these situations; The experiments mainly validate the sample efficiency of the RL training process, there is no sufficient ablation study on the learned context information; Compared to the vanilla WM baseline, the performance of RL training on the Meta-world benchmark seems not very impressive; The predicted video frames seem not very promising; There is no discussion of limitations and failure cases.

Questions

The authors propose to randomly select a frame for predicting the context feature, have the authors tried to use a fixed frame (e.g., the first frame or the last frame of the video clip)? Would the context feature extracted from the different frames be consistent with each other? And how would choosing different frames affect the final performance of the RL training? It is assumed that the learned context feature captures the static properties of objects (e.g., colors, shapes), would modifying the context feature allow us to generate diverse video frames (e.g., changing the color of the object)? It would be interesting to show more temporal consistent video frames by editing the context feature.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

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

2 fair

Presentation

3 good

Contribution

1 poor

Limitations

There is no discussion of limitations and failure cases.

Summary

Learning a world-model that can generalize to different domains and tasks is difficult. The authors enhanced an existing framework for pre-training world models using in-the-wild videos, which can be fine-tuned on downstream tasks. In particular, the authors introduce a contextual encoder which helps in disentangling temporal dynamics from static contexts. Additionally, they include a cross-attention mechanism and dual reward predictors to improve the learning of task-relevant representations.

Strengths

The manuscript is very well written and structured. ● The idea of using a context to encode static information is novel, well-motivated and it shows good results in the DMC remastered task ● Learning world models on in the wild videos is hard (and so far does not help downstream tasks, as shown by Seo et al. [49]) - The extension of the Action-Free Pre-training from Videos approach to in-the-wild videos is convincing and the context modulation via a U-Net for reconstruction is innovative. ● The authors perform several relevant ablations to illustrate the role of the different proposed components.

Weaknesses

Overall, the context is only used together with the latent dynamic to decode the image. Is there a clear reason why the context is not taken into account for predicting the reward (in Fig. 3b the arrow of the context that goes to the reward seems to be misleading)? As the authors stated, the context might implicitly encode some important information about the task, e.g., the static position of an object, which can be helpful during the fine-tuning phase.

Questions

● What does the context encode? The motivation and the qualitative analysis are convincing, but it would also be useful to test this further. This could be done via a decoding analysis. In particular, it would be very interesting to see the difference between the context and the dynamic in the case of the DMC remastered task, where the contextWM provides a significant advantage. ● Why do you think is the gap in performance to prior methods particularly strong for DMCR? ● Looking at figure 6b (bottom), the performance of the pre-trained ContextWM on the SSv2 dataset for the CARLA driving task seems to outperform the one that is not-pretrained. However, the effects of the dataset domain of figure 7b (right) shows that there is almost no difference between the one pre-trained on SSv2 and the one without pre-training. What am I missing? ● Does the contextWM help to achieve better generalization? Could one, e.g., purposefully change the color, size or shape of the object for the meta-world and achieve better performance than the WM?

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

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

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Fine

Summary

This paper studies whether large-scale in-the-wild datasets can be used to pre-train world models for efficient downstream reinforcement learning. Specifically, they introduce Contextualized World Models (ContextWM), an architecture specifically designed to learn to separate context and dynamics modeling. Their experimentation reveals that the proposed methodology outperforms DreamerV2 on a variety of downstream tasks.

Strengths

- Well-written. The paper is well-written, and the figures aid in the understanding of the methodology. - Variety of experiments. The authors conduct a variety of experiments, comparing not just to DreamerV2, but also analyzing the effects of pre-training, architecture choices, dataset domain, etc. Furthermore, both quantitative and qualitative comparisons are included. - Strong experimental. Relative to DreamerV2, the proposed methodology achieves a strong performance -- the gap seems to be particularly large for the DMC Remastered tasks.

Weaknesses

- Missing comparison to prior work. APV [49] is the most similar prior work which this work builds off of (and seems to be the SOTA in this space), and yet the proposed method is not compared to APV. Many of the tasks used, hyper-parameters, and evaluation protocols adopted are from APV, allowing for a comparison, yet somehow, this comparison is omitted. Looking at the results plots in the APV paper, visually, the performance of ContextWM seems similar to that of APV (and in some cases clearly worse, e.g. in the dial turn task). While one may claim that a comparison to APV may be unfair because both the data and the model would be different, this would still reveal whether the ability to use a larger amount of in-the-wild data as well as the changes to the model architecture are actually beneficial. Furthermore, training APV on this large-scale data and seeing whether ContextWM outperforms it would be an experiment which would reveal whether the architectural changes proposed are significant. Finally, APV primarily compares to DreamerV2 as this was the SOTA when APV was published -- now, DreamerV3 seems to be the SOTA, so a comparison to DreamerV3 should be conducted. - Incomprehensive ablations. It is not clear which one environment the ablation study is conducted on. It would be a lot more convincing if the ablation study was done across tasks, and if the trends held true across tasks as this would alleviate concerns of cherry-picking a task.

Questions

- Is it possible to qualitatively demonstrate that ContextWM is able to separate context and dynamics modeling? - For some tasks, pre-training and choice of data to pre-train with makes a huge difference, e.g. DMC Re-mastered. For others, not so much, e.g. CARLA. Is there a sense of why this is the case?

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

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

2 fair

Presentation

4 excellent

Contribution

2 fair

Limitations

There is no discussion on the limitations of the work.

Summary

This paper proposes a Contexturelized World Model with In-the-wild Video Pretraining, which extends recently proposed action-free pre-training from videos (APV) to the case of contextualized video-prediction models. Specifically, they propose to sample a randomly chosen frame and use it as "contextualized information" for the prediction model. The "contextualized" information is combined into the prediction model via multi-scale cross-attention mechanisms. Besides, during the model-based RL phase, they propose to predict both pure task reward and the sum of the task reward plus a weighted intrinsic reward. Empirical validations on several well-known benchmarks (Meta-world, DMC, and CARLA) show the superior performance of the proposed method.

Strengths

1. The paper is generally well-written and easy to follow. 2. While the proposed method is simple, experimental results show a stable performance gain by the proposed methods. 3. Besides the main contribution, the paper also provides an interesting analysis of the choice of pre-training datasets.

Weaknesses

1. Technical novelty of the proposed method is not high. The main proposal is to add contextualized information to facilitate a better prediction. Besides, the paper does not discuss how the proposed approach (random sample selection) incorporates contextualized information well. Especially, I'm not sure we can call the random variable $c$ a contextualized vector since it is computed via only single frames rather than the full context of the trajectories. In this sense, I'm wondering what happens if we predict/reconstruct observations using the same multi-scale cross-attention architecture but with the image from the same time step for each frame. In the setup, the model does not use contextualized information but has similar architecture to the proposed method, and thus more appropriate for the baseline. 2. Lack of discussion with transformer-based world models. As shortly discussed, several studies incorporate transformer architecture in the prediction model. Since the transformer architecture predicts the future via autoregressive fashion, I think it is more natural to handle the contextualized information. However, the current manuscripts lack discussion on this point and comparison with transformer-based world models

Questions

1. Did you try methods other than random sample selection to compute contextualized information? 2. Did you try a comparison with transformer-based architecture? 3. The logic behind several sentences is hard to capture. Could you please explain more about these sentences? - In line 182, "Nevertheless, ..... Therefore, we propose a dual reward predictor" => Why does the dual predictor resolve the issue? Why not just balance lambda depending on the task? - In line 234, "indicating that the performance gain of IPV with vanilla WM is primarily due to the intrinsic exploration bonus. " => The logic is unclear. [Minor] - In eq. 6, How do you compute KL when t=1? - In Fig. 5-b, Pre: O should be Pre: ✓ or with Pre or something like that. - If I correctly understood, Fig. 3 is a bit misleading as a context variable $c$ is directly input into each frame, rather then recurrent prediction.

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

A limitation section should be added.

Summary

This submission presents the contextualized world model (ContextWM), a framework for leveraging in-the-wild videos for pre-training of a world model to be used in model-based reinforcement learning. Following the work from Seo et al. (2022), the authors pre-train an action-free version of the recurrent state-space model (RSSM) with two important modifications: 1) a context encoder processes a randomly sampled frame of the input video to provide context features, which the decoder can directly access to better reconstruct static visual details, enabling the dynamics model, on which the decoder is also conditioned, to focus on temporally varying information. 2) the authors also opt for using a dual-reward predictor during fine-tuning, which predicts the pure task reward in addition to the combined task and video-based intrinsic novelty reward proposed by Seo et al. (2022). This facilitates task-relevant representation learning. The ContextWM is evaluated on Meta-world, the remastered DeepMind Control Suite, and a task with varying weather conditions in the CARLA driving simulator. In most benchmarks, ContextWM shows significant improvements in terms of sample efficiency or final performance. An ablation study and other analytical experiments further show the effectiveness of ContextWM and its design decisions.

Strengths

The proposed contextualized world model is novel and the design decisions are well motivated. For the most part the description of the method and experimental setup is very clear. The performance improvements are significant and in some cases very impressive. The qualitative analysis provides some interesting insights, for instance the clear separation of video representations for two videos with contrastive labels in Figure 8b.

Weaknesses

1. As the main contribution the cross-attention to the context should be explained in a bit more detail. Since the authors mention U-Nets, I wonder whether the decoder attends to context features at each of the corresponding three stages shown in Table 1 in the supplementary material or only at the decoder input level?

Questions

2. Why did you choose BatchNorm instead of LayerNorm? Did you experiment with both? 1. maybe use a check mark in the legend of Figure 6b instead of the circle to indicate "with pretraining". 1. It'd be very interesting to see more examples with other contrastive labels for the Video representations experiment in Section 5.5. ## Acknowledgement of rebuttal I have read the rebuttal, other reviews. My relatively minor concerns have been addressed, given that the authors have provided some requested clarifications, a discussion of limitations, and additional insightful experiments. I strongly believe this paper should be accepted.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and 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

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The discussion of limitations is very *limited*. The discussion section mostly mentions future work on scaling and exploring other pre-training objectives.

POST-REBUTTAL

I acknowledge the authors' dedication to incorporating the updates. Following a meticulous examination of the rebuttal materials, my reservations regarding the modest advancement in performance and the absence of profound technical contributions persist. Given these considerations, I am leaning toward upholding my initial evaluation.