Learning Generalizable Agents via Saliency-guided Features Decorrelation
Authors: Sili Huang
Yanchao Sun
Jifeng Hu
Siyuan Guo
Hechang Chen
Yi Chang
Lichao Sun
Bo Yang
Venue: NeurIPS 2023 spotlight (2023)
Keywords: reinforcement learning
generalization

Abstract: In visual-based Reinforcement Learning (RL), agents often struggle to generalize well to environmental variations in the state space that were not observed during training. The variations can arise in both task-irrelevant features, such as background noise, and task-relevant features, such as robot configurations, that are related to the optimal decisions. To achieve generalization in both situations, agents are required to accurately understand the impact of changed features on the decisions, i.e., establishing the true associations between changed features and decisions in the policy model. However, due to the inherent correlations among features in the state space, the associations between features and decisions become entangled, making it difficult for the policy to distinguish them. To this end, we propose Saliency-Guided Features Decorrelation (SGFD) to eliminate these correlations through sample reweighting. Concretely, SGFD consists of two core techniques: Random Fourier Functions (RFF) and the saliency map. RFF is utilized to estimate the complex non-linear correlations in high-dimensional images, while the saliency map is designed to identify the changed features. Under the guidance of the saliency map, SGFD employs sample reweighting to minimize the estimated correlations related to changed features, thereby achieving decorrelation in visual RL tasks. Our experimental results demonstrate that SGFD can generalize well on a wide range of test environments and significantly outperforms state-of-the-art methods in handling both task-irrelevant variations and task-relevant variations.

Learning Generalizable Agents via Saliency-Guided Features Decorrelation
Sili Huang
Yanchao Sun?
Jifeng Hu
Siyuan Guo Bo Yang
Hechang Chen
Yi Chang
Lichao Sun?
1,8 Key Laboratory of Symbolic Computation and Knowledge Engineering of Ministry of Education
6.446,,SShool of Artificial Intelligence, Jilin University, China 2 Department of Computer Science, University of Maryland, College Park, USA Lehigg University, Bethlehem, Pennsylvania, USA
Abstract
In visual-based Reinforcement Learning (RL), agents often struggle to general- ize well to environmental variations in the state space that not observed during training. The variations can arise in both task-irrelevant features, such as background noise, and task-relevant features, such as robot configurations, that are related to the optimal decisions. To achieve generalization in both situations, agents required accurately understand the impact of changed features the decisions, i.c., establishing the true associations between changed features and decisions in the policy model. However, due to the inherent correlations among features in the state space, the associations between features and decisions be- come entangled, making difficult for the policy to distinguish them. To this end, we propose Saliency- Guided Features Decorrelation (SGFD) to eliminate these correlations through sample reweighting. Concretely, SGFD consists of two core techniques: Random Fourier Functions (RFF) and the saliency map. RFF is utilized to estimate the complex non-linear correlations in high-dimensional images, while the saliency map is designed to identify the changed features. Under the guidance of the saliency map, SGFD employs sample reweighting to minimize the estimated correlations related to changed features, thereby achieving decorrelation in visual RL tasks. Our experimental results demonstrate that SGFD can generalize well on wide range of test environments and significantly outperforms state-of-the-e art methods in handling both task-irrelevant variations and task-relevant variations.
Introduction
Learning control from high-dimensional visual observations is critical requirement for real-world applications and has received significant attention in recent years 43
18 10 Reinforcement
Learning (RL) has demonstrated its effectiveness in solving visual tasks by learning from compact representations of sensory inputs 122 [3]. However, current algorithms exhibit limited reliability when deployed in unseen environments, despite achieving satisfactory performance during training 7 488 It remains an open challenge to learn policies with good generalization across both changed lask-irreeevant and relevant features (24) Consider a simple example during the evaluation stage:
well-trained robot may fail to accomplish the task when confronted with iesk-irreeevant visual disturbances or changes to the task-relevant robotic arm, as illustrated in Figure I
Generalization to iask-irreeevant features in vision RL tasks has been widely discussed 49
Since a well-established consensus that task-irrelevant features should not affect the decision, agents are typically designed to maintain an invariant policy when these features are changed
"Corresponding Author. Bo Yang, Hechang Chen and Yi Chang.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Testing tasks with
(b) Training tasks
[c] Testing tasks with changed task relevntt feahures
changed task invelevann feahures
Figure Visualizing tasks across different variations. denotes the optimal decision, which should invariant task-irrelevant features but different changed task-relevant features. Note that "changed feature' refers to feature whose value has changed across tasks.
To training the invariant policy, previous works have proposed many solutions, such as bisimulation metric 2 33 data augmentation [46), and contrastive learning [29], to learn invariant representations that filter out the environmental variations. However, the invariant representations may ignore the changed task-relevant features that lead different optimal decisions, shown in Figure C Although Dunion et al. [15] proposed disentangled representation to preserve task-relevant features, this approach could only recover the sub-optimal decisions in the unseen environment. Therefore, a desired policy should not only remain invariant to task-irrelevant features but also adapt to changed task- rese-releannt features.
25 28 31.
Fundamentally, generalization fails because the agent does not understand how the changed features, whether they are task-irrelevant or relevant, affect its decisions. From the perspective of causal inference [14] 335 33], the associations between features and decisions are often confused with each other, which be attributed to the inherent correlations among input features in the training data. For instance, in a grasping task, heavier objects may more frequently be placed at lower horizontal heights, which could mislead the robot into adjusting its force based height. Therefore. promising approach is to climinate these correlations, enabling the agent to distinguish the true association between each input feature and decisions [5]] 26]. However, recent studies have indicated that achieving perfect decorrelation between all pairs of input features is challenging
44 This difficulty is especially pronounced in visual RL tasks, the high-dimensional images that maintain sequential relationships result in complex non-lineat and strong correlations among features [40].
45
To this end, we propose the Guided Features Decorrelation (SGFD) method to eliminate the inherent correlations among input features through sample reweighting. Concretely, SGFD consists of two core components designed to address the challenges presented by non-lineat and strong correlations in visual RL tasks. The first component is Random Fourier Functions (RFF), which has been demonstrated to assist in approximating the non-linear correlations with linear computational complexity [5I]. SGFD employs the cross-covarianc, matrix in conjunction with RFF to estimate the non-linear correlations in high-dimensional images. To address the strong correlations, SGFD integrates the second component, a classification model specifically designed to distinguish the sources of the images. Given that the values of changed features are unique each environment, these features can be pinpointed by computing saliency map 5 of the classification model. Guided by this saliency map, SGFD reweights the training samples to minimize the estimated correlations, with particular focus on the identified features. To evaluate SGFD, we conduct various experiments on visual RL tasks from DeepMind Control Suite 41 and Causal World [4, which range from scenes with different task-irrelevant features, such distracting background videos, to task-relevant features, such as the length of robotic arms. The experimental results demonstrate that our SGFD method can significantly improve the generalization across various types of environmental variations.
Summary Contributions. (1) We propose the SGFD model, which achieves improved general- ization in visual RL tasks, covering both task-relevant and task-irrelevan situations. (2) We design a sample reweighting method for RL tasks that encourages the agent to understand the impact of changed features on its decisions. (3) We validate the performance of our algorithm through vari- ous experiments and demonstrate that SGFD significantly outperforms state-of oate-of-theeare in handling changed task-relevant features.

---

2 Related Work
Generalization in Image-Based Reinforcement Learning. Numerous efforts have been made to achieve generalization in image-based RL tasks, employing strategies such as image augmentation
19 20), encoding inductive biases 3 29 36], and learning invariant representations 49 g among
other approaches [52] [ The fundamental idea behind image augmentation is to bolster the robustness of representations by introducing image perturbations 7 30 21]. Kostrikov et al. [25] applied an array of augmentation techniques, including cutouts, translation, and cropping, to the RL model. These augmented images can also be utilized to construct an auxiliary task for encoding inductive biases [32]. Laskin et al. (28) sought to learn representation that maximizes the similarity between different augmentations of the same observation. However, augmentation techniques may not alter the task-relevant features because they can impact the rewards signal [2) and optimal decision- making. Invariance techniques strive to learn a representation that disregards distractors in the image, thereby efficiently generalizing to unseen backgrounds [48 31] Recently, there has been an increasing interest in using varying rask-relevant situations as an additional experimental setup to evaluate model's generalization ability [47. To address this situation, Dunion al. (15) developed a representation that compresses changed task-relevant features into a small dimension, ensuring that the learned behaviors exhibit minimal changes in the test environment. However, this method requires slight fine-tumin fine-tuning to recover the optimal decisions in the testing environment.
Features Decorrelation in Generalization. A promising direction to achieve generalization is through features decorrelation, which climinates the correlation between inputs features [1] [66]. Shen et al. (37], Kuang et al. [7 attempted to achieve global decorrelation on linear regression tasks by using sample reweighting. Following this, Zhang et al. [51] empirically extended the reweighting method to deep learning models, noting its great potential. Recently, Xu et al. 44 conducted theoretical analysis of feature decorrelation and proved that perfect reweighting could help select task-relevant features regardless of whether the data generation mechanism is linear nonlinear. However, it is not easy to eliminate the correlation between all input features under finite training samples[45) 45) This is especially serious in RL tasks because data often exhibit a sequential relationship [4]. resulting in strong correlations between features. To compensate for the imperfection of sample reweighting, Yu et al. [45] introduce a sparsity constraint under the assumption that the correlation between the changed features and other features is significantly smaller. However, this assumption not applicable RL tasks because the changed features could be task-relevant and, therefore, significantly correlated with other features in the task. In contrast, we introduce a saliency- guided model identify the changed features without additional assumptions.
Preliminaries
Markov Decision Process. We define an environment Markov Decision Process (MDP) [4]. represented by a tuple M = (S, A, P, 7), where S denotes the high-dimensional state space, A is the action space, P(s'Is, a) is the state-transition function, R S x A - IR. is the reward function, and y E [0, 1) is the discount factor. A decision-making policy, parameterized by 9, is a function (a]ss) that maps state distributions over actions. The objective of RL agents that maximizes the expected cumulative return, expressed as Es, 1E (s,, at)].
find a policy
Soft Actor-Critic We utilize Soft Actor-Critic (SAC) [17 as base RL algorithm. SAC aims to maximize a variant of the RL objective augmented with an entropy term: EIE, 7",(s o.) aHH(((-(s))) To optimize this objective, SAC learns two models:
state-action value function Qo(s, a) and stochastic policy Tt (a(s), where and & are the parameters for Qo a) and &e (als). respectively. The parameters 0 of the state-value function are trained by minimizing the soft Bellman residual, given as:
Jo(0) E(n, (Qo(s,,
(C az) (r(se, at) + TE PP (V
where D is the replay buffer storing the data, Vi denotes the value function, and 0 are the target parameters of 0. The stochastic policy (a(s)s) is trained by maximizing the following objective
(Ø) Eg,ND Ear la
a.)],
(2)
where a is temperature coefficient.
3

---

Section 4.2: Sallency-guided model
Mulliple training eavironments
Hightrrababilll
Saliency map
Clessllirr
Probabinty vector
Replay buffer
multiplication
extractor
butch samples
Sample weights
Latent features
REF Section 4.1: Decorrelation with REF
The architecture our SGFD. SGFD aims reduce correlations in image features by reweighting samples. This involves five steps: (1) We fetch a sample batch from the replay buffer, which may come from multiple environments with different backgrounds or robot configurations. (2) The image in the sample is compressed by encoder into latent features. (3) Then, augment these features using multiple Random Fourier Functions to capture nonlinear correlations. (4) Concurrently, we train a classifier and apply saliency maps to detect features that shift across environments. (5) Finally, SGFD reweights samples to eliminate correlations between identified features and other features.
Figure 2:
Notations. For existing visual RL methods [1O 18 433, the state is typically compressed into compact vector representation z. Therefore, use z as the default in our calculations unless stated otherwise. For clarity, we define Zj as the j-th feature in the d-dimensional vector, and Z as the value of Z, in sample For instance, if Z2 corresponds to the angle of a robot arm, ži,j would represent the encoded value of the angle in image i.
4 Method
In this section, we introduce the SGFD method, sample reweighting technique designed to improve generalization across broad spectrum of visual RL tasks, as depicted in Figure Specifically, SGFD utilizes RFF to enrich image features, thereby enabling the estimation of non-linear correlations in high-dimensional images, as detailed in Section [.1] Then, SGFD integrates a classifier model and saliency map to detect changed features across environments. Based on the estimated correlation between identified features and other input features, SGFD adjusts the weights of the samples to achieve decorrelation, as elaborated upon in Section 4.2
4.1 Decorrelation with RFF
To eliminate the correlation between any pair of features, the first step is to quantify their correlation. Given a pair of features ZZ, in the RL states, we gather batch of samples from the replay buffer D, denoted as (Zi,5: 21.5). (Z2,i, z2.j)... iZn,i, 2..j). ). The primary challenge lies accurately
estimating the correlation between these features based on the available samples.
renowned method for evaluating independence is the Hilbert-Schmidt Independence Criterion (HSIC), which calculates cross-covariance operator in the Reproducing Kernel Hilbert Space([6). However, the computational cost of HSIC can be considerable, especially when the batch size compute the
is large. To overcome this issue and apply HSIC to modern RL models, independence using the Frobenius norm, shown to be equivalent to the Schhmitt norm in Euclidean space (38] The cross-covariance matrix can be formulated as follows:
E2,z,
(u(zz,) - Eiu(Z,)))* (v(Za,)) 1=1
E(v(Z,))),,
(3)
where E(u(Z,)) = 1 Ei=1 u(z,,), u and v are vectors, and each component is a function sampled from the RFF space, formulated as follows:
uss (z. um (Z.,) E HRFF, Vm, vzz,s) = (vi (2.,)) v2(2.,) ..., uM (2.,)),; um (Z..) E HRFF, Vm,
u(z..)) (uy (z.
(4)

---

where HRFF denotes the RFF space and has been demonstrated to approximate non- non-linear correlation between image features [3I] Specifically, we sample M functions from the RFF space, defined as HRFF {h x + (2cos(wr-v0) W ~ NV 1), vo Uniform(( (0, 2T where is sampled from standard normal distribution and v is sampled from a uniform distribution. The independence between Z, and Z, can be calculated by the Frobenius norm of the cross- cross-covariance matrix ||Σ 112 If ıı,,, 11 / is zero, the two features Z; and Z are considered independent.
We reformulate the task of correlation elimination as an optimization problem aimed minimizing the Frobenius norm of the cross- cosa-covre matrix ı| 1| To this end, we introduce set of learnable sample weights, denoted by IR, subject the constraint that Σk=1 W/ Based on the weighted samples, the independence can be defined as:
R (w (Z Elwu(Z.)))* (wEV(ZE)) (Z,)l,
Ezizyiw
Σ
(5)
n1
where [E[ [L(wu(Z,) [Z)] 1 Σm= (zz The correlation among all features in RL states can be eliminated by solving the following optimization problem:
n.
arg min W
Σ Igisjdd
1|S..Z,
where d is the dimension of Z. When the number of samples is not limited, Equation 6 is proved to have multiple solutions that achieve perfect decorrelation 27
4.2 Saliency-Guided Decorrelation with RFF
In practice, achieving desirable weights that ensure perfect independence between all pairs of features with finite samples is significant challenge [45]. Therefore, relying solcly on the features decorrelation approach introduced in Section 1.1 may not be sufficient. Furthermore, in RL tasks, the sequential nature of the data creates strong dependencies between features, further complicating the sample reweighting. To overcome this challenge, we propose a sesieency-giider model that prioritizes the most crucial correlations between changed features and other image features.
The saliency map [3). well-known computer vision method, uses the backpropagation algorithm to compute the attributions on neural networks. It is widely used to interpret models by visualizing important for decision-making. Inspired by the work of Bertoin et al. (7]. we incorporate this functionality into the training phase by introducing a classifier model that distinguishes the environmental source of states. The contribution of each feature can be determined by computing a saliency map. For example, the contribution of the i-th feature in the state is M,(f, Z) (z( /82, where f denotes the classifier. If the i-th feature changes across the environments, it will provide critical information for the classifier's decisions, leading to a large value of M,(f,)).
which pixels are
Given our focus on the correlations related to changed features, it is anticipated that when a pair of features is more likely to be part of the changed set, they should be prioritized for decorrelation. Accordingly, the probability that each feature belongs to the changed part can be calculated by the saliency map, i.e., p( (4) M. Eisisa M, We incorporate this term into the Equation ( revising as:
arg min Σ (Z,0)(Z)))1
15i<jdd
Based on the weighted samples, we can reformulate the Equation 2 as:
(d) = Σ w W.Earrr lo Qp(se, at]]
(8)
With the reformulated Equation g in which the changed features are considered independently from other features, the agent can learn the actual association with output decisions. More details of SGFD are described in Appendix A and Appendix B
5

---

SGED
TED
AMBS
SGQN
1-3 Iaste-rrelevant variations
10
10: taslerelevant variations
1: Walker-wall
6: Mr-Walkes-walk
3
2: Cheetah-run
MT-Chestal-uu
3: spigenspin
8: MT Fingerspin
8
4: Walker-nun
9: MI-Wälker-run
5: Fingestum
10: MI-Fingertumm
7
DBC
DR
Figure 3: The normalized performance generalization under changed task-irrelevant and relevant features. Each polygon represents one algorithm across 10 tasks. Each vertex of the polygon denotes the normalized performance, which matches counterclockwise from walker-walk to MT-finger-turm Note that the results related to task-relevant features are the average of the two settings presented in Table
Experiments
To carry out comprehensive investigation of SGFD, we have performed series of experiments on visual RL tasks, utilizing the DeepMind Control Suite [41] and Causal World [4]. We provide a brief introduction to the tasks, baselines, and evaluation metrics in Section 5.1 In Section 5.2 we verify the generalization capabilities SGFD across a variety of tasks, ranging from scenarios with diverse fusk-irrelevant features to those with task-relevant features. To gain deeper insight into the workings of SGFD, we evaluated its decorrelation potential and visualized the weighted samples, as detailed in Section 5.3 Then, conducted several ablation experiments RFF and saliency-guider model in Section 5.4 Finally, we investigated the interpretability of RL models in Appendix C
5.1 Experiment Setting
Environments.
The DeepMind Control Suite [1 is widely-used visual control toolkit that offers
a wealth of simulated animal behavior tasks. Drawing on previous works 448 50) we constructed environments with both background noises and varying robot settings test the generalization capabilities the SGFD. The backgrounds were sampled from the Kinetics dataset (23], while the robot settings followed those proposed by Zhang et al. [50]. Furthermore,
utilized the Causal World [4) which comprises series of robot manipulation tasks. In these manipulation tasks, we employed perception states to evaluate the potential of SGFD for decorrelation. More implement
details are provided in Appendix
Baselines. We compared our method against several recently established baselines, which we outline as follows. Our baseline for disentanglement representation is TED [15]. TED compresses changed features into small dimension, thereby ensuring that learned behaviors exhibit minimal change in testing environment. We used DBC [48] as a baseline method that learns invariant representations based on the bisimulation metric. We also compared with AMBS [1], an enhanced version of DBC that integrates a contrastive method. In addition, we compared with SGQN [7]. a novel method that uses data augmentation and the saliency map to ignore the task-in muskiirreevvant features in the images. Lastly, we included Domain Randomisation (DR) 42 in our comparison, a technique commonly utilized in control tasks.
The primary metric under consideration the cumulative reward from the testing environments, which reflects the generalization capabilities of the learned policy models. Additionally, we employed the Pearson correlation coefficient to assess correlation based on samples, both with and without weighting. Lastly, we utilized the saliency map an interpretability metric to visualize the learned policy models. All experimental assessments calculate the mean and standard deviation of the results across five seeds unless otherwise stated.
Evaluation Metrics.

---

Table 1: Generalization changed lask-irrelevant features. For each lask, cost 1ef steps in training environments and evaluate in the lesting environment under background noises.
Tasks
SGFD
TED
AMBS
SGQN
DBC
DR
walker-walk
959.1 26.3
871.5= 60.6
926.7 53.2
815.1=530 53,9
800.9 41.4
712.4 93.7
cheetah-run
599.6: 47.2
544.7=229 22,9
517.7= 73.4
332
312.1= 20.3
340.0= 44.0
iineer-spin
965.7::
932.1=71.0
50.
943.3=462 46.2
663.7: 68.7
860.8= 42.1
walker-run
420.7: 39.2
387.8=27.1
398.7=
317.2=3455 34.5
332.4 37.1
231.3+
finger-turn
984.3.1 11.5
963.9 94.5
966.71 37.0
971.3 26.0
931.2.1 41.6
947.2 21.7
Total
3929.5
3700.0
3734.9
3379.7
3040.3
3091.7
Table 2: Generalization to changed task-relevant features. For each task, we cost 1e6 steps in training envi- ronments and evaluate in the testing environment with different robot parameters. We consider two setups for evaluation: interpolation setup and extrapolation setup where the variations in the task-relevant features are interpolations and extrapolations between the variations the training environment, respectively.
Tasks
SGED
TED
AMBS
SGQN
DBC
DR
Interpolation
wawwalker-waik
549.4 42.5
471.9+
532.5+ 81.7
287.1 34.5
245.7+
343.8
MT-cheetah-rum
395.9 35.2
367.8+
298.64 53.5
40.9
191.7+ 23.5
216.3+ 33.1
MT-linger-spin
234.1:1 11.9
201.7=
161.3 17.3
135.7 16.9
21.
207.1 28.2
MT-walker-run
170.3.1 05.7
125.6 13.2
161.4 06.7
126.3 18.6
97.2.t 19.0
160.3 ::: 16.6
MT-finger-turn
923.7.1 26.1
748.9
821.81:
786.6 65.6
358.51
704.7::
Total
2273.4
1910.6
1975.6
1561.4
1114.7
1632.2
Extrapolation
MT-walker-walk
541.7:1 65.4
365.9 17.7
467.5+ 91.7
271.24 75.4
229.8 89.9
307 8 58.9
MT-cheetah-run
392.3:1 32.1
311.91 52.7
270.2+ 35.5
167.2.t
174.0#
196.6:: 6.::
MT-finger-spin
231.8:1 11.5
199.7 18.0
160.2.t 17.6
135.6: 11.3
221.41 43.0
197.1 21.5
MT-walker-run
170.0:1 07.2
126.7 13.2
156.2.1 07.5
118.9 118.9410 18.2
89.7 19.7
156.94 12.7
MT MT-finger-turn
917.3+ 22.6
743.6# 58.3
803.5 57.4
653.3 56.6
335.6# 56.5
611.7= 53.6
Total
2253.1
1747.8
1857.6
1346.2
1050.5
1470.1
5.2 Evaluation on the Generalization of SGFD
Evaluation on Changed Task-Irrelevant Features. In this experiment, we assessed the gener- alization capability of SGFD model with respect different task-irrelevant features. During the training phase, the policy model was allowed to interact with the environments million times. During the evaluation, we assessed the cumulative rewards in testing environment with previously unseen backgrounds. As indicated in Table and vertices 1-5 of the polygons in Figure 3 our model outperformed other baselines under the background noises. Although AMBS and SGQN aimed to direct the policy model to focus on task-relevant features, the learned representations may still contain partial background features that were effectively filtered out. Consequently, the policy model established associations with these tasks-irrelevant features, leading to incorrect gacision-making when the background changed. contrast, SGFD enabled the policy to understand that background features should not affect decision-makinn The results indicate that decorrelation is effective RL tasks with background noises and that SGFD exhibits strong generalization capabilities to new values of task-irrelevant features.
thereby achieving superior generalization performance.
Evaluation on Changed Task-Relevant Features To evaluate the generalization ability of SGFD with respect to task-relevant features, we conducted evaluations under two experimental setups: interpolation and extrapolation. In the interpolation setup, task-relevant features underwent changes that were interpolated between those in the training environments. In contrast, in the extrapolation setup, changes exceeded the range observed in the training environments. As depicted in Table 2 and at vertices 6- 10 of the polygons in Figure 3 SGFD demonstrated superior performance under both setups. Generally, extrapolation poses more significant challenge for generalization, as it necessitates the model to comprehend how changed features impact optimal decision-making. Owing the associations between changed features and decisions were not confused, our model outperformed other baselines, with more substantial performance gaps observed in the extrapolation
7

---

110
1.0
$10
1.0
$10
1.0
650
0.00
0.5
0.5
0.5
6:11
S3T 033
5g
$2
$2
$1
0.0
0.0
0.0
5,5255
$10
$2
$10
5,52
$10
(a) Original states
(bb) Deconelation
(c) SGFD
Figure 4: Pearson correlation coefficients among features: a) raw data, b) on weighted data, c) on saliency- guided weighted data. The features 81 to 81a represent the state information of the environment and act as conditional inputs for decision- making (a(a)), such as the coordinates of the object or the angle of the manipulator. The feature $10 is the one that varies across the environments; thus, the correlations between other features play crucial role generalizations, The strength of correlation reflected by the color intensity- darker the color, the stronger the correlation.
States Env
Reweighting
1.05
1.03
1.11
1,23
1.13 >1.0
1.21
0.66
0.86
0.98
0.84
L10
<10"
States in Env b
Reweighting
0.94
0.91
$1.
0.96
0.99
0.73
0.78
$1.0%
0.75
0.80
1.46
>1.
1.82
(a) 20 states from environments with different background
States Env c
Reweighting
0.97
0.91
-1.0
0.89
0.93
0.98
0.97
<1.0
0.94
0.96
1.21
>10
.115
States in Env c
Reweighting
1.03
1.08
1.07
1.08
1.11
0.97
0.84
0.93
0.83
$10
10
states from environmentss with differeett robot parameters
Figure 5: Visualization of weighted samples obtained from different environments with background noises and robot parameters. States that belong to the same dotted box with a specific color are considered semantically similar. SGFD balances the weights of samples to eliminate the correlations hetween changed features and other features. As shown in (a), due to the interference of the background noises, the robot fell more often (states the green dotted box) than did without the background. SGFD reduces the weighting of states with correlation
backgrounds and enhances the weighting states without backgrounds, effectively neutralizing between background and body posture. A comparable result was observed in (b), where the reweighting process aimed to equalize the number of states in each group across environments with different body lengths.
setting. These results demonstrate that SGFD can effectively adapt to new values of task-relevant features, showing its robust generalization capability in a variety of scenarios.
5.3 Case Studies on the SGFD
Correlation of the Weighted Samples. To evaluate the decorrelation capability of our model, we implemented 'Picking' task in Causal World. This task involves robotic arm grasping objects, enabling the model to access the state with physical meaning. We recorded the Pearson correlation coefficients among features under three conditions: a) on raw data, b) on weighted data, and c) on saliency-guidee weighted data. Among these features, we singled out $1u. a feature that represents mass, as the focal feature for environmental variation. Consequently, when object with new mass, it becomes critical to carefully examine the correlations
model given the
task of handling between the mass feature and other features. As depicted in Figure 4 due to limited samples and strong correlations, the effectiveness of direct decorrelation of all features is significantly diminished.

---

RFF-decomelation
doxomiation
rD dexommiiiion
sulleny--uiddd decomrlation
doramsstiom docommiiiion
620
500
400
reaand
reward
840
250
200
460
160 (b)
Cheetalt-run
MI--chetah-nun
Chestah-run
MT-chootah-uun 1.0
(a)
Figure 6: Ablation studies. (a) The performance difference of ablation studies the RFF. (b) The performance difference of ablation studies on the saliency-guided model. The x-axis denotes steps the training environments. The y-axis denotes the cumulative reward recorded in the testing environment with background noises (Chectah-run) or varied robot parameters (MT-chectahh-rum).
In contrast, our SGFD model can effectively reduce the correlations between the mass feature (810) and other features. We also conducted an experiment involving changes in multiple features and obtained similar results. Detailed information about this experiment is described in Appendix C.1
Visualization of the Weighted Samples. To delve deeper into the implications of weighted samples from our SGFD method, we collected 20 state instances during training, each exhibits distinct backgrounds or robot parameters from various environments. The identified states were sorted into three groups in each environment, each consisting of multiple similar state instances. As illustrated Figure 5 (a), the green dotted boxes indicate several states with "fall" posture. Due the impact of background noises on the policy, we identified 4 similar states in the environment a and similar states the environment b. After reweighting process, both environments a and had approximately 3 similar states, effectively neutralizing the correlation between background and body posture. A similar outcome was observed F wherein the reweighting process endeavored to balance the number of states in each group across environments with varied body lengths. As result of the correlation removal, the policy model was better equipped to accurately capture the association between the changed features and decision-making.
5.4 Ablation Studies
Ablation on RFF. To examine the role of RFF, we compared the complete SGFD algorithm, the decorrelation method without RFF, and the algorithm without any decorrelation. The testing en- vironments comprise varied backgrounds and robot parameters. As depicted in Figure 6 (a), the decorrelation without RFF outperforms the approach without decorrelation but is significantly inferior SGFD. This discrepancy arises due to the complex nonlinear correlations of features within the image, which RFF can approximate. Without RFF, the basic decorrelation method can only eliminate linear correlations, thus achieving limited improvement.
Ablation on Sauiency-Gudeed Model. We conduct a separate ablation study on the saliency-guides model to evaluate its impact on the generalization capability of the overall model. Analogous to the RFF study, we test the performance of three distinct approaches: no decorrelation, decorrelation, and saliency- guided decorrelation. As shown in Figure (b). the performance significantly deteriorates when the saliency-guin guided model is removed from SGFD. This outcome is predictable because directly decorrelating all pairs of features is challenging, thereby complicating the model's 's ability to discern the associations between the variant features and decisions.
We conduct additional ablation studies for further insights into the roles of RFF and the saliency- guided model. Detailed analyses and results of these studies are available Appendix C.2 Moreover, report on the learning curve of the classifier, as shown in Appendix C.3
Conclusion, Limitations, and Broader Impact
In this work, we introduce SGFD, sample reweighting method designed enhance generalization in visual reinforcement learning tasks across environments with unseen task-irrelevant and task-relevant features. SGFD composed of two core components: RFF and a saliency-guidet model, which empower the RI. agent to understand the impact of changed features on its decisions. We assess SGFD using the DMControl benchmark, where it demonstrates significant improvements in generalization across various environmental variations. Furthermore, our case studies conducted on the Causal

---

World benchmark, as well as the visualization of the weighting, highlight the superior decorrelation achieved by our model. The results from SGFD underscore the potential value of sample reweighting in generalization RI. tasks.
In terms of limitations, consecutive frames to approximate a fully observable condition. This might not be optimal for all scenarios, particularly when the changed features cannot directly observed. potential solution is to train an inference model that uses a small amount of data from the deployment environment to rapidly predict the potential variations. Regarding societal impacts, we do not anticipate any negative consequences stemming from the practical application of our method.
13 relies stacking note that our work, like many others 48
Acknowledgments
This work was supported by the National Key R&D Program of China under Grant Nos. 2021ZD01 12501 and 2021ZD01 12502: the National Natural Science Foundation of China under Grant Nos. U19A2065, U22A2098, 61976102, 62172185, 62206105 and 62202200; the International Cooperation Project of Jilin Province under Grant Nos. N0s.202200000099G
References
[1] Julius Adebayo, Justin Gilmer, Michael Muelly, Ian Goodfellow, Moritz Hardt, and Been Kim. Sanity checks for saliency maps. Advances in neural information processing systems, 31, 2018.
[2] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[3] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Con- trastive behavioral similarity embeddings for generalization in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[4] O Ahmed, F Träuble, A Goyal, A Neitz, M Wüthrich, Y Bengio, B Schölkopf, and S. Bauer. Causalworld: robotic manipulation benchmark for causal structure and transfer learning. arXiv peprint arXiv:2010.04296 2020.
[5] Akanksha Atrey, Kaleigh Clary, and David Jensen. Exploratory not explanatory: Counterfactual analysis of saliency maps for deep reinforcement learning. In International Conference on Learning Representations, 2020.
[6] David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization in reinforee- ment learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
[7] David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[8] Ioana Bica, Daniel Jarrett, and Mihaela van der Schaar. Invariant causal imitation learning for generalizable policies. In Advances in Neural Information Processing Systems (NIPS), 2021.
[9] Yihan Cao, Siyu Li, Yixin Liu, Zhiling Yan, Yutong Dai, Philip S Yu, and Lichao Sun. A comprehensive survey of ai-generated content (aigc): A history of generative ai from gan to chatgpt. arXiv preprint arXiv:2303.042: 2023.
[10] Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In International Conference on Robotics and Automation (ICRA), 2019.
[11] Jianda Chen and Sinno Jialin Pan. Learning generalizable representations for reinforcement learning via adaptive meta-learner behavioral similarities. Proceedings of the International Conference Learning Representations (ICLR), 2022.
10

---

[12] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICML), 2020.
[13] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. Decoupling represen- tation learning from reinforcement learning. In International Conference on Machine Learning (ICML), 2020.
[14] Peng Cui and Susan Athey. Stable learning establishes common ground between causal inference and machine learning. Nature Machine Intelligence (12):100-115, 2022.
[15] Mhairi Dunion, Trevor Mclnmee, Kevin Luck, Josiah Hanna, and Stefano Albrecht. Temporal disentanglement of representations for improved generalisation in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[16] Arthur Gretton, Kenji Fukumizu, Choon Teo, Le Song, Bernhard Schölkopf. and Alex Smola. kernel statistical test of independence. Advances in neural information processing systems, 20, 2007.
[17] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning (PMLR), page 1861-1870, 2018.
[18] Aleksi Hamalainen, Karol Arndt, Ali Ghadirzadeh, and Ville Kyrki. Affordance learning for end-toOend visuomotor robot control. In International Conference on Intelligent Robots and Systems (IROS), 2019.
[19] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 13611-13617. IEEE, 2021.
[20] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learnin with convnets and vision transformers under data augmentation. Advances in neural information processing systems, 34:3680-3693, 2021.
[21] 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.
[22] Jifeng Hu, Yanchao Sun, Hechang Chen, Sili Huang, Yi Chang, Lichao Sun, et al. Distributional reward estimation for effective multi-agent deep reinforcement learning. Advances in Newral Information Processing Systems,
2022.
[23] Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint 1rXiv:1705.06955 2017.
[24] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint arXhv:2111.09771 2021.
[25] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[26] Kun Kuang, Peng Cui. Susan Athey, Ruoxuan Xiong, and Bo Li. Stable prediction across unknown environments. In proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pages 1617-1626, 2018.
[27] Kun Kuang, Ruoxuan Xiong, Peng Cui. Susan Athey, and Bo Li. Stable prediction with model misspecification and agnostic distribution shift. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2020.
[28] Michael Laskin, Aravind Srinivas, and Pieter Abbeel, Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 5659-5650. PMLR, 2020.

---

[29] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
[30] 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-1998 19895, 2020.
[31] Bonnie Li, Vincent François-Lavet, Thang Doan, and Joelle Pineau. Domain adversarial reinforcement learning. arXiv preprint arXiv:2102.07097, 2021.
[32] Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 369868 2020.
[33] Bogdan Mazoure, Ahmed M. Ahmed, Patrick MacAlpine, and R Devon Hjelm. Cross- trajectry representation learning for zero-shot generalization in rl. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
[34] Martin I. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[35] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning, pages 5389-5400 PMLR, 2019.
[36] Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data-efficient reinforcement learning. Advances Neural Information Processing Systems, 34: 12699- 2021.
[37] Zheyan Shen, Peng Cui, Tong Zhang, and Kun Kuang. Stable learning via sample reweighting. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2020.
[38] Eric V Strobl, Kun Zhang, and Shyam Visweswaran. Approximate kernel-based conditional independence tests for fast non con-parametric causal discovery. Journal of Causal Inference, 7(1), 2019.
[39] Jiawei Su, Danilo Vasconcellos Vargas, and Kouichi Sakurai. One pixel attack for fooling deep neural networks. IEEE Transactions on Evolutionary Computation, 23(5):828 841, 2019.
[40] Yanchao Sun, Shuang Ma, Rainesh Madaan, Rogerio Bonatti, Furong Huang, and Ashish Kapoor. Smart: supervuper multi-task pretraining with control transformers. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[41] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, and Andrew Lefrancq. Deepmind control suite. arXiv peprint arXiv:1801.0080 2018
[42] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation the real world. In 2017 IEEE/RSJ international conference intelligent robots and systems (IROS), pages 23-30. IEEE, 2017.
[43] Mel Vecerik, Jean- Baptiste Regli, Oleg Sushkov, David Barker, Rugile Pevceviciute, Thomas Rothorl, Raia Hadsell, Lourdes Agapito, and Jonathan Scholz. S3k: Self-supervised semantic keypoints for robotic manipulation via multi-view consistency. In In Conference on Robot Learning (PMLR), 2021.
[44] Renzhe Xu, Xingxuan Zhang, Zheyan Shen, Tong Zhang, and Peng Cui. A theoretical analysis on independence-ddiver importance weighting for covariate-shift generalization. In Interna- tional Conference on Machine Learning, pages 24803-24829 PMLR, 2022.
12

---

[45] Han Yu, Peng Cui, Yue He, Zheyan Shen, Yong Lin, Renzhe Xu, and Xingxuan Zhang. Stable learning via sparse variable independence. Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2023.
[46] Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: aware lipschitz data augmentationfor visual reinforcement learning. arXiv peprint wwxhv:2222.099822 2022.
[47] Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang. Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning, arXiv preprint urXiv:2212.08866 2022.
[48] Zhang. R Mcallister, R Calandra, Y Gal, and Levine. Learning invariant representations for reinforcement learning without reconstruction. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[49] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning (PMLR), 2020.
[50] Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pincau. Learning robust state abstractions for hidden-parameetr block mdps. arXiv preprint urXiv:2007.07266, 2020.
[51] Xingxuan Zhang, Peng Cui, Renzhe Xu, Linjun Zhou, Yue He, and Zheyan Shen. Deep stable learning for out-of-distribution generalization. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, 2021.
[52] Ce Zhou, Qian Li, Chen Li, Jun Yu, Yixin Liu, Guangjing Wang. Kai Zhang, Cheng Ji, Qiben Yan, Lifang He, et al. A comprehensive survey pretrained foundation models: A history from bert chatgpt. arXiv preprint arXiv:2302.094 2023.
13

---

Appendix
A Pseudocode of Saliency-Guided Features Decorrelation
In Algorithm[] | we introduce the training process of SGFD, At each step, we apply the policy model To to interact with training environments and sample batch of transitions along with environment labels, as described in lines 3-5. Before the sample reweighting, we assess the current classifier's accuracy. If the accuracy exceeds 0.9, we proceed with sample reweighting, as described in lines 6-9. In contrast, if the accuracy is lower and > 1e5, we update the classifier based on the cross-entropy loss, described in lines 9-13. In experiments, the classifier could be quickly converged. Based on the converged classifier, SGFD uses Equation (7) to learn a set of sample weights, as described in lines 14-17. This process also converged quickly, as only 128 parameters (the batch size) needed to be updated. Finally, SGFD uses Equation (T) and Equation (8) to update state-action value function Qo and policy model T based the weighted data, described in lines 18-20.
Algorithm 1 Saliency-Guided Features Decorrelation
1: Input: Initialize policy model TT. state-action value function Qo, classifier model fø+ and set of training environments with environment labels fe', ,e,...,,e* e where ex is K- T-dimensional one-hot vector and the k-th component is
2: Output: The learnt policy T
3: for (= 1 T
4: Applying the policy model T to interact with training environments
5: Sampling a batch of transitions along with environment labels (s;, ag Ti, S4, et) ~ D
6: for iler to 10 do
7:
if / < 1e5 or the accuracy of ((s) > 0.9 then
8:
Break this cycle
9:
else
10:
Sampling batch of transitions along with environment labels (s,, aj, Ti, Tr, ct) ~ Update the classifier model fo by minimizing the following loss
11:
((b) = - Σm Σ et,, (og()=(s))));
where cx denotes the j-th component of vector et, fo((s); denotes the predicted probability of the j-th class end if end for
12:
13:
14:
Initialize set of sample weights that can be learned
15:
for iter = to 10 do
16:
Update the sample weights by minimizing Equation end for
17:
18:
Update the state-action value function Qs by minimizing Equation 1I
19:
Update Te by minimizing Equation (
20:
end for
B Implement Details
Experimental Setting Details. For task-irrelevant variations, follow the prior work 330 sample several images from Kinetics (23) In training, we create four parallel environments with different image backgrounds. In testing, we create new environment with an unseen image background. For task-relevant variations, we follow the prior work (50) to revise the robot configures. The controllable configuration contains 10 different settings for each task, e.g., the length of the robot torso. In training, we create four parallel environments with different robot configures, where the values of configures are sampled from a predefined distribution, i.e., a uniform distribution between and In testing, we exhibit evaluations under two setups: interpolation and extrapolation. In the interpolation setup, manner interpolated between those in the training environments, whereas in the extrapolation setup, changes exceeded the range of those in the training environments, i... the tenth setting with the farthest training distribution. The visualizations are illustrated in
task-relevant features changed
Figure
14

---

Table 3: Architectures of SGFD.
Hyperparameter
Value
Convolutional layers
4
Latent representation dimension
50
Image size
(84,84)
Stacked frames
3
Encoder architecture
kernel size
3 x
Channels
3
stride
2 for the first layer, 4 otherswise
Activation
ReLU
MLP layers
2
Policy and value function
Hidden dimension
1024
Activation
ReLU
MLP layers
2
Classifier architecture
Hidden dimension
128
Activation
ReLU
Table 4: Hyperparameters SGFD.
Hyperparameter
Value
Optimizer
Adam
Learning rate
1e-3
Action repeat
Training for policy
Replay buffer capacity
1000000
Batch size
128
Target soft-update rate (r)
0.01
Actor update frequency
Optimizer
SGD
momentum
0.9
Training for sample reweighting
Learning rate
le-2
weight decay
le-4
The number of RFFs m
5
For each experiment, we report the mean and standard deviation over 10 random seeds.
Compute. Experiments are carried out on NVIDIA GeForce RTX 3090 GPUs and NVIDIA A10 GPUs. Besides, the CPU type is Intel(R) Xeon(R) Gold 6230 CPU 2. 2.10HHz. Each run of the experiments spanned about 12-24 hours, depending on the algorithm and the complexity of the task.
the same encoder architecture, which consists of convolutional layers. The encoder weights are shared between the policy and the state-action value function. Each convolutional layer has 2, 2, and all other layer has a stride 1. There is ReL.U activation between each convolutional layer. The convolutional layers are followed by trunk network with linear layer, layer normalization, and finally, tanh activation. As we focus on the sample reweighting, the encoder is updated by referring the AMBS method.
Architectures. Following AMBS [1.
3 kernel size and 32 channels. The first layer has stride
The policy Te and the state-action value function 2o are both 2-layer MLPs with hidden dimension of 1024. We apply ReLU activations after each layer except the last layer. The classifier is imple- mented with 2-layer MLPs with a hidden dimension of 128. We also apply ReLU activations after each layer except the last layer. Table shows the values of the architectures for the encoder, the policy, the value function, and the classifier.
Hy Hlpperparameters Table 4 shows the hyperparameters used in our SGFD model.
15

---

cheelah
finger
walker
Training enviromments with different robot par parameters
tesling with unsenn robot parameters
environments
testing environment
with difffrent backgrouad
with unseen robot parameters
visualizations rohot configurations
The visualizations background noise
Figure 7: The visualization of training environments and testing environments.
$10
1.0 $
1.0
$0
1.0
5g
MC
0.5
0.5
0.5
SP 50 CM
53
$2
6:1
52
$q 5g
5g
Sg
S10
0.0
51
53
5g $10
0.0
52
Sg Sg $10
0,0
(a) Original states
(b) Docordation
SGFD
Figure 8: Pearson correlation coefficients among variables: a) on raw data, b) on weighted data, and c) on saliency-guided weighted data. The variables Ss, Sp, and S10 are the ones that have been changed among the environments.
C Additional Experimental Results
C.1 Additional Correlation Studies on the Weighted Samples
To further evaluate the decorrelation capability of our model, we implement the 'Pushing' task in Causal World and allow the model to access state features that have physical meaning. We record the Pearson correlation coefficients among features under three conditions: a) on raw data, b) on weighted data, and c) saliency- guided weighted data (SGFD). Among these features, the Ss. Sg. and Sjo. which represent the length, width, and height, as the features of interest for environmental variations. Therefore, when the model is required push an object with new size, the correlation between size and other features needs to be carefully considered. As shown in Figure g due limited samples and strong correlation between features. the effectiveness of directly decorrelating all features is greatly reduced. In contrast, SGFD is capable of effectively reducing the correlation between the size features (Sg, Sg, and S10) and other features.
designate
C.2 Additional Ablation Studies
We performed additional ablation experiments test the effect of RFF and the saliency-guided
model. As depicted in Figure / (a) and Figure 9 (b), when RFF is removed, the generalization performance drops significantly.
contrast, the global decorrelation can cause over-adjusting of samples weighting and impair the performance, as demonstrated by the reward curves in Figure (c). Due to the need to eliminate correlations between variables using limited samples, the model may assign very low weights to certain training samples. However, this can potentially harm the efficiency of training. Thanks to the caliency-guides model, SGFD balances both efficiency and decorrelation.
C.3 Evaluation on the Classification Model
We recorded the converge curve of the classifier during training. The classifier starts to update when the replay buffer collects le5 steps. To avoid overfitting, stopped to update the classifier when the accuracy exceeded 0.9. As shown in Figure the classifier converged quickly and achieved the accuracy of 0.9, which means that the classifier does not take up many computing resources.
16

---

RFF-ddxomdaation doconelation no decomrllatin
suliexy--uiddd decomalaiion
doondaatoon- doconelation
430
560
500
reward
rewend
406
290
250
300
330
420
1.0% 0
(a) walker-num
(b) MT-walka-walk
(c) walkornun
(c)) MT-walkoewlll
Figure 9: Ablation studies. (a)-(b) The performance difference of ablation studies the RFF. (c)-(d) The guided model. The x x-axis denotes steps in training environments. The y-axis denotes the cumulative reward recorded in the testing environment with hackground noises (walker-run) or varied rohot parameters (M (MT-walker-walk)
performance difference of ablation studies on
Accura
0.9
0.9
0.9
0.9
0.9
le6
16
e6
106
(a) walker walk
(b) chestah
fingerspin
(cD walkernu
(e) finger tum
All
Accutay
0.9
0.9
0.9
0.9
Je6
1e6
8 MT MT-walleewwal (g) MT-cheetah-run
&n MT-finger spin (i) MTwwaaleernm
1.0 0
(i) MT-finger-tum
Figure 10: The accuracy of the classifier during training. The x-axis denotes the number of training steps taken training environments, while the y-axis indicates the accuracy. The classifier starts to update when the replay buffer collects le5 steps. The line corresponds to the 0.9 threshold, which indicates that the classifier stops updating if its accuracy exceeds 0.9.
C.4 Saliency Maps on the Policy Models
An intuitive approach to explaining visually-based models involves identifying pixels that have significant impact the final decision [1] To determine whether the model focuses on the object or background during decisions, we visualize
gradient the action with respect to the input pixels. Experiments conducted in environments with varying backgrounds and robot parameters, as depicted in Figure IT] Saliency maps of the baseline models reveal that various backgrounds draw considerable attention from the policy model while failing make decisive contributions to our model. Moreover, SGFD can notice the changing aspects of the robot's 's body that other bascline models overlook. As result, the saliency maps demonstrate that the decorrelation benefits the model's generalization in both task-relevant and fask-irrelevant situations.
C.5 Additional Visualization of the Weighted Samples
In this experiment, we significantly expanded the dataset to include 300 state instances. which is substantially larger than the 20 states mentioned in Section The identified states were also sorted into three groups in each environment, each consisting of multiple similar state instances. Moreover, we carefully analyzed the distribution of these states within each group. and the corresponding proportions are visually depicted in Figure T2 illustrated in Figure [2 (a), the green dotted boxes indicate that 50% of the states in environment A and 35.4% of the states in environment B exhibited similarities. After the reweighting process, both environments and had approximately 42% similar states, effectively neutralizing the correlation between background and body posture. similar outcome was observed Figure 12 (b), wherein the reweighting process endeavored to balance the number of states in each group across environments with varied body lengths.
17

---

Observation
SGFD
TED
AMBS
SGQN
DBC
DR
Figure 11: Visualization of salieney maps produced by the SGFD and other baselines when the rohot with different backgrounds and robot parameters. The brighter the pixel is, the more contributions it makes to the decision.
C.6 Evaluation on the Generalization of SGFD
Evaluation Changed Task-Irrelevant Features. We report the training curves of SGFD in environments with different task-irrelevant features, where the end points of the curves correspond to the results in Table | As illustrated in Figure [3 our model achieved superior performance compared to other baselines in previously unseen background environments. As the task becomes difficult (Figure 13 (b) and Figure 13 (d)), the advantage SGFD achieves more significance. The results demonstrate that the decorrelation is suitable for RL tasks with different mask-irreevaan features, and SGFD exhibits strong generalization capabilities to new values task- iere-ireeeevat features.
Evaluation on Changed Eask-Relevan Features. We also report the training curves SGFD in environments with different task-relevant features, where the end points of the curves correspond the results Table 2 As shown in Figure 14 the reward curve exhibited a consistent positive correlation between the number of training steps and the average reward in the testing environment under both interpolation and extrapolation setups. Generally speaking, the extrapolation setting is more challenging for generalization because requires the model fully understand how changed features affect optimal decisions. Benefiting from decorrelation, our model achieved superior performance compared other baselines, where the gaps were more significant under the extrapolation setting, as illustrated in Figure 14 (f). Our results demonstrate that the decorrelation facilitates our model's s understanding the true associations between changed features and decisions, leading to strong generalization on new values of task-relevant features.
18

---

25.3%
50%
States Env
24.7%
0.86
1,23
0.88
Mean reweighting
11.3%
35. 35.4%
States Env B
53.3%
1.26
1.16
Mean reweighting
300 states environments with different background
44%
10%
States in Env
16%
0.91
1.33
0.83
Mean reweighting
21.4%
31
11.3%
States in Env
47.3%
1.08
0.82
.21
Mean reweighting
(b) 300 states from environments with different robot parameters
Figure Visualization weighted samples obtained from different environments with background noises and robot parameters. States that belong to the same dotted box with a specific color are considered semantically similar. SGFD balances the weights of samples to climinate the correlations between changed features and other features. As shown in (a), the walker agents in Env A and Env B have different backgrounds, entangled with different distributions of the body posture. SGFD reduces the weighting of states with backgrounds and enhances the weighting states without backgrounds, effectively neutralizing the correlation between background and body posture. comparable result was observed in (b), where the reweighting process aimed to equalize the number of states in each group across environments with different body lengths.
SGPD
TED
AMBS
9GQN
DBC
DR
1000
1000
500
1000
renyard
300
500
300
05
10
0. (b) chestal--un
16
0.
0.
1.0
166
055
10G
(a) walker-wall
(c) singersin
(cD walker nun
(e) finger
Figure 13: Generalization changed task ineeevrteleantt features. The x-axis denotes the number of training steps taken training environments, while y-axis indicates average reward in testing environment under different backgrounds. The reward curve demonstrates that our SGFD model generalizes well to new values of lask-irrelevant features.
SGFD
TED
AMBS
SGON
DBC
Interpelttonn
400
200
750
100
300
(6) wT-wwells--was walk
1/66
(b) A-T-cheedah--mr
Lo
(c) MT-Oreer spin
(d) MT wal-wallennn
1d0
(e) Mi-finge
05
120
200
200
100
100
500
12 105
(LS
05
055
0.5
( wi-vallarwalk
(g) MT-chastal--um
(b) Mi-lineessiin
MT-waller-m
M..-fligee-baam
Figure 14: Generalization to changed task-relevant features. The x-axis denotes the number of training steps taken in training environments, while the y-axis represents the average reward in the testing environment with different robot parameters. We consider setups for evaluation: an interpolation setup ((a)-(e)) and extrapolation setup ((f)(j), where the changes the task-relevant features are interpolations extrapolations between the changes in the training env environments, respectively. The reward curve shows that our SGFD generalizes well to the new values of task-relevant features.
19


Reviews:
[{"review_id": "OfOMSe1Pgv", "reviewer": "Reviewer_tXZg", "summary": "This paper presents a method to weight experiences in reinforcement learning so that the state features they depict are maximally decorrelated. The assumption is that learning with decorrelated features leads to better generalization to unseen task instances because the policy does not confound the role of different features when making decisions. The idea is novel and interesting and the experimental evaluation includes several simulated embodied AI tasks and comparisons to other baselines.\n", "questions": "- Explain the features per task and where they come from. Discuss the applicability of the method if it requires pre-specifying optimal features. How would the method work with learned latent features? Could you include an evaluation on that?\n- Can you include completely distracting features to evaluate the robustness of the solution?\n- Please, discuss what would be necessary for this method to work directly with images as input. It feels strange to only go “the other way”: from features to images (saliency) in order to be able to check when they are co-ocurring, but there is no discussion about the normal direction: from images to features.\n- How does the method scale with the number of features? And with the number of samples? What is the computational cost? Can you experiment with that and provide some evaluation?\n", "limitations": "- The limitations section is acceptable, but several significant limitations for the method to be truly impactful are not discussed (see my comments above). Please, include them.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The presented method is novel. The ideas of the paper are well-grounded.\n- The method seems mathematically sound. The appendix includes all relevant proofs.\n- The experimental evaluation, including several baselines and different tasks, demonstrates the strengths of the algorithm\n- The problem addressed in the paper is critical for the robot learning community: improving the generalization capabilities of trained policies.\n\n", "weaknesses": "- The writing could be improved. The text focuses on describing the steps but would benefit from explaining more the reasoning behind the design decisions. \n- The assumptions on the problem structure are unclear. What type of relationship between features and actions and features-to-features should be present in the problem? What happens if features are not (cannot be) decorrelated?\n- It is unclear what the features are. I can’t find that information either in the text or in the appendix, maybe I’m missing it. Please, clearly discuss what the features are and whey they come from. Include it in the main text, e.g., in the experimental setup description. Explain and discuss how they are selected/created. In these times of representation learning, especially for vision-based control, is this a strong limitation of the method?\n- Fig. 2 is quite complex. What are the colors of the “enriched features” indicating? What are the “G”s indicating? This figure could serve as a good summary of the method but it needs to be improved for that and include a more complete caption.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly encouraged that the reviewer finds our method novel and effective.\n\n**Reply to the weakness**\n>**W1. Explaining more the reasoning behind the design decisions.**\n\nAs suggested, we introduce an example to explain why feature decorrelation can promote generalization, as shown in Figure 1 in the \"global\" response PDF file. \n\n**Figure 1: the motivation for feature decorrelation.** Consider a classifier trained on images labeled 'cat': three with sofas and one with grass in the background. The model may falsely link sofas to cats due to frequent pairing. To mitigate this issue, we can reweight the images to decorrelate the features between backgrounds and animals. In RL, features like varied background noise can correlate with the robot's state, as in Figures 5 (main text) and 11 (Appendix). Such correlations may lead agents to form spurious associations, altering optimal actions with background shifts.\n\nThis example can be first proposed in the method section and then fluently lead to the motivation of feature decorrelation.\n\n>**W2. What type of relationship between features and actions and features-to-features should be present? What happens if features cannot be decorrelated?**\n\n**Features-to-actions and inter-feature relationships.**\nWe aim to recover the true association between features and actions, and for inter-feature relationships, we target their decorrelation. In W1's example, sample reweighting decorrelate the robot's state from the background, facilitating the model to recognize that the true association between background features and actions is 0.\n\n**What happens if features cannot be decorrelated?** In Figure 6 (a), the generalization will be hurt if the features are not decorrelated. In practice, due to the limited data, the sample reweighting often yields approximately independent results, not completely independent, as shown in Figure 4.\n\n>**W3. What the features are and where they come from.**\n\nOur method takes images as inputs, extracting latent features via an encoder trained alongside RL models. In our method, we leverage saliency map to identify features that vary across environments. In the main text, latent features are first mentioned on line 124. The structure of the encoder is mentioned in line 478 of the appendix.\n\n>**W4. Fig. 2 is quite complex. What are the colors of the “enriched features” indicating? What are the “G”s indicating?**\n\nThe term \"enriched features\" refers to latent features mapped by multiple RRFs to assess nonlinear correlation. The symbol \"G,\" initially representing the robot's center of gravity, was removed to prevent confusion. This is clarified in the updated Figure 2 in the \"global\" response PDF. We've replaced \"enriched features\" with \"RFF map\" for clarity, referencing Equation (4). The architecture now aligns better with the method's two subsections.\n\n**Figure 2: the architecture of SGFD.** Essentially, our SGFD aims to reduce correlations in image features by reweighting samples. This involves five steps: (1) We fetch a sample batch from the replay buffer, which may come from multiple environments with different backgrounds or robot configurations. (2) The image in the sample is compressed by an encoder into latent features. (3) Then, we augment these features using multiple RRFs to capture nonlinear correlations. (4) Concurrently, we train a classifier and apply saliency maps to detect features that shift across environments. (5) Finally, SGFD reweights samples to eliminate the correlations between identified features and other features.\n\n**Reply to questions**\n>**Q1. How would the method work with learned latent features? Could you include an evaluation on that?**\n\nAs mentioned in W3, our method processes images directly, with the encoder utilizing the AMBS [1] update method alongside RL models. We'll detail the experimental setup in the main text, not just the appendix.\n\n>**Q2. Can you include completely distracting features to evaluate the robustness of the solution?**\n\nIn Table 1, test environments have all distracting features. Inputs are 3x84x84 RGB pixel values, perturbed by real photos. We visualized the training and test environments in Figure 3 in the \"global\" response PDF file.\n\n>**Q3. What would be necessary for this method to work directly with images as input.**\n\nOur method takes images as input and uses an encoder to transform pixels into latent features. This encoder serves both the classifier and RL models. In experiments, the gradients from saliency maps are only propagated to the latent features, not the initial image layer.\n\n>**Q4. How does the method scale with the number of features? And with the number of samples? What is the computational cost?**\n\n**Number of features** As the inputs are images, the number of features is assumed to be the output dimension of the encoder. When accessing the ground-truth state of the environment, we demonstrate the decorrelation ability of the model for different numbers of changed features in Figure 4 of the main text and Figure 7 of the Appendix. \n\n**Number of samples and computational cost** The computational cost grows linearly with batch size. Using an off-policy RL algorithm, the cost is O(N x B), with N as environment steps and B as sample batch size.\n\nFor clarity on computational cost, we've timed training with and without SGFD for batch sizes 64, 128, and 256. The total number of steps is 1 million, tested on an NVIDIA RTX 3090 GPU and Intel Xeon Platinum 8255c CPU at 2.50 GHz.\n\n|batch size|64|128|256|\n|-|-|-|-|\n|SGFD|13 h|16 h|22 h| \n|no SGFD|12 h|15 h|20 h|\n\nThe results show that SGFD indeed brings extra computational cost due to the sample reweighting for decorrelation. However, compared to the RL algorithm, SGFD brings significant generalization at a cost of less than 10%.\n\n**References**  \n[1] Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. ICLR 2022."}, {"review_id": "kw6z069EdP", "reviewer": "Reviewer_NDX8", "summary": "This paper propose SGFD, an novel approach aimed at improving the generalization ability of distinguishing task-irrelevant and task-relevant situations. SGFD leverages two core techniques: Random Fourier Functions (RFF) and the saliency map, to estimate the complex non-linear correlations in high-dimensional images, and identify the changed features to achieve decorrelation. Furthermore, it conducts experiments on two benchmarks to exhibit its improved generalization ability.", "questions": "1. Can you visualize the examples of training enviroments and test environments?  I don't know the gap between these two environments and how difficult to generalize to the novel scenes. \n\n2. Can you illustrate the Figure 4 to me? I think you need to decorrelate the association between features and the policy, not among different features? But the Figure 4 seems to decorrelate the features.\n\n3. Why the calculated $w*$ only apply to actor loss? Why not critic?\n\n4. Is it possible to apply SGFD on the environments which contains both task-relevant and task-irrelevant features?\n\nIf you can address my questions, I would consider increasing the score.", "limitations": "Mentioned in weakness and questions.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Its proposed method has the capability to address problems in two distinct settings.\n\n2. I think utilizing saliency map to further achieve decorrelation is an interesting idea.", "weaknesses": "1. Figure 2 is confusing. This figure is meant to illustrate the core content of the paper, its unclear presentation undermines the overall comprehensibility. For instance, it would be clearer if the authors had depicted only one type of environment. The inclusion of two types of environments creates ambiguity as to whether the training was performed in one or both environments.\n\n2. This paper lacks a clear logical flow in the writing, particularly when it comes to the technical details. \n\n3. The performance of the proposed SGFD method does not seem to significantly surpass that of other algorithms. The overlap within the standard deviation intervals between SGFD and other algorithms suggests that the advantages of SGFD are not as pronounced as might be expected.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThanks for the reviewer's positive appraisal, insightful comment, and criticism of our paper.\n\n**Reply to the weakness**\n>**W1. Figure 2 is confusing.**\n\nWe split the original Figure 2 into two figures in the global response PDF: Figure 1 discusses the motivation for feature decorrelation, while Figure 2 details the technique flow.\n\n**Figure 1: the motivation for feature decorrelation.** Consider a classifier trained on images labeled 'cat': three with sofas and one with grass in the background. The model may falsely link sofas to cats due to frequent pairing. To mitigate this issue, we can reweight the images to decorrelate the features between backgrounds and animals. In RL, features like varied background noise can correlate with the robot's state, as in Figures 5 (main text) and 11 (Appendix). Such correlations may lead agents to form spurious associations, altering optimal actions with background shifts.\n\n**Figure 2: the architecture of SGFD.** Essentially, our SGFD aims to reduce correlations in image features by reweighting samples. This involves five steps: (1) We fetch a sample batch from the replay buffer, which may come from multiple environments with different backgrounds or robot configurations. (2) The image in the sample is compressed by an encoder into latent features. (3) Then, we augment these features using multiple Random Fourier Functions to capture nonlinear correlations. (4) Concurrently, we train a classifier and apply saliency maps to detect features that shift across environments. (5) Finally, SGFD reweights samples to eliminate the correlations between identified features and other features.\n>**W2. This paper lacks a clear logical flow in the writing, particularly when it comes to the technical details.**\n\nOur writing logic is as follows. In the Introduction, we first discuss the generalization when task-relevant and task-irrelevant features are shifted and point out that it can be achieved by feature decorrelation. Then, the Preliminary Section indicates that the learned encoder obtains the image features. The Method Section starts with using the Random Fourier Function for assessing feature correlation and elaborates on decorrelation through sample reweighting. It then delves into the challenge of decorrelating all feature pairs, introducing the saliency map to focus on shifted features. Finally, Our experiments in Sections 5.2 to 5.4 assess generalization, feature decorrelation, and conduct ablation studies. This logic flow will be clearer in the updated version with the revised Figure 2.\n>**W3. SGFD does not seem to significantly surpass other algorithms.**\n\nThe advantages of our SGFD method are reflected in two aspects. Firstly, it improves generalization for both task-relevant and task-irrelevant cases, while previous work usually only focuses on one. Secondly, it shows clear superiority in task-relevant cases, outperforming the state-of-the-art method by 23% on average, as shown in Table 2. This is because feature decorrelation encourages the model to form true associations with the features and thus be able to adjust actions in unseen environments correctly.\n\n**Reply to questions**\n>**Q1. Can you visualize the training and test environments?**\n\nWe visualize the environments in the \"global\" response PDF file. As shown in Figure 3, the model is trained in 4 environments with different task-irrelevant features or task-relevant features. During testing, the model is deployed in unseen environments without extra training. Differential settings between training and testing follow the baseline algorithms.\n>**Q2. Can you illustrate the Figure 4? Why eliminate the correlation between features？**\n\n**Illustration Figure 4:** Figure 4 depicts a 10x10 heatmap, visualizing the correlations between ten features ($s_1-s_{10}$) in a robotic arm grasping task. To compute the correlation of weighted data, we fetch a sample batch from the weighted distribution and then calculate the Pearson correlation coefficient. The model could access true states ($s_1-s_{10}$) with physical meaning. Notably, $s_{10}$ denotes the object's size, which varies across environments. For optimal RL model generalization to objects of unobserved sizes, it's crucial to assess correlations between $s_{10}$ and $s_1-s_9$. As highlighted in Figure 4, our method successfully reduces this correlation.\n\n**Why decorrelation:** Due to the word limit, please refer to the motivation for feature decorrelation in the W1 response.\n\n**Decorrelate associations between features and policies:** In fact, feature decorrelation leads the policy to decorrelate with some features. As the example in reply W1, the policy will eliminate the dependence of the background on decision-making when the semantics of images are independent of the background.\n>**Q3. Why the calculated $w$ only apply to actor? Why not critic?**\n\nWith the actor loss under the $w^*$ effect, the model could recover the true associations between the changed features and the optimal actions. In the testing phase, we only deploy the policy network to the new environment without transferring the critic; thus, $w^*$ does not need to change the critic loss.\n>**Q4. Can you apply SGFD on both task-relevant and irrelevant features?**\n\nAs suggested, we test the environment where both task-relevant and task-irrelevant features are changed. Specifically, in the MT-cheetah-run task, both torso length and background noise varied during testing. Similarly, in the MT-finger-turn task, the robot arm's length and background noise changed.\n\n|Tasks|SGFD|TED|AMBS|SGQN|DBC|DR|\n|-|-|-|-|-|-|-|\n|Mix-cheetah-run|$378.6±36.2$|308.6±47.6|265.6±37.4|161.2± 47.5|171.2±36.2|189.4±46.3|\n|Mix-finger-turn|$899.3±29.5$|726.2±52.2|802.5±53.1|632.9± 52.8|321.3±61.6|605.6±57.8|\n\nFrom the results, the advantage of feature decorrelation is further amplified and outperforms the baseline by 28% when the changes in both cases co-occur. We will add this experiment to the latest version."}, {"review_id": "tMDLuHXNfZ", "reviewer": "Reviewer_jEbK", "summary": "This paper deals with a novel sample re-weighting method designed to enhance generalization in visual reinforcement learning tasks across environments with unseen task-irrelevant and task-relevant features. SGFD is composed of two core components: RFF and a saliency-guided model. The paper puts down results comparison with standard datasets to establish the claim. This is a very important problem w.r.t. many downstream tasks of agents where main reason for failures is due to lack of generalization in learning.", "questions": "Fig. 2 weighted samples in figure, why those weights and what value it adds?\nNot clear why HSIC was discussed in first place, line 142.\nIs not Figure 3 too perfect wrt other approaches? Why so? I will also like to see worst results in supplementary to see variance.\nTable 2 - in term of metric results there exists 2 clusters (walker-walk, finger-spin,finger-turn) and (cheetah-run, walker-run). Is the are any analysis for this?", "limitations": "The work is limited to the specific datasets tried on in terms of training, otherwise adaptable.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Code and supplementary material content are good.\nThe paper is well written with example in intro.\nAppendix A,B are the crux of this paper in terms of contributions.", "weaknesses": "Related work needs an overhaul in terms of gaps and only putting exactly related prior art.\nContribution section please re-write and make it crisp to two main claims.\nApart from Pearson 218, other correlation checks could be put into function.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly encouraged that the Reviewer jEbK finds our method novel and effective.\n\n\n**Reply to the weakness**\n>**W1. Related work needs an overhaul in terms of gaps.**\n\nBuilding on the work discussed in Section 2, we further complement related work that generalizes to novel environment configurations. These works are typically studied with procedurally generated environments [1][2]. Some approaches to this problem leverage techniques from supervised learning, such as regularization, curriculum strategies, hyper-parameter tuning, and using self-supervised objectives [3-5]. Recurrent Independent Mechanisms leverage modules to learn a state function improved out-of-distribution generalization [6]. Feature-Attending Recurrent Modules showed that a modified attention mechanism led to strong generalization improvements with RL [7]. Our approach differs from these approaches in that we make no assumptions about the structure of task rewards or states, exploiting feature decorrelation to improve the generalization. \n\n>**W2. Contribution section please re-write and make it crisp to two main claims.**\n\nAccording to the suggestion, we revised the contribution part and made it crisp to two main claims. Summary of Contributions: (1) We propose the SGFD model that utilizes a sample reweighting method to improve generalization in visual RL tasks, covering both task-relevant and task-irrelevant situations. (2) Experimental results demonstrate that SGFD generalizes well on a wide range of novel environments and significantly outperforms state-of-the-art methods in handling task-relevant variations.\n\n>**W3. Apart from Pearson, other correlation checks could be put.**\n\nWe additionally introduce Spearman's rank correlation coefficient to test our method. Since Spearman's rank correlation coefficient cannot be directly applied in the case of weighted data, we first randomly select 128 samples based on the distribution of weighted values and calculate the standard Spearman's rank correlation coefficient based on these 128 samples. To save space, we take Figure 7 in the appendix as an example and calculate the correlation between the changed features ($s_8-s_{10}$) with other features ($s_1-s_7$). Similar to the Pearson correlation coefficient, we control the value range between 0 and 1, where 1 means complete correlation and 0 means independent of each other.\n\n|Unweighted samples|$s_1$|$s_2$|$s_3$|$s_4$|$s_5$|$s_6$|$s_7$|\n|-|-|-|-|-|-|-|-|\n|$s_8$|0.81|0.72|0.78 |0.79|0.74|0.74|0.92| \n|$s_9$|0.81|0.82|0.89 |0.78|0.86|0.82|0.84|\n|$s_{10}$|0.79|0.75|0.80|0.91|0.87|0.89|0.74|\n\n|Weighted samples|$s_1$|$s_2$|$s_3$|$s_4$|$s_5$|$s_6$|$s_7$|\n|-|-|-|-|-|-|-|-|\n|$s_8$|0.43|0.16|0.31|0.24|0.35|0.14|0.20| \n|$s_9$|0.28|0.27|0.31|0.18|0.28|0.24|0.39|\n|$s_{10}$|0.28|0.32|0.31|0.36|0.31|0.38|0.24|\n\nFrom the results, $s_8-s_{10}$ and $s_1-s_7$ show a strong correlation in the unweighted data, which is significantly reduced after reweighting. \n\n\n**Reply to questions**\n>**Q1. Fig. 2 weighted samples in figure, why those weights and what value it adds?**\n\nThe weights in Figure 2 could eliminate the correlation of features in the image, which is beneficial to improve the generalization. For ease of understanding, we discuss the motivation for feature decorrelation in Figure 1 in the general response PDF file. Consider a classifier trained on images labeled 'cat': three with sofas and one with grass in the background. The model may falsely link sofas to cats due to frequent pairing. To mitigate this issue, we can reweight the images to decorrelate the features between backgrounds and animals. In RL, features like varied background noise can correlate with the robot's state, as in Figures 5 (main text) and 11 (Appendix). Such correlations may lead agents to form spurious associations, altering optimal actions with background shifts.\n\n>**Q2. Not clear why HSIC was discussed in first place, line 142.**\n\nThe HSIC is the theoretical motivation behind our approach. Therefore, we introduce HSIC, leading to our correlation evaluation method based on the Random Fourier Function.\n\n>**Q2. Is not Figure 3 too perfect wrt other approaches? Why so?**\n\nFigure 3 does not represent that our method achieves perfect results on every task. Specifically, each polygon vertex denotes a task, and we normalize all algorithms' performances against the state-of-the-art result. Figure 3 is a visualization of Table 1 and Table 2, so the variance of each task can be seen in Table 1 and Table 2. In addition, we also record the learning curve, and the shaded part of the curve also shows the variance of the performance, as shown in Fig. 12 and Fig. 13 in the appendix.\n\n>**Q3. Table 2 - there exists 2 clusters (walker-walk, finger-spin,finger-turn) and (cheetah-run, walker-run). Is the are any analysis for this?**\n\nBased on the description, the reviewer may be talking about the results in Table 1. Each task in Table 1 is derived from DeepMind Control Suite [8], which is a widely used benchmark. In this benchmark, each task has an independent reward setting, e.g., 600 is a bad score in finger-spin, but it is close to the optimal solution in cheetah-run. Overall, our method achieves decent results on every task. Therefore, two clusters are observed because they have similar maximum reward settings.\n\n\n**References**  \n[1] Leveraging procedural generation to benchmark reinforcement learning. ICML 2020.  \n[2] MiniHack the Planet: A Sandbox for Open-Ended Reinforcement Learning Research. NeurIPS 2021.  \n[3] Reinforcement learning with augmented data. NeurIPS 2020.  \n[4] Generalization in reinforcement learning with selective noise injection and information bottleneck. NeurIPS 2019.   \n[5] Procedural generalization by planning with self-supervised world models. ICLR 2022.  \n[6] Recurrent independent mechanisms. 2019.  \n[7] Feature-attending recurrent modules for generalization in reinforcement learning. ICLR 2022.  \n[8] Deepmind control suite.2018."}, {"review_id": "NMO9xZgxVc", "reviewer": "Reviewer_SPDk", "summary": "In visual-based Reinforcement Learning (RL), agents often struggle with generalization to environmental variations that were not observed during training that have changed task-irrelevant features and changed task-relevant features.\nTo achieve generalization in environmental variations, The authors popose a sample reweighting method for RL tasks that encourages the agent to understand the impact of changed features on its decisions, called Saliency-Guided Features Decorrelation (SGFD).  The authors demonstrated that SGFD significantly outperforms state-of-the-arts in handling changed task-relevant features.", "questions": "Additional experiments and analysis of mixture of changes of task-irrelevant feature cases and task-releeant feature cases would be included.\nAdditional analysis of policy dependency would be helpful to support the authos claims.", "limitations": "The authors describe limitations, but more insights about limitations in real situations that include mixture of changes of task-irrelevant feature cases and task-relevant feature cases would be helpful.", "rating": 7, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This paper is well organized and described. Details of the proposed method are well described with supplementary materials. The performance advantages are not huge, but the authors have detailed experiments and analysis of changed feartures that support their claims.", "weaknesses": "The proposed method use the saliency to calculate and discriminate changed features under environmental changes.\nTheir experimental results and analysis include only one of task-irrelevant feature cases or task-relevant feature cases.\nThe proposed method is dependent on the saliency of the learned policy. Additional analysis of policy dependency would be helpful to support the authos' claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are encouraged that Reviewer SPDk finds our method novel and the idea well-grounded.\n\n**Reply to the weakness**\n>**W1. The proposed method use the saliency to calculate and discriminate changed features under environmental changes. Their experimental results and analysis include only one of task-irrelevant feature cases or task-relevant feature cases.**\n\nFirstly, we wish to clarify that variations in task-relevant or task-irrelevant features are widespread in the real world. Therefore, many works focus on one of these situations and propose novel ideas [1-4]. Second, our work improves generalization in both cases and achieves comparable performance to other algorithms in their focus scenarios. Lastly, as suggested, we evaluate the most challenging environment where both task-relevant and task-irrelevant features are changed. Specifically, in the MT-cheetah-run task, both torso length and background noise varied during testing. Similarly, in the MT-finger-turn task, the robot arm's length and background noise changed. Neither the changed background nor the robot's configuration was present in the training environments.\n\n| Tasks | SGFD | TED | AMBS | SGQN | DBC | DR |\n|-------|-------|-------|-------|-------|-------|-------|\n| Mix-cheetah-run | $\\mathbf{378.6 \\pm 36.2}$ | 308.6 $\\pm$ 47.6 | 265.6 $\\pm$ 37.4 |  161.2 $\\pm$ 47.5 |  171.2 $\\pm$ 36.2 |  189.4 $\\pm$ 46.3 | \n| Mix-finger-turn | $\\mathbf{899.3 \\pm 29.5}$ | 726.2 $\\pm$ 52.2 | 802.5 $\\pm$ 53.1 | 632.9 $\\pm$ 52.8 |  321.3 $\\pm$ 61.6 |  605.6 $\\pm$ 57.8 | \n\nFrom the results, the advantage of feature decorrelation is further amplified and outperforms the baseline by 28% when the changes in both cases co-occur. This is because our method can handle the generalization problem of task-relevant and task-irrelevant cases, while previous work usually only focuses on one. In addition, feature decorrelation encourages the model to recover the true associations between the changed features and the optimal behavior and thus be able to adjust actions in unseen environments correctly. We will add this experiment to the latest version.\n\n>**W2. The proposed method is dependent on the saliency of the learned policy. Additional analysis of policy dependency would be helpful to support the authors' claims.**\n\nBased on the reviewer's comment on the weakness, we speculate that the \"policy dependency\" mentioned here means that the policy's generalization depends on the accuracy of the saliency map. In the paper, we performed two types of experiments directly related to the saliency map. \n\nThe first class experiments are Figures 4 and 6 in the main text, which test the performance of our method without the assistance of saliency maps. As shown in Figure 4 (b), without the guidance of the saliency map, the decorrelation ability of the model is obviously limited, and this result also corresponds to the generalization of Figure 6 (b).\n\nAnother type of experiment is Figure 9 in the appendix, which focuses on testing the accuracy of the classifier. From the results, the classifier achieves an accuracy of over 0.9 within 1e5 steps in each task, which demonstrates that saliency maps can accurately identify features that shift across environments.\n\n**Reply to questions**\n>**Q1. Additional experiments and analysis of mixture of changes of task-irrelevant feature cases and task-relevant feature cases would be included.**\n\nAccording to the suggestions, we evaluate our method on the mixture of changes of task-irrelevant feature cases and task-relevant feature cases. The experiment results can be find in the response of W1.\n\n>**Q2. Additional analysis of policy dependency would be helpful to support the authors claims.**\n\nSimilarly, we discussed the \"policy dependency\" in our response to W2.\n\n**References**  \n[1] Temporal disentanglement of representations for improved generalisation in reinforcement learning. ICLR 2023.  \n[2] Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. NeurIPS 2022.  \n[3] Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. ICLR 2022.  \n[4] Learning robust state abstractions for hidden-parameter block mdps. ICLR 2021."}, {"review_id": "3yg2aUnmCe", "reviewer": "Reviewer_NdCS", "summary": "This paper introduces an original approach named Stochastic Gradient Feature Decorrelation (SGFD), which aims to amplify the generalization capabilities of reinforcement learning (RL) agents across various environmental variations. These variations can encompass task-irrelevant visual attributes such as backgrounds, as well as task-related factors like physical configurations. The authors propose achieving this through a decorrelation of features, executed via a resampling technique intended to minimize the Frobenius norm of the cross-covariance matrix, derived from Random Fourier features. Recognizing the inherent challenges of complete decorrelation, the authors shift focus towards effectively decorrelating the most variable features. This is facilitated by leveraging the saliencies from an environment classification model, which, under ideal circumstances, makes its decisions based on distinct features that are not common to different environments.", "questions": "* Could the authors elaborate on the inference procedure employed during testing? This could help elucidate the practical applicability of the methodology.  \n* Does the encoder also benefit from the gradients produced through resampling? Understanding this aspect could contribute to a more comprehensive understanding of the process.  \n* Given the necessity to decorrelate features that vary across environments from those that remain constant — and assuming that the variable features have a higher saliency — it appears that the $p\\left(\\mathbf{Z}_i\\right) p\\left(\\mathbf{Z}_j\\right)$ in equation 7 tends to assign more importance to tuples of varying features. Would it potentially be more efficient to replace it with $|p\\left(\\mathbf{Z}_i\\right) - p\\left(\\mathbf{Z}_j\\right)|$to accentuate the decorrelation between the varying and consistent features?\n", "limitations": "While the authors have reasonably addressed the methodological limitations of their approach, they assume no potential negative societal impacts stemming from their method. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is commendably well-written and coherent, effectively explaining complex ideas in an accessible manner. The authors demonstrate a strong theoretical grounding, with well-motivated intuitions supporting their methodology. Their saliency-guided optimization is an interesting approach, backed by a robust ablation study. SGFD successfully tackles generalization issues relating to both task-irrelevant and task-relevant features, demonstrating a broad scope of applicability. The proposed method do present noticeable enhancements in generalization performance, more so in the case of task-relevant features.\n", "weaknesses": "Despite the paper's strengths, there are some areas where it could be improved. Firstly, the methodology requires several environments with variations to train their environment classifier, introducing an element of manual supervision into the learning process, which may not be ideal in all scenarios. Furthermore, the full algorithm can be challenging to comprehend without first referring to Appendix A, suggesting that the main body of the text might benefit from additional clarification. It might also be beneficial to introduce the general objective - namely, the reweighting of the batch sampled from the buffer - earlier in the paper to give readers a clearer understanding of the process. Lastly, while the authors promote their method as an improvement over Soft Actor-Critic (SAC), the achieved results actually utilize the same updates for the encoder as employed by Adaptive Meta-learner of Behavioral Similarities (AMBS). This could lead to misunderstandings, as readers may infer that the saliency-guided resampling alone yields these performances. It would therefore be beneficial for clarity and fairness if the authors explicitly acknowledged this in the experimental section.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly encouraged that the Reviewer NdCs finds our method novel and effective.\n\n**Reply to the weakness**\n\n>**W1. Firstly, the methodology requires several environments with variations to train their environment classifier, introducing an element of manual supervision into the learning process, which may not be ideal in all scenarios.**\n\nWe agree with the point of view. From a causal perspective, invariance needs to sum up from changing data. It is an intuitive way to train general policies from different environments. In the experiments, we set up four training environments for each task, which puts the cost in an acceptable range.\n\n>**W2. Furthermore, the full algorithm can be challenging to comprehend without first referring to Appendix A, suggesting that the main body of the text might benefit from additional clarification. It might also be beneficial to introduce the general objective - namely, the reweighting of the batch sampled from the buffer - earlier in the paper to give readers a clearer understanding of the process.**\n\nAccording to the suggestion, we split the original Figure 2 into two figures in the global response PDF. We clearly illustrate the motivation for feature decorrelation in Figure 1, visualizing the difference before and after sample reweighting. In Figure 2, we explicitly indicate that the model aims to reweight a batch of samples from the replay buffer. Based on the revised Figure 2, we give a more precise technical flow description.\n\n**Figure 2: the architecture of SGFD.** Essentially, our SGFD aims to reduce correlations in image features by reweighting samples. This involves five steps: (1) We fetch a sample batch from the replay buffer, which may come from multiple environments with different backgrounds or robot configurations. (2) The image in the sample is compressed by an encoder into latent features. (3) Then, we augment these features using multiple Random Fourier Functions to capture nonlinear correlations. (4) Concurrently, we train a classifier and apply saliency maps to detect features that shift across environments. (5) Finally, SGFD reweights samples to eliminate the correlations between identified features and other features.\n\n>**W3. Lastly, while the authors promote their method as an improvement over SAC, the achieved results actually utilize the same updates for the encoder as employed by Adaptive Meta-learner of Behavioral Similarities (AMBS). This could lead to misunderstandings, as readers may infer that the saliency-guided resampling alone yields these performances. It would therefore be beneficial for clarity and fairness if the authors explicitly acknowledged this in the experimental section.**\n\nThanks for the suggestion; we will explicitly describe the details of the encoder in the experimental section, not just in the appendix.\n\n**Reply to questions**\n>**Q1. Could the authors elaborate on the inference procedure employed during testing?**\n\nOur method consists of four neural network models: an actor model (policy), a critic model, an encoder, and a classifier model. Among these four models, the critic model is used to improve the actor's performance during the training phase. At the same time, the classifier is used to assist the sample reweighting process. During testing, we deployed the encoder and actor model to the new environment. During inference, the encoder compresses the observed image into a latent representation, and the actor model predicts actions until the environment is terminated or the task is completed. This process will not involve the calculation of the classifier and the critic model.\n\n>**Q2. Does the encoder also benefit from the gradients produced through resampling?**\n\nThe role of the encoder is to compress high-dimensional images into compact representations, which are the basis of other models. In our experiments, the encoder is only updated along with the critic model, which is not affected by sample reweighting. The sample reweighting is only applied to the actor model to avoid confusing which part of the model brings the final result improvement. We will clarify this in the final version to make it easier for readers to understand the process.\n\n>**Q3. Would it potentially be more efficient to replace $p(\\mathbf{Z}_i)p(\\mathbf{Z}_j)$ with $|p(\\mathbf{Z}_i)-p(\\mathbf{Z}_j)|$ to accentuate the decorrelation between the varying and consistent features?**\n\nIntuitively, $|p(\\mathbf{Z}_i)-p(\\mathbf{Z}_j)|$ can identify pairs of varying features and invariant features with high probability. According to your suggestion, we conducted the related experiment and compared it with Equation (7). Due to the time limit, we tested several scenarios with task-irrelevant situations.\n\n|  | walker-walk | cheetah-run | finger-spin | walker-run | finger-turn |\n|-------|-------|-------|-------|-------|-------|\n| $p(\\mathbf{Z}_i)p(\\mathbf{Z}_j)$ | $\\mathbf{959.1 \\pm 26.3}$ | 599.6 $\\pm$ 47.2 | 965.7 $\\pm$ 45.9 | $\\mathbf{420.7 \\pm 39.2}$ | $\\mathbf{984.3 \\pm 11.5}$ | \n| $\\mathbf{abs}(p(\\mathbf{Z}_i)-p(\\mathbf{Z}_j))$ | 951.2 $\\pm$ 27.2 | $\\mathbf{605.6 \\pm 46.6}$ | $\\mathbf{975.5 \\pm 39.8}$ | 406.3 $\\pm$ 46.1 | 974.2 $\\pm$ 13.5 | \n\nFrom the preliminary results, $|(p(\\mathbf{Z}_i)-p(\\mathbf{Z}_j)|$ achieved a comparable performance as $p(\\mathbf{Z}_i)p(\\mathbf{Z}_j)$. Two points may cause the fluctuation in performance: (1) Due to the data limit, sample reweighting only achieves approximate perfect decorrelation. (2) The classifier's error may affect the recognition of changed features. Overall, $|p(\\mathbf{Z}_i)-p(\\mathbf{Z}_j))|$ is a promising approach and will complement experiments in task-relevant situations. Supplementary experiments will be updated in the latest version."}]

Author rebuttal:
Dear Reviewers,

We are very grateful to the reviewers for their valuable suggestions, which further improved our work. We provide three visualizations about our motivation, technique, and environmental setting with a submitted 1-page pdf.

- Figure 1: The motivation for feature decorrelation used in our method.
- Figure 2: The technique flow of our method.
- Figure 3: The visualization of training and testing environments.

Thank you again for your careful review and helpful comments.

Kind regards,

Paper1844 Authors

Meta-review: This paper proposes a method to improvement generalization in visual-based reinforcement learning by training the policy-network to ignore the spurious correlation of input features present in the training data. This is achieved by using gradient-based saliency to guide the decorrelation of features extracted from input images. The authors adequately addressed the concerns raised by reviewers and the overall consensus is that this paper meets the bar for acceptance.
Decision: Accept (spotlight)