Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions

Authors: Kyungmin Kim, JB Lanier, Pierre Baldi, Charless Fowlkes, Roy Fox
Primary area: reinforcement learning
Keywords: Visual Control, Robust Representation Learning, Model-Based Reinforcement Learning

Abstract:
Recent advancements in Model-Based Reinforcement Learning (MBRL) have made it a powerful tool for visual control tasks. Despite improved data efficiency, it remains challenging to train MBRL agents with generalizable perception. Training in the presence of visual distractions is particularly difficult due to the high variation they introduce to representation learning. Building on DREAMER, a popular MBRL method, we propose a simple yet effective auxiliary task to facilitate representation learning in distracting environments. Under the assumption that task-relevant components of image observations are straightforward to identify with prior knowledge in a given task, we use a segmentation mask on image observations to only reconstruct task-relevant components. In doing so, we greatly reduce the complexity of representation learning by removing the need to encode task-irrelevant objects in the latent representation. Our method, Segmentation Dreamer (SD), can be used either with ground-truth masks easily accessible in simulation or by leveraging potentially imperfect segmentation foundation models. The latter is further improved by selectively applying the reconstruction loss to avoid providing misleading learning signals due to mask prediction errors. In modified DeepMind Control suite (DMC) and Meta-World tasks with added visual distractions, SD achieves significantly better sample efficiency and greater final performance than prior work. We find that SD is especially helpful in sparse reward tasks otherwise unsolvable by prior work, enabling the training of visually robust agents without the need for extensive reward engineering.

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

Reviews and discussion:
<thread>
creation date:
1730231999242

thread_type:
official_review

<thread_object>
creator: reviewer_l53c

<content>
summary:
The paper presents Segmentation Dreamer (SD), which improves Model-Based Reinforcement Learning (MBRL) in visually distracting environments. It focuses on task-relevant components using segmentation masks, and achieves better sample efficiency and performance, especially in sparse reward settings and when exposed to significant visual distractions.

strengths:
Originality: Propose selective L2 loss to handle not perfect segmentation input
Quality: A good approach to introduce segmentation model to help MBRL 
Clarity: Visualizations easy to read
Significance: Deal with distraction in RL policy deployment, which is important

weaknesses:
The main weakness is that proposed algorithm mainly tested with visual distraction with swapped 2D background shift. This set of benchmarks are surrogate to real domain shifts but the segmentation model directly aims to remove that. Thereby, it would be hard to generalize to other cases such as change of viewing angle, color, lighting, texture etc.

questions:
Can you also compare against other methods also utilizing  segmentation model? e.g., using SAM to get relevant part and feed that to some RL algorithm?  

Can you discuss how would some extension of the proposed method handle cases such as change of viewing angle, color, lighting, texture etc.?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730634048832

thread_type:
official_review

<thread_object>
creator: reviewer_qm91

<content>
summary:
The paper introduces an approach to enhance MBRL in environments with visual distractions. Building upon the DREAMER framework, the authors propose a method called Segmentation Dreamer (SD). This method utilizes segmentation masks to focus on task-relevant components of the visual input, reducing the complexity of representation learning. To address the limitations of predicted masks, the authors introduce a selective L2 loss function.

strengths:
To address the limitations of predicted masks, the authors propose a selective L2 loss function.
This is a thoughtful addition to improve the robustness of the model to segmentation errors.

weaknesses:
1. While the paper presents an integration of segmentation masks into MBRL, it lacks a thorough comparison with existing works that also combine foundation models with RL. A more detailed analysis against these related works would strengthen the paper's claims and provide a clearer picture of the method's advantages and potential limitations[1, 2].
2. The generalizability of the approach across a wider range of real-world scenarios beyond the controlled environments should be considered. In real-world scenarios, distractions may closely resemble task-relevant features, potentially occluding critical information. The paper could benefit from evaluating the method in more realistic settings with complex distractions[3, 4].
3. The paper addresses some limitations, particularly regarding the dependence on accurate segmentation masks. However, a more detailed discussion on potential model biases, especially how segmentation errors might affect performance in untested environments, would be beneficial.



Reference:
[1] "Generalizable Visual Reinforcement Learning with Segment Anything Model." arXiv, 2024
[2] "Empowering Embodied Visual Tracking with Visual Foundation Models and Offline RL." European Conference on Computer Vision, 2024
[3] "Towards distraction-robust active visual tracking." International Conference on Machine Learning. PMLR, 2021.
[4] "Anti-distractor active object tracking in 3D environments." IEEE Transactions on Circuits and Systems for Video Technology 32.6 (2021): 3697-3707.

questions:
1. Could the authors elaborate on how their approach of integrating segmentation masks into MBRL differs from and improves upon existing methods that combine foundation models with RL?
2. What are the authors' thoughts on the generalizability of their method to real-world tasks where distractions are inherently complex and closely intertwined with task-relevant features? 
3. Can the authors discuss potential biases in their model that may arise from errors in segmentation masks? What strategies could mitigate the impact of segmentation errors in untested environments?
4. In environments where distractions and task-relevant features are dynamic, how does the method ensure consistent performance?
5. Is there a mechanism for the model to adapt to new environments or changing task dynamics without complete retraining?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649387064

thread_type:
official_review

<thread_object>
creator: reviewer_yfcf

<content>
summary:
This paper proposes Segmentation Dreamer (SD), an approach to learn task relevant features in model-based RL framework by using segmentation mask from prior domain knowledge as an auxiliary objective. The proposed approach only reconstruct task relevant components instead of entire original image to avoid noisy learning signals from distractions. Additionally a selective l2 loss is proposed to handle error predictions from segmentation masks. Experiments are performed on visual control benchmarks including DMC and Meta-World with visual distractions by replacing them with diverse background scenes. Experiments demonstrate that SD outperforms Dreamer and other baselines in better sample efficiency and convergence performance. Ablations show effectiveness of each component, specifically focusing on segmentation quality and usefulness of proposed selective l2 loss.

strengths:
1. Paper is well-written, offering clarity and easy to understand
2. The motivation of only reconstructing task-relevant parts in image is reasonable. The proposed approach can handle potential error of adapting pre-trained segmentation model to visual RL domains
3. Extensive experiments including baseline comparisons and ablation analysis demonstrate benefits of proposed approach in improving sample-efficiency

weaknesses:
1. This framework requires significant prior knowledge about which part is task-relevant in the image. This may limit the scope of potential usage of the proposed approach, where domain knowledge may be ambiguous or hard to obtain.
2. Since only parts of image is reconstructed specified by mask, this problem setting is simpler than the original, which may be hard to directly compare to other baselines
3. Experiments are limited to performance of visual control tasks; we do not know how good pre-trained segmentation models are adapted to downstream domains.

questions:
1. Since segmentation quality is important to task performance (Figure 3b), what will happen if more samples, than 1 or 5-10, are used in fine-tuning pre-trained segmentation models to downstream domains? Will the proposed selective l2 loss become less important if segmentation model improves with more training samples?
2. How is binary mask predictor head trained, i.e. what is loss function objective here? What will happen if binary mask decoder is trained but selective l2 loss is not used in learning world model, i.e. is the mask prediction task or addressing segmentation model errors more important in this framework?
3. In terms of $mask_\mathbf{SD}$ and $mask_\mathbf{FM}$, could authors visualize how these segmentations look like in order to understand what are similarity and differences in predictions from adapting pre-trained models and those from the process of learning world models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688004008

thread_type:
official_review

<thread_object>
creator: reviewer_rwsa

<content>
summary:
This paper addresses the challenge of training Model-Based Reinforcement Learning (MBRL) agents in visually distracting environments, where standard approaches struggle with representation learning due to high visual variability. Building on the DREAMER framework, the authors propose Segmentation Dreamer (SD), an auxiliary task that uses segmentation masks to focus representation learning on task-relevant elements, thereby simplifying latent encoding by excluding irrelevant objects. In experiments with visually complex tasks, SD significantly improves sample efficiency, performance, and robustness, particularly in sparse reward scenarios, enabling effective agent training without extensive reward engineering.

strengths:
* This paper is well-written and easy to follow. The idea of using the segmentation to capture the task-relavent information is reasonable.

* The literature review is sufficient and the authors have adequately discussed the major difference between the proposed method and existing works.

* The evaluations of the proposed method are comprehensive. The authors compare the proposed methods with several state-of-the-art algorithms and evaluations are conducted on two benchmarks.

weaknesses:
* Although the authors conduct experiments on several benchmarks, I still have concerns about the generalization ability. In DMControl tasks, the agent’s visual state consistently appears at the center of the image, and the fixed camera makes segmenting the agent's body relatively straightforward. Could the authors provide additional visualized results for Meta-World and other more complex tasks? For instance, how would segmentation handle the robotic arm and objects in these more challenging scenarios?

* The segmentation process will inevitably introduce additional computational overhead during the model training. For example, the proposed method need masks of the agent body, which is obtained from a fine-tuned SAM model.

questions:
1. How does the quality of the mask affect the results of the proposed method? The authors propose a selective L2 loss during training. Can the authors provide results about using native L2 loss and selective L2 loss for training respectively?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731082747464

thread_type:
official_review

<thread_object>
creator: reviewer_3s6k

<content>
summary:
This paper proposes Segmentation Dreamer, a model-based reinforcement learning method that extends the Dreamer framework by introducing an auxiliary task focused on predicting task-related segmentation masks. Specifically, the reconstruction task in Dreamer is modified to work with images that contain task-relevant parts, rather than reconstructing the entire image. The authors evaluate the method on two RL benchmarks: DeepMind Control and Meta-World, demonstrating its effectiveness in these environments.

strengths:
1. Clear Presentation: The paper is well-written, with clear and easy-to-understand figures, text, and mathematical formulations. The explanation of the model and its components is thorough, making it accessible to both newcomers and experienced researchers in the field.

2. Intuitive Idea: The core idea of the paper—introducing task-relevant segmentation as an auxiliary task—is intuitively appealing. It draws inspiration from human perception, where individuals tend to focus on the most relevant parts of their environment when performing a specific task. This concept feels natural and aligns with cognitive science.

3. Solid Experimental Comparison: The authors conduct experiments for comparing with some strong baselines, and provide detailed ablations for SD's designs.

weaknesses:
1. Limited Scope of Contribution: One of the main weaknesses of the paper is that the contribution feels somewhat narrow. While the introduction of task-relevant segmentation masks is an interesting idea, it could potentially have a broader application beyond Dreamer or visual distraction tasks. Other model-based RL frameworks, such as transformer-based models (e.g., Trajectory Transformer), also use reconstruction as an auxiliary task, meaning the innovation may not be entirely novel. The applicability of this approach to other RL models and tasks remains underexplored, which limits its perceived impact.

2. Stability Concerns in Meta-World: The experimental results on Meta-World show large reward variance, which raises concerns about the stability of the model. This suggests that the proposed approach may struggle with robustness, particularly in more complex or dynamic environments. A deeper analysis of the reasons for this instability and potential solutions would strengthen the paper.

3. Limited Scope of Experiments: The experiments are conducted on relatively simple and controlled environments (DMC and Meta-World), which might not fully capture the complexity and variability of real-world tasks. The lack of experiments in more complex or high-dimensional settings limits the generalizability of the findings. It would be helpful to see how the method performs on more challenging tasks or in real-world applications where the environment is less predictable.

questions:
How to ensure that the segmentation masks correctly identify task-relevant parts of the image, especially when the foundation model (used for segmentation) does not inherently understand what is important for the task. In some cases, the background details may be crucial for the task, and it is not clear how the model can distinguish between what is task-relevant and what is not. How can you ensure that the masking process does not remove essential information, especially in environments where the task might require nuanced understanding of the background or context? In other words, how can the sparse rewards guide good segmentation masks prediction, when they are not enough at guiding good policies?

The variance in the rewards observed in the Meta-World experiments could be indicative of instability in the model. Can the authors provide more insights into why this happens and how this issue might be mitigated? For example, are there any hyperparameters or training strategies that could be tuned to improve stability across environments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732575026106

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear Reviewers,

We sincerely thank you for your valuable feedback and insightful comments. We are particularly grateful for recognizing our work’s contribution to addressing background distractions in visual control tasks, a well-known and challenging problem.

- Reviewer 3S6K’s positive feedback on the core idea of introducing task-relevant segmentation as an auxiliary task is especially encouraging. 

- We appreciate Reviewer RWSA and Reviewer yfcf’s acknowledgment of our efforts in providing a comprehensive evaluation of our method and baselines. 

- We are also grateful for Reviewer qM91 and Reviewer L53c’s positive comments on the selective L2 loss function and its role in increasing robustness.

We hope this discussion addresses all your concerns and further strengthens the paper. If the reviewers find our responses satisfactory, we kindly request that they consider increasing their scores by the end of the rebuttal period.


## 1. Generalization to more complex, realistic environments

While our current experiments primarily focus on background distractions, the proposed method has the potential to be extended to handle other types of distractions, including foreground occlusions, color variations, and changes in camera angle.

To better reflect real-world scenarios, we have added evaluations for our method on three additional types of visual control distractions:

1. **Foreground distractions**: Potentially occlude task-relevant objects.

2. **Color variations in foreground objects**: Simulating changes due to environmental factors like lighting.

3. **Camera angle changes**: Altering the perspective of the scene.

For sample images illustrating perturbations and detailed experimental setups, please refer to **Appendix Figures 11–14**. To handle a variety of distractions at test time, both the segmentation models and SD are trained and evaluated in domain-randomized environments.

The success of our method hinges on two key aspects, which we analyze below:

1. **Robustness of Visual Foundation Models (Segmentation Models)**: Can segmentation models remain robust to these perturbations and distractions?

- Answer: Yes, our results demonstrate that segmentation models are highly effective at identifying objects of interest, even in the presence of these distractions. This aligns with observations from our primary experiments on background distractions. Leveraging segmentation models remains a robust strategy for training RL agents under these conditions. For segmentation quality, please refer to **Figures 11–14 in the Appendix**. Additionally, we plan to include quantitative evaluations of segmentation model performance in the revision to provide further insights.

2. **Capacity of SD to Handle Perturbations**: Can SD manage perturbations that segmentation models cannot filter?

- Answer: Unlike background distractions, perturbations in color and camera angle cannot be fully filtered out and are retained in the target. While this might raise concerns about wasted capacity due to spurious information, our results (see Table below) show that SD effectively handles these perturbations. The table summarizes the test returns (mean and standard error) over four runs, where 100 data points were used for fine-tuning. Notably, our method, by focusing on agent-centric features, outperforms the standard Dreamer trained in the standard environment. Additionally, the selective L2 loss proves highly effective in handling occlusions, enabling the recovery of occluded foreground agents. This underscores the versatility and value of the selective L2 loss across multiple scenarios. While we designed SD primarily to mitigate visual distractions outside of task-relevant objects, these new results show that it is also robust to other realistic forms of distractions.

|                |                                        |                             |                                                     |           |                  |                             |
| -------------- | -------------------------------------- | --------------------------- | --------------------------------------------------- | --------- | ---------------- | --------------------------- |
|                | Standard environment (Standard Dreamer) | **Foreground distractions** | **Foreground distractions** (**No** selective loss) | **Color** | **Camera angle** | **Background with GT mask** |
| Walker         | 752 ± 9                                | 740 ± 4                     | 688 ± 37                                            | 761 ± 3   | 752 ± 10         | 719 ± 19                    |
| Cartpole Swing | 818 ± 52                               | 860 ± 2                     | 852 ± 7                                             | 870 ± 4   | 863 ± 4          | 869 ± 4                     |
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces Segmentation Dreamer (SD), an extension of the DREAMER framework designed to improve Model-Based Reinforcement Learning (MBRL) in visually distracting environments. SD focuses representation learning on task-relevant components of the visual input by using segmentation masks, simplifying latent encoding and avoiding noisy signals from irrelevant objects. The method is evaluated on tasks from the DeepMind Control and Meta-World benchmarks, demonstrating improved sample efficiency and performance.

The strengths of this paper include a clear presentation with easy-to-understand explanations and visualizations, as well as an intuitive idea of using segmentation masks to focus on task-relevant components. The main weaknesses include limited generalizability, as the experiments are primarily conducted in controlled environments, the method's dependence on prior knowledge to identify task-relevant components, and its computational cost.

This paper is decided to be rejected, with an average score of 5.4, including two borderline positive scores and three borderline negative scores.