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

[Q1] Comparison with Other Methods Utilizing Segmentation Models

• Prior works [1, 2] have indeed used segmentation models with task-specific prior knowledge in RL. However, these approaches utilize segmentation for preprocessing input observations, which can be unstable at test time. If the segmentation model fails during deployment—such as when encountering unfamiliar scenarios—it disrupts the input and causes the agent to malfunction. In contrast, our method leverages segmentation models as an auxiliary task, significantly improving robustness to segmentation errors. This advantage is evident in Table 1 and Figure 9, where the “As Input”  method represents these prior approaches, highlighting their vulnerability in untested environments. By using segmentation as an auxiliary task rather than a preprocessing step, our design avoids these pitfalls and ensures greater reliability and efficiency while addressing the same core problem.

• References:

• [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 (ECCV), 2024.

[W1] Generalization to other domain shifts

• Thank you for this thoughtful observation. We conducted additional experiments with three types of distractions: changes in camera angles, foreground occlusions, and color variations. Please refer to General Response 1 for further details.

[Q2] Extension of the Proposed Method to Handle Cases Like Viewing Angle, Color, Lighting, and Texture Changes

• While our method demonstrates robustness under domain randomizations, such as variations in camera angles, foreground occlusions, and color changes, we acknowledge that more rigorous perturbations or combinations of distractions require further exploration. Enhancing invariance in these scenarios is a promising direction for future work. For example, decoding canonical representations [7] as an additional auxiliary task could promote invariance to foreground perturbations. Similarly, incorporating a contrastive loss to map similar behaviors of different foreground objects closer together in the feature space could improve robustness against foreground variations. We plan to explore these approaches to further enhance the generalizability and applicability of our method to more complex, real-world scenarios.

• [7] James, Stephen, et al. "Sim-to-real via sim-to-sim: Data-efficient robotic grasping via randomized-to-canonical adaptation networks." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.

[W1, Q1] Comparison with existing works that combine foundation models with RL

• Thank you for suggesting these related works. Below, we summarize the key distinctions and advantages of our approach:

  • Use of Segmentation Models: Existing works [1, 2] utilize segmentation models for preprocessing input observations, meaning their success is heavily dependent on the quality of the segmentation model at test time. If the segmentation model fails during deployment (e.g., when encountering unfamiliar scenarios), it can disrupt the input and cause the agent to malfunction. In contrast, our method uses segmentation models for an auxiliary task, making the agent more robust to segmentation errors. The instability of using processed inputs is demonstrated in Table 1 of our paper, where “As Input” represents this family of prior approaches.

  • Training Overhead: Because the quality of segmentation models is critical, methods like [2, 5] require substantial effort to fine-tune segmentation models, often involving building simulations to generate synthetic datasets for fine-tuning. This introduces considerable time and resource overhead. In contrast, our approach can effectively use out-of-the-box segmentation models with minimal data (e.g., 1, 5, or 10 examples). Because we use segmentation masks for an auxiliary target, our method remains resilient to segmentation errors and eliminates the need for extensive fine-tuning during training.

  • Inference Efficiency: Existing methods that rely on segmentation models for preprocessing add computational overhead during test time. This is undesirable for real-world scenarios where lightweight, efficient inference is critical, particularly in resource-constrained or on-device settings. Our method avoids this issue by limiting segmentation only to training, ensuring no additional overhead during deployment while retaining robustness.

• While prior works have strengths such as simplifying input observations when the segmentation is accurate, our method uniquely balances robustness, simplicity, and efficiency. Additionally, we acknowledge that incorporating ideas from related works, such as instance-level segmentation and tracking (e.g., [3, 4]), could enhance segmentation quality and make our method even more robust. Using canonical representation as an input can be especially robust to foreground variations which we can also use as another auxiliary task to be invariant to foreground perturbations. We plan to explore these directions in future work.

• We have updated the paper to include additional prior work and provide a clearer picture of how our method compares. Please refer to Section 2 (marked in red).

• [5] John So, Amber Xie, Sunggoo Jung, Jeffrey Edlund, Rohan Thakker, Ali Agha-mohammadi, Pieter Abbeel, and Stephen James. Sim-to-real via sim-to-seg: End-to-end off-road autonomous driving without real data. In Conference on Robot Learning, 2022.

[W2, Q2] Generalizability to Real-World Tasks with Complex Distractions

• Thank you for raising this insightful point. To demonstrate the promise of our method in more realistic scenarios, we evaluated our method against three additional types of distractions, including occlusions. Please refer to General Response 1 for further details.

• While our method effectively handles a wide range of distractions, we acknowledge its limitations in scenarios where distractions closely resemble task-relevant features, as the segmentation model lacks knowledge of controllability. A promising direction for future work is to incorporate controllability knowledge into either the segmentation model or the RL agent, enabling them to differentiate between controllable and uncontrollable features. Another viable direction is to introduce a flag as an additional auxiliary task to explicitly identify controllable foreground objects, combined with video segmentation models capable of object tracking, as demonstrated in [2] in the context of input preprocessing. These extensions represent tangible directions to improve our method’s generalizability and applicability to tasks with more intertwined distractions and task-relevant features.

[W3, Q3] Model bias 

• Thank you for raising this important point. Figure 9 in the Appendix (“SD^PerSAM_1”) provides insight into how test-time segmentation errors might affect performance in untested environments. In this figure, the x-axis represents episodic segmentation errors at test-time, and the y-axis shows test returns. While our method does not rely on segmentation during testing, we performed segmentation ad hoc to assess its impact. Data points on the left (low test IoU) indicate high test-time segmentation errors from the model used during training. Despite these errors, our method exhibits stable performance, even in scenarios where the segmentation would have failed. This robustness highlights the method's resilience to unseen environments.

• In this paper, we also propose two strategies to mitigate the impact of segmentation errors: 

• 1) By only using segmentation predictions as training targets, we prevent failures that could be caused by deployment time prediction errors. This increase in robustness is made evident by comparing our SD^PerSAM_1 method to the “As Input” variation in Table 1 and Figure 9. “As Input” applies segmentation mask predictions to deployment time observation inputs. Our proposed method avoids using segmentation masks in deployment and is thus less sensitive to segmentation accuracy.

• 2) We leverage Selective L2 loss to attempt to skip training on regions of reconstruction targets where the segmentation model is likely incorrect. The positive effect of reducing incorrect training examples on performance is shown in Table 1, where SD^PerSAM_1 with Selective L2 loss shows a clear improvement over our Naive L2 loss ablation.

• We hope that this answers your question concerning the impact of segmentation errors on test-time performance and how the effects of poor accuracy can be mitigated. Did we address your question fully?

[Q4] Handling Dynamic Distractions and Task-Relevant Features

• Dynamic Task-Relevant Features: When task-relevant features, such as the position, shape, or appearance of foreground objects, shift during execution, our method can remain effective. If these changes are anticipated and incorporated as prior knowledge, segmentation models can be fine-tuned with larger datasets to account for such dynamics, ensuring consistent performance.

• Dynamic Distractions: In cases where distractions, such as background or foreground elements, change unpredictably or frequently, our experiments show that our method is robust to significant background variations. Furthermore, our additional experimental results highlight the method’s potential in managing dynamic foreground elements effectively.

• Looking forward, a promising direction for future work is to autonomously learn the distinction between distractions and task-relevant segmentation features directly from data, particularly in environments where these features change dynamically over time.

[Q5] Mechanism for Adapting to New Environments

• Currently, our framework does not include a mechanism to adapt to new environments or changing task dynamics without complete retraining. This limitation aligns with the broader challenge in MBRL methods, which typically require fine-tuning on the target task or pre-training with multi-task learning to adapt effectively.

• However, in our framework, if the source and target tasks share the same task-relevant objects, transfer can be facilitated. Developing a mechanism for seamless adaptation to new environments without retraining remains an intriguing direction for future work.

[Q1] The effect of the proposed l2 loss if the segmentation model improves

• As the segmentation model improves with more training samples, the performance gap between the naive L2 loss and the selective L2 loss is expected to decrease. However, this improvement comes at the cost of additional computational overhead during training, including the need for larger datasets and extended fine-tuning of the segmentation model.

• Furthermore, our additional experiments added during this discussion phase show that the selective L2 loss remains effective in scenarios involving occlusions. In these cases, SD with the selective L2 loss can recover occluded parts of foreground objects and maintain robustness. This underscores the versatility and value of selective L2 loss across various challenging conditions. We plan to include these findings in the revision to provide additional evidence of its effectiveness in different scenarios.

[Q2] The effect of mask prediction

• The binary mask predictor is trained using the standard L2 loss.

• When both the RGB and binary mask decoders are trained with the standard L2 loss, the performance significantly drops compared to our method. This is partially because there is no mechanism to correct errors in the RGB decoder, allowing incorrect signals to propagate throughout training. Additionally, binary mask prediction does not serve as an effective auxiliary task, as it tends to generate overly abstract and coarse features that fail to support downstream tasks properly.

• These findings highlights two important implications:

  • Addressing Segmentation Errors: Correcting segmentation model errors is critical for achieving robust performance. This highlights the value of the selective L2 loss, which effectively mitigates segmentation inaccuracies and prevents error propagation during training.

  • Careful Design of Auxiliary Losses: The auxiliary loss must be thoughtfully designed to balance feature abstraction and task relevance. Selective L2 loss provides this balance by selectively focusing on regions where predictions are more reliable, enhancing overall learning stability.

• We will include these insights in the revision to provide a deeper understanding of the importance of addressing segmentation errors and designing auxiliary tasks aligned with the needs of our framework.

[W1] Prior Knowledge Requirement and Its Potential Limitations

• Yes. Our method indeed assumes that task-relevant parts of image observations can be identified using prior knowledge. This aligns with how humans leverage prior knowledge, such as affordances [6], to efficiently solve tasks. We believe there are many domains where obtaining prior knowledge is straightforward, such as robotics tasks involving object manipulation or navigation, where task-relevant features are naturally evident. In such cases, our method suggests an effective and efficient interface to incorporate this prior knowledge, significantly simplifying the learning process while maintaining minimal computational overhead.

• That said, we recognize that in domains where prior knowledge is ambiguous or difficult to obtain, this could limit the applicability of our approach. A promising direction for future work is to automatically learn task-relevant knowledge from data during training, using segmentation maps as an intermediate representation. This would expand the applicability of our method to more complex environments where task-relevant components are not explicitly known.

• [6] Shikhar Bahl, Russell Mendonca, Lili Chen, Unnat Jain, Deepak Pathak. Affordances from Human Videos as a Versatile Representation for Robotics. In  Computer Vision and Pattern Recognition, 2023.

[W2] Simpler problem setup?

• By only reconstructing task-relevant portions of images we are indeed simplifying the functionality that the world model has to learn. This is an intended benefit of our proposed method that we see as a strength. Our comparisons with baselines like standard Dreamer show the benefit of leveraging prior knowledge to create a purposefully more tractable learning task via task-relevant reconstruction. By leveraging task-relevant reconstruction to only model the necessary components of observations, we can train a world model and agent more robust to task-irrelevant visual perturbations.

• Moreover, prior works [1, 2] have also leveraged segmentation models with task-specific prior knowledge to inform segmentation. However, these methods use segmentation for preprocessing input observations, which has been shown to be unstable at test time. This instability is evident in Table 1, where “As Input” represents this family of methods, demonstrating their vulnerability in untested environments. Our careful design eliminates these limitations, ensuring robustness and efficiency while solving the same underlying problem.

[W3] Limited to visual control task

• While our method leverages segmentation models specifically for visual control tasks, these tasks encompass a wide range of applications, including object manipulation and navigation. Visual observations are a commonly used multimodal representation for perceiving the world, making our approach valuable for shaping features to focus on task-relevant image components in such scenarios.

• To evaluate the performance of adapted segmentation models in downstream domains, we suggest two approaches:

  • Qualitative Evaluation: By inferring the fine-tuned segmentation model on test scenarios, one can visually inspect segmentation quality. This provides insights into how well the RL agent might perform, given the observed correlation between segmentation quality and RL test-time performance (Fig. 3b).

  • Quantitative Evaluation: Creating an evaluation set to quantify segmentation model performance offers a practical way to analyze its quality and assess its adaptation to downstream tasks as part of ad-hoc analysis, as illustrated in Fig. 9 of the Appendix.

[W1] Segmentation Quality in More Challenging Scenarios

• To address concerns about generalization, we conducted additional experiments on DMC with three different distraction types, including camera view changes. Both the segmentation models and the SD RL algorithm demonstrated robustness to these perturbations. Please refer to General Response 1 for details.

• Additionally, Figures 16–18 in the Appendix show segmentation results on more complex tasks in Meta-World, involving multiple objects and occlusions. The segmentation models were fine-tuned using only 10 data points, yet they performed robustly under challenging conditions such as unseen poses and occlusions (e.g., robotic arm occluding itself or objects). This robustness can be attributed to the models’ pre-training on large-scale datasets. Naturally, some segmentation errors still occur, but these are effectively managed in our framework using selective L2 loss.

[W2] Training time computational overhead

• Yes. Incorporating segmentation models during training introduces additional computational overhead. To quantify this, we conducted an analysis on DMC-1M using an NVIDIA RTX A4500 and SegFormer with the MiT-b0 backbone, as used in our main experiments:

  • Training Time: Standard Dreamer required 539 minutes (≈9 hours), while our method took 813 minutes (≈13.5 hours), reflecting a ~50% increase. These numbers are averaged over 8 runs. One important point to mention is that our current implementation is relatively unoptimized, where images from experience collection are each sequentially processed individually by the segmentation model. With further code optimizations, such as batch inference and asynchronous image processing, we expect a significant acceleration in training time. For example, using batch inference with a batch size of four (matching the number of simulators running), the additional training time would reduce to approximately 68.5 minutes, corresponding to only a ~12% increase instead of the current ~50%.

  • VRAM Usage: Standard Dreamer utilized 3200MiB VRAM, while our method required an additional 1113MiB, a ~35% increase.

• We believe this computational overhead during training is a reasonable trade-off for achieving a significantly better policy. Importantly, our method introduces no additional overhead during test time, unlike approaches that incorporate segmentation modules for preprocessing in the inference pipeline. Test-time efficiency is particularly critical in scenarios when delays are unacceptable or on-device inference is required, where additional computational demands are highly undesirable. By using segmentation masks for an auxiliary target during training rather than for preprocessing inputs, our approach avoids these issues while maintaining effectiveness.

[Q1] How does the quality of the mask affect the results of the proposed method?

• Figure 3b plots the training-time segmentation quality against the RL agent’s test-time performance. We observe that better segmentation tends to lead to higher RL performance, as accurate targets better highlight task-relevant components. This suggests that improved segmentation models can enhance agent performance without ground-truth masks.

[Q2] Results about using native L2 loss and selective L2 loss for training respectively?

• Table 1 demonstrates that using the selective L2 loss (SD^PerSAM_1) consistently outperforms the naive L2 variants, particularly in complex tasks such as Cheetah Run and Walker Run. Segmentation models often miss embodiment components. With naive L2 loss, the model replicates these errors, leading to incomplete latent representations and harming dynamics learning. In contrast, the selective L2 loss addresses this issue by skipping L2 computation in regions where PerSAM targets are likely incorrect. Additionally, our experiments on foreground distractions further reinforce the effectiveness of the selective L2 loss in enhancing robustness when occlusions affect the foreground agent.

[W3] Evaluation on more challenging tasks or in real-world applications

• While DMC is a simulated and controlled environment, solving tasks with distractions in these settings remains a challenging problem. Many prior works on DMC struggle significantly, especially when rewards are sparse, whereas our method demonstrates strong performance under such conditions.

• Beyond DMC, we evaluate our method on Meta-World, a widely recognized and challenging benchmark for RL. Meta-World involves tasks requiring object manipulation with multiple objects, frequent occlusions, and varying object poses, all of which introduce significant complexity. Meta-World resembles real-world scenarios, where agents must handle dynamic interactions and multi-step processes to achieve their goals.

• To further evaluate our method’s robustness, we also present preliminary experiments on three additional types of distractions to increase environmental complexity. Please refer to General Response 1 for details. Expanding to real-world applications, where unpredictability and variability are even greater, is a promising direction that we plan to explore in future work.

[Q1] Ensuring Segmentation Masks Identify Task-Relevant Parts

• In our current framework, the foundation model does not inherently understand task relevance. Instead, it is informed about the task using a few labeled examples provided before the RL training loop begins. Our paper focuses on designing the most effective way to integrate MBRL with segmentation models, given examples of a few correct task-relevance masks as prior knowledge.

• By reconstructing agent-centric parts of the observation with the help of the segmentation model, our method successfully solves tasks even in scenarios with sparse rewards, where other baselines fail completely (see Fig. 3a, Cartpole Swing Sparse). 

• Making the segmentation model inherently aware of task relevance and enabling it to learn and adjust task-relevant features directly from rewards during RL training remains an exciting direction for future research.

[W1] Broader application beyond Dreamer or visual distraction tasks

• The focus of our paper is on addressing distractions in visual control tasks, a highly challenging problem that reflects the complexities of real-world scenarios. Numerous prior works (e.g., Deng et al., 2022; Nguyen et al., 2021; Zhang et al., 2021) have also concentrated exclusively on this specific issue.

• That said, we agree that demonstrating the versatility of our method would make it even more impactful.

  • Applicability to Other MBRL Frameworks: Indeed, our auxiliary task is not limited to Dreamer and could potentially be integrated into other reconstruction-based model-based RL (MBRL) frameworks, including transformer-based models such as Trajectory Transformers (Janner et al., 2021) or STORM (Zhang et al., 2023). While these frameworks rely on reconstruction as an auxiliary task, they likely struggle in distracting environments because naive reconstruction preserves unnecessary information, hindering agent learning. Our auxiliary task can address this issue by enabling those methods to focus on agent-centric features and mitigate distractions through feature shaping with segmentation models. Since our approach is orthogonal to advancements in MBRL architectures, it serves as a complementary enhancement. We see great potential for future work to apply our method to other architectures.

  • Use Beyond Visual Distraction Tasks: The idea of using segmented images as auxiliary targets extends beyond addressing visual distractions. For instance, it can facilitate object-centric representation learning by reconstructing the segmentation of individual objects. This concept has been explored by FOCUS (Ferraro et al., 2023), which demonstrated its effectiveness using ground-truth masks. We believe that by leveraging segmentation models, our approach can support disentangled representation learning, even when ground-truth masks are unavailable. Additionally, it has the potential to address occlusions often encountered in object manipulation tasks through the use of selective L2 loss.

• In summary, while this work specifically focuses on visual distractions, its underlying concepts are versatile and applicable to a wider range of tasks and frameworks, offering promising directions for future research.

[W2, Q2] Reward variance in Meta-world

• Thanks, this is a good point that we will elaborate on better in subsequent revisions. In figure 5, while “SD^SegFormer_10” exhibits higher performance than other baselines, it also shows slightly more variance. The primary contributing factor to this variance is the accuracy of the segmentation model we pair with Segmentation Dreamer. We can improve the variance by A) using a better segmentation foundation model, B) using more finetuning image/mask examples (we used 10 in this case), and C) using a more principled approach to select a diverse set of representative finetuning image/mask pairs. In our current experiments, we selected finetuning data points by rolling out random behavior in each environment and selecting random frames from the trajectories. This selection method already effectively showcases the benefit of Segmentation Dreamer over other approaches; however, refining the data selection process could further enhance performance stability which we leave for future work.

• Importantly, a demonstration of Segmentation Dreamer performance and variance with a better (perfect) segmentation mask model is shown “SD^GT”.

Dear reviewer 3S6K,

Thanks for your efforts in reviewing this paper. Could you please let us know if the authors have addressed your comments? Or if you still have any other concerns?

Best regards, Your AC

Thanks for the responses. I see the authors provide some justifications for the scope of this paper, but this does not convince me. The paper's contribution is still limited when only focusing on Dreamer and current simulated distraction environments. Besides, simply introducing pre-trained segmentation models without letting the model be aware of task information also makes the method less insightful and too straightforward. This does not address the fundamental problem and challenges in visual control tasks with distractions. Therefore, I decide to keep my score.