Leveraging Separated World Model for Exploration in Visually Distracted Environments

Thank you for your review and feedback. The term "task-irrelevant" on L51 should be "task-relevant". Here are our responses to your questions:

• Strong assumption

  • Firstly, many Visual RL works use separation assumptions, such as Denoised MDP[3] and TIA[2], which have natural assumptions and practical significance. SeeX is inspired by this line of research. We believe no single assumption is universally applicable (e.g., CNNs for images and RNNs for text). So whether an assumption is "strong" is subjective and context-dependent.
  • Secondly, we focus on URL problems with task-irrelevant and action-independent noise and we model URL as minmax problems through theoretical analysis. A practical scenario fitting our assumption is a book-finding robot in a library. The robot must differentiate between task-relevant items (e.g., books) and task-irrelevant items (e.g., posters or people) to find the specific book. Thus, our work is grounded in both theory and practical application.
  • Finally, the robot manipulation you mentioned is discussed in L331-334.

• Inseparable noises

  • We emphasized in the paper (caption of Figure 1 and the conclusion) that we only focus on handling moving distractors. Moving distractors and inseparable noises represent distinct research directions. According to the No Free Lunch theorem, no single method can be effective for all tasks.

• Discussion on whether EX-BMDP and latent disagreement are orthogonal

  • Separated world model
    • Firstly, we adopted(L78-79) rather than proposed EX-BMDP[1]. The similarity you see between EX-BMDP and TIA is due to their shared use of the separation concept, but separation is not a specific method but rather a general idea[1,2,3,4].
    • Additionally, EX-BMDP differs significantly from TiMDP[2]. TiMDP applies action to the exogenous branch, while EX-BMDP does not. Did the reviewer notice this? The agent can only influence the endogenous branch under our assumption, which is why EX-BMDP is suitable.
  • Intrinsic reward
    • As we all know, during pretraining phase of unsupervised RL, task-specific rewards are typically avoided in favor of intrinsic rewards, which come in various forms[5,6,7,8]. Any intrinsic reward can support SeeX's exploration, and we chose the classic disagreement[5].
  • Our method formalizes URL as a minmax problem, where the outer layer minimizes REGRET’s upper bound using a separation assumption (Eq 2), leading us to adopt EX-BMDP. Inner layer maximize state uncertainty, so we approximate Eq 7 with Eq 11, making disagreement a reasonable choice. Reviewer considers the methods to be orthogonal, but we believe that the underlying ideas abstracted (e.g., separation and intrinsic reward) can inspire cross-field applications. Our step-by-step design and component choices are based on necessity, and our results validate the framework's effectiveness.

• L193, the relationship between Eq 7 and Eq 11

  • In Eq 7, the goal of $\pi_{expl}$ is to maximize the total variation (TV) distance between the learned model and the true latent dynamics in the task-relevant part. Practically, we estimate $D_{TV}$ using the disagreement among an ensemble of neural networks, inspired by [9].

• How to ensure $s^-$ does not contain information about the reward?

  • Firstly, our Exo-Rec design follows [2,3], which states that task-relevant information is a small part of the observation. The Exo-Rec term aims to encode as much information from $o$ into $s^-$ (L164-166), not to remove reward-related information from $s^-$ as you mentioned.
  • Secondly, We assume that only the task-relevant part of $o$ is affected by the action. Based on EX-BMDP, we design a separated world model where the action limits the encoding of task-relevant information into $s^+$.
  • In summary, we believe there is a misunderstanding of our method and the Exo-Rec design. In visual settings, reward-relevant and reward-irrelevant information are inherently mixed and cannot be perfectly separated. Are you asking us to ensure a 100% separation between the two? This is not feasible. Our results(Figure 3, Table 1) show that our method attempts to separate $s^-$ and $s^+$ as much as possible, benefiting the task of handling moving distractors.

• The implementation of baseline methods

  • We have mentioned in the paper (L47-49, L73-77, L201) that our method builds on URLB[10] and URLB-pixel[11], using the same implementation for the baseline. The specific implementation of the baseline can be found in the URLB-pixel GitHub repository.

• The REGRET bound in Eq (6)

  • According to your statement, if we extract precise task-irrelevant information ($\hat{T}^-$) from the observation, it wouldn’t help in reducing REGRET? This seems unreasonable because reward-relevant and reward-irrelevant information are complementary and together constitute the observation. Knowing one allows us to infer the other.
  • Although $s^-$ is unrelated to REGRET, we use losses of $s^-$ and $s^+$ to bound REGRET (Eq 6). If you believe that $s^-$ is insignificant, you might try deriving a tighter REGRET bound based solely on $s^+$.

We hope this clarifies our work. Please feel free to ask any further questions.

[1] Provably Filtering Exogenous Distractors using Multistep Inverse Dynamics

[2] Learning Task Informed Abstractions

[3] Denoised MDPs: Learning World Models Better Than the World Itself

[4] SeMAIL: Eliminating Distractors in Visual Imitation via Separated Models

[5] Planning to Explore via Self-Supervised World Models

[6] Curiosity-driven exploration by self-supervised prediction

[7] Exploration by random network distillation

[8] Curiosity-driven exploration via latent bayesian surprise

[9] Reward-free curricula for training robust world models

[10] URLB: Unsupervised Reinforcement Learning Benchmark

[11] Mastering the unsupervised reinforcement learning benchmark from pixels

Dear reviewer NjAE,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors

Thank you for reviewing our rebuttal. We are glad to hear that you do not have any new questions. If you have any further concerns or need clarification, please feel free to reach out at any time.

Thank you for your detailed review. We have noted the typographical and formatting issues mentioned by the reviewer. BA on Line 157 refers to the Barber-Agakov bound[1], named using the authors' initials, which is a commonly used bound for mutual information. The second "=" on Line 164 should indeed be a "+". We believe these adjustments do not affect the main content and conclusions of the paper. Below are our responses to the issues you raised.

• Regarding conclusions of Figure 4
  • Our original intention was to show SeeX's performance stability between 100k and 500k steps; we sampled and tested the performance of other methods between 100k~500k, but none exceeded 90% expert performance, so we did not report them in the paper. We show the performance of all steps for the other methods in Pdf-Table-B and our refined conclusion is that SeeX achieves 90% expert performance with fewer pre-training steps than other methods, demonstrating the efficiency of our bi-level framework and the separated world model during exploration.
• Other trajectory plots except for Figure 9
  • In Figure 9, we have shown the reconstruction trajectories of SeeX after fine-tuning in the Walker environment, but there was a lack of trajectory visualization of other environment during the exploration phase. Therefore, Pdf-Figure-C displays the reconstruction trajectories of SeeX during exploration in the Quadruped and Jaco environments. The trajectories are organized into five columns: from top to bottom, they represent the original observation $o$; the reconstruction $\hat{o}$ obtained from the joint use of $s^+$ and $s^-$ with the help of the mask $m$; the endogenous reconstruction $\hat{o}^+$ obtained using only $s^+$, representing the task-relevant part; the exogenous reconstruction $\hat{o}^-$ obtained using only $s^-$, representing the task-irrelevant moving distractors; and the mask $m$ used to reconstruct the entire image. By analyzing the reconstruction trajectories, we can draw the following conclusions:
    • Although moving distractors were not completely removed, the reconstruction results show that $s^+$ contains most of the task-relevant information (including the agent's torso), significantly reducing the impact of background noise.
    • During the exploration phase, we use intrinsic rewards to encourage the agent to explore a more diverse range of states, aiming to build as comprehensive a world model as possible. This foundation supports imagination during policy training in the fine-tuning phase. The reconstruction results indicate that the agent does not stick to fixed actions but attempts to explore diverse states, demonstrating the effectiveness of the intrinsic rewards used.
    • The Jaco agent is a multi-joint robot with high flexibility, making it more complex than both the Walker and Quadruped. The reconstruction quality for Jaco is noticeably worse than for Quadruped, and the removal of moving distractors is less effective, which aligns with the experimental results in our paper. However, some interesting observations can be made: (1) The red ball, while not part of the agent's torso, is related to the reward function. It appears in $\hat{o}^+$ but not in $\hat{o}^-$, demonstrating SeeX's ability to automatically identify task-relevant information. (2) The base of Jaco, though fixed and not controlled by actions, is part of the agent's torso but has little relation to the reward. Therefore, SeeX classifies it as task-irrelevant, corresponding to the reconstruction image $\hat{o}^-$.
• Regarding the ablation studies in Figure 5
  • When formatting Figure 5, we only considered aesthetic factors and overlooked the potential for misleading the reviewers. Therefore, we have re-drawn the conclusions of the ablation study and improved the readability of the charts, and added the ablation experiments on the walker (Pdf-Table-C) as well as the performance with $\alpha=0$ (Pdf-Figure-B). By analyzing the charts, we can draw the following conclusions:
    • The choice of weight for the Exo-Rec term significantly impacts SeeX's performance. For Walker, selecting $\alpha=3$ is appropriate, while for Quadruped and Jaco, $\alpha=1$ is more suitable. Removing the Exo-Rec term (equivalent to setting $\alpha=0$) leads to varying degrees of performance degradation across different environments, indicating that the Exo-Rec design contributes positively to performance.
    • Observing Pdf-Table-C, SeeX's sensitivity to different hyperparameters shows similar conclusions for Walker and Quadruped: (1) Removing the separated world model significantly impacts performance; (2) Removing the Exo-Rec term aligns with the conclusions from Pdf-Figure-B; (3) SeeX's performance is generally positively correlated with the number of predictive heads $K$, but considering the computational cost associated with increasing $K$, $K=5$ is a suitable choice.

Thank you again for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper.

[1] The IM algorithm: a variational approach to information maximization

Dear reviewer HYxU,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors

Thank you for your prompt and thoughtful response to our rebuttal. We appreciate the time and effort you’ve taken to review our rebuttal and provide additional feedback. Below are our responses to the new questions and concerns you raised:

• Additional ablation results on Walker
  • As you mentioned in your review, “there are no data points at 200k/300k/400k for baselines,” we understand your concern. To address this, we have included the results in a table, providing more specific quantitative values to better illustrate the performance comparison of different methods at various time steps. Considering your preference for the same type of qualitative line charts as before, we will update the new data points in the line charts of the later formal version.
• Ablation of $\alpha$​
  • Firstly, thank you for pointing out in your review that Figure 5 was somewhat confusing. In Figure 5, the top-left, top-right, and bottom-left plots show ablations with different values of $\alpha$ in various environments; since their content is consistent, we present them as bar charts in the PDF for a clearer comparison of their normalized returns. The bottom-right plot depicts ablation experiments of different modules, and we provide a more detailed analysis of these experiments in the table, which shows the specific numerical improvements.
  • Secondly, to address your and future readers' confusion, in the later formal version, we will replace the previous inconsistent Figure 5 with a unified and clearer set of figures. Additionally, we will include the original line plots in the appendix for reference, so readers can still gain a qualitative understanding of the performance of the ablation experiments.
• Reservations about the novelty
  • First, thank you for considering that our approach of using only task-related information to measure disagreement is an improvement. Below, I will provide a more detailed explanation of our novelty and contributions.
  • In real-world scenarios, there are many sources of action-independent and task-irrelevant noise. For example, when a drone is patrolling, surrounding buildings, or distant vehicles do not directly affect the drone's flight path or the completion of its patrol task. Many current works overlook this issue, and we are the first to address such noise in the URL setting. Additionally, we propose a bi-level optimization framework centered around task-related information and design our methods based on this theoretical foundation. Finally, our experimental results demonstrate the effectiveness of our approach, and ablation studies confirm the utility of different components.
• Statements without experimental results
  • Considering the URL problem setting, our model-based framework requires intrinsic reward to assist in early exploration. Intrinsic reward is a modular and can be implemented in various ways. For example, the disagreement measure $Var(\hat{h}^+_{t+1,i})$ used in our work and the variance of the log of predicted probabilities $Var(\log \hat{T}_i(z_t))$ used in [1], both of the intrinsic rewards can be used within our framework.
  • Our main contribution is a theoretically-based framework for solving the URL problem with visual disturbances. In principle, intrinsic reward can be implemented based on the ideas from different methods in URLB[2], as well as the two specific approaches mentioned earlier, meaning our framework is intrinsic-reward-agnostic. As you mentioned, the experimental performance of these intrinsic rewards has indeed not been tested. We will consider conducting more extensive and comprehensive comparisons in future work.

Thank you for your continued effort in providing multiple rounds of feedback to improve our work. Your questions have been very valuable. Please feel free to reach out with any further questions or concerns you may have in the future.

[1] Offline Reinforcement Learning from Images with Latent Space Models

[2] URLB: Unsupervised Reinforcement Learning Benchmark

Thank you for the detailed review and valuable suggestions regarding our paper. We are pleased to hear that you find our work to have clear technical contributions and to perform well in experiments. In response to the questions and suggestions you have raised, we provide the following answers:

• Comparison between our research direction and data augmentation (DA)
  • Our designed bi-level separation framework extracts task-relevant information, while the DA method extracts task-relevant representations through data augmentation. As you mentioned, there are indeed similarities between the two approaches. However, since there are few methods that directly use DA for exploration, we incorporated the DA method into both Plan2Explore and SeeX to investigate whether it enhances performance. This is to verify whether the DA method provides benefits in the moving distractor setting.
  • We used the classic data augmentation technique of random shift (pad by 4 pixels) proposed in DrQ[1], applying it 4 times to the same image. To thoroughly evaluate the effectiveness of DA, we tested its impact during both the pretraining (p-DA) and finetuning (f-DA) stages on the walker task. The performance results are shown in Pdf-Table-A.
  • By analyzing the data in the table, we can draw the following conclusions:
    • Applying data augmentation (f-DA) on specific tasks is effective, as it can mitigate the impact of distractors and improve performance, which supports the reviewer's insight.
    • However, for tasks with moving distractors, the separation design in SeeX performs better than data augmentation.
    • Using data augmentation during the pretraining stage leads to a performance drop, possibly because it interferes with the learning of the world model. The specific reasons will be explored in future work.
• Further clarification on the content from L331-334
  • Based on whether they affect the agent's reward function and action, distractors can be categorized into four types: task-irrelevant + action-independent (the setting addressed in our work); task-relevant + action-independent (which I understand to be the multi-agent setting you mentioned); task-irrelevant + action-dependent; and task-relevant + action-dependent. The latter two are the settings I referred to on L331-334, where distractors can interact with agents and even influence their reward function.
  • We believe the reviewer's understanding is correct and we appreciate you pointing this out. We will include the settings you mentioned in the camera-ready version.
• Ablation in which the policy is optimized with access to both branches
  • Firstly, by formalizing the URL problem as a minmax problem, we further implement it as a separated world model consisting of two branches: endogenous and exogenous. The exogenous branch extracts task-irrelevant information, which is unrelated to the task. Therefore, it is reasonable to use only the endogenous branch to optimize the policy.
  • Secondly, I understand that you are referring to the ablation experiments for using both branches as inputs to the policy. To address your concerns comprehensively, we consider the following two scenarios: (1) using both exogenous and endogenous states simultaneously for policy optimization, and (2) merging the endogenous and exogenous branches into one, using a single encoder to extract the state $s$ for policy optimization.
    • (1) We added the following comparative experiments: both $s^+$ and $s^-$ are simultaneously input to both the actor $\pi(a|s^+,s^-)$ and the critic $v(s^+,s^-)$, with the action still applied to the endogenous branch. We compared the performance of this setup (SeeX-both-branch) with SeeX in the walker environment, and the specific performance curves are shown in Pdf-Figure-A. In all tasks, SeeX outperformed SeeX-both-branch, especially in the last three challenging tasks. This indicates that using only endogenous information is advantageous for tasks with moving distractors.
    • (2) Using a single encoder to encode observations for training the policy $\pi(a|z)$, which is the approach of Plan2Explore[2]. The superiority of SeeX over Plan2Explore is demonstrated in Figures 3 and 4 of our paper, showcasing the effectiveness of the separated world model.
  • In summary, we believe that using only the endogenous branch for policy optimization is reasonable, and our experiments and the content in the paper have demonstrated this. Although using information from two branches under our assumption doesn't offer advantages, as noted by the reviewer, there are multi-agent settings where noise, despite action independence, can affect the agents' reward functions. In such cases, incorporating part of $s^-$ into policy optimization might be effective. This is an interesting research direction that we are glad to explore in future work.

Thank you again for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper. We will ensure that the additional experimental results and any necessary changes are included in the camera-ready version of the paper.

[1] Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels

[2] Planning to Explore via Self-Supervised World Models

Thank you for your detailed review and valuable feedback on our work. We are very pleased that you found our algorithm theoretically motivated and experimentally successful. In response to your questions and suggestions, we have made the following points:

• Exploration of Real-World Applications
  • Our work focuses on handling moving distractors, where we need to separate task-relevant states from task-irrelevant states. In this setup, there are many common downstream tasks. For example, consider a book-finding robot in a library that needs to scan the shelves to find a specific book. The robot will encounter task-relevant objects (such as books) and task-irrelevant objects (such as posters on the wall or passing people), so the robot needs to distinguish which information is beneficial for the task.
  • We agree with your view that real-world applications are crucial for validating the practicality of the algorithm. However, the self-driving and navigation challenges you mentioned are difficult to analyze and computationally expensive in practice. We would be willing to address these in future work. And this will involve designing experiments closer to real-world conditions to verify the algorithm's effectiveness in handling complex real-world scenarios.
• Discussion on Method Limitations
  • The assumption of completely separating endogenous and exogenous states that you mentioned is indeed a limitation of our method. We have discussed this in the paper, and our work has preliminarily validated the effectiveness of the bi-level optimization framework and the separated world model under the setting where endogenous and exogenous information is fully separable. In future work, we will explore how to introduce interactions between states to better understand and improve the applicability of the algorithm.

Once again, thank you for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper.

Dear reviewer SZqk,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors