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 2024 poster (2024)
Keywords: Reinforcement Learning
Visual-based RL
Generalization
Primary area: reinforcement_learning
TL;DR: We propose SMG, which utilizes a reconstruction-based auxiliary task to extract task-relevant representations from visual observations and further strengths the generalization ability of RL agents with the help of two consistency losses.

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

Focus On What Matters: Separated Models For Visual-Based RL Generalization
Hai Zhang Feifan Yang Hang Yu Chang Huang Hongtu Zhou Chen Ye Changjun Jiang
Di Zhang Bowen Lv
Junqiao Zhao
Department Computer Science, Tongji University. Shanghai, China MOE Key Lab of Embedded System and Service Computing. Tongji University, Shanghai, China (2331922, 2151769, zhanghai12138, 2153299, chaojungiao) @tongji edu. (2053881, 2130790, zhouhongtu, yechen, cijiang]@tongj i edu. cn
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 rask-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (Separated Models for Generalization), a novel ap- proach 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 thereby achieving free from overfitting. Extensive experiments in DMC demon- strate the SOTA performance of SMG in generalization, particularly excelling in rideo-backgroun settings. Evaluations robotic manipulation tasks further confirm the robustness of SMG real-world applications. Source code is available at hattps://anonynn 4open. science/r/MMG/
scenarios,
Introduction
Visual-based Reinforcement Learning (RL) has demonstrated remarkable success across various tasks, including Atari games [27] II 18), robotic manipulation [23] 9), and autonomous navigation (26] (6). However, deploying visual-based RL algorithms in real-world applications requires high generalization ability due to numerous factors that can induce distribution shifts between training and deployment scenarios, such as variations lighting conditions, camera viewpoints, and backgrounds. Many visual-based RI. algorithms are prone to overfitting to the training observations [5] 34 limiting their applicability scenarios where fine-tuning with deployment observations not allowed.
To address the generalization gap in visual-based RL, current studies primarily focus utilizing data augmentation techniques [19
20 33 and exploring various auxiliary tasks [2 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 RL agents [4 [0] Z 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
"Corresponding author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 acconstruction-based auxiliary task 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 ss 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- 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, 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- off-policy RL algorithms.
SMG emphasizes the significance of task tesk-relevant features in visual-based RI. generalization and successfully incorporates a reconstruction loss into this setting.
Extensive experimental results demonstrate that SMG achieves tftte-o--thearr performance across various visual-based RL tasks, particularly excelling in video-background settings and robotic manipulation tasks.
2 Background
A, p,r,,,),, where S is the state space. .A is the action space, : S x A x S -> [0, 1] is the state transition probability function, S x A x S + R is the reward function, and 7 E [0, is the discount factor. At each time step f, the agent receives a state 8, E S, selects an action ag E A, and then receives a reward rz R. The agent's goal is learn a optimal policy (a, (s)) that maximizes the expected return E(e, ar) (2 - 7'r,). where Pr defines the discounted state-action visitation of n.
A Markov Decision Process (MDP) can be defined tuple (S,
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 thing, at each timestep t, the visual observation o can only capture partial information about the true state 8t, as certain elements may be obscured in the image. For another, the dimension of o is much higher than that of st, which makes it difficult utilize o, directly for policy learning.
To infer the true underlying state from visual observations, existing methods usually employ parameterized encoder f to map a stacked frame sequence 2; = =Or, Ov+1, 0z) to a compact low-dimensional latent vector Rt, 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 [4I). 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 77 21 2. By adding another parameterized image decoder g, the reconstruction loss is defined by maximizing the likelihood function:
Liccn = -DD Exr [(o
(I)
2

---

Fore
Fore
Foreground
Fore
Fore
Angmented
Attribution
Raw Observation
Reinforcement
Observation
Learning
Miok
Reconstruction
Back
Back
Fore
Background
Attribution
Raw Observation
Augmentation
Reconstruction
Foreground Consistency Loss
Background Loss
Prodicted
GrounddTuut
Q-vaine Consistency Loss
Background
Backgrround
Strop Gradieet
(a) Learning Task-Relevant Representations With SMG
(b) Improving Generalization With SMG
Figure I: Architecture of SMG. One-way arrows represent different types of data flows with the same input. Two-way arrows represent different types of loss.
3 Approach
3.1 What Matters in a Reinforcement Learning Task?
Learning to generalize is hard for RL agents, particularly when utilizing an image reconstruc- tion 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 new task. Even when colors or backgrounds are changed, hu- mans still leverage the prior knowledge to complete the task by focusing on task-relevant features. Considering robotic manipulation task where the agent must move the arm to the red target (Figure[)) 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 this task. We aim for our RI. agent to learn an optimal policy that solely relies on these task-relevant features while disregarding irrelevant regions.
Control Relevant
Reward Relevant
Figure 2: A robotic manipulation task explanation for task-relevant parts in the environment.
Formally, we decompose the latent representation 2, into task-relevant part 2," and task-irrelevant part 2, These two representations are independent, as p(2glo;) = p(2 lo,)p(2; lo.). 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 associated with the reward signal (the arm and the target), both are crucial for policy learning.
3.2 Learning Task-Relevan Representations with Separated Models
3.2.1 Separated Models and Reconstruction
The representation learning objective SMG to maximize the mutual information I (or; 2x) be- tween the observation and the latent representation and we further derive an image reconstruc- tion objective incorporating the combination of task-relevant representation 2 and task-irrelevant representation 2₁ follows:
Lyconn T(o;; 24) < E.,,
D E - (f (o.) log ((0,1= 2, )]]
Inspired by previous works
30 that explore how to mitigate background distractions, we
plement the reconstruction process by introducing the foreground encoder ft and the background

---

encoder f to extract different types of representations simultaneously, which forms separated models architecture. We also incorporate two decoders. The foreground decoder g employed to reconstruct the foreground image of and predict mask M, with values between (0, 1). The background decoder is employed to reconstruct the background image 0 The full image or is then reconstructed by of_ 9r and the mask M, via of of @ M, + of @ (1- Mc) (O denotes the Hadamard product), the reconstruction process is illustrated by the black arrows in Figure Ia Notably, the area where the agent is focusing can be visualized as of @ Mr, which we term the "attribution" of the agent, formally defined Aftrib(or).
3.2.2 Additional Loss Terms
Based on the separated models architecture, we define four additional loss terms to enhance the model's ability to distinguish between two types 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 hyperparameter p, termed the mask ratio, to constrain the proportion of the foreground part in the mask. As shown explicit form of an information bottleneck, as the percentage determines the number of pixels of of retained in the final reconstruction. This constraint forces ft to prioritize the task-relevant parts of the observation during encoding. Empirical results in Section demonstraa that Lmakk facilitates learning a more precise mask.
in Equation
3 we regard mask
Timask Ens Mc(i,j)
p)²
image size?
Background reconstruction loss. Improving the precision of background prediction can conse- quently 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;- To achieve this, we propose new type of data augmentation called attribution augmentation tailored for SMG, as illustrated in Figure 3b This augmentation involves augmenting the raw observation with its corresponding predicted mask M, via Tauib (oz) = or Me + c (1 MM), where e represents randomly sampled image. This simulates the video- baceobbackgroundd setting in deployment scenarios. We define the background reconstruction loss Loack follows:
Louck Bo [E 2- (o:) [lo (c(=
Q-value loss. Recall that the task-relevant representation =, has two key properties: reward-relevant and control-relevant. Satisfying the former is relatively straightforward, the representation 2, used for policy learning. Through the Bellman residual update objective [35] outlined in Equation
zt will progressively enhance its correlation with the reward signal.
Lg = EE-DD((QC a.) (re + (V
(5)
Empowerment loss For the control-relevant property, integrate an empowerment term T(a,, 12,"" (28 based on conditional mutual information, which quantifies the relevance be- tween the action and latent representation. Maximizing the empowerment term further leads to
maximizing a variational lower bound [ 2x+) as shown in Equation 6 This objective neces- sitates that a, is predictable when neighboring representations are known. We implement this objective by incorporating an inverse dynamic model.
Laction = (a 2 (xt) S -E p(a.. 1=/=) [log (aa 2 xt
(6)
The whole separated models architecture is shown in figure Ta
3.3 Generalize Task-Relevant Representations with Separated Models
Utilizing the separated models architecture, SMG can successfully extract tesk-relevant representa- tions 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 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 consisteney loss. To force the agent focus on the same task-relevant area in trans- formed scenarios, we train the foreground models to predict the attribution under augmented obser- vation Attrab(r(o:)) with the supervision of the ground truth attribution Attrib(or) (as Attrib(o, is relatively easier to converge to an accurate value, and we discuss it in detail in Appendix[F] (F). The foreground consistency loss consisc is defined as Equation (where sg means the stop-gratd gradient operation).
Liore_comsee
(btrD(AAtrri((((a)) - sg( [
Q-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 Equation [8 This approach also
regularizes the agent to learn an accurate task-relevant representation, as the gradient of La comsiist is back pack-propagated to the latent space.
bg consist Eor =D (10(((/000))))))(((((())))))))))
(8)
The above two consistency losses illustrated in Figure Tb
3.4 Overall Objective
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 Leiitic as the sum of the Q-value loss Lg and the
Q-value consistency loss Lypcomiist
Leritec Lg + Log.onniit
(9)
Additionally, the auxiliary loss Laux comprises five previously mentioned loss terms:
Traux Arecon Freeon + Amuus Tmmakk +
Aback T.back +
Aaction Frccion + Xtore
>fooreccomiia consist I-fore
(10)
Although Laux contains five loss terms, experimental results show that using average weights for the first four terms and 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 are provided in Appendix A
4
Experimental Results
4.1
Setup
We benchmark SMG against the following baselines: (1) SAC 8). serving as the foundational algorithm for all other base- lines: (2) DrQ [19], utilizing random shift augmentation; (3)
SODA 2 incorporating consistency loss on latent repre-
sentations: (4) SVEA [14, focusing on stabilizing Q-values; (5) SRM [5], 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 seeds.
(a) Overlay
(b) Attribution
Figure 3: Two types of data augmen-
tations using in SMG.
results are calculated by four random
5

---

Training Observation
Evaluation
Observation
Mask
Background
Attribution
Reconstruction
Color- Hard
Video-
Hard
Color- Jitter
Figure 4: Visualizing the reconstruction process of SMG in different tasks (from top bottom: walker-walk, cheetah-run, peg in box).
To achieve stable performance across various evaluation settings, we train SMG using hybrid data augmentation approach for To:), involving random overlay [14] and attribution augmentation for all tasks (each time we randomly select type of data augmentation, as shown in Figure 3p. The network design for SMG and more detailed experiment settings are reported in Appendix
4.2 DMControl Results
We first conduct experiments on five selected tasks from DMControl 36 and adopt the same evaluation setting DMControl Generalization Benchmark 112 (DMC-GB) used, which contains random-colors and video-hackground modifications across four different levels: color-easy, color- hard, video-easy and video-hard. Figure 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 the four evaluation settings.
(a) Training
(b) Color-easy
(c) Color-hard
(d) Video-eaxy
(e) Video-hard
Figure 5: Example training and testing observation for DMC-GB (walker-walk), (a) is the training observation. (b-c) indicates different degrees of color change; (d-e) replaces the background with random videos, with (c) additionally removing the floor and the walker's 's shadow.
To provide a clear explanation of how SMG reconstructs images, we present the image outputs of walker-walk and cheetah-run after 500k training steps of training in the first two rows of Figure 4 The last four columns illustrate the model outputs necessary for reconstructing the evaluation observations. 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 iise-itrelevant elements such as the skybox, the floor, and even the random color variation. This indicates that the task-relevant representation 2t contains only the information required to accomplish the task, which is crucial for generalization. Note that we aim to maintain the similarity between Allrib(T(e)) and Atrr)((o), even random-color settings. As shown by the first row of color-hard setting, SMG predicts yellow attribution despite the input evaluation observation being orange.

---

Table 1: DMControl results video-background settings. We evaluate each seed five times and calculate the mean value. Then, we calculate the mean and standard deviation with four random seeds. Red indicates the best and blue indicates the second-best A improvement SMG over the second best.
DMControl
SAC
DrQ
SODA
SVEA (overlay)
SRM
SGQN
SMG
(video-easy)
(ours)
cartpole, swingup
175
606
617
718
615
717
839
+121
1:1
176
1108
=16
17%
finger, spin
171
511
615
817
6422
860
952
+92
=99
1:56
191
1101
282
1:18
11%
walker, stand
908
924
947
961
1155
1:3
1:28
1:50
#14
12:10
119
1%
walker, walk
325
720
518
691
662
830
904
74
4:26
1:69
1:92
#120
#75
#:58
234
9%
cheetah, run
179
211
215
278
253
308
348
+-00
13%
DMControl
SAC
DrQ
SODA
SVEA (overlay)
SRM
SGQN
SMG
(video-hard)
(ours)
cartpole,
156
168
346
510
254
599
764
+ +165
swihgup
159
28%
finger, spin
22
54
310
353
131
710
910
+00
110
144
1:72
1:71
2:89
1:159
261
28%
walker, stand
278
406
558
955
::41
179
1:88
1:57
1:9
778
:19
10%
walker, walx
132
110
175
348
165
634
814
+180
+26
4:3
#31
2:80
199
1:33
=51
28%
cheetah, run
56
38
118
105
87
135
303
+168
#24
124%
Table reports the generalization performance of SMG and baseline methods with the video- background modification, which is the most challenging evaluation setting. The table shows that SMG outperforms all baselines all ten tasks. Particularly impressive is SMG's superiority in video-hard; when removing the floor and the walker's 's shadow, the performance of all baseline methods drops significantly. However, SMG is less affected by this substantial distribution shift and maintains stable performance across all tasks, with episode returns boosted more than 160 over the second-best in four out five tasks (as walker-stand is much easier task to train), showcasing its exceptional generalization capability.
4.3 Robotic Manipulation Results
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 previously mentioned.
(a) Training
(b) Test
(c) Test 2
(d) Test 3
(c) Test 4
(f) Test 5
Figure Examples training and testing observation for the robotic environment (Peg-in-box). (b-f) indicates five different evaluation settings varying in background colors and table textures.

---

Table 2 presents the evaluation results for peg- peg-in--ox, a task where robot must insert a peg tied to its arm into 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 bbaed auxiliary loss can capture more detailed features in the image, which is hard for methods that mainly rely on data augmentation techniques.
Table 2: Robotic manipulation results in peg-in-box. Red indicates the best and blue indicates the second-best. improvement of SMG over the second best. The last row reports the average performance over all six evaluation settings.
Cobtic-Manipulatton
SAC
DrQ
SODA
SVEA (overlay)
SRM
SGQN
SMG (ours)
A
(peg-in-box)
train
31
233
232
212
227
232
237
+44
473
1:20
1:09
115
1:19
16
2%
test1
33
63
34
18
55
67
237
1174
+25
=30
114
+:0
#:98
#26
18
2767
teat2
76
11
194
219
+-25
119
2.1
test3
-8
15
66
67
147
198
2,37
+33
+:16
107
147
#:7
1:11
#:4
15
20%
teat4
-42
72
80
109
112
-51
237
+125
131
1:28
1:122
1:08
1:123
1:46
2:17
112%
test5
-2
544
104
26
143
108
237
+-44
Average
-24
48
64
72
116
66
234
+118
102%
4.4 Ablation Study
In order to explore the role played by different loss terms in SMG, we conduct an ablation study in DMControl tasks. Table B 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, Lygconiirr exhibits the most substantial impact on performance, highlighting the importance of maintaining stable Q-value estimation in generalization tasks. Moreover, the performance drop without Liuct or Lmaak around 20% 30%, underlining the importance of attribution augmentation in enhancing SMG's generalization in video-backgroun settings, as the two loss terms directly affect the quality of the attribution augmentation. Additionally, Laction aids in learning better task-relevant representation. As for Ljjor_comsiit it also contributes to improving generalization ability, particularly relatively challenging tasks where the performance improvement ranges from 15% to 25%.
Table 3: Ablation study in DMControl (video-hard). Red indicates the performance drop of the ablated model compared to the full model.
DMControl SMG (full)
w/o comsis
w/o Lystion
w/o Loack
w/o Liusk
w/o Lg consis
(video hard)
cartpole, swingup
764 = 32
720 + 100
631 = 92
763 + 44
590 + 84
302 + 30 -62 (00%)
(6%)
133 (17%)
(0%)
177 (23%)
f finger, spin
910 = 61
695 + 103
609 + 352
412 + 170
731 = 130
509 + 83
-215 (24%)
(3396)
(00%)
177 (20%)
4401 (44%)
walker, stand'
955
885 + 45
855 +
775 + 144
836 127
432 + 210
-70 (7%)
100 (10%)
-180 (19%)
1119 (12%)
-323 (55%)
walker, walk
814 51
642 + 63
670
657 + 103
416 + 98
282 + 34
-172 (21%)
- 144 (188)
-157 (19%)
-398 (49%)
-532 (65%)
cheetah, run
303 + 16
247 + 40
212 + 52
233 4 110
162 - 100
130 + 37
-66 (18%)
-1 (30%)
-70 (23%)
144 (47%)
-173 (579))
Average
-15%
-22%
-3%
30%
56%
8

---

To better grasp the significance of Lmask and Liack in SMG, we showcase the predicted masks and their corresponding attribution augmentations in Figure 7 When musk 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 Back causes the backgroung model to learn all features excessively, resulting in attribution augmentation images devoid of task-relevan information. The ablation results underscore that both Lmakk and Ltack 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 reveal that maskk serves guiding factor in mask learning and SMG is not significantly influenced by variations in the hyperparameter mask ratio p.
(a) SMG
(b) without T.misk
(c) without Track
Figure 7: Predicted masks and corresponding attribution augmentations. (a) is the full model, (b) and (c) are the models without Lmakk and Lmask respectively.
5 Related Work
Improving generalization ability of RL agents has drawn increasing attention in recent years. Researchers primarily explore aspects: using data augmentation techniques inject useful priors when training (20
13
16
22
14
38 and employing various auxiliary tasks to guide the learning For example, Hansen and Wang [12] regularize the representations
32
process 13 3
42
40
12
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 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 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 novel reconstruction-baal 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
Bharadhwaj al. use empowerment term to distill control-relevant features from the task; Zhu et al.
45 43
45
bolster the resilience of RL agents learns compact representations by
by regularizing the posterior predictability; Zhang et al.
4
bisimulation metrics. Additionally, methods utilizing separated model architectures to extract different types of features simultaneously have been proposed [6]
37). For instance, Wang al. 39 decompose the latent state into four parts based on their interaction with actions and rewards; Pan 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 model-free setting and propose novel loss terms to enhance the accuracy of image predictions.
39
30
25
A detailed comparison between SMG and other methods is provided in Appendix F2
6 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-uue 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 RI agent's focus under deployment scenarios. We believe that the proposed method can be applied to 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 mask increases. This can lead to decline in SMG's performance. For instance, in an autonomous navigation task, the presence 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 real applications.
Acknowledgments and Disclosure of Funding
This work is supported by the National Key Research and Development Program of China (No. 2020YFA071 1402).
References
[1] Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.0522 055265, 2021.
[2] Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pages 6-20. PMLR, 2021.
[3] David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! saliency-guidee q-networks for generalization in visual reinforcement learning. Advances in Neural Information Processing Systems, 35:30693-30700 2022.
[4] Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. In- formation prioritization through empowerment in visual model-based rl. arXiv preprint arXiv:2204.08585 2022.
[5] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International conference on machine learning, pages 1282-1289. PMLR, 2019.
[6] Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
[7] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 2018.
2018.
[8] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[9] Tuomas Haarnoja, Ben Moran, Guy Lever, Sandy H Huang, Dhruva Tirumala, Markus Wulfmeier, Jan Humplik, Saran Tunyasuvunakool, Noah Y Siegel, Roland Hafner, et al. Learn- ing agile soccer skills for bipedal robot with deep reinforcement learning. arXiv preprint ar0iv:2304.1 13653, 2023.
[10] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 191.::11,, 2019.
[11] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193. 2020.
[12] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 1-61171 IEEE, 2021.
10

---

[13] Nicklas Hansen, Rishabh Jangir. Yu Sun, Guillem Alenyå, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. arXiv preprint arXiv:2007.043009 04330, 2020.
[14] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. Advances in neural information processing systems, 34:3680-3693,, 2021.
[15] Yangru Huang, Peixi Peng, Yifan Zhao, Guangyao Chen, and Yonghong Tian. Spectrum random masking for generalization in image-based reinforcement learning. Advances in Neural Information Processing Systems, 35:20393-20406, 2022.
[16] Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to- sim: Data-efficient robotic grasping via randomized-to-caooiica adaptation networks. In P Proceedings of the IEEE/CVF conference on computer vision and paltern recognition, pages 12627-12637, 2019.
[17] Rishabh Jangir, Nicklas Hansen, Sambaran Ghosal, Mohit Jain, and Xiaolong Wang. Look closer:
Bridging egocentric and third-person views with transformers for robotic manipulation. IEEE Robotics and Automation Letters, 7(2):3046 3053, 2022.
[18] Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model- based reinforcement learning for atari arXiv preprint 2XXv:19933.00344, 2019.
[19] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.
[20] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems, 33:19884 19895, 2020.
[21] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. Advances in Neural Information Processing Systems, 33:741-7555 2020.
[22] Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. arXiv preprint 2019.
[23] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 17(39): 2016.
[24] Langing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Junqiao Zhao, and Pheng- Heng. Towards an information theoretic framework of context-based offline meta-reinforcement learning. arXiv preprint arXin:2402.024299 2024.
[25] Yuren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. Advances in Neural Information Processing Systems, 36, 2024.
[26) Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andrew Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre, Koray Kavukcuoglu, et al. Learning to navigate in complex environments. arXiv preprint arXiv: I.06673, 2016.
[27] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Joannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 132.5513 2013.
[28] Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. Advances in neural information processing systems, 28, 2015.

---

[29] Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning, pages 81 8130-8139. PMLR, 2021.
[30] Minting Pan, Xiangming Zhu, Yunbo Wang. and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. Advances in neural information processing systems, 35:23178-23191, 2022.
[31] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. International Conference on Machine Learning, pages 5171-5180. PMLR, 2019.
[32] Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus, Automatic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34:5402-5415, 2021.
[33] Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmentation for deep learning. Journal of big data, 6(1):1-48, 2019.
[34] Xingyou Song. Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. arXiv preprint arXiv:1912.02995. 2019.
[35] Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9 44, 1988.
[36] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 18010006900 2018
[37] Shenghua Wan, Yucen Wang, Minghao Shao, Ruying Chen, and De-Chuan Zhan. Semail: eliminating distractors in visual imitation via separated models. In International Conference Machine Learning, pages
PMLR, 2023.
[38] Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforce- ment learning with mixture regularization. Advances in Neural Information Processing Systems, 33:7968-7998 2020.
[39] Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206. 15477, 2022.
[40] Ziyu Wang, Yanjie Ze, Yifei Sun, Zhecheng Yuan, and Huazhe Xu. Generalizable visual reinforcement learning with segment anything model. arXiv preprint arXiv:2312.17116, 2023,
[41] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10674-10681, 2021.
[42] Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: Task-aware lipschitz data augmentation for visual reinforcement learning. arXiv preprint arXiv:2202. 09982, 2022.
[43] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint htth:20000.004422 2020.
[44] Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study on overfitting in deep reinforcement learning. arXiv preprint arXiv: /004. 1804.0689, 2018.
[45] Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta. Repo: Resilient model- based reinforcement learning by regularizing posterior predictability. Advances in Neural Information P Processing Systems, 36, 2024.
[46] Yuke Zhu, Roozbeh Mottaghi, Erie Kolve, Joseph J Lim, Abhinav Gupta, Li Fei-Fei, and Ali Farhadi. Target-driven visual navigation in indoor scenes using deep reinforcement learning. In 2017 IEEE international conference on robotics and automation (ICRA), pages 3357-3364. IEEE, 2017.
12

---

A Derivations
We formulate the representation learning objective variational lower bound of the mutual information [3T 24 between the observation Og and the representation 2;- By considering the
independence between the task-relevant and task-irrelevant representations, we can decompose the mutual information as:
I(on m)) = 2)) logp(oe(<)) log p(:)]
2 E s.) [lgp(()/=z)
2 s() [ogp(or
(11)
Bostslo.ppoo)llg ((o])2)))
(o (o)) lo))((.)) (o) [llg <(0rl=)
z[
ED [B.] ~f (0r). ~f (o.) log (orl2(
We use the empowerment term [aa, =++))=+)) introduced in 28 to quantify the information con- tained in the representation 2/4,1 about the selected action in goal of enhance the control-relevant property of the rele-relevant representation 2," We derive the variational lower bound of the empow- erment term as:
Maa 2++11=+))) = E play + flog
p(arl2)))
a(arl2t1110 2:
plarlz 2it)
Epton flog
log
p(arlzit)
g(arlztt1,
a101l2111 27)
E Ep
Blog
p(arlztt)
)p(= [lo ((: at ++1] / p(z )paell=) ( + lltt a:) log p(a)22))
E Pl
E Eplac [lo ((a - z + E ppc) bete .a) (4) (p(ael=) +
E plac )ot [log g(a: 22 x[]]
(12)
In practice, we integrate parameterized inverse dynamic model to predict the action based on the two continuous representations 2, and = We employ the Mean Squared Error (MSE) loss to guide the training of the inverse dynamic model.
13

---

B Pseudocode
Algorithm 1 SAC with Separated Models
Denote network parameters 0, mask ratio p. batch size N, replay buffer B Denote policy network тe. foreground encoder fg. background encoder fg foreach iteration time step do
a. a. To(fō (o)), P(o, a), R(
+ B (o, a. o.))
foreach update time step do
for, ai ri, o'hie,NN ~
of, maakk; ~ fo((0) (o:)
of -0((0))
0r +- t +mask; + + (1 - mask;) // is sampled from image dataset
Lireon - -(oo,, of5 * maski o (1 -maak))) // Equation 2
Iifore_consist +- Mot 10 (o)=²)) // Equation 7
Libckk + Lie, fo (o(=0)) // Equation 4
Liaction - L(oi, a,))) // Eqqutionn 6
Limask L(maski, p) // Eqquation 3
Lig Lig_eonsist +- [[Qn(SØ (oi), Qo J (o; a)) // Equation 8 Q
T.aux T.pecon + "'fore_constir + L.back + Lraction + T'mas // auxiilirry loss
Leritie Ly connsiit |l critic loss
update with Lactor; Locitier Lawx
end for end for
Lg Locor are defined by SAC
C More Experiment Details
C.1 Computing Hardware
We conduct all experiments on single machine equipped with an AMD EPYC 7B12 12 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
Table 4: Wall time comparison of different methods on DMControl tasks.
Algorithm
Wall Time (DMControl, 500k)
SAC
10 hours
DrQ
~ 13 hours
SODA
12 hours
SVEA
~ 12 hours
SRM
hours
SGQN
12 hours
SMG (ours)
22 hours
C.2 Network Architecture
reproduce baseline methods with the official code of DMC-GB (https://github. com/ published by Nicklas Hansen, and we build our model on top of the SAC implementation. We use the same encoder and decoder architecture the baseline methods to ensure a fair comparison.
14

---

Figure 8 provides detailed view of the encoder and decoder architecture. The input observation shape 9 84 x 84, achieved by stacking three continuous frames. The encoder network contains 12 stacked convolutional layers, each with 32 filters of size 3 x 3. The stride is set to for the first layer and for the subsequent ones, facilitating down-sampling of the visual input. Then, after flatten operation and fully connected layer, embedding of size embedding size x is obtained. Before decoding, SMG first expands the embedding into triples of the same size, aiming to decode three stacked imput images separately. These three embeddings decoder network, which consists of two groups of convolutional and upsampling layers to reconstruct the observation. The foreground decoder outputs the reconstructed foreground and 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.
then individually fed into the same
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 primarily due to the novel model architecture rather than the increase in the number of parameters, we use encoder and decoder networks similar to those in the baseline methods.
Embedding
Conv2D 3x3 padding + Sigmoid
Conv2D 3x3 stride I
Conv2D 3x3 padding
Upsample IReLU
Conv2D 3x3
Flattenei Lincar
Embedding
Linear RReUU
Embedding
Linear R Reshape
ReLU
Embedding
x11
x2
Figure 8: SMG network architecture (foreground encoder + foreground decoder).
C.3 Hyperparameters
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 Laux in SMG comprises five loss terms, which seems challenging to balance the weights. However, through experiments, we found that setting average weights for Tercoo: Temask, Lraetio T-hack is sufficient to achieve good performance (except the Atack is set to since the background model should train to fit more complex images). Regarding the Ljore. 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 set it to 0.1.
D
More Experiment Results
D.1 Training Curves
We present the training curves for all seven tasks Figure 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-basee auxiliary task in enhancing sample efficiency. SMG exhibits superiority, particularly in the video-hard setting 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.
15

---

Table Hyperparameters.
Hyperparameter
Value
Observation size
84
Frume stack
Discount factor "
0.99
Batch
64
Embeddiing size
256
Action repeat
(carspoe-swvnguu)) (walker-waff, walker run)
2 (finger-spin). (meach, peg-in-box)
Train stops
250k (walker-stand, reach, peg-te-box), 500k (others)
Replay buffer
500k
Actor optimizer
Adam (rr le 3. = 0.9, 00 0.999)
optimizer
Adam 3, 8,
0.999)
Auxiliary lusk optimizer
Adam (tr le 3, 0.9, .g 0.999)
Auxiliary update frequency
Reconstruction loss weight Aucee
Background roconstruction loss weight Abak
Mask ratio loss weight Auot
Empowerment loss weight Awaion
Q-value consistency weight bn cosair
Foreground consistency loss weight Afve p
0.1
(reach, peg-in-box), 0.06 (walker-wull,, walker-stand, cheetab-run) 0.04 (cartpole-wtingpp... htgp:e/e-owiiggu.....e-----p (inger-spnn)
Moreover, SMG shows more stable performance across different evaluation settings, which is crucial for real-world applications.
D.2
More Table Results
Table 6: DMControl results random-color settings.
DMControl-GB
SAC
DrQ
SODA
SVEA (overlay)
SRM (SAC)
SGQN
SMG (ours)
(color-easy)
cartpole.
178
845
720 +109
800
856
764
854
-2
swingup
0%
finger,
296
827
761
919
916
852
957
+38
422
1174
187
1:43
1:34
112
2:52
4%
walker, stand
592
827
929
957
953
906
965
+88
#274
#:97
#23
+:4
+5
=000
#:13
156
walker, walk
430
669
539
705
632
805
915
+110
433
1:08
4:11
#124
1:9
4:47
1:36
14%
cheetash, run
253
237
219
289
272
312
346
+34
#27
#:74
#46
4:4
#:44
#:34
#227
11%
DMControl-GB SAC
DrQ
SODA
SVEA (overlay)
SRM (SAC)
SGQN
SMG
(color-hard)
(ours)
cartpole,
717
385
752
swingup
126
2133
190
1:86
1:10
111
2662
35
finger, spin
271
663
834
27
#23
#214
+106
1:74
+:00
219
#113
3%
walker,
719
807
788
878
+71
1182
+:139
1:122
935
walker,
379
456
396
571
483
632
739
+107
#:192
cheetash
208
147
199
238
203
210
299
+61
26%
Table @/oows the generalization performance of SMG and all baseline methods with the random- color modification in DMControl tasks. SMG outperforms all baselines in 7 out of 10 tasks, with the performance gap within 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 achieved because overlaying the observation with random images can also introduce color shift.
16

---

Table Training and average performance in DMControl.
DMControl-GB
SAC
DrQ
SODA
SVEA (overlay)
SRM (SAC)
SGQN
SMG
A
(training)
(ours)
cartpole, swingup
186
872
687
809
871
805
858
14
#6
=10
#173
#:42
#10
#58
=9
2%
finger, spin
306
884
801
923
922
961
+36
=12
+116
266
#:66
+33
#61
24
4%
walker, stand
630
955
959
952
964
#:24
=18
#01
=0
266
#:17
=18
walker, walk
422
827
581
753
715
876
924
-442
-61
L129
1143
174
1445
1.31
cheetah,
311
333
225
300
298
343
357
+14
DMControl-GB
SAC
DrQ
SODA
SVEA (overlay)
SRM (SAC)
SGQN
SMG
A
(average)
(ours)
cartpole, swingup
176
642
591
720
676
704
808
+-88
+:26
=110
finger, , spin
213
586
630
776
690
809
924
+115
walker, stand
489
747
772
891
8445
893
945
+ 52
#147
walker, walk
338
556
442
614
531
755
859
+ 104
1:109
1:254
1:146
1199
1:0
8:72
14%
cheetah, run
201
199
195
242
223
262
331
+-69
=46
+110
#:40
#:72
#75
#:77
=24
26%
For more direct measurement of the generalization ability in DMControl, we further calculate the average performance across five evaluation settings (including performance under training observa- tion) and report the results in Table 7 As shown in the table, SMG achieves state-of-the fi-thh--rt art zero-shot generalization capability in all five D'Control tasks, surpassing all baseline methods by margin of up to 26%. The results also demonstrate SMG's stability across different evaluation settings, with standard deviations less than 80 all tasks. contrast, the standard deviations of other methods range from 100 to 250.
Table 8: Robotic manipulation results in reach.
Robtic-Manipulation
SAC
DrQ
SODA
SVEA (overlay)
SRM (SAC)
SGQN
SMG (ours)
A
(Reack)
train
32
11
33
30
33
30
test1
-16
-26
-22
33
19
30
+33
#23
#9
#10
#25
#:
58%
teat2
-21
24
26
test3
-32
38
20
-13
24
33
30
+14
+20
231
#10
222
4:2
9%
test4
-19
-21
24
29
150
1:25
116
121
130
16
11
21%
test5
54
33
50
37
16
29
1:77
Average
21
-6
20
-10
6
21
29
+-8
The experiment results of robotic manipulation reach are reported in Table SMG also shows stable
and superior performance this task, with an average improvement 38% over the second-best method.
17

---

E More Ablation Study
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 Ly 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 video- background settings, so the Q-value estimation is less affected. A similar phenomenon is observed in Liuck and Lmask, indicating that attribution augmentation is more crucial in video-base bileo-background settings.
The mask ratio p is a hyperparameter that controls the expected proportion of the foreground area. However, this parameter is 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 o values in the walker-walk task of the video-hard setting. We select p 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 P values between 0.04 and 0.08 achieve similar performance. Moreover, when p is too small (0.02) or too large (0.1), the performance drops around 6% compared to the optimal p value (0.06). We also report the predicted masks of different p values in the figure. As p increases, the predicted masks start to include background areas, so too high value leads decreased performance. Conversely, when p is too small, the mask depicts an inaccurate foreground area (e.g. the legs of the walker with p 0.02), resulting in a performance drop as well.
Table 9; Ablation study in DMControl (average performance).
DMControl
SMG (full)
w/o Lior comss
w/o Lystiom
w/o Luuck
w/o Linisk
w/o Ly consis
(average)
cartpole, swihgup
808 + 53
763 + 28
762 + 71
795 32
758 + 97
646 + 191
-45 (6%)
-46 (0%)
133 (2%)
50 (8%)
-662 (20%)
finger, spin
924 = 45
815 = 66
791 112
640 + 115
866 + 73
773 + 151
1009 (12%)
- 133 (14%)
284 (31%)
-58 (8%)
-151 (10%)
walker, stand"
945 + 33
918 + 26
874 + 17
915 + 70
930 + 17
598 + 114
-27 (3%)
-71 (8%)
-30 (39))
-15 (2%)
-317 (37%)
walker, walk
859 - 72
727 = 54
757 - 61
756 1: 103
693 = 145
613 + 227
-132 (15%)
(1296)
(12%)
166 (19%)
-246 (29%))
cheetah, run
331 24
304 + 29
325 - 58
270 + 25
319 + 82
269 + 110
27 (8%)
(2%)
-61 (18%)
-12 (4%)
62 (19%)
Average
-9%
8%
-13%
-7%
24%
(a) P = 0.02 (800) (b)p= = 0.04 (857) (c) p = 0.06 (859) (d) = 0.08 (823) (e) p = 0.1 (793)
Figure Ablation study of mask ratio p in walker-walk of average performance across five evaluation settings. The images and numbers in parentheses indicate the predicted masks and the corresponding performance, respectively.
F
More Discussion
F1
Bootstrapping Process in SMG
The 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 the early stages, which means it cannot satisfy the two requirements immediately. We dig into the training process of SMG
18

---

by experiments and provide the model outputs in different training stages in Figure [0] In the very early stage ( 1000 steps), the model has difficulty predicting accurate masks, leading the attribution augmentation more likes overlay augmentation. However, the model rapidly learns to predict relatively accurate masks and generate meaningful attribution augmentation images that can help optimize Brack and Fere consist (after 2000 steps), aided by the constraint of T... Subsequently, with the inclusion of Ltack and Lioe consist, L 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.
F.2 Comparison with Related Work
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-basee method focusing on eliminating task-irrelevant distractors in training observations, while SMG aims to utilize task-relevant features across diverse deployment scenarios 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 background-one reconstruction loss and requires the background model to reconstruct the full observation, which may cause the background branch to overly fit task 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 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 33. 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- gradienttintt operation.
(a) Training Steps: 1000
(b) Training Steps: 2000
(c) Training Steps: 6000
(d) Training Steps: 30000
Figure 10: Masks, attributions, and corresponding attribution augmentation images in different training stages.
19

---

Traiiing
Color-Easy
Color-Hard
Video-Easy
Video-Hard
300
150
Cortpole-swinge
690
(00
Episode
450
300
450
300
433
199
6.
2.0
3.0
4.0
5.0
5.0
1.0
2.0
3.0
3.0
1.0
2.8
3.0
0.5
1.0
2.0
9.0
30
4.0
5.0
1000
1000
110
Episods
400
410
1.3
2.0
4.
5.0
0.
1.0
2.0
5.0
3.0
1.0
2.0
1.0
4.0
2.0
3.0
40
5.0
800
600
903
830
Epdcade
400
200
0.5
3.9
600
Epicade Return
690
Walker-walk
150
150
650
420
500
150
150
0.0
1.0
2.0
5.0
4.0
5.0
82
20
02
1.0
0.0
1.8
20
30
5.0
600
Epiands Rotui
750
750
793
Choctah--un
-000
200
300
300
150
150
190
130
2.0
3.0
2.0
30
4.0
5.0
0.
3.0
2.0
30
4.0
5.0
00
1.0
2.0
30
1.0
2.0
30
5.0
Training
Testl
Test2
Test3
Test4
Episode Return
199
143
140
101
100
100
-23
-63
-60
0.0
3.0
10
Return
Reach
-15
-15
-55
-30
0.0 05
14
15
20
25
65
05
1.
15
25
25
0.
0.5
10
1.
20
24
00
55
18
13
20
25
50
05
340
15
20
25
Number (ram))
(x10f)
(x10f)
framnes(xtoe)
(x10P)
SMG (ours)
SGQN
SRM
SVEA
DrQ
SODA
Figure 11: Training curves in all seven tasks. We evaluate each seed three times and then calculate the mean episode return for every 10k training steps, and the variance shown as the shaded area by calculating four random seeds.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly state the contributions of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations and future work are discussed in the Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
21

---

The derivations of the representation learning objective and the empowerment term are provided in the Appendix A
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Source code is available at httts:///aoonmonut .trrmouuusooo......o.ooo science/r/SMG/ and the experimental setting and details are described in the Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open Access to Data and Code
22

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Source code is available at httts::/anonymoouut aa/nymoonnmmmmuo 4open. science/r/SMG/
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipss.cc/pill wttpcccpppppppppeeeeeeeeeeeeeeee httaa//gggieee//eeuuebiissssssss /Cc////oddiiisssseeeeeee ssiiiiiiccyyyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi- dence intervals, statistical significance tests, 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 with given experimental conditions).
23

---

The method for calculating the error bars should be explained (closed form formula, call 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 the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The paper provides the details of the compute resources in the Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We have reviewed the NeurlPS Code Ethics and ensured that our research conforms to it.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
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 why the paper does not address societal impact.

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have 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 high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 best faith effort.
12. 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]
The paper properly credits the original owners the assets and mentions the license and terms of use. The official code of DMC-GB (https://4 //gg:///bbtt com/ ni nicklashanenn/ wwtpee//cmmoooooooooeee hatpee/eeeeennoooooooeoaaaaaaaaa ben.aiiiiiiiiiii-bbaaaaa.a.a benchmark uses the MIT license.
Justification:
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 dataset.
The authors should state which version of the asset used and, possible, include 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.
25

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
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 available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We publish source code, and the new assets well documented in this paper.
Guidelines:
The answer NA means that the paper does not release 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
26

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
27


Reviews:
[{"review_id": "wOjwxilLJq", "reviewer": "Reviewer_C7ia", "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: \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.", "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": "The authors have adequately described limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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;", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for your constructive comments and suggestions. We address each of your comments as follows.\n\n### Q1: The writing is a bit sloppy, with many typos and confusing sentences\n---\nA1:\nThank you for your careful review. We have thoroughly reviewed the paper multiple times and assure you that the errors listed in __Author Rebuttal 1.1__ will be corrected in the final version of the paper.\n\nWe hope that the current version is clear enough for readers to understand our key ideas. If there are any sentences that are still prone to misunderstanding, please let us know.\n\n### Q2: The resulting objective is too complex and has too many terms\n---\nA2: \nThis phenomenon is common in methods that utilize separate models. For example, Iso-dreamer [4] has six loss terms, TIA [2] has seven loss terms, and IFactor [5] has eight loss terms, yet all achieve excellent performance. Therefore, we believe the focus should be on the difficulty of adjusting the loss weights for different tasks rather than on the number of loss terms.\n\n**In all seven tasks, we used the same weights for the five auxiliary loss terms.** The results demonstrate that weighting four of the five terms equally is sufficient to achieve outstanding performance. This indicates that our method is robust enough across different tasks and does not require additional time for adjusting the weights.\n\nThe main concern may stem from why we use a smaller weight for $L_{fore\\ consist}$ than others. This is because 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), and cause the foreground encoder cost more time steps to learn an accurate task-relevant representation. **We have added a new ablation study to provide readers with more insight into our loss weight setting.** By setting $\\lambda_{fore\\ consist}$ to 1 (refer to __Author Rebuttal 2.4__), the performance of SMG drops around 10~20% across different tasks.\n\n### Q3: No comparison to TIA, although the presented method is quite similar. Was that because you only compare to model-free methods?\n---\nA3:\nTo the best of our knowledge, there are still no model-based methods that effectively address the generalization problem in visual-based RL. Therefore, in our original paper, we did not include comparisons with model-based methods.\n\nAlthough a series of model-based methods (e.g., TIA[2], DenoisedMDPs[3]) use a similar separated models structure like SMG, they address a fundamentally different problem. **These algorithms primarily aims to enhance the agent’s robustness to distractors during training.** Specifically, TIA trains and tests both in video background environments, **whereas our method trains without a background and tests in the video background environments (just like the difference between I.I.D and OOD).** Methods like TIA do not incorporate specific designs, such as data augmentation techniques, to bridge the gap between training and testing scenarios, and therefore lack generalization capabilities. \n\nNevertheless, we have included TIA as a baseline to support our point; please refer to the __Author Rebuttal 2.1__. As shown in Table 1, TIA fails to generalize in video-background settings in all seven tasks. This further demonstrates the effectiveness of our improvements (two consistency loss and data augmentation techniques) to the separated models architecture under the generalization settings.\n\n\n### Q4: Have you tried minimizing MI between background representations and actions and or rewards?\n---\nA4:\nThis is a very interesting question. Why we aim to maximize MI between foreground representation and actions or rewards is because this is the most straightforward way to learn a good task-relevant representation. We did consider minimizing the MI between background representation and actions or rewards as well, but we abandoned this for the following reasons:\n\nTIA[2] attempts to minimize the mutual information between the background representation and rewards with the loss term $L_{Radv}=-max_qq(r_t|s_t^-)$ . TIA implements this by maximizing the prediction loss of the background reward model (https://github.com/kyonofx/tia/blob/main/Dreamer/dreamers.py#L368). This approach appears problematic because the reward model can easily learn parameters that satisfy this condition, such as outputting 0 for any state, which would make the prediction loss consistently high. We have also done ablation experiments on TIA before, and the results show that this loss term does not improve the model performance significantly. Given that incorporating this term would further complicate the optimization objective, we chose not to include it.\n\nAs for the MI between the background representation and actions, since the background encoder does not include actions as its input, the background representation inherently lacks action-relevant features. Therefore, there is no need to minimize this MI term.\n\n\n\n### Reference\n---\n[1] Generalizable visual reinforcement learning with segment anything model.\n\n[2] Learning task informed abstractions.\n\n[3] Denoised mdps: Learning world models better than the world itself.\n\n[4] Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models.\n\n[5] Learning world models with identifiable factorization.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "15tbehfjAC", "reviewer": "Reviewer_jV8P", "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.\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.", "questions": "See 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.", "limitations": "There is no separate section for limitations in the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This 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.", "weaknesses": "This 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for your constructive comments and suggestions. We address each of your comments as follows.\n\n### Q1: Some testing scenarios appear to be overlooked in this paper.\n---\nA1:\nWe admit that real-world deployment scenarios can be more diverse and complex. **However, the four testing scenarios used in our paper are widely accepted benchmarks in visual-based RL [3,4,5].** For example, the color variation settings simulate different illumination conditions, while the video variation settings simulate training a robot in different environments, such as training indoors and testing outdoors. In addition, the testing scenarios in the two robotic manipulation tasks simulate texture changes, Therefore, we don’t think the real scenarios are overlooked.\n\nMoreover, no algorithm has yet effectively solved the video-hard setting except ours, which is the closest scenario to realistic applications. We hope our paper can promote more research in this challenging setting.\n\n\n### Q2: SMG introduces a lot of loss terms.\n---\nA2:\nThis phenomenon is common in methods that utilize separate models. For example, DenoisedMDPs [3] has four loss terms, Iso-dreamer [4] has six loss terms, TIA [2] has seven loss terms, and IFactor [5] has eight loss terms, yet all achieve excellent performance. Therefore, we believe the focus should be on the difficulty of adjusting the loss weights for different tasks rather than on the number of loss terms.\n\n**In all seven tasks, we used the same weights for the five auxiliary loss terms.** The results demonstrate that weighting four of the five terms equally is sufficient to achieve outstanding performance. This indicates that our method is robust enough across different tasks and does not require additional time for adjusting the weights.\n\n### Q3: About mask ratio $\\rho$\n---\nA3:\nPlease refer to __Author Rebuttal 2.6__\n\n### Q4: Although the authors describe this as a bootstrapping process, further analysis regarding its stability and potential failure modes would be beneficial.\n---\nA4: \nThanks for rising this concern.\n**The key guarantee that ensures the models learns accurate foreground images and masks is the utilization of $L_{q}$ and $L_{action}$ to extract reward-relevant and action-relevant features.** The bootstrapping process then starts based on these features.\n\nWe have added an experiment to show a failure case by removing $L_{action}$ and stopping the gradient from $L_{q}$ to $z^+$; please refer to __Author Rebuttal 2.7__. As shown in Figure 3, the foreground models failed to learn meaningful images, and the background models learned everything from the input observations. This failure case occurs approximately once every five seeds. However, when we add the two loss terms back, we have never observed such failures.\n\n\n### Q5: The paper could be strengthened by considering a broader range of baselines.\n---\nA5:\nThanks for your suggestion. We have added two more baselines:\n- SAM-G [3] is a method that utilizes both pre-trained image encoders and the Segment Anything Model. **Although the results reported in the paper use four times the batch size and twice the training steps compared to SMG, SMG still achieves comparable performance.** This is because methods with pre-trained models heavily rely on the similarity between pre-training datasets and the RL task scenarios. Utilizing Large vision foundation models also increasing parameters and training time (SAM-G takes around 2d to train). Despite being trained from scratch, SMG can quickly distinguish foreground and background parts from environments and is more flexible than these pre-trained methods. Please refer to __Author Rebuttal 2.2__.\n\n- TIA[1] is a model-based method that uses a similar separated models methods to us. However, as TIA is not designed for visual generalization, it fails in our testing scenarios. Please refer to __Author Rebuttal 2.1__.\n\n\n### Q6: The additional computation overhead introduced by the extra modules is concerning.\n---\nA6: \nWe agree that algorithms are required to be time efficient to train in realistic tasks, but we do think the time consumption of SMG is reasonable.\n\n- Although SMG requires 22 hours to train for 500k time steps, **in almost all tasks, SMG already outperforms the baselines’ 500k performance at 200k time steps (please refer to Figure 11 in our original paper).** We trained for 500k time steps to ensure a fair comparison. Therefore, it is possible to train for only half the time reported in our paper.\n\n- Algorithms that perform well in more complex tasks inevitably lead to more complicated model structures and an increase in the number of parameters. Despite the fact that baselines training consume less time, they simply can't generalize at all in the robotic manipulation tasks. And we think the training time for SMG is acceptable as methods like SAM-G cost much more training time.\n\n\n### Q7: \"foreground\" actually refers to the task-relevant parts of the image, while \"background\" refers to the task-irrelevant parts, correct?\n---\nA7:\nYes, you are right. We think the terminologies “foreground” and “background” are more accessible for readers to understand our key idea. Therefore, we will add a clear definition in Section 3.1:\n\n“In the following paragraphs, we will use ‘foreground’ to refer to task-relevant parts and ‘background’ to refer to task-irrelevant parts.”\n\n### Q8: The necessity for background reconstruction is unclear.\n---\nA8:\nPlease refer to __Author Rebuttal 1.5__\n\n### Q9: There is no separate section for limitations in the paper.\n---\nA9:\nPlease refer to __Author Rebuttal 1.2__\n\n### Reference\n---\n[1] Learning task informed abstractions.\n\n[2] Learning world models with identifiable factorization.\n\n[3] Generalizable visual reinforcement learning with segment anything model.\n\n[4] Stabilizing deep q-learning with convnets and vision transformers under data augmentation.\n\n[5] Generalization in Reinforcement Learning by Soft Data Augmentation", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xc1ivg0InA", "reviewer": "Reviewer_ZRY4", "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.", "questions": "Please see the weaknesses section.", "limitations": "This paper discusses the limitations, but that is not enough.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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).", "weaknesses": "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.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for your constructive comments and suggestions. We address each of your comments as follows.\n\n### Q1: It would be beneficial to also include a comparison with model-based RL methods.\n---\nA1:\nTo the best of our knowledge, there are still no model-based methods that effectively address the generalization problem in visual-based RL. Therefore, in our original paper, we did not include comparisons with model-based methods.\n\nAlthough a series of model-based methods (e.g., TIA[1], DenoisedMDPs[2]) use a similar separated models structure like SMG, they address a fundamentally different problem. These algorithms primarily aims to enhance the agent’s robustness to distractors during training. Specifically, TIA trains and tests both in video background environments, whereas our method trains without background and tests in the video background environments (just like the difference between I.I.D and OOD). Methods like TIA do not incorporate specific designs, such as data augmentation techniques, to bridge the gap between training and testing scenarios, and therefore lack generalization capabilities. \n\nNevertheless, we have included TIA as a baseline to support our point; please refer to the __Author Rebuttal 2.1__. As shown in table 1, TIA fails to generalize in video-background settings in all seven tasks. This further demonstrates the effectiveness of our improvements (two consistency loss and data augmentation techniques) to the separated models architecture under the generalization settings.\n\n### Q2: The paper lacks sufficient discussion and analysis of its limitations.\n---\nA2:\nThank you very much for the reminder, we have added a paragraph to Section 6 that describes the limitation of our work. Please refer to the __Author Rebuttal 1.2__.\n\n### Q3: The serial numbers in some figures appear to be somewhat disorganized.\n---\nA3: Thank you for your feedback. The order of Figures 3 and 4 is indeed incorrect due to LaTeX’s automatic formatting. We will ensure this is corrected in the final version of the paper.\n\n### Reference\n---\n\n[1] Learning task informed abstractions.\n\n[2] Denoised mdps: Learning world models better than the world itself.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u7SuIhCH92", "reviewer": "Reviewer_iwu6", "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.", "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$?\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.**", "limitations": "Yes, this work has discussed its limitations in Sec. 6.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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).", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for your constructive comments and suggestions. We address each of your comments as follows.\n\n### Q1: I think the novelty of learning mask models to distinguish noise from the environment is limited. \n---\nA1:\nThank you for your professional analysis. However, we would like to elaborate further on the novelty of our method:\n- **The key idea behind SMG is to focus on the task-relevant features across training and testing scenarios.** The separated models structure is employed as an effective tool for extracting these task-relevant representations from observations.\n- Compared to previous approaches, SMG introduces several key innovations to enhance performance in generalization settings. **We are the first to extract reward-relevant features via the Q-network in a model-free setting, provide the background encoder with accurate supervision, use the learned foreground to bootstrap the training process, and more fully utilize the learned mask through attribution augmentation.** We believe these improvements can inspire future research to advance the separated models structure into broader fields.\n- We have added a experiment to test TIA[5] in generalization setting; please refer to the __Author Rebuttal 2.1__. As shown in Table 1, TIA fails to generalize in video-background settings in all seven tasks. This further demonstrates the effectiveness of our improvements (two consistency loss and data augmentation techniques) to the separated models architecture under the generalization settings.\n\n\n### Q2: I'm curious about the performance of the proposed method in some more challenging settings, like RL-Vigen.\n---\nA2:\nWe have attempted to conduct experiments on three Adroit tasks introduced by [RL-Vigen](https://github.com/gemcollector/RL-ViGen/tree/ViGen-adroit?tab=readme-ov-file)[3]. However, we found that SAC-based algorithms (including SVEA, SGQN, and our SMG) perform poorly in these challenging tasks. The results reported in RL-ViGen are based on using VRL3 [4] as the backbone.\n\nVRL3 relies Adroit offline data to pre-train the encoder. When we attempted to build SMG upon VRL3, we discovered that the [download link](https://drive.google.com/drive/folders/14rH_QyigJLDWsacQsrSNV7b0PjXOGWwD?usp=sharing) for the offline data is currently invalid. As a result, we are unable to conduct experiments at this time. We have emailed the author to report this issue and will proceed with the experiments once we receive a response.\n\nNevertheless, we report the reconstruction process of SMG in Adroit-Pen; please refer to __Author Rebuttal 2.5__. SMG successfully outputs an accurate foreground image in this challenging task, and we believe VRL3 would benefit from the learned task-relevant representation.\n\n### Q3: 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---\nA3:\nThank you for your advice, the pseudo code is showed in __Author Rebuttal 1.4__\n\n\n### Q4: I'm curious about its performance with other methods.\n---\nA4:\nThis is a good suggestion. We have added experiments that use SVEA[2] as a backbone; please refer to __Author Rebuttal 2.3__. The performance of SMG improves in most tasks, which is reasonable since SVEA is a stronger algorithm than SAC.\n\n\n### Q5: Some other types of generalization also should be discussed.\n---\nA5: \nThank you for your reminder. We have revised Section 5; please refer to __Author Rebuttal 1.3__.\n\n\n### Q6: How do you determine the hyperparameter $\\rho$ ?\n---\nA6:\nWe invite the reviewer to refer to __Author Rebuttal 2.6__.\n\n### Q7: Why do reconstruction-based methods benefit generalization? \n---\nA7:\nThere might be some misunderstanding. **We do not think that directly using a reconstruction-based loss benefits generalization; rather, it may exacerbate the overfitting problem to task-irrelevant features. The advantage of reconstruction-based methods lies in their ability to improve sampling efficiency [1].** Our key idea is to enhance the agent’s generalization ability while utilizing the reconstruction-based loss to improve sampling efficiency. To achieve this, we introduce the separated models (please refer to lines 39–43 of our paper).\n\n### Reference\n---\n[1] Improving sample efficiency in model-free reinforcement learning from images.\n\n[2] Stabilizing deep q-learning with convnets and vision transformers under data augmentation.\n\n[3] RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization.\n\n[4] Vrl3: A data-driven framework for visual deep reinforcement learning.\n\n[5] Learning task informed abstractions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We revised the paper and added suggested experiments according to the reviewer’s comments. The detailed revisions are described as follows. The additional figures and a table are attached in the pdf file.

# 1. Revisions

### 1.1. Revise some typos and sentence

line 13: achieving free from overfitting -> achieving free from overfitting to task-irrelevant features.

line 38: further strengths -> further strengthens 

line 100: focused in -> focused on

line 458: perceptrons -> perceptions

line 523: Bootstraping -> Bootstrapping

line 536: Comparision -> comparison

### 1.2. Added a paragraph describing limitations in section 6

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.

### 1.3. Introduce more related works in section 5

In addition to view generalization, considerable research has focused on dynamic generalization [3,4,5] to develop a global model capable of generalizing across different dynamics. Additionally, several studies [6,7,8] have explored task generalization, which aims to enable learned agents to generalize to new tasks.

### 1.4. Add the pseudo code of SMG

| __Algorithm 1__ SAC with separated models |
| :----|
|__Denote__ network parameters $\theta$, mask ratio ρ, batch size $N$, replay buffer $\mathcal{B}$|
|__Denote__ policy network $\pi_{\theta}$, foreground encoder $f^+_{\theta}$,background encoder $f^-\_{\theta}$|
|__for__ each iteration time step __do__|
|$\qquad a,o',r\sim\pi_{\theta}(f^+_{\theta}(o)),\mathcal{P}(o,a),\mathcal{R}(o,a)$|
|$\qquad \mathcal{B}\leftarrow \mathcal{B}\ ∪\ (o,a,r,o')$|
|$\qquad$ __for__ each update time step __do__|
|$\qquad \qquad \\{o_i,a_i,r_i,o'\_i\\}_{i\in[1,N]}\sim \mathcal{B}$|
|$\qquad \qquad o^+\_i,mask_i\sim f^+_{\theta}(o_i)$|
|$\qquad \qquad o^-\_i\sim f^-_\theta(o_i)$|
|$\qquad \qquad o^{aug}_i \leftarrow o^+\_i\ast mask_i+\epsilon\ast(1-mask_i)$ // $\epsilon$ is sampled from image dataset|
|$\qquad \qquad L_{recon}\leftarrow L(o_i,\ o^+_i\ast mask_i+o^-_i\ast(1-mask_i))$ // Equation 2| 
|$\qquad \qquad L_{fore\ consist}\leftarrow L(o^+\_i,f^+_\theta(o^{aug}_i))$ // Equation 7|
|$\qquad \qquad L_{back}\leftarrow L(\epsilon,f^-_\theta(o^{aug}_i))$ // Equation 4|
|$\qquad \qquad L_{action}\leftarrow L(o_i,o'_i,a)$ // Equation 6|
|$\qquad \qquad L_{mask}\leftarrow L(mask_i,\rho)$ // Equation 3|
|$\qquad \qquad L_{q\ consist}\leftarrow L(Q_\theta(f^+\_\theta(o_i),a),Q\_\theta(f^+_\theta(o^{aug}_i),a))$ // Equation 8|
|$\qquad \qquad L_{aux}\leftarrow L_{recon}+L_{fore\ consist}+L_{back}+L_{action}+L_{mask}$ // auxiliary loss|
|$\qquad \qquad L_{critic}\leftarrow L_{q}+L_{q\ consist}$ // critic loss|
|$\qquad \qquad$ __update__ $\theta$ with $L_{actor},L_{critic},L_{aux}$|
|$\qquad$ __end for__|
|__end for__|
|$L_{q},L_{actor}$ is defined by SAC|
||

### 1.5 Revise line 136

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

# 2. Experiments

### 2.1. Add TIA [2] as a new baseline

The results are shown in the 4th column of __Table 1__. 

### 2.2 Add SAM-G [1] as a new baseline

The results are shown in the 5th column of __Table 1__.

Training SAM-G for one seed on our RTX3090 graphics card took nearly two days. Therefore, we are currently unable to complete all the training during the rebuttal period. Instead, we report the results from the SAM-G paper, noting that they used four times the batch size and twice the training time steps compared to SMG. Despite this, SMG still achieves comparable performance to SAM-G.

### 2.3. Add SVEA as a new backbone method

The results are shown in the 8th column of __Table 1__.

### 2.4. Add a new ablation of $L_{fore}$

The results are shown in the 6th column of __Table 1__.


### 2.5. Visualize the reconstruction process in adroit-pen

The results are shown in the __Figure 1__.

### 2.6. Further study the impact of ρ in peg-in-box

ρ does need to be set for different tasks but the choice of ρ can be roughly estimated based on the percentage of the task-relevant area. It is also possible to directly reuse settings from similar tasks. For instance, we used ρ=0.06 for three different tasks and ρ=0.12 for two different tasks.

In addition, an imprecise ρ only slightly affects the performance. We have added a new ablation for ρ in peg-in-box task with a much larger interval; The results in Figure 2 indicate that variations do not significantly influence the performance. As the optimal ρ being 0.12, even when setting ρ to 0, SMG still accurately masks out the task-relevant area and the performance drops by only 8%.

### 2.7. Add a new ablation of removing $L_{q}$ and $L_{action}$

The results are shown in the __Figure 3__.


# Reference

[1] Generalizable visual reinforcement learning with segment anything model.

[2] Learning task informed abstractions.

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

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

[5] Graph networks as learnable physics engines for inference and control.

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

[7] Task Aware Dreamer for Task Generalization in Reinforcement Learning.

[8] Learning modular neural network policies for multi-task and multi-robot transfer.

Meta-review: 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!
Decision: Accept (poster)