Title: Focus On What Matters: Separated Models For Visual-Based RL Generalization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Visual-based RL

Generalization
Abstract: A primary challenge for visual-based Reinforcement Learning (RL) is to generalize effectively across unseen environments. Although previous studies have explored different auxiliary tasks to enhance generalization, few adopt image reconstruction due to concerns about exacerbating overfitting to task-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (\blue{S}eparated \blue{M}odels for \blue{G}eneralization), a novel approach that exploits image reconstruction for generalization. SMG introduces two model branches to extract task-relevant and task-irrelevant representations separately from visual observations via cooperatively reconstruction. Built upon this architecture, we further emphasize the importance of task-relevant features for generalization. Specifically, SMG incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting. Extensive experiments in DMC demonstrate the SOTA performance of SMG in generalization, particularly excelling in video-background settings. Evaluations on robotic manipulation tasks further confirm the robustness of SMG in real-world applications. Source code is available at \url{https://anonymous.4open.science/r/SMG/}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nVisual-based Reinforcement Learning (RL) often fails to generalize across unseen environments. This work proposes SMG (Separated Models for Generalization) to improve the generalization in VRL by introducing two models to separately extract task-relevant and task-irrelevant representations through image reconstruction. Specifically, SMG proposes two additional consistency losses on relevant features, improving generalization. Extensive experiments, including video-hard DMC, color-hard DMC, and manipulation tasks, show SMG excels in diverse settings and tasks, demonstrating robust performance.\n\nSTRENGTHS:\n- Separating foreground and background for reconstruction makes sense for improving the generalization in VRL.\n\n- Extensive experiments in various experimental settings demonstrate the effectiveness of SMG.\n\n- The learned mask looks very effective (Fig. 3 and Fig. 7).\n\nWEAKNESSES:\n- Distinguishing between controllable and uncontrollable parts for learning a mask model has been widely discussed in the community, like TIA [1], Denoised MDP [2], ISO-Dream [3] and so on. Although I appreciate authors' efforts to discuss its difference against TIA (appendix E.2), I think the novelty of learning mask models to distinguish noise from the environment is limited. Nevertheless, I believe that this paper has made contributions in applying mask models to the field of visual RL generalization.\n\n- I'm curious about the performance of the proposed method in some more challenging settings, like RL-Vigen [4].\n\n- As there are many losses, it is better to add a detailed pseudo code about how to calculate all these losses, which can make the paper more readable.\n\n- This proposed SGM is considered to be seamlessly combined with any existing off-policy RL algorithms. As the experiments mainly consider SAC as the RL backbone, I'm curious about its performance with other methods, like DrQ or SVEA.\n\n- The related work part only discusses observation generalization in RL and some other types of generalization also should be discussed, like dynamic generalization [5,6] and task generalization [7,8].\n\nOverall, I lean toward boardline of this work. I will participate in subsequent discussions and would like to adjust my scores, especially for the response to my concerns about experiments.\n\n[1] Learning Task Informed Abstractions\n\n[2] Denoised MDPs: Learning World Models Better Than the World Itself\n\n[3] Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models\n\n[4] RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization\n\n[5] Context-aware dynamics model for generalization in model-based reinforcement learning\n\n[6] Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability\n\n[7] Zero-shot task generalization with multi-task deep reinforcement learning\n\n[8] Task Aware Dreamer for Task Generalization in Reinforcement Learning\n\nQUESTIONS:\n- How do you determine the hyperparameter $\\rho$? In Fig.9, this paper shows different results about $\\rho$ in walker-walk. Are there more results to show the relationship between SGM and $\\rho$?\n\n- Why do reconstruction-based methods benefit generalization? Are there any explanations?\n\n\n---------------\n\n**After reading the authors' response and other reviewers' comments, I have raised my scores from 4 to 5.**\n\nLIMITATIONS:\nYes, this work has discussed its limitations in Sec. 6.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel method that utilizes two model branches to extract task-relevant and task-irrelevant representations separately from visual observations, aiming to enhance the zero-shot generalization ability of RL agents. The approach introduces four additional loss terms and two consistency losses to guide the agent's focus towards task-relevant areas across different scenarios. The proposed method can be seamlessly integrated into existing standard off-policy RL algorithms as a plug-and-play module. Experimental results demonstrate the effectiveness of the proposed model on two environments, surpassing previous benchmarks such as SAC and DrQ.\n\nSTRENGTHS:\n1. This paper is clearly written and easy to follow.\n2. Based on the separated models architecture, this paper proposes multiple effective loss functions to focus on task-relevant features in visual-based RL generalization.\n3. The authors provide detailed validations on the DMC environment and robotic manipulation tasks. They demonstrate the advantages of the proposed loss terms across multiple tasks in DMC (Table 3) and showcase the state-of-the-art performance of SMG (Table 1, 2).\n\nWEAKNESSES:\n1. While the paper compares the performance with model-free RL methods, it would be beneficial to also include a comparison with model-based RL methods. Previous works such as DreamerPro [1], Iso-Dream [2], and Denoised-MDP [3] have addressed visual distractions to enhance the generalization ability of RL agents.\n\n[1] Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations.\n\n[2] Iso-Dream: Isolating Noncontrollable Visual Dynamics in World Models.\n\n[3] Denoised mdps: Learning world models better than the world itself.\n\n2. The paper lacks sufficient discussion and analysis of its limitations. \n3. The serial numbers in some figures appear to be somewhat disorganized.\n\nQUESTIONS:\nPlease see the weaknesses section.\n\nLIMITATIONS:\nThis paper discusses the limitations, but that is not enough.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel approach called SMG (Separated Models for Generalization) to improve generalization in visual-based reinforcement learning (RL). The approach works by using separate foreground and background encoders/decoders and employing a mask to isolate task-relevant regions. In addition, it also applies four additional losses(mask ratio, background, Q-value and empowerment losses) to to enhance the model’s ability to distinguish between two types of representations. To make the learned models generalize to different visual styles, it introduces attribution augmentation and consistency losses. The authors position this as a plug-and-play method that can enhance existing RL algorithms' generalization capabilities.\n\nExperiments show SMG outperforms baseline methods, particularly in video-background settings where it maintains performance even with significant visual changes. Ablation studies validate the importance of each component.\n\nThe main contributions are:\n\n- SMG: A separated model architecture with two branches to extract task-relevant and task-irrelevant representations from visual observations.\n\n- Two consistency losses to guide the agent's focus on task-relevant areas across different scenarios.\n\n- Strong performance on DMControl benchmark tasks, especially in video-background settings.\n\nSTRENGTHS:\nThis paper has several strengths:\n\n- SMG achieves state-of-the-art performance on the DMControl Generalization Benchmark, particularly excelling in the challenging video-background settings. This demonstrates the practical effectiveness of the approach.\n\n- SMG is a plug-and-play method that can enhance existing RL algorithms' generalization capabilities. It is designed to be easily integrated with existing off-policy RL algorithms, enhancing its practical value and potential for wide adoption.\n\n- This paper includes detailed ablation studies that validate the importance of each component in the SMG architecture, providing insights into the method's workings.\n\n- This paper is well-written. And it also provides clear visualizations of the reconstruction process, helping readers understand how SMG extracts and utilizes task-relevant features.\n\nWEAKNESSES:\nThis paper has several weaknesses:\n\n- My major concern is the overclaim made by this paper. While it claims to address the generalization gap in visual-based reinforcement learning, the method proposed primarily tackles scenarios where only the backgrounds differ. However, visual generalization challenges are more diverse and include variations such as different lighting conditions and textures, which are common in real-world robotics applications. These scenarios appear to be overlooked in this paper.\n\n- SMG introduces a lot of loss terms and associated hyperparameters, which could complicate tuning in practical applications.\n    - Specifically, the mask ratio $\\rho$ appears to be crucial for performance, as it is the sole factor preventing the model from classifying everything as foreground. Given that $\\rho$ represents the ratio between the foreground and the entire image, it likely necessitates per-task tuning, which could prove to be challenging and not scalable.\n\n- The foreground consistency loss, as discussed in Section 3.3, heavily depends on the predicted mask to construct the augmented observation. During the initial stages of training, this process relies on potentially inaccurate mask predictions and attributions. Although the authors describe this as a bootstrapping process, further analysis regarding its stability and potential failure modes would be beneficial.\n\n- The paper could be strengthened by considering a broader range of baselines. For example:\n    - Recent studies [1] suggest that visual encoders pre-trained on large-scale image datasets can improve the visual robustness of a policy. This paper does not make any comparisons with visual pre-training methods.\n    - Large vision foundation models like SAM [2] could potentially be utilized to provide supervision for generating foreground masks. Would this approach be more effective than training a mask predictor from scratch?\n\n\n- The additional computation overhead introduced by the extra modules is concerning.\n    - The architecture, which involves separate models, essentially doubles the number of parameters compared to baseline methods. Although the authors argue that the performance improvements are due to the novel architecture rather than the increased number of parameters, this could still be problematic for practical applications with limited computational resources.\n    - Training time: The reported wall time for SMG is significantly longer than that of the baseline methods (22 hours versus 8-13 hours for 500,000 steps).\n\n[1] Hansen, Nicklas, et al. \"On pre-training for visuo-motor control: Revisiting a learning-from-scratch baseline.\" arXiv preprint arXiv:2212.05749 (2022).\n\n[2] Kirillov, Alexander, et al. \"Segment anything.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.\n\nQUESTIONS:\nSee the \"Weaknesses\" section. Some additional questions are noted below:\n\n- The terminology used for \"foreground\" and \"background\" is somewhat confusing. To clarify, \"foreground\" actually refers to the task-relevant parts of the image, while \"background\" refers to the task-irrelevant parts, correct?\n- The necessity for background reconstruction is unclear. The authors claim that \"improving the precision of background prediction can consequently enhance the foreground as well,\" but a more detailed explanation of this assertion would be beneficial.\n\nLIMITATIONS:\nThere is no separate section for limitations in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel objective to improve robustness of the visual encoder in RL to background noise and to color perturbations. First, the authors split the visual encoder into two models: background encoder/decoder and foreground encoder/decoder. The proposed training objective contains multiple components: \n- overall reconstruction loss that combines outputs of the background and the foreground decoders modulated by a mask;\n- mask ratio loss that prevents the foreground mask from taking up too much of the image;\n- background reconstruction loss that uses the learned mask to generate a new data sample;\n- q-value loss that makes the foreground representation capture value information;\n- empowerment loss that makes the foreground representations capture the agent actions;\n- foreground and q-value consistency losses that make sure that changing the background (using the learned mask) doesn't change the foreground features and q-values\n\nThe method is tested on DMC generalization benchmark and on robotic manipulation tasks.\n\nSTRENGTHS:\n- The method performs really well with various distractors;\n- The idea of re-using the learned masks for augmentations is interesting and, as far as I can tell, novel;\n\nWEAKNESSES:\n- The writing is a bit sloppy, with many typos and confusing sentences;\n- The resulting objective is too complex and has too many terms;\n- No comparison to TIA, although the presented method is quite similar. Was that because you only compare to model-free methods?\n\nTypos (some of them, I didn't write down all of them, please run a spell checker on the text):\nline 13: achieving free from overfitting : not clear what this means\nline 38: further strengths -> further strengthens\nline 100: focused in -> focused on\nline 536: Comparision -> comparison\n\nQUESTIONS:\n- In your objectives, you're maximizing the mutual information between foreground representations of two consecutive states and the action that was taken between them. Have you tried minimizing MI between background representations and actions and or rewards? This could be done with information bottleneck method for example. If you haven't tried this, do you this can help?\n\nLIMITATIONS:\nThe authors have adequately described limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}