Focus On What Matters: Separated Models For Visual-Based RL Generalization

Authors: di  zhang
bowen  lv
hai  zhang
feifan  yang
junqiao  zhao
hang  yu
chang  huang
hongtu  zhou
chen  ye
changjun  jiang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 05:16:54

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/}.

Main paper: 
 introduction: 
 Visual-based Reinforcement Learning (RL) has demonstrated remarkable success across various tasks, including Atari games [27, 11, 18] , robotic manipulation [23, 9] , and autonomous navigation [26, 46] . However, deploying visual-based RL algorithms in real-world applications requires a high generalization ability due to numerous factors that can induce distribution shifts between training and deployment scenarios, such as variations in lighting conditions, camera viewpoints, and backgrounds. Many visual-based RL algorithms are prone to overfitting to the training observations [5, 34, 44] , limiting their applicability in scenarios where fine-tuning with deployment observations is not allowed.To address the generalization gap in visual-based RL, current studies primarily focus on utilizing data augmentation techniques [19, 20, 33] and exploring various auxiliary tasks [12, 3, 13] . However, few of the previous works successfully incorporate reconstruction loss to this field, which is commonly adopted in standard visual-based RL settings and has been demonstrated to improve the sample efficiency of RL agents [41, 10, 7] . This is because reconstructing the entire input observation can exacerbate the overfitting problem to task-irrelevant features and thus weaken the generalization ability. Although several works also explored extracting task-relevant features from visual observations [6, 39, 45] , little attention has been paid to the potential of leveraging these features in improving generalization.In this paper, we propose SMG (Separated Models for Generalization), a method that utilizes a reconstruction-based auxiliary task to extract task-relevant representations from visual observations and further strengthens the generalization ability of RL agents with the help of two consistency losses. The core mechanisms behind SMG can be summarized in two parts: First, we introduce two model branches to disentangle foreground and background representations underlying in the visual observations. This separated model framework circumvents the risk of overfitting task-irrelevant features inherent in a single model structure by prudently designing the reconstruction paths, allowing our model to benefit from reconstruction loss without sacrificing generalization ability. Second, we introduce two consistency losses to align the agent's focus on the task-relevant features between raw and augmented observations. This approach enables the foreground model to extract more robust task-relevant representations, which substantially boost the generalization capability of RL agents across diverse deployment scenarios.We evaluate SMG's effectiveness across a range of challenging visual-based RL tasks, including five tasks from DMControl [36] and two more realistic robotic manipulation tasks [17] . We also adapt different evaluation settings with random-color and video-background modifications. Through comparisons with strong baseline methods, SMG demonstrates state-of-the-art performance in terms of generalization, particularly showcasing superiority in video-background settings and robotic manipulation tasks.In summary, the main contributions of this paper are as follows:• We present SMG, a novel approach that aims to enhance the zero-shot generalization ability of RL agents. SMG is designed as a plug-and-play method that seamlessly integrates with existing standard off-policy RL algorithms.• SMG emphasizes the significance of task-relevant features in visual-based RL generalization and successfully incorporates a reconstruction loss into this setting.• Extensive experimental results demonstrate that SMG achieves state-of-the-art performance across various visual-based RL tasks, particularly excelling in video-background settings and robotic manipulation tasks. 
background: 
 A Markov Decision Process (MDP) can be defined as a tuple (S, A, p, r, γ), where S is the state space, A is the action space, p : S × A × S → [0, 1] is the state transition probability function, r : S × A × S → R is the reward function, and γ ∈ [0, 1] is the discount factor. At each time step t, the agent receives a state s t ∈ S, selects an action a t ∈ A, and then receives a reward r t ∈ R. The agent's goal is to learn a optimal policy π(a t |s t ) that maximizes the expected returnE (st,at)∼ρπ [ ∞ t=0 γ t r t ],where ρ π defines the discounted state-action visitation of π. Learning an optimal policy from visual observations poses a substantial challenge for RL agents due to the inherent partial observability of the environment, a characteristic of POMDPs (Partially Observed MDP). For one thing, at each timestep t, the visual observation o t can only capture partial information about the true state s t , as certain elements may be obscured in the image. For another, the dimension of o t is much higher than that of s t , which makes it difficult to utilize o t directly for policy learning.To infer the true underlying state from visual observations, existing methods usually employ a parameterized encoder f to map a stacked frame sequence x t = (o t ′ , o t ′ +1 , ..., o t ) to a compact low-dimensional latent vector z t , which is then used as input by policy and value function. However, training the encoder solely to rely on the reward signal is demonstrated to sample inefficiency and may lead to suboptimal performance [41] . To tackle this issue, various auxiliary tasks have been proposed to enhance encoder training, with one common choice being to extract features from pixels via image reconstruction loss [7, 21, 2] . By adding another parameterized image decoder g, the reconstruction loss is defined by maximizing the likelihood function: Learning to generalize is hard for RL agents, particularly when utilizing an image reconstruction loss. While images are rich in information, requiring the agent to reconstruct the entire input observation can lead the autoencoder network to overfit to features that are unrelated to the task (e.g. colors, textures, and backgrounds). In contrast, humans can accurately figure out what matters visually when learning a new task. Even when colors or backgrounds are changed, humans can still leverage the prior knowledge to complete the task by focusing on task-relevant features. Considering a robotic manipulation task where the agent must move the arm to the red target (Figure 2 ), despite variations in background colors and textures across four test scenarios on the left, only the arm's orientation and the target position should be focused on this task. We aim for our RL agent to learn an optimal policy that solely relies on these task-relevant features while disregarding irrelevant regions.EQUATIONFormally, we decompose the latent representation z t into task-relevant part z + t and task-irrelevant part z - t . These two representations are independent, as p(zt |o t ) = p(z + t |o t )p(z - t |o t ).The task-relevant representation can be further subdivided into the "control-relevant" part, which is directly affected by the agent's actions (the arm); and the "reward-relevant" part, which is associated with the reward signal (the arm and the target), both are crucial for policy learning. 
separated models and reconstruction: 
 The representation learning objective of SMG is to maximize the mutual information I(o t ; z t ) between the observation o t and the latent representation z t , and we further derive an image reconstruction objective incorporating the combination of task-relevant representation z + t and task-irrelevant representation z - t as follows:EQUATIONInspired by previous works [6, 30] that explore how to mitigate background distractions, we implement the reconstruction process by introducing the foreground encoder f + and the background encoder f -to extract different types of representations simultaneously, which forms a separated models architecture. We also incorporate two decoders. The foreground decoder g + is employed to reconstruct the foreground image o + t and predict a mask M t with values between (0, 1). ′ t = o + t ⊙ M t + o - t ⊙ (1 -M t ) (⊙denotes the Hadamard product), the reconstruction process is illustrated by the black arrows in Figure 1a . Notably, the area where the agent is focusing can be visualized as o + t ⊙ M t , which we term the "attribution" of the agent, formally defined as Attrib(o t ).Based on the separated models architecture, we define four additional loss terms to enhance the model's ability to distinguish between two types of representations. These include the mask ratio loss and background reconstruction loss, which supervise the model's pixel outputs; along with the Q-value loss and empowerment loss, designed to consider the two properties of task-relevant representation.Mask ratio loss. To further refine the accuracy of mask prediction, we introduce a hyperparameter ρ, termed the mask ratio, to constrain the proportion of the foreground part in the mask. As shown in Equation 3, we regard L mask as an explicit form of an information bottleneck, as the percentage ρ determines the number of pixels of o + t retained in the final reconstruction. This constraint forces f + to prioritize the task-relevant parts of the observation during encoding. Empirical results in Section 4.4 demonstrate that L mask facilitates learning a more precise mask.EQUATIONBackground reconstruction loss. Improving the precision of background prediction can consequently enhance the foreground as well. Since the foreground and background are complementary, providing supervision for the background prevents the foreground from learning all parts of the observation. Therefore, we add additional supervision to the task-irrelevant representation z - t . To achieve this, we propose a new type of data augmentation called attribution augmentation tailored for SMG, as illustrated in Figure 3b . This augmentation involves augmenting the raw observation o t with its corresponding predicted mask M t via τ attrib (o t ) = o t ⊙ M t + ϵ ⊙ (1 -M t ), where ϵ represents a randomly sampled image. This simulates the video-background setting in deployment scenarios. We define the background reconstruction loss L back as follows:EQUATIONQ-value loss. Recall that the task-relevant representation z + t has two key properties: reward-relevant and control-relevant. Satisfying the former is relatively straightforward, as the representation z + t is used for policy learning. Through the Bellman residual update objective [35] outlined in Equation 5, z + t will progressively enhance its correlation with the reward signal.EQUATIONEmpowerment loss. For the control-relevant property, we integrate an empowerment term I(a t , z + t+1 |z + t ) [28] based on conditional mutual information, which quantifies the relevance between the action and latent representation. Maximizing the empowerment term further leads to maximizing a variational lower bound q(a t |z + t+1 , z + t ) as shown in Equation 6. This objective necessitates that a t is predictable when two neighboring representations are known. We implement this objective by incorporating an inverse dynamic model.EQUATIONThe whole separated models architecture is shown in figure 1a .Utilizing the separated models architecture, SMG can successfully extract task-relevant representations from raw observations. Nevertheless, the agent still lacks the ability to generalize effectively and may struggle to extract meaningful features from scenarios with transformed styles. To address this issue, we treat the task-relevant representation under raw observations as the ground truth and train SMG on more diversely augmented samples. Instead of directly optimizing the distance between the representations under raw and augmented observations, we introduce two types of consistency losses, considering both attribution and Q-values for more explainable supervision. By doing so, the foreground model can learn to extract task-relevant representations across different deployment scenarios.Foreground consistency loss. To force the agent to focus on the same task-relevant area in transformed scenarios, we train the foreground models to predict the attribution under augmented observation Attrib(τ (o t )) with the supervision of the ground truth attribution Attrib(o t ) (as Attrib(o t ) is relatively easier to converge to an accurate value, and we discuss it in detail in Appendix F). The foreground consistency loss L fore_consist is defined as Equation 7(where sg means the stop-gradient operation).EQUATIONQ-value consistency loss. In addition to the attributions, the Q-values obtained from transformed observations also exhibit high variance [14] , indicating instability in both the extracted representations and the Q function. To address this, we regularize the Q-values under augmented observations to be consistent with those under raw observations, as shown in Equation 8. This approach also regularizes the agent to learn an accurate task-relevant representation, as the gradient of L q_consist is back-propagated to the latent space.EQUATIONThe above two consistency losses are illustrated in Figure 1b .Our proposed separated models architecture can seamlessly integrate as a plug-and-play module into any existing off-policy RL algorithms. In this work, we leverage SAC [8] as the base algorithm. Throughout the training phase, SMG iteratively performs exploration, critic update, policy update, and auxiliary task update. We define the critic loss L critic as the sum of the Q-value loss L q and the Q-value consistency loss L q_consist :EQUATIONAdditionally, the auxiliary loss L aux comprises five previously mentioned loss terms:EQUATION)Although L aux contains five loss terms, experimental results show that using average weights for the first four terms and a smaller weight for the last term can achieve satisfactory performance. Detailed information about hyperparameters tuning is provided in Appendix C.3. The detailed derivation of Equation 2 and Equation 6 are provided in Appendix A.4 Experimental Results We benchmark SMG against the following baselines: (1) SAC [8] , serving as the foundational algorithm for all other baselines; 
setup: 
 (2) DrQ [19] , utilizing random shift augmentation; (3) SODA [12] , incorporating a consistency loss on latent representations; (4) SVEA [14] , focusing on stabilizing Q-values;(5) SRM [15] , proposing a novel data augmentation technique; (6) SGQN [3] , the previous SOTA method integrating saliency maps into RL tasks. We reproduce the results using the same settings reported in the original papers, with the exception of setting the batch size to 64 for all methods. Additionally, all results are calculated by four random seeds. To achieve stable performance across various evaluation settings, we train SMG using a hybrid data augmentation approach for τ (o t ), involving random overlay [14] and attribution augmentation for all tasks (each time we randomly select a type of data augmentation, as shown in Figure 3 ). The network design for SMG and more detailed experiment settings are reported in Appendix C.Evaluation ObservationWe first conduct experiments on five selected tasks from DMControl [36] and adopt the same evaluation setting as DMControl Generalization Benchmark [12] (DMC-GB) used, which contains random-colors and video-background modifications across four different levels: color-easy, colorhard, video-easy and video-hard. Figure 5 shows an example in walker-walk task. We train all methods for 500k steps (except walker-stand for 250k, as it converges faster) on the training setting and evaluate the zero-shot generalization performance on the four evaluation settings. The predicted attribution (the fifth column) highlights the extracted task-relevant area, which shows SMG accurately depicts the attribution of the input observation while omitting the task-irrelevant elements such as the skybox, the floor, and even the random color variation. This indicates that the task-relevant representation z + t contains only the information required to accomplish the task, which is crucial for generalization. Note that we aim to maintain the similarity between Attrib(τ (o t )) and Attrib(o t ), even in random-color settings. As shown by the first row of color-hard setting, SMG predicts a yellow attribution despite the input evaluation observation being orange. ±30 ±26 ±40 ±13 ±24 ±44 ±46 124%Table 1 reports the generalization performance of SMG and all baseline methods with the videobackground modification, which is the most challenging evaluation setting. The table shows that SMG outperforms all baselines in all ten tasks. Particularly impressive is SMG's superiority in video-hard; when removing the floor and the walker's shadow, the performance of all baseline methods drops significantly. However, SMG is less affected by this substantial distribution shift and maintains a stable performance across all tasks, with episode returns boosted more than 160 over the second-best in four out of five tasks (as walker-stand is a much easier task to train), showcasing its exceptional generalization capability.To further validate SMG's applicability to more realistic tasks, we conduct experiments on two goal-reaching robotic manipulation tasks [17] , including peg-in-box and reach, and following similar generalization settings used in [3] . As illustrated in Figure 6 , there are five different testing settings with different colors and textures for the background and the table. We train all methods for 250k steps and use random convolutions [22] as the data augmentation for baseline methods, as it aligns better with the testing scenarios. SMG continued to use hybrid augmentation as previously mentioned. Table 2 presents the evaluation results for peg-in-box, a task where a robot must insert a peg tied to its arm into a box. SMG achieves dominant performance across all evaluation settings, boosting an average improvement of 102% over the second-best method. Impressively, SMG exhibits remarkable stability across the six evaluation settings, with a standard deviation of only 7, while baseline methods all fail in some evaluation settings. This underscores SMG's generalization capability. These results also highlight SMG's superiority in realistic tasks, as its reconstruction-based auxiliary loss can capture more detailed features in the image, which is hard for methods that mainly rely on data augmentation techniques. In order to explore the role played by different loss terms in SMG, we conduct an ablation study in DMControl tasks. Table 3 presents the performance drop without each loss term compared to the full model in the video-hard setting. The results indicate that every loss term contributes significantly to the final performance. Notably, L q_consist exhibits the most substantial impact on performance, highlighting the importance of maintaining stable Q-value estimation in generalization tasks. Moreover, the performance drop without L back or L mask is around 20% to 30%, underlining the importance of attribution augmentation in enhancing SMG's generalization in video-background settings, as the two loss terms directly affect the quality of the attribution augmentation. Additionally, L action aids in learning a better task-relevant representation. As for L fore_consist , it also contributes to improving generalization ability, particularly in relatively challenging tasks where the performance improvement ranges from 15% to 25%. To better grasp the significance of L mask and L back in SMG, we showcase the predicted masks and their corresponding attribution augmentations in Figure 7 . When L mask is removed, the model generates an almost white mask, indicating that the foreground model overly captures irrelevant features without the constraint of mask ratio loss. Consequently, only a few parts are replaced by a random image in the attribution augmentation. In contrast, removing L back causes the background model to learn all features excessively, resulting in attribution augmentation images devoid of task-relevant information.The ablation results underscore that both L mask and L back are vital in crafting meaningful attribution augmentations, which in turn are utilized by the two consistency losses and impact the representation learning process. We conduct more experiments in Appendix E to reveal that L mask serves as a guiding factor in mask learning and SMG is not significantly influenced by variations in the hyperparameter mask ratio ρ.  
related work: 
 Improving generalization ability of RL agents has drawn increasing attention in recent years.Researchers primarily explore two aspects: using data augmentation techniques to inject useful priors when training [20, 15, 16, 22, 14, 32, 38] and employing various auxiliary tasks to guide the learning process [13, 3, 1, 42, 40, 12] . For example, Hansen and Wang [12] regularize the representations between observations with its augmented view through an auxiliary prediction task; Hansen et al. [14] stabilize Q-values via delicately design the data augmentation process; Bertoin et al. [3] introduce saliency maps to visualize the focus of Q-functions; Wang et al. [40] extract the foreground objects by employing a segment anything model. Orthogonal to existing works, we argue that focusing the RL agent on task-relevant features across diverse deployment scenarios can substantially boost the generalization capability. We propose a novel reconstruction-based auxiliary task to achieve this goal.Decision-making based on task-relevant features can substantially enhance the performance and robustness of RL agents [4, 45, 43, 29] . Bharadhwaj et al. [4] use an empowerment term to distill control-relevant features from the task; Zhu et al. [45] bolster the resilience of RL agents by regularizing the posterior predictability; Zhang et al. [43] learns compact representations by bisimulation metrics. Additionally, methods utilizing separated model architectures to extract different types of features simultaneously have been proposed [6, 39, 30, 25, 37] . For instance, Wang et al. [39] decompose the latent state into four parts based on their interaction with actions and rewards; Pan et al. [30] leverage both controllable and non-controllable states in policy learning; Wan et al. [37] apply task-relevant features to imitation learning. Our work also employs separated models. However, we prudently design this architecture in a model-free setting and propose novel loss terms to enhance the accuracy of image predictions.A detailed comparison between SMG and other methods is provided in Appendix F.2. 
conclusion and future work: 
 In this paper, we propose SMG for visual-based RL generalization and show its superiority in sample efficiency, stability, and generalization through extensive experiments. The success of SMG can be attributed to two key factors: (i) a delicately designed reconstruction-based auxiliary task with separated models architecture, which enables the RL agent to extract task-relevant and task-irrelevant representations from visual observations simultaneously; (ii) two consistency losses to further guide the RL agent's focus under deployment scenarios. We believe that the proposed method can be applied to a wide range of tasks.SMG is particularly well-suited for robotic manipulation tasks in realistic scenarios. However, when the observation contains too many task-relevant objects, the complexity of accurately learning a mask increases. This can lead to a decline in SMG's performance. For instance, in an autonomous navigation task, the presence of numerous pedestrians in the view makes it challenging to accurately mask all of them.The future work includes exploring more advanced backbones for task-relevant feature extraction, taking into account the generalization on non-static camera viewpoints and the test of SMG on realistic tasks to verify its generalization ability in real applications.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: The experimental setting and details are described in the Appendix C.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The paper reports the mean and standard deviation of the results in the tables and figures.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not release data or models that have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for existing assets: 
 Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: The paper properly credits the original owners of the assets and mentions the license and terms of use. The official code of DMC-GB (https://github.com/ nicklashansen/dmcontrol-generalization-benchmark) uses the MIT license.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We publish our source code, and the new assets are well documented in this paper.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We formulate the representation learning objective as a variational lower bound of the mutual information [31, 24] between the observation o t and the representation z t . By considering the independence between the task-relevant and task-irrelevant representations, we can decompose the mutual information as:EQUATIONWe use the empowerment term I(a t , z + t+1 |z + t ) introduced in [28] to quantify the information contained in the representation z + t+1 about the selected action a t , in goal of enhance the control-relevant property of the task-relevant representation z + t . We derive the variational lower bound of the empowerment term as:EQUATIONIn practice, we integrate a parameterized inverse dynamic model to predict the action a t based on the two continuous representations z + t and z + t+1 . We employ the Mean Squared Error (MSE) loss to guide the training of the inverse dynamic model.Algorithm 1 SAC with Separated Models Denote network parameters θ, mask ratio ρ, batch size N , replay buffer B Denote policy network π θ , foreground encoder f + θ , background encoder f - We conduct all experiments on a single machine equipped with an AMD EPYC 7B12 CPU (64 cores), 512GB RAM, and eight NVIDIA GeForce RTX 3090 GPUs (24 GB memory). We report the training wall time of different methods on DMControl tasks in Table 4 . θ foreach iteration time step do a, o ′ , r ∼ π θ (f + θ (o)), P(o, a), R(o, a) B ← B ∪ (o, a, r, o ′ ) foreach update time step do {o i , a i , r i , o ′ i } i∈[1,N ] ∼ B o + i , mask i ∼ f + θ (o i ) o - i ∼ f - θ (o i ) o aug i ← o + i * mask i + ϵ * (1 -mask i ) // ϵ is sampled from image dataset L recon ← L(o i , o + i * mask i + o - i * (1 -mask i )) // Equation 2 L f ore_consist ← L(o + i , f + θ (o aug i )) // Equation 7 L back ← L(ϵ, f - θ (o aug i )) // Equation 4 L action ← L(o i , o ′ i , a) // Equation 6 L mask ← L(mask i , ρ) // Equation 3 L q_consist ← L(Q θ (f + θ (o i ), a), Q θ (f + θ (o aug i ), a)) // Equation 8 L aux ← L recon + L f ore_consist + L back + L action + L mask // auxiliary loss L critic ← L q + L q_consist //We reproduce all baseline methods with the official code of DMC-GB (https://github.com/ nicklashansen/dmcontrol-generalization-benchmark) published by Nicklas Hansen, and we build our model on top of the SAC implementation. We use the same encoder and decoder architecture as the baseline methods to ensure a fair comparison.Figure 8 provides a detailed view of the encoder and decoder architecture. The input observation shape is 9 × 84 × 84, achieved by stacking three continuous frames. The encoder network contains 12 stacked convolutional layers, each with 32 filters of size 3 × 3. The stride is set to 1 for the first layer and 2 for the subsequent ones, facilitating down-sampling of the visual input. Then, after a flatten operation and a fully connected layer, an embedding of size embedding_size × 1 is obtained. Before decoding, SMG first expands the embedding into triples of the same size, aiming to decode three stacked input images separately. These three embeddings are then individually fed into the same decoder network, which consists of two groups of convolutional and upsampling layers to reconstruct the observation. The foreground decoder outputs the reconstructed foreground and a mask, while the background decoder outputs only the reconstructed background. For the inverse dynamic model, we adopt the architecture from [13] , which utilizes multi-layer perceptions to project the concatenation of two embeddings into the action space.The number of parameters in SMG is approximately double that of the baseline methods due to the use of two model branches. However, the performance improvement is primarily due to the novel model architecture rather than the increase in the number of parameters, as we use encoder and decoder networks similar to those in the baseline methods. We report the hyperparameters used in our experiments in Table 5 . We use the same hyperparameters for all seven tasks, except the action repeat and the mask ratio ρ. The L aux in SMG comprises five loss terms, which seems challenging to balance the weights. However, through experiments, we found that setting average weights for L recon , L mask , L action , L back is sufficient to achieve good performance (except the λ back is set to 2 since the background model should train to fit more complex images).Regarding the L fore , a too-large weight would lead to the model overfitting the inaccurate attribution predictions in the early stage (as we use the model output under raw observation as ground truth), so we set it to 0.1.We present the training curves for all seven tasks in Figure 11 , including four evaluation settings of DMControl and Robotic Manipulation tasks. As depicted in the figure, SMG demonstrates notably faster convergence and higher asymptotic performance across nearly all training and evaluation settings, showcasing the effectiveness of the reconstruction-based auxiliary task in enhancing sample efficiency. SMG exhibits superiority, particularly in the video-hard setting of DMControl tasks, where the performance of other methods drops evidently when random videos replace the background. Additionally, the figure underscores the considerable challenge posed by Robotic Manipulation tasks, with only SMG and SGQN successfully achieving zero-shot generalization in evaluation settings. Moreover, SMG shows more stable performance across different evaluation settings, which is crucial for real-world applications. Table 6 shows the generalization performance of SMG and all baseline methods with the randomcolor modification in DMControl tasks. SMG outperforms all baselines in 7 out of 10 tasks, with the performance gap within 5% in the other three tasks. The results indicate that SMG not only performs well in video-background settings but also exhibits superior generalization capability in random-color settings. This is achieved because overlaying the observation with random images can also introduce color shift. For a more direct measurement of the generalization ability in DMControl, we further calculate the average performance across five evaluation settings (including performance under training observation) and report the results in Table 7 . As shown in the table, SMG achieves state-of-the-art zero-shot generalization capability in all five DMControl tasks, surpassing all baseline methods by a margin of up to 26%. The results also demonstrate SMG's stability across different evaluation settings, with standard deviations less than 80 in all tasks. In contrast, the standard deviations of other methods range from 100 to 250. The experiment results of robotic manipulation reach are reported in Table 8 . SMG also shows a stable and superior performance in this task, with an average improvement of 38% over the second-best method.We report the effect of removing each loss term to the average performance across five evaluation settings in DMControl tasks in Table 9 . Compared with Table 3 , L q_consist still exhibits the most substantial impact on performance, though the performance drop is slightly smaller. This may be because the random-color settings do not shift the observations heavily compared to the videobackground settings, so the Q-value estimation is less affected. A similar phenomenon is observed in L back and L mask , indicating that attribution augmentation is more crucial in video-background settings.The mask ratio ρ is a hyperparameter that controls the expected proportion of the foreground area. However, this parameter is an empirical choice and may not precisely match the actual proportion of a given task. To investigate the sensitivity of SMG to the mask ratio, we conduct experiments with different ρ values in the walker-walk task of the video-hard setting. We select ρ values ranging from 0.02 to 0.1 with an interval of 0.02 and report the average performance across five evaluation settings in Figure 9 . The results indicate that variations do not significantly influence SMG in the mask ratio, as ρ values between 0.04 and 0.08 achieve similar performance. Moreover, when ρ is too small (0.02) or too large (0.1), the performance drops around 6% compared to the optimal ρ value (0.06). We also report the predicted masks of different ρ values in the figure. As ρ increases, the predicted masks start to include background areas, so a too high value leads to decreased performance. Conversely, when ρ is too small, the mask depicts an inaccurate foreground area (e.g. the legs of the walker with ρ = 0.02), resulting in a performance drop as well. EQUATIONThe attribution augmentation utilized in SMG requires the model to predict an accurate mask, and the foreground consistency loss also requires a precise attribution prediction of the model. This might seem contradictory, as the model struggles to make meaningful predictions in the early stages, which means it cannot satisfy the two requirements immediately. We dig into the training process of SMG by experiments and provide the model outputs in different training stages in Figure 10 . In the very early stage (≤ 1000 steps), the model has difficulty predicting accurate masks, leading the attribution augmentation more likes an overlay augmentation. However, the model rapidly learns to predict relatively accurate masks and generate meaningful attribution augmentation images that can help optimize L back and L fore_consist (after 2000 steps), aided by the constraint of L q . Subsequently, with the inclusion of L back and L fore_consist , the network begins to focus more on task-relevant areas in the observation, thereby in turn comes back to enhance the accuracy of Q-values and foreground predictions. Consequently, we view the training of SMG as a bootstrapping process.TIA [6] also designs two model branches to capture task and distractor features, similar to our separated models architecture. However, SMG differs from TIA in several essential aspects: (i) TIA is a model-based method focusing on eliminating task-irrelevant distractors in training observations, while SMG aims to utilize task-relevant features across diverse deployment scenarios to enhance the generalization capability of RL agents; (ii) SMG operates in a model-free setting, which can be more efficient to train and more flexible for applying data augmentation techniques; (iii) TIA uses a background-only reconstruction loss and requires the background model to reconstruct the full observation, which may cause the background branch to overly fit task-relevant features. In contrast, SMG addresses this issue by introducing attribution augmentation images to supervise the background model; (iv) SMG utilizes mask ratio loss to learn a more precise mask, while the masks in TIA are prone to containing distractors, as reported in its original paper.SODA [12] also improves the generalization ability of RL agents by regularizing the representations between observations and their augmented views, similar to the consistency losses in SMG. However, SODA implements this by simply minimizing the L2 distance between the two representations, which imposes a too rigid constraint and lacks interpretability. We achieve this by introducing Q-value consistency loss and foreground consistency loss, which provide more explainable supervision and additionally improve the stability of Q-values and predicted attributions.Note that the core idea underlying the Q-value loss in Equation 8 differs significantly from the consistency regulation objective proposed by SGQN [3] . SGQN focuses on prioritizing pixels that belong to the saliency map during encoding, primarily to enhance the accuracy of Q-value estimation under raw observations. In contrast, SMG treats the Q-values under raw observations as the ground truth and aims to achieve consistency between these Q-values and those obtained under augmented observations. Thus, we additionally use a stop-gradient operation. Justification: Source code is available at https://anonymous.4open.science/r/SMG/.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Justification: There is no societal impact of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: This submission proposes a new method to enhance the generalization capabilities of visual-based reinforcement learning (RL). The proposed method is composed of two branches to independently extract task-relevant and task-irrelevant features from visual inputs, utilizing a mask to isolate critical areas. It then combines multiple loss values to increase the robustness across different environments. By its simple design, it can be complementary to standard RL methods. The authors showed that the proposed method outperforms baseline methods, even on challenging cases with visual perturbations.

All authors agreed that this is a good submission to be shared with the community, and the meta reviewer also agrees with the reviewers. It is a clear accept!

Reviews: review: paper_summary: paper_summary: The 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: 
- overall reconstruction loss that combines outputs of the background and the foreground decoders modulated by a mask;
- mask ratio loss that prevents the foreground mask from taking up too much of the image;
- background reconstruction loss that uses the learned mask to generate a new data sample;
- q-value loss that makes the foreground representation capture value information;
- empowerment loss that makes the foreground representations capture the agent actions;
- 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

The method is tested on DMC generalization benchmark and on robotic manipulation tasks.
strength_weakness: strength_weakness: - The writing is a bit sloppy, with many typos and confusing sentences;
- The resulting objective is too complex and has too many terms;
- No comparison to TIA, although the presented method is quite similar. Was that because you only compare to model-free methods?

Typos (some of them, I didn't write down all of them, please run a spell checker on the text):
line 13: achieving free from overfitting : not clear what this means
line 38: further strengths -> further strengthens
line 100: focused in -> focused on
line 536: Comparision -> comparison
questions: questions: - 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?
limitations: limitations: The authors have adequately described limitations and potential negative societal impact of their work.
paper_summary: paper_summary: This 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.

Experiments 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.

The main contributions are:

- SMG: A separated model architecture with two branches to extract task-relevant and task-irrelevant representations from visual observations.

- Two consistency losses to guide the agent's focus on task-relevant areas across different scenarios.

- Strong performance on DMControl benchmark tasks, especially in video-background settings.
strength_weakness: strength_weakness: This paper has several weaknesses:

- 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.

- SMG introduces a lot of loss terms and associated hyperparameters, which could complicate tuning in practical applications.
    - 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.

- 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.

- The paper could be strengthened by considering a broader range of baselines. For example:
    - 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.
    - 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?


- The additional computation overhead introduced by the extra modules is concerning.
    - 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.
    - 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).

[1] Hansen, Nicklas, et al. "On pre-training for visuo-motor control: Revisiting a learning-from-scratch baseline." arXiv preprint arXiv:2212.05749 (2022).

[2] Kirillov, Alexander, et al. "Segment anything." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
questions: questions: See the "Weaknesses" section. Some additional questions are noted below:

- 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?
- 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.
limitations: limitations: There is no separate section for limitations in the paper.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

[1] Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations.

[2] Iso-Dream: Isolating Noncontrollable Visual Dynamics in World Models.

[3] Denoised mdps: Learning world models better than the world itself.

2. The paper lacks sufficient discussion and analysis of its limitations. 
3. The serial numbers in some figures appear to be somewhat disorganized.
questions: questions: Please see the weaknesses section.
limitations: limitations: This paper discusses the limitations, but that is not enough.
paper_summary: paper_summary: Visual-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.
strength_weakness: strength_weakness: - 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.

- I'm curious about the performance of the proposed method in some more challenging settings, like RL-Vigen [4].

- 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.

- 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.

- 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].

Overall, 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.

[1] Learning Task Informed Abstractions

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

[3] Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models

[4] RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization

[5] Context-aware dynamics model for generalization in model-based reinforcement learning

[6] Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability

[7] Zero-shot task generalization with multi-task deep reinforcement learning

[8] Task Aware Dreamer for Task Generalization in Reinforcement Learning
questions: questions: - 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$?

- Why do reconstruction-based methods benefit generalization? Are there any explanations?


---------------

**After reading the authors' response and other reviewers' comments, I have raised my scores from 4 to 5.**
limitations: limitations: Yes, this work has discussed its limitations in Sec. 6.

Comments: title: Discussion reply to reviewer C7ia
Discussion reply to reviewer jV8P
Discussion reply to reviewer jV8P
Thanks for your response
Discussion reply to the reviewer iwu6
Response to the rebuttal
comment: Yes, we will revise some paragraphs and add the additional experiments to the appendix. Thank you again for helping us improve the paper and update the score, we really appreciate your valuable comments!
Thank you for clarifying the difference with TIA and running the additional experiments.
Also thank you for clarifying that you use the same hyperparameters across all tasks. Please emphasize that in the text, I think that strengthens your claims.

I increase my score to 7.
Thank you for helping us improve the paper and update the score, we really appreciate your valuable comments!
Thank you for the explanation. I have updated my rating.
Thanks for your response. There seems to be some misunderstanding regarding our method and the experimental settings. SMG is definitely not tailored specifically for visual background settings, and we’d like to elaborate further:

- SMG is not primarily tackles scenarios where only the backgrounds differ. The key idea behind SMG is to focus on the task-relevant features across training and testing scenarios, such idea is applicable to any generalization tasks. Although more diverse generalization scenarios exist in the real world, we have tried to include a sufficient and widely accepted range of test settings in the paper to demonstrate the generalization ability of SMG. 

- We adopt a widely used evaluation setting similar to previous works [1,2,3,4], which includes color-easy, color-hard, video-easy, video-hard, and five different testing scenarios in robotic manipulation tasks. Only two of the nine settings are related to visual backgrounds: the other two color settings in DMC introduce randomized colors (both the foreground and background), while the five settings in robotic manipulation tasks involve changes to both color and textures (please refer to figure 5 and figure 6 in our paper).

- Despite SMG’s impressive performance in the video background settings, we’d like to highlight the improvements SMG delivers in other settings as well. In the random-color settings (Table 6 in our paper), SMG outperforms all baselines in 7 out of 10 tasks, with the performance gap within 5% in the other 3 tasks. In robotic manipulation tasks (Table 2 in our paper), SMG is the only method to maintain stable performance across 5 random-texture settings. These results indicate that SMG not only performs well in video-background settings but also exhibits superior generalization capability in random-color and random-texture settings.


---

Other than that, I hope the answers to your other questions have addressed your queries. If there are any additional concerns, please let us know!

---
Reference

[1] Generalization in Reinforcement Learning by Soft Data Augmentation.

[2] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation.

[3] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning.

[4] Spectrum Random Masking for Generalization in Image-based Reinforcement Learning.
Thank you to the authors for their response.

I still feel that the proposed method appears to be tailored primarily for generalization across varying visual backgrounds. I would recommend that the authors consider moderating the scope of their claims accordingly.
I'd like to thank the authors for the detailed and thorough response. I remain positive about this paper and believe it will be of interest to researchers of generalization in visual RL.
Thank you for helping us improve the paper and update the score. We really appreciate your valuable comments!

We will address your concerns as follows:

- Q1: Thank you for recognizing our work! We also hope our paper can promote more research in visual-based generalization problems.

- Q2: Thanks for your suggestion, we’d like to add the experiments in the camera-ready version.

- Q3-6: Yes, we will revise some paragraphs and add the additional experiments to the appendix.

- Q7: Overfitting to task-irrelevant features typically occurs when a single autoencoder is used to reconstruct the entire observation, as the encoder must learn all the features present in the observation. However, since SMG employs two autoencoders to separately fit task-relevant and task-irrelevant features, our model avoids such overfitting problem and can also take the advantage of high sample efficiency. But for typical methods using a single autoencoder, it's definitely a disadvantage in generalization tasks.

- Other comments: We'd like to discuss these settings separately：
  -  Foreground noise: SMG is particularly suited for such setting, as we use the extracted foreground for data augmentation. With a small modification, such as adding some noises to the foreground image (e.g., color jitter) before synthesizing it with the background image during the data augmentation stage, we can successfully simulate the testing scenario during training.
  - Camera views changing: This setting is relatively more challenging due to the involvement of spatial transformations. Directly applying SMG to this setting may not work, as SMG is trained with a fixed camera view. However, this issue could be addressed by incorporating an additional camera during training and adding STN blocks[1] to our model, adopting a similar approach to MoVie[2].
  - Overall, when applying SMG to different generalization settings, we should first consider the similarity between the data augmentation types and the testing scenarios, and then adjust the augmentation methods accordingly.


---
__Reference__

[1] Spatial transformer networks.

[2] Movie: Visual model-based policy adaptation for view generalization.
I'd like to appreciate the authors' efforts in addressing my concerns and supplementing extra experiments.

- Q1: Overall, I think the concept novelty of this work is limited as there are many works learning a mask model to distinguish between controllable and uncontrollable parts for handling distracting environments. However, I still believe this work has enough technique novelty to apply this insight to improving the generalization in visual RL, which is also beneficial for the community by extending relatively mature technologies to new settings. 

- Q2: Thanks for the reconstruction results in Adroit. In RL-Vigen, several settings can be handled by SAC-based algorithms like SGQN. For example, in Table-top Manipulation, RL-Vigen also considers visual appearances and lighting changes, and SGQN can handle this setting to some degree.

- Q3-6: Thanks for the response and supplemented results, which should be added to the paper for better highlighting the paper's contribution.

- Q7: Thanks for your clarification. There is still a small concern. As the authors have mentioned, reconstruction-based losses may exacerbate the overfitting problem but it can improve sampling efficiency. So is it more advantageous or disadvantageous for this task?

- Other comments: I'm curious about whether this method can be extended to more general visual generalization tasks, like foreground noises, camera views changing, and so on.

Overall, most of my concerns are addressed. Although there are several minor points, I have raised my scores into 5.