Visual Representation Learning for World Models by Predicting Fine-Grained Motion

Authors: Zhao-Han Peng, Shaohui Li, Zhi Li, Yu LIU, You He
Primary area: reinforcement learning
Keywords: world models, model-based reinforcement learning, visual representation learning

Abstract:
Originating from model-based reinforcement learning (MBRL) methods, algorithms based on world models have been widely applied to boost sample efficiency in visual environments. However, existing world models often struggle with irrelevant background information and omit moving tiny objects that can be essential to tasks. To solve this problem, we introduce the Motion-Aware World Model (MAWM), which incorporates a fine-grained motion predictor and entails action-conditional video prediction with a motion-aware mechanism. The mechanism yields compact and robust representations of environments, filters out extraneous backgrounds, and keeps track of the pixel-level motion of objects. Moreover, we demonstrate that a world model with action-conditional video prediction can be interpreted as a variational autoencoder (VAE) for the whole video. Experiments on the Atari 100k benchmark show that the proposed MAWM outperforms current prevailing MBRL methods. We further show its state-of-the-art performance across challenging tasks from the DeepMind Control Suite.

Decision: reject
Reviewer scores: 5, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1730610691019

thread_type:
official_review

<thread_object>
creator: reviewer_klro

<content>
summary:
This paper tackles the problem that model-based visual RL methods tend to ignore small objects that are essential for tasks while learning irrelevant background information. The proposed approach MAWM introduces a pixel-level loss function based on motion information together with adaptive motion-aware scheduler based on timesteps and video prediction loss instead of image reconstruction as representation learning objective. Authors perform experiments on Atari and DM control benchmarks. Authors compared MAWM to several model-based and model-free RL baselines, and demonstrated improvements averaging across multiple tasks. Ablation studies show contributions of each component in this framework.

strengths:
1. The proposed approach shows improvements over baselines in benchmark evaluations
2. Authors perform various ablations to demonstrate effectiveness of each component
3. Paper is well structured and experiments are well organized

weaknesses:
1. The main difference from existing model-based RL methods is an auxiliary objective based on pixel-wise motion information and replacing image reconstruction loss with video prediction objective, both of are not significant changes.
2. Since the motivation of MAWM is to avoid distractions from irrelevant background information, in addition to standard control benchmarks, environments such as distracted DM control [1] and DM control generalization [2] can better demonstrate benefits of the proposed approach than baselines.
3. Current experiments do not quantitatively establish correlation whether proposed approach has more benefits in environments or tasks where objects tend to be smaller, which is a main motivation in method design.

[1] Zhang, Amy, et al. "Learning invariant representations for reinforcement learning without reconstruction." arXiv preprint arXiv:2006.10742 (2020).

[2] Hansen, Nicklas, and Xiaolong Wang. "Generalization in reinforcement learning by soft data augmentation." 2021 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2021.

questions:
1. In ablation study (L486-L491), why video prediction loss instead of image reconstruction loss in MAWM framework is important for DM control experiments but not Atari experiments?
2. What is the intuition and rationale to use an adaptive Gaussian mixture model to extract ground truth label of whether each pixel belongs to foreground or background?
3. How does MAWM compare to MWM [3] that uses MAE objective on convolution features to better learn representations, which is also shown specifically more suitable for small objects?

[3] Seo, Younggyo, et al. "Masked world models for visual control." Conference on Robot Learning. PMLR, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730620256174

thread_type:
official_review

<thread_object>
creator: reviewer_pz6p

<content>
summary:
This paper introduces a new method for learning the model in model-based RL with special considerations in motion prediction modeling. The authors introduces motion awareness by adding foreground motion prediction and adaptive motion-blur losses over traditional video generative pipelines. The resulting model achieves comparable results against existing MBRL methods.

strengths:
The proposed motion awareness in model learning is intuitive and the proposed method achieves comparative results with state-of-the-art methods.

weaknesses:
One concern about this paper is the significance of the proposed method. First, except in curves and Fig.3 from the ablation studies, the authors might want to provide more visualizations on the effect of the motion awareness introduced, especially when considering the current experiment setting is only in atari and dm-control-suite where data are all synthetic and should potentially be simpler compared to real-world videos. Second, I wonder if the same pipeline could be applied to real-world videos as currently there is increasing trend in leveraging video generative models as "world models" to facilitate various tasks. It would be better if the authors could find a proper way to compare or address such line of methods. Lastly, the paper organization could be improved for better clarity.

questions:
See the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692901051

thread_type:
official_review

<thread_object>
creator: reviewer_qy1e

<content>
summary:
The paper presents a novel approach in terms of model-based reinforcement learning (MBRL), where an image-based world model is utilized to improve sample efficiency by mimicking a scalable digital copy of the environment, such as DreamerV3. The authors claims that traditional world models tend to ignore moving tiny objects and their connections with tasks, thereby they propose a Motion-Aware World Model (MBRL) to account for them. Specifically, MAWM 1. focuses on small objects via pixel-level attention mechanisms and 2. deals with rapid changes of them via an adaptive control scheduler. The results on Atari 100k and DeepMind Control Suite (DMC) depict the superiority of the proposed method against DreamV3.

strengths:
The paper presents several strengths:
1. The experimental validation of the proposed model on two standard datasets, Atari 100k and DMC, demonstrates its effectiveness and generalizability.
2. The writing of the paper is clear, making it easy for readers to understand the proposed method and its experimental validation.
3. The details of the model in the appendix are valuable contributions to the community, as it enables other researchers to reproduce the results and build upon the proposed method.

weaknesses:
1. One concern is regarding the novelty of the proposed two techniques. The motion-aware auxiliary loss is not a novel topic [1]. Additionally, the technical contributions of the proposed pixel-level attention and scheduler are not clear enough to me.
[1] 3D Motion Decomposition for RGBD Future Dynamic Scene Synthesis, CVPR 2019. 

2. Is it possible to compare MAWM with pretrained video generation models like stable video diffusion, to figure out whether they can capture the patterns of moving targets or not?

3. I also have some concerns about the experimental results. As an example, Table 1 is confusing. In the last row, the Median score of a counterpart SimPLe is 1.34, which seems to obtain the best results without being marked bold. Besides, the different components proposed in this paper don’t demonstrate convincing improvement in Figure 4. A more comprehensive and significant ablation study could help address this.

questions:
Please see the weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699437797

thread_type:
official_review

<thread_object>
creator: reviewer_vtbc

<content>
summary:
This paper presents the Motion-Aware World Model. It integrates a fine-grained motion predictor with an adaptive motion-aware scheduler and involves action-conditional video prediction to filter out backgrounds and track object motion at the pixel level.

strengths:
1. The attention that the authors give to relevant foreground information and moving tiny objects is meaningful for constructing world models.

2. The performance gain of the proposed method is considerable on the Atari 100K benchmark and DeepMind Control Suite.

weaknesses:
1. The author elaborates on some details in the method section; however, many technical aspects, like the Convolutional Block Attention Module, were not employed in previous approaches. Nevertheless, the relevant experiments have not undergone ablation, which could readily prompt people to doubt the fairness of the experiments.

2. The results of the ablation experiments are rather confusing. It appears that MAWM is not consistently the best choice. Moreover, for the majority of the experiments, it doesn't seem that they have reached convergence.

3. The motivation of this paper is to solve the problem of struggling with irrelevant background information and omitting moving tiny objects. Is there any corresponding visualization to verify this?

4. Many experimental details are not clearly stated. For example, what are the values of $\alpha$ and $r_{dizzy}$?  

5. Discussions on limitations are necessary.

questions:
Please see weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732969784693

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their detailed reviews. Based on their comments, we have revised our paper as listed below:
- We have run experiments on an additional benchmark to demonstrate the generalization ability and benefits of our proposed approach (in Appendix L), following suggestions of Reviewer [KLRo](https://openreview.net/forum?id=8BJl6LQgW5&noteId=LKr56goOPB).
- We have added several ablation studies on CBAM, harmonizers, and choice of the autoencoder to demonstrate the effectiveness of our motion-aware mechanism and the importance of video prediction (in Appendix F), according to Reviewers [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH), [qY1E](https://openreview.net/forum?id=8BJl6LQgW5&noteId=sqISwzHENv), and [KLRo](https://openreview.net/forum?id=8BJl6LQgW5&noteId=LKr56goOPB).
- We have illustrated that our proposed methods can capture the motion of tiny objects and make almost perfect predictions, in comparison to pretrained video generation models, revealing that our motion-aware mechanism is the key to efficient performance for most tasks (in Appendix M), for Reviewers [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH), [qY1E](https://openreview.net/forum?id=8BJl6LQgW5&noteId=sqISwzHENv), [pZ6p](https://openreview.net/forum?id=8BJl6LQgW5&noteId=4dLWWd9qsd), and [KLRo](https://openreview.net/forum?id=8BJl6LQgW5&noteId=LKr56goOPB) concerns.
- We have discussed differences between our motion-aware mechanism and existing methods in the response to Reviewers [qY1E](https://openreview.net/forum?id=8BJl6LQgW5&noteId=sqISwzHENv), and extended our related work (in Appendix N).
- We have conducted an experiment to ensure that MAWM improves beyond standard settings of 100k interactions on Atari Games, implying that MAWM may have better efficiency and effectiveness than existing methods, in answer to Reviewers [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH) and [pZ6p](https://openreview.net/forum?id=8BJl6LQgW5&noteId=4dLWWd9qsd).
- We have included a discussion of the limitations of our work (in Conclusion), according to Reviewer [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH).
- We have extended experiments to show MAWM consistently satisfactory performance on all the 20 tasks from DeepMind Control Suite, setting a state-of-the-art on the whole benchmark, as showcased in Table 12, which can be a reply to [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH) and [qY1E](https://openreview.net/forum?id=8BJl6LQgW5&noteId=sqISwzHENv).
- We have reflected feedbacks on revisions of our manuscript from Reviewers [vTbC](https://openreview.net/forum?id=8BJl6LQgW5&noteId=snUvzjSHNH), [qY1E](https://openreview.net/forum?id=8BJl6LQgW5&noteId=sqISwzHENv), [pZ6p](https://openreview.net/forum?id=8BJl6LQgW5&noteId=4dLWWd9qsd), and [KLRo](https://openreview.net/forum?id=8BJl6LQgW5&noteId=LKr56goOPB).
- We showcase visualization [here](https://anonymous.4open.science/r/mawm-C555)  to help understand our intuition and rationale for our motion-aware mechanism. We will release our code to promote future work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission addresses the problem of model-based reinforcement learning in visual environments. It aims to introduce a "world model" that is better at capturing the visual dynamics of tiny objects that are relevant to the target tasks. This is achieved by adding foreground motion prediction and adaptive motion-blur losses. Evaluations are performed on Atari and Deepmind Control Suite. The submission received four borderline reject (5) ratings initially, and the authors provided their rebuttal. Unfortunately, none of the reviewers engage in the post-rebuttal discussion. After reading through the rebuttals and the submission, the AC believes that the submission should be revised to better explain its contributions with respect to prior work, along with the generalizability and limitations of the approach when applied to visually more complex scenarios. Overall, the AC finds no ground to overturn the consensus among all four reviewers.