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

We sincerely appreciate your recognition of our work, as it means so much to us. We hope the following response can help clarify concerns you may have about our methods:

• W1,Q3: We sincerely appreciate your valuable feedback. We are currently incorporating comparisons with APV, a straightforward yet effective action-free pretraining method based on DreamerV2. To ensure fairness in the comparison, we have implemented APV on the foundation of DreamerV3 to eliminate any confounding factors. This experiment will be included in subsequent revisions of the paper.
• W2: RCWM does have certain limitations. It cannot handle more complex scenarios that may arise in real-world settings, such as a tilted table or viewpoint changes, which could introduce errors in the robot dynamics learned by the world model. Addressing robustness in such cases would require incorporating corresponding disturbances during training. However, this is not the focus of this paper. In future work, we will explore more robust methods for world model learning.
• Q1: Yes, we used the same set of trajectories for evaluation. Specifically, we selected 100 successful trajectories collected by a third-party policy that is independent of the evaluated policies as our evaluation dataset. This ensures fairness. Moreover, successful trajectories contain rich dynamic information, making them highly suitable for evaluating dynamics learning.
• Q2: In the stick-push task, the robot needs to first grasp the stick and then use it to push the kettle. During the stick-grasping phase, the robot may collide with the stick and the table, altering the dynamics and increasing dynamic errors. This highlights a limitation of the RCWM approach: it may struggle to effectively infer the robot's dynamics in scenarios involving significant collisions or resistance. Nevertheless, our method is still more accurate than vanilla world model that uses a single model to learn global dynamics.
• Q4: Thank you very much for your valuable feedback. RCWM is the first step toward decoupling robot dynamics from environmental dynamics. We firmly believe that this decoupling and explicit interaction modeling represent a valuable research direction. Our goal is to enable robots to first learn to predict their own motion and then progressively learn to interact with the external world. However, explicitly modeling interaction relationships can be challenging, especially for complex interactions, such as those involving flexible objects. Currently, RCWM is limited to learning relatively simple interactions, such as pushing objects or opening doors. Furthermore, when factors in real-world scenarios cause changes in robot dynamics, RCWM may struggle to robustly provide accurate dynamic predictions. In future work, we plan to explore methods for learning more robust robot dynamics and more accurately modeling interactions.

Thank you once again for the time and effort you have dedicated to our work!

We hope that the following responses will help reviewers better evaluate our approach:

• Q1: We use SAM2 to obtain robot segmentation masks, a technique that is well established and works well. We describe it in detail in A.4. Only manual labeling of the first frame of each trajectory is required to automatically obtain masks for subsequent frames, which results in lower manual costs. Furthermore, it is important to emphasize that our approach requires only a few data with robot masks to effectively decouple the dynamics of the robot and the environment, and doesn't require additional mask data for subsequent fine-tuning. In the paper, we only used 16 trajectories with a total of 4000 steps with masks for warm-up, which means that even if each trajectory needs to be manually labeled in order to obtain masks using SAM2, only 16 images need to be manually labeled. And further reducing the amount of data the model may still be able to function properly. In fact, even without mask data, our model still exhibits the ability to decouple structured information. We illustrate this in A5.1 as well as in Figure 12. Nevertheless, it is still difficult to stably decouple robot and environment dynamics without prior masks. It is not clear to us if there is a way to do so.

• Q2: Learning world models in dynamic stochastic environments is still an open problem. For DreamerV3, its ability to model random parts of the environment, such as random walks of monsters in Atari games, through stochastic representations is still limited. It is quite difficult to accurately model the physical interactions of unpredictable objects. Existing world models are all data-driven in nature and are deficient in modeling dynamic stochastic environments. However, for simple physical interactions such as pushing a cube, opening a door, etc., our approach can effectively learn this interaction and generate realistic imaginary trajectories.

• Q3: If only the replay buffer from a single task is used as the dataset for pre-training, it may lead to insufficient learning of the robot dynamics, as the covered state space might not be comprehensive enough. This could prevent the learned robot model from providing accurate dynamic predictions for new tasks. We focus on enabling the same robot to learn a variety of tasks and aim to better utilize the data collected from previous tasks to pre-train the world model. This is useful in real applications. In fact, as long as the state and action spaces of the robot in the pre-training dataset are sufficiently covered, even data collected through unsupervised exploration rewards in a single task can lead to a good learning of robot dynamics.

We sincerely thank the reviewers for their valuable comments to make this paper better!

We hope that the following responses will help reviewers better evaluate our approach:

• W1: TD-MPC2 is currently a leading model-based RL algorithm and has garnered significant attention from the community. However, it is primarily focused on locomotion control tasks that rely on state inputs, rather than on visual robot manipulation tasks that depend on image inputs. Additionally, we believe that directly comparing our method with TD-MPC2 would not provide a valid assessment of our approach, as our method is an enhancement based on DreamerV3, and performance differences between TD-MPC2 and DreamerV3 could obscure our method's effectiveness. While our approach could theoretically be adapted to build world models in TD-MPC2, doing so would require substantial modifications. We chose DreamerV3 as the foundation for our algorithm because it demonstrates exceptional potential in visual robot manipulation tasks. Our goal is to investigate whether there exists a more suitable method for constructing world models specifically for visual robot manipulation tasks, while using TD-MPC or Actor-Critic for policy learning simply represents an alternative way to leverage this world model. And for some model-free methods that focus on sample efficiency, the same can be disruptive in evaluating our methods. DrM, for example, uses neural network perturbations to enhance exploration, which is important in Meta-world tasks. There also exist Model-based rl methods for studying efficient exploration that could be combined with our approach to improve exploration efficiency, but this is not the dimension we wish to evaluate. We believe that our comparison with the vanilla world model in DreamerV3 is sufficient to illustrate the effectiveness of our approach and is in the dimension of world model construction.

• W2: For dexterous manipulation tasks in Adroit, we believe that state-based control algorithms such as TD-MPC2 are more suitable. While visual model-based rl methods may have difficulty in effectively learning high-dimensional hand dynamics from only images, which remains a current area of endeavor. And the interaction between the dexterous hand and external objects is more complex compared to the gripper, and it is quite difficult to make the world model learn this complex interaction accurately. We believe that our chosen tasks are also challenging, which typically require more than a million interactions to learn effective policies and can sufficiently validate the effectiveness of our approach. To the best of our knowledge, we are the first to separate robot dynamics and environment dynamics and simultaneously model their interactions, which we believe is a valuable and successful attempt. However, for learning more complex robot dynamics and interactions, we will try it in our future work.

• W3: RL training in the real world is often costly, and although the method is already sample efficient, it still requires a significant amount of environment interaction. We believe that experiments in simulated environments have sufficiently demonstrated the effectiveness of our approach. We will try real-world experiments in future work if conditions permit.

We sincerely hope that the reviewers will re-evaluate our approach in light of our responses. We greatly value any feedback on our responses. Once again, thank you for your valuable suggestions despite your busy schedule!

Thanks for your reply! I have changed my scores to Rating: 3 and Confidence: 3.

The authors did not include any of the experiments I proposed in the original review and responded with incorrect information: TD-MPC2 is not a state-based RL algorithm and has been validated on diverse manipulation tasks. As a leading model-based RL algorithm, it is reasonable to expect the authors to compare its performance with the proposed method. I hope the authors will include more comparisons in the future.

We have updated the comparison results with TD-MPC2 on our website and hope you will take a look. Experimentally, our method is competitive with TDMPC2. Since TDMPC2 has significant advantages over DreamerV3 in terms of both policy learning and decision making, our approach focuses on the construction of the world model without modifying the policy learning, and therefore does not show a consistent advantage over TDMPC2. We found that TDMPC2 still has excellent performance even without pre-training. However, since the learning of representations and dynamics relies heavily on task-relevant reward information, it is difficult for TDMPC2 to obtain valid information from pre-training. While our method can efficiently learn representations and robot dynamics that are useful for new task learning. Although our method is not as lightweight as TDMPC2 and requires the use of pre-training, robot masks, etc. to perform well, we believe that task-independent pre-trained world models will be one of the important research directions in the future. And our approach offers a new way of thinking about this.

Once again, we sincerely appreciate the time and patience you have dedicated to reviewing our work.

We sincerely hope that the reviewers will re-evaluate our approach in light of our responses. We greatly value any feedback on our responses. Once again, thank you for your valuable suggestions despite your busy schedule!

We hope that the following responses will help reviewers better evaluate our approach:

• Q1: Our approach is not an improvement for TD-MPC2, nor does it address the problem of multi-task offline learning. We focus on online training for visual robot manipulation. And, to the best of our knowledge, TD-MPC2 does not perform multi-task learning with images as input either. This remains a difficult problem for the current research and we believe it is beyond the scope of this paper.
• Q2: We detailed the computation of the attention map in A.5.3 and supplemented the visualization results in Figures 15 and 16. The attention map consists of two parts, one for $\hat{z}_t^{R}$ and $z _ {t-1}^E$ with a shape of 3232 and one for $\hat{h}_t^R$ and $z _ {t-1}^E$ with a shape of 832. We concatenate both as 4032 attention maps and resize to 8064 to align the size of the image observations for visualization. The parts of the attentional map without significant activation are the attentional maps of $\hat{z}_t^{R}$ and $z _ {t-1}^E$, and the parts showing significant activation are the attentional maps of $\hat{h}_t^R$ and $z _ {t-1}^E$. This is because $\hat{z}_t^{R}$ is a stochastic representation generated by $\hat{h}_t^R$, which contains richer information, such as action information. Although the attention maps for $\hat{z}_t^{R}$ and $z _ {t-1}^E$ are not significant, they are still computed to ensure consistency between the robot and environment representations. The activations are not meaningless either. We observe significant activation when the robot touches an object, suggesting that the interaction model is capable of learning the interaction relationship between the robot and the object.
• Q3: The gripper may lose some details when reconstructed due to its small size, but the opening and closing state of the gripper can also be clearly represented, which does not affect the imagination of the operation. In Figure 5, in the case of occlusion of the gripper, it may appear to disappear due to its small size, but it still does not affect the operation of the door. Methods based on reconstruction usually lose some details when reconstructing, which is due to the nature of the autoencoder. Even though the reconstructed image may be missing some details, the hidden state may not ignore this information. We think reconstruction is just a way to better help us understand what's going on in our world model, because state transitions and decisions happen in the latent space.

We sincerely thank the reviewers for their valuable comments to make this paper better!

We hope that the following responses will help reviewers better evaluate our approach:

• W1: We considered the potential performance gains that could result from the use of robot masks, so we added the same mask-guided decoder as RCWM when training DreamerV3 and used the same amount of mask data for pre-training to eliminate the interference. We explain this in line 417 of the paper. For the interaction model, which is an integral part of RCWM, the same is part of the proposed improvements. Our RCWM, by decoupling robot dynamics from environment dynamics, is able to allow us to model the interaction between the two more explicitly, which is one of the advantages of RCWM over the vanilla world model.

• W2: Our approach is not an improvement aimed at improving generalization. The improvement in generalization for environment disturbances is a result of our separation of predictions for robot dynamics and environment dynamics. When the environment changes, our robot branch is able to provide accurate robot representations and dynamics predictions almost unaffected. Thus even if the environment representation changes, the policy can still generate reasonable actions based on the robot representation. Figure 8 shows that our approach is more robust to environment disturbances compared to DreamerV3, which strongly supports our idea. As to why both our approach and DreamerV3 produce significant performance degradation on the sweep-into task, we believe that it is because the sweep-into task needs to determine where the holes and objects are by the desktop texture, and thus is very sensitive to the desktop texture. For a description and visualization of the sweep-into task check out https://meta-world.github.io/. Furthermore, since we focus on using the same robot for multiple tasks, there is no need to be robust to robot appearance. And when position and viewpoint change, the dynamics in the 2d image-based world model will be severely affected, which is still lacking research.

• W3: The tasks we chose are very challenging tasks in the Meta-world benchmark, which usually require more than a million steps of interaction with the environment to learn effective policy. And these tasks mostly involve interactions with objects in the environment, which can effectively evaluate whether our method is able to learn such complex interactions. Our focus is on visual robot manipulation tasks rather than locomotion control tasks. As for the dexterous hand in MyoSuite evaluated in TD-MPC2 or the humanoid robot control task in Humanoidbench, it is very difficult to use visual inputs for control, and the current approach is still attempting on state inputs. Moreover, it is also very difficult to accurately learn the complex dynamics of the dexterous hand from vision, which we believe is beyond the scope of this paper.

• W4: Since every component of RCWM is indispensable, we think that ablation experiments are not necessary. For example, without the warm-up phase, it is difficult for our model to stably disentangle representations of the robot and the environment without prior robot masks. For the choice of warm-up data, which is not critical for this method, it is sufficient that the data covers the operational space relatively adequately and thus allows our model to learn the robot dynamics, or even trajectories collected with unsupervised stochastic exploration. As for the components in RCWM, such as mask-guided decoder, interaction model, etc., RCWM does not work properly after removal.

• W5: Although reconstruction error is not a good metric to measure the ability to model robot dynamics, it also reflects the ability to some extent. In order to avoid the influence of environment on evaluating robot dynamics modeling, we multiply the imagined reconstructed image with the ground-truth robot mask and compute the mean square error with the ground-truth robot image to serve as the robot dynamics prediction error. We mentioned in line 447 of the paper. With this evaluation method, we were able to consider only the reconstruction results for the robot part. We perform a full evaluation in Figure 13 in Appendix A.5.2.

• W6: TD-MPC2 mainly focuses on locomotion control tasks using state as input rather than images, and does not use reconstruction loss to learn representations. In contrast, our approach focuses on visual robot manipulation tasks and uses reconstruction loss to learn representations. Applying RCWM to TD-MPC2 would require large modifications, which may be beyond the scope of this paper and may lead to new approaches. We believe that the performance improvement over the vanilla world model is sufficient to illustrate the effectiveness of our approach.

Apologies for the delay in response.

I particularly want to thank the authors for taking so much time to explain their settings and answer all the questions I had. Here is a summary of the main concerns that were not addressed by the authors:

• I wish I could make my statement about 'TD-MPC2 should be included as a baseline' clearer. I know that the authors aim to convey a promising direction, which I agree with and appreciate: decoupling robot dynamics and environments is helpful. However, from my perspective, the current submission only verifies the feasibility of this interesting idea, but is far from the conclusion on its generalizable effectiveness. Specifically, the training setup is the same as the one in TD-MPC2 and DreamverV3 but with an additional pre-training phase. In other words, the goal of this manuscript is to improve the accuracy of world modeling. If TD-MPC2 or other SOTA MBRL algorithms, which are not included in this manuscript, can already outperform RCWM by learning from scratch, then it is hard to see the motivation of the community to use the proposed RCWM that even require additional efforts (e.g., SAM 2 annotation) to design the pre-training stage. In this case, using other SOTA MBRL algorithms will be enough and burden-easy.
• That is also the reason why I wish to see the results on more tasks. Possibly, RCWM may perform badly in tasks with complex interactions. If the authors could show the superiority of RCWM on diverse task domains, even compared with only DreamerV3, the concern above could also be alleviated.
• That is also the reason why I wish to see additional emergent properties, like generalization ability and robustness to visual disturbance that are not discovered in previous methods. These properties could also greatly motivate the use of RCWM.

Regarding other responses:

• TD-MPC2 is not a solely state-based algorithm, recognized by other reviewers as well. There already exists an extension of TD-MPC2 to visual humanoid control tasks [1], which are likely more complex than manipulation tasks. Based on their results, DreamerV3 significantly underperforms TD-MPC2. Meanwhile, the authors of TD-MPC2 also open-source their code with high reusability to the tasks considered in this manuscript.
• ManiSkill2 does not include only simple tasks like grasping. Tasks like insertion, stacking, and assembling are also included. I do not think they are simpler than the selected ones. ManiSkill2 also supports more realistic rendering, which may be a better choice to show the promise of RCWM.
• Regarding MyoSuite tasks or Adroit mentioned by other reviewers, it is true that they are hard to model. Informally speaking, based on my personal experience, these tasks could be successfully learned with DreamerV3 with ~33 PSNR. It would be better to give the conclusion after conducting the results.
• I also carefully read the discussions with other reviewers. The authors said that "in Adroit, the viewpoints are generally inconsistent between tasks, making it nearly impossible to extract reusable dynamics." I would like to say that the views can simply be consistent across all Adroid tasks, based on my past research. The authors also said that "the joints of the dexterous hand often experience severe occlusion, which presents a significant challenge for learning the robot's dynamics." This is true. But, Adroid tasks are inherently partially observable and can be solved by DreamerV3 without the additional tuning, based on my research as well. This is not the reason to exclude the possibility of doing experiments in this domain.
• Ablation studies are necessary. I understand this paper, and I know that some components cannot be removed. What I wish to see are some other design choices, like model size, pre-training dataset, model choice, etc. Doing more experiments on design choices will help to guide the community to extend RCWM in the future. Meanwhile, the authors tell me some conclusions (in both manuscript and rebuttal page), but do not show the numerical results to verify them. I would suggest providing more empirical numbers to strengthen some words.

In summary, I would maintain my score and genuinely suggest the authors address the concerns above with more empirical results.

[1] Hansen, N., SV, J., Sobal, V., LeCun, Y., Wang, X., & Su, H. (2024). Hierarchical World Models as Visual Whole-Body Humanoid Controllers. arXiv preprint arXiv:2405.18418.

We sincerely appreciate your recognition of our work, as it means so much to us. We hope the following response can help clarify concerns you may have about our methods:

• W1: To eliminate interference caused by model capacity, we reduced the latent space dimension when constructing the two branches in RCWM. The latent space dimension of the dynamics model in DreamerV3 is 512, with 5.7M parameters, whereas the latent space dimension of the dynamics model in RCWM is 256, with a total of 3.7M parameters for both branches and the interaction model combined. In fact, our method uses fewer parameters in the dynamics model. This also indirectly demonstrates that RCWM is an efficient approach to constructing world models. We believe that expanding the capacity would likely result in performance improvements, as validated in DreamerV3.
• W2: Our aim is to learn generalized robot dynamics that can be directly used for new tasks through pre-training, therefore hoping that it can inference independently of the environment. This prevents us from designing the information exchange between branches, even though doing so might lead to better performance. We explicitly models the effect of the robot on the environment while implicitly modeling the effect of the environment on the robot. If the interaction model is disabled, the environment branch will not be able to access the action information and the details of the robot's movement and thus it will not work properly. If we simply use two identically constructed branches to predict the dynamics of the robot and the environment separately, we believe that this may not allow for effective dynamic separation and consistent interaction prediction. Since the learning of policies requires multi-step rollout in the world model, if there is no information exchange between two branches completely, the prediction results of the two branches may produce significant differences during rollout due to the presence of compounding errors thus failing to produce consistent interaction prediction. We therefore believe that constructing an interaction model is necessary, and this is the key to the success of our approach in modeling the manipulation process.

Thank you once again for the time and effort you have dedicated to our work!

Dear Reviewer pYBH, Could you please read the authors' rebuttal and give them feedback at your earliest convenience? Thanks. AC