Pre-Training Robo-Centric World Models For Efficient Visual Control

Authors: Long Qian, Ziru Wang, Sizhe Wang, Lipeng Wan, Zeyang Liu, Xingyu Chen, Xuguang Lan
Primary area: reinforcement learning
Keywords: Pretraining World Models, Model-based Reinforcement Learning, Visual Robot Control

Abstract:
Humans can accurately anticipate their movements to behave as expected in various manipulation tasks. We are inspired to propose that integrating prior knowledge of robot dynamics into world models can effectively improve the sample efficiency of model-based reinforcement learning (MBRL) in visual robot control tasks. In this paper, we introduce the Robo-Centric World Model (RCWM), which explicitly decouples the robot dynamics from the environment and enables pre-training to learn generalized and robust robot dynamics as prior knowledge to accelerate learning new tasks. Specifically, we construct respective dynamics models for the robot and the environment and learn their interactions through cross-attention mechanism. With the mask-guided reconfiguration mechanism, we only need a few prior robot segmentation masks to guide the RCWM to disentangle the robot and environment features and learn their respective dynamics. Our approach enables independent inference of robot dynamics from the environment, allowing accurate prediction of robot movement across various unseen tasks without being distracted by environmental variations. Our results in Meta-world demonstrate that RCWM is able to efficiently learn robot dynamics, improving sample efficiency for downstream tasks and enhancing policy robustness against environmental disturbances compared to the vanilla world model in DreamerV3. Code and visualizations are available on the project website: https://robo-centric-wm.github.io.

Decision: reject
Reviewer scores: 6, 3, 5, 8

Reviews and discussion:
<thread>
creation date:
1730209588533

thread_type:
official_review

<thread_object>
creator: reviewer_tbdh

<content>
summary:
This paper introduces the Robo-Centric World Model (RCWM), which decouples robot dynamics from environmental dynamics, and employs an interaction model to evaluate the impact of the robot's actions on the environment. The authors use SAM2 to create segment masks for the robot, which are then used to warmup the mask reconstruction process. The subsequent pre-training is conducted without mask usage. Experiments are conducted in the Meta-world domain, demonstrating through quantitative and visual results that RCWM can learn distinct dynamics and show robustness against disturbances.

strengths:
- The decomposition of robot and environmental dynamics is a straightforward method that enables transfer of shared knowledge.
- The distillation of SAM into world models that enables the agent to distinguish robot and background is an interesting practice.
- The experimental results show improvement over the base model, with the rich visualization results offering valuable insights.

weaknesses:
- Discussion and comparison to related works are missing. The authors state that action-free pre-training methods "are of limited help when confronted with robot manipulation tasks that require accurate predictions". However, the paper only carries out experiments in the Meta-world domain, where other action-free pre-training methods also demonstrate good performance.
- Limitations are not addressed adequately. RCWM limits the interaction between the robot and the environment from two-way to one-way, which is acceptable for Meta-world since pre-training and fine-tuning are carried out on the same platform. However, RCWM may encounter difficulties when applied to certain types of real-world visual robotics tasks, such as dealing with a different friction coefficient or a different tilting angle for the table.

questions:
- In Figure 7(a), do RCWM and vanilla WM use the same set of trajectories for evaluation, or are the trajectories sampled independently for each model? If the models share the same trajectories, how are they sampled?
- In Figure 7(b), a sharp turn can be observed between 10 and 20 rollout steps for the stick-push task. Why would this happen?
- A discussion section comparing RCWM to previous action-free pre-training methods should be added. The authors should compare RCWM to previous action-free pre-training methods specifically on robot manipulation tasks that requires accurate predictions.
- A limitation section about the applicability and generality scope of RCWM should be added. The authors should explicitly discuss the potential limitations in applying RCWM to real-world scenarios with varying physical properties.

Minor comments:
- The notation for the baseline method is a bit inconsistent. I suggest changing all "vanilla WM" notations to "DreamerV3" or vice versa for better readability.
- Typo: "dynamic model" should be "dynamics model"
</content>
</thread_object>

</thread>
<thread>
creation date:
1730527070277

thread_type:
official_review

<thread_object>
creator: reviewer_neog

<content>
summary:
The paper introduces the Robo-Centric World Model (RCWM), designed to improve sample efficiency and robustness in model-based reinforcement learning (MBRL) for visual robot control. RCWM achieves this by decoupling robot dynamics from environment dynamics, allowing each component to be modeled independently while an interaction model, based on cross-attention, captures the effect of robot actions on the environment. The model’s training pipeline includes a mask-guided warmup using robot segmentation masks, followed by mask-free pre-training and fine-tuning, ensuring RCWM can robustly handle novel tasks and disturbances.

strengths:
* The paper is well-written and well-organized.
* The idea of decoupling robot and environment dynamics is very novel and holds potential value for improving the efficiency and robustness of learning algorithms.
* The proposed method demonstrates resilience against visual disturbances and changing backgrounds.

weaknesses:
* More baselines should be included, such as a leading model-based RL algorithm TD-MPC2 (https://arxiv.org/abs/2310.16828). Meanwhile, it is important to report the sample-efficiency comparison with leading model-free algorithms like DrM (https://arxiv.org/pdf/2310.19668).
* More experiments should be conducted in more challenging tasks such as dexterous manipulations in Adroit (https://arxiv.org/pdf/1709.10087) to validate the effectiveness of the proposed method.
* It is also essential to validate the proposed method in real-world experiment (such as Box Close task).

questions:
* How might RCWM handle scenarios where prior robot segmentation masks are unavailable or difficult to obtain? Could an alternative approach to disentangling robot and environment dynamics be feasible in these cases?
* Given RCWM's reliance on cross-attention for modeling interactions, how well would it handle environments with more unpredictable or dynamic elements, such as deformable objects or non-static obstacles?
* What will happen if we only use the replay buffer from 1 task trained by DreamerV3 as the dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730623031632

thread_type:
official_review

<thread_object>
creator: reviewer_qufp

<content>
summary:
This paper presents robot-centric world models that distinctively model robot and environmental dynamics to enhance visual model-based reinforcement learning. Results from eight MetaWorld tasks demonstrate improved learning efficiency compared to the vanilla baseline, DreamerV3.

strengths:
- The paper is well-organized, featuring visualizations that aid in understanding the proposed robot-centric world models.
- The approach appears to effectively learn robot-centric dynamics, though I have some questions about it.
- Experimental results show performance gains over the state-of-the-art DreamerV3 across 8 MetaWorld tasks.

weaknesses:
- A primary concern is that the performance improvement may not solely be attributed to the proposed idea. Several additional implementations, such as the transformer in the interaction model, mask-guided encoder, and predictor head, confuse the attribution of improvements.
- The generalization ability of the proposed model seems limited. Significant performance drops occur under environmental disturbances, particularly in the Sweep-Into task (Figure 8(b)). Additionally, the focus on disturbance testing should extend to other variables like robot arm appearance, base position, and camera view.
- The selected tasks appear relatively simple. It would be beneficial to evaluate the modeling effectiveness on more complex embodiments, such as dexterous hands or humanoids, as evaluated in [1].
- The absence of ablation studies. An analysis of each design's impact—such as the warm-up stage, warm-up data choice, or removal of components—would assist the understanding.
- Construction error may not be a good metric to measure the ability to model robot dynamics, because the error may mainly caused by the environmental dynamics error. As an extreme case, the vanilla world model may be good at robot dynamics but very poor at environmental modeling. The authors should address this concern.
- Testing the extension capability of the proposed method on other state-of-the-art algorithms, like TD-MPC2 [1], would significantly strengthen this paper.

questions:
- What is the performance in a multi-task setup, following [1]?
- What do the attention maps signify in all figures? The highlighted regions seem meaningless.
- The gripper reconstruction in Figure 5 appears inadequate. Could the authors clarify this?

[1] Hansen, N., Su, H., & Wang, X. TD-MPC2: Scalable, Robust World Models for Continuous Control. In *The Twelfth International Conference on Learning Representations*.

[2] Sferrazza, C., Huang, D. M., Lin, X., Lee, Y., & Abbeel, P. (2024). Humanoidbench: Simulated humanoid benchmark for whole-body locomotion and manipulation. *arXiv preprint arXiv:2403.10506*.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730853475168

thread_type:
official_review

<thread_object>
creator: reviewer_pybh

<content>
summary:
This paper considers the problem of world model pre-training in robotics applications. Different from previous approaches that use a single model to train the scene evolution, the key idea of this paper is to decouple world dynamics into robot dynamics and environment dynamics. The proposed model uses a dedicated robot branch to predict the future images of the robot, an environment branch to predict the future images of the environment, and an interaction module to inject the information from the robot branch to the environment branch.

strengths:
This paper is well-written and the logic is clear. The idea of decupling robot dynamics and environment dynamics is interesting. The claim is supported by promising experiment results.

weaknesses:
It is unclear if the gains are from the additional capacity introduced by the additional branch. (2x model capacity?) To verify this, the authors could expand the dreamer model capacity and compare the results.

It is strange that the robot branch can learn effects of the environment on the robot, even there is no information exchange from environment branch to the robot branch.
This makes me wonder if the true technical idea is simply using two models to predict the evolution of the robot and the rest of the scenes separately. This can be validated by disabling the interaction part and compare with the proposed model in terms of dynamics prediction and policy learning.

questions:
See my questions above
</content>
</thread_object>

</thread>
<thread>
creation date:
1732192851375

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
## **We hereby summarize some common misunderstandings and concerns raised by the reviewers regarding our work.**

First, we would like to emphasize that our work is oriented towards robot manipulation tasks with purely visual inputs. Our approach is improved based on DreamerV3, since the Dreamer series of methods has been widely used to solve visual robot manipulation tasks.

* **The reviewers mentioned the comparison with TD-MPC2 and the expansion based on it, and we would like to clarify the difference and comparability of our method with TD-MPC2**. The main contribution of our approach is to improve the construction of the world model in Dreamerv3 to make it more suitable for learning robot manipulation tasks. Although DreamerV3 and TD-MPC2 are both advanced model-based algorithms, there are significant differences between them. First, there are significant differences in the construction of the world models. DreamerV3 uses reconstruction loss to learn the hidden space representation and constructs RSSM to learn the dynamics. In contrast, TD-MPC2 uses contrast learning to learn the hidden space and uses MLP to learn dynamics with frame stacking as input. Second, the two use completely different approaches for policy learning. Moreover, TD-MPC2 focuses more on control tasks with state as input rather than visual control tasks, while DreamerV3 focuses more on visual control tasks. Therefore, we believe that the comparison with TD-MPC2 does not reflect the improvement of our approach because our improvement focuses on the world model, while the two are quite different both in terms of world modeling and policy learning. Moreover, there are obvious problems in extending our approach to TD-MPC2 due to the different construction forms of the world model. **We would like to clarify that our approach and TD-MPC2 are not directly comparable due to fundamental differences in world modeling and policy learning.**. 
* The reviewers refer to evaluations on more complex control tasks, such as dexterous hand or humanoid robot control tasks. We would like to emphasize that our approach makes a valuable attempt to effectively decouple robot dynamics as well as interaction learning. However, learning complex robot dynamics such as dexterous hands from visual input is still an open problem. The fact that dexterous hands have many joints and are often occluded when performing manipulation tasks which leads to the difficulty of learning such high-dimensional dynamics simply from 2D images. Moreover, the interaction between dexterous hands and objects is much more complex compared to gripper-object interaction, and trying to learn this interaction accurately may require more complex models and more data. Although our approach may not be effective in solving such complex tasks, we believe that we have performed a valuable exploration and provided a new way of thinking about world model construction. Learning about complex robot dynamics is our future endeavor.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a world-model-based framework for robot learning. However, the paper does not have enough experiments and ablative studies to prove the effectiveness of the proposed idea. Moreover, the generalization ability is not tested in the paper. MIssing baselines give another issue. Overall, the paper is below the bar of ICLR.