Title: Learning Generalizable Agents via Saliency-guided Features Decorrelation
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

Learning control from high-dimensional visual observations is a 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 [12, 13]. However, current algorithms exhibit limited reliability when deployed in unseen environments, despite achieving satisfactory performance during training [7, 48]. It remains an open challenge to learn policies with good generalization across both changed task-irrelevant and relevant features [24]. Consider a simple example during the evaluation stage: a well-trained robot may fail to accomplish the task when confronted with task-irrelevant visual disturbances or changes to the task-relevant robotic arm, as illustrated in Figure 1.

Generalization to task-irrelevant features in vision RL tasks has been widely discussed [49, 11, 6, 3]. Since it is 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

(a) Testing tasks with

changed task-irrelevant features

(b) Training tasks

(c) Testing tasks with

changed task-relevant features

𝑎 ∗

𝑎 ∗

立

L

[25, 28, 31]. 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 to different optimal decisions, as shown in Figure 1(c). Although Dunion et al. [15] proposed a 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-relevant features.

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, 35, 39], the associations between features and decisions are often confused with each other, which can 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 on height. Therefore, a promising approach is to eliminate these correlations, enabling the agent to distinguish the true association between each input feature and decisions [51, 26]. However, recent studies have indicated that achieving perfect decorrelation between all pairs of input features is challenging [45, 44]. This difficulty is especially pronounced in visual RL tasks, as the high-dimensional images that maintain sequential relationships result in complex non-linear and strong correlations among features [40].

To this end, we propose the Saliency-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-linear 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 [51]. SGFD employs the cross-covariance 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 to each environment, these features can be pinpointed by computing a saliency map [5] of the classification model. Guided by this saliency map, SGFD reweights the training samples to minimize the estimated correlations, with a 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 as 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 of Contributions. (1) We propose the SGFD model, which achieves improved generalization in visual RL tasks, covering both task-relevant and task-irrelevant 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 various experiments and demonstrate that SGFD significantly outperforms state-of-the-arts 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, 8], among other approaches [52, 9]. 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 a 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 [22] and optimal decisionmaking. Invariance techniques strive to learn a representation that disregards distractors in the image, thereby efficiently generalizing to unseen backgrounds [48, 11, 31]. Recently, there has been an increasing interest in using varying task-relevant situations as an additional experimental setup to evaluate a model's generalization ability [47]. To address this situation, Dunion et 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-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 eliminates the correlation between inputs features[14, 26]. Shen et al. [37], Kuang et al. [27] 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 a 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 or nonlinear. However, it is not easy to eliminate the correlation between all input features under finite training samples[45]. This is especially serious in RL tasks because data often exhibit a sequential relationship [40], 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 is not applicable to 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 to identify the changed features without additional assumptions.

3 Preliminaries

Markov Decision Process. We define an environment as a Markov Decision Process (MDP) [34], represented by a tuple M = ( S , A , P , R , γ ) , where S denotes the high-dimensional state space, A is the action space, P ( s ′ | s, a ) is the state-transition function, R : S × A → R is the reward function, and γ ∈ [0 , 1) is the discount factor. A decision-making policy, parameterized by ϕ , is a function π ϕ ( a | s ) that maps a state to distributions over actions. The objective of RL agents is to find a policy that maximizes the expected cumulative return, expressed as E π ϕ [ ∑ ∞ t =0 γ t R ( s t , a t )] .

Soft Actor-Critic. We utilize Soft Actor-Critic (SAC) [17] as our base RL algorithm. SAC aims to maximize a variant of the RL objective augmented with an entropy term: E [ ∑ t γ t R ( s t , a t ) + α H ( π ( ·| s t ))] . To optimize this objective, SAC learns two models: a state-action value function Q θ ( s, a ) and a stochastic policy π ϕ ( a | s ) , where θ and ϕ are the parameters for Q θ ( s, a ) and π ϕ ( a | s ) , respectively. The parameters θ of the state-value function are trained by minimizing the soft Bellman residual, given as:

where D is the replay buffer storing the data, V ¯ θ denotes the value function, and ¯ θ are the target parameters of θ . The stochastic policy π ϕ ( a | s ) is trained by maximizing the following objective

where α is a temperature coefficient.

Classifier

RFF

extractor

Saliency map

Sample weights

A batch of samples

Encoder

Latent  features

⋮

RFF map

⋮

Probability vector

Replay buffer

Multiple training

environments

Section 4.2: Saliency-guided model

Section 4.1: Decorrelation with RFF

Elementwise

multiplication

High probability

Notations. For existing visual RL methods [10, 18, 43], the state s is typically compressed into a compact vector representation z . Therefore, we use z as the default in our calculations unless stated otherwise. For clarity, we define Z j as the j -th feature in the d -dimensional vector, and z i,j as the value of Z j in sample i . For instance, if Z j corresponds to the angle of a robot arm, z i,j would represent the encoded value of the angle in image i .

4 Method

In this section, we introduce the SGFD method, a sample reweighting technique designed to improve generalization across a broad spectrum of visual RL tasks, as depicted in Figure 2. Specifically, SGFD utilizes RFF to enrich image features, thereby enabling the estimation of non-linear correlations in high-dimensional images, as detailed in Section 4.1. Then, SGFD integrates a classifier model and a 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 ( Z i , Z j ) in the RL states, we gather a batch of n samples from the replay buffer D , denoted as ( z 1 ,i , z 1 ,j ) , ( z 2 ,i , z 2 ,j ) . . . ( z n,i , z n,j ) . The primary challenge lies in accurately estimating the correlation between these two features based on the available samples.

A renowned method for evaluating independence is the Hilbert-Schmidt Independence Criterion (HSIC), which calculates a cross-covariance operator in the Reproducing Kernel Hilbert Space[16]. However, the computational cost of HSIC can be considerable, especially when the batch size n is large. To overcome this issue and apply HSIC to modern RL models, we can compute the independence using the Frobenius norm, shown to be equivalent to the Hilbert-Schmidt norm in Euclidean space [38]. The cross-covariance matrix can be formulated as follows:

where E [ u ( Z i )] = 1 n ∑ n k =1 u ( z k,i ) , u and v are vectors, and each component is a function sampled from the RFF space, formulated as follows:

where H RFF denotes the RFF space and has been demonstrated to approximate non-linear correlation between image features [51]. Specifically, we sample M functions from the RFF space, defined as H RFF = { h : x -→ √ 2cos( ωx + ψ ) | ω ∼ N (0 , 1) , ψ ∼ Uniform(0 , 2 π ) } , where ω is sampled from a standard normal distribution and ψ is sampled from a uniform distribution. The independence between Z i and Z j can be calculated by the Frobenius norm of the cross-covariance matrix || ˆ ∑ Z i Z j || 2 F . If || ˆ ∑ Z i Z j || 2 F is zero, the two features Z i and Z j are considered independent.

We reformulate the task of correlation elimination as an optimization problem aimed at minimizing the Frobenius norm of the cross-covariance matrix || ˆ ∑ Z i Z j || 2 F . To this end, we introduce a set of learnable sample weights, denoted by w ∈ R , subject to the constraint that ∑ n k =1 w k = n . Based on the weighted samples, the independence can be defined as:

where E [ wu ( Z i )] = 1 n ∑ n k =1 w k u ( z k,i ) . The correlation among all features in RL states can be eliminated by solving the following optimization problem:

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 a significant challenge [45]. Therefore, relying solely on the features decorrelation approach introduced in Section 4.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 saliency-guided model that prioritizes the most crucial correlations between changed features and other image features.

The saliency map [5], a 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 which pixels are more 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 i ( f, Z ) = ∂f ( Z ) /∂ Z i , 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 i ( f, Z ) .

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 ( Z i ) = e M i / ∑ 1 ≤ j ≤ d e M j . We incorporate this term into the Equation (6), revising it as:

Based on the weighted samples, we can reformulate the Equation (2) as:

With the reformulated Equation 8, 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.

SGFD

AMBS

SGQN

DBC

DR

1

2

3

4

5

6

7

8

9

10

1-5: task-irrelevant variations

1: Walker-walk

2: Cheetah-run

3: Finger-spin

4: Walker-run

5: Finger-turn

6-10: task-relevant variations

6: MT-Walker-walk

7: MT-Cheetah-run

8: MT-Finger-spin

9: MT-Walker-run

10: MT-Finger-turn

5 Experiments

To carry out a comprehensive investigation of SGFD, we have performed a 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 of SGFD across a variety of tasks, ranging from scenarios with diverse task-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, we conducted several ablation experiments on RFF and saliency-guided 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 [41] is a widely-used visual control toolkit that offers a wealth of simulated animal behavior tasks. Drawing on previous works [48, 50], we constructed environments with both background noises and varying robot settings to test the generalization capabilities of the SGFD. The backgrounds were sampled from the Kinetics dataset [23], while the robot settings followed those proposed by Zhang et al. [50]. Furthermore, we utilized the Causal World [4], which comprises a 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 B.

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 a small dimension, thereby ensuring that learned behaviors exhibit a minimal change in the testing environment. We used DBC [48] as a baseline method that learns invariant representations based on the bisimulation metric. We also compared with AMBS [11], 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-irrelevant features in the images. Lastly, we included Domain Randomisation (DR) [42] in our comparison, a technique commonly utilized in control tasks.

Evaluation Metrics. The primary metric under consideration is 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 the correlation based on samples, both with and without weighting. Lastly, we utilized the saliency map as 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.

5.2 Evaluation on the Generalization of SGFD

Evaluation on Changed Task-Irrelevant Features. In this experiment, we assessed the generalization capability of our SGFD model with respect to different task-irrelevant features. During the training phase, the policy model was allowed to interact with the environments 1 million times. During the evaluation, we assessed the cumulative rewards in a testing environment with previously unseen backgrounds. As indicated in Table 1 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 not effectively filtered out. Consequently, the policy model established associations with these tasks-irrelevant features, leading to incorrect decision-making when the background changed. In contrast, SGFD enabled the policy to understand that background features should not affect decision-making, thereby achieving superior generalization performance. The results indicate that decorrelation is effective in RL tasks with background noises and that SGFD exhibits strong generalization capabilities to new values of task-irrelevant features.

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 a more significant challenge for generalization, as it necessitates the model to comprehend how changed features impact optimal decision-making. Owing to the associations between changed features and decisions were not confused, our model outperformed other baselines, with more substantial performance gaps observed in the extrapolation

(a) Original states

(b) Decorrelation

(c) SGFD

s

1

s

1

s

2

s

2

s

3

s

3

s

10

s

10

…

…

s

1

s

1

s

2

s

2

s

3

s

3

…

…

s

1

s

1

s

2

s

2

s

3

s

3

…

…

s

10

s

10

s

10

s

10

0.0

1.0

0.5

0.0

1.0

0.5

0.0

1.0

0.5

0.67

0.740.650.650.680.72

0.67

0.68

0.68

0.460.46

0.560.52

0.53

0.49

0.49

0.4

0.4

0.3

0.140.140.140.140.140.270.21

0.2

0.67

0.740.79

0.7

0.72

0.72

0.67

0.68

1

0.68

0.460.46

0.47

0.66

0.56

0.4

0.49

0.4

0.4

0.67

0.64

0.7

0.65

0.72

0.67

0.680.68

1

0.2

0.67

0.740.73

0.8

0.65

0.72

0.67

1

0.680.68

0.46

0.46

0.44

0.52

0.41

0.44

0.49

0.4

0.4

0.63

0.65

0.7

0.650.68

0.67

0.68

1

0.68

0.21

0.67

0.65

0.87

0.71

0.65

0.67

1

0.67

0.670.67

0.59

0.46

0.41

0.47

0.44

0.49

0.490.49

0.49

0.67

0.65

0.65

0.65

0.67

0.67

1

0.680.680.27

0.63

0.65

0.740.81

0.65

1

0.67

0.72

0.72

0.72

0.44

0.46

0.41

0.52

0.55

0.49

0.44

0.4

0.49

0.67

0.65

0.6

0.7

0.65

1

0.67

0.67

0.67

0.14

0.67

0.65

0.74

0.6

1

0.65

0.65

0.65

0.72

0.68

0.51

0.47

0.41

0.41

0.55

0.44

0.41

0.56

0.53

0.67

0.65

0.71

0.7

1

0.65

0.67

0.680.72

0.14

0.9

0.65

0.74

1

0.6

0.81

0.71

0.8

0.7

0.65

0.51

0.46

0.41

0.41

0.52

0.47

0.52

0.66

0.52

0.67

0.6

0.73

1

0.7

0.7

0.65

0.650.65

0.14

0.8

0.64

1

0.740.740

0.74

0.87

0.73

0.79

0.65

0.51

0.46

0.41

0.41

0.41

0.41

0.44

0.47

0.56

0.8

0.64

1

0.73

0.71

0.6

0.65

0.7

0.7

0.14

0.69

0.69

1

0.640.65

0.8

0.9

0.65

0.67

0.65

0.63

0.65

0.67

0.74

0.67

0.740.74

0.67

0.67

0.46

0.46

0.51

0.46

0.46

0.51

0.47

0.51

0.46

0.44

0.46

0.59

0.46

0.46

0.46

0.46

0.46

0.46

0.69

0.69

0.64

0.8

0.67

0.6

0.65

0.67

0.65

0.67

0.650.650.64

0.67

0.630.67

0.14

0.3

1.05

1.03

1.11

1.23

1.13

1.21

0.66

0.86

0.98

0.84

0.94

0.91

0.96

0.99

0.73

0.78

0.75

0.80

1.46

1.82

States in Env

a

States in Env

b

Reweighting

> 1.0

> 1.0

< 1.0

< 1.0

< 1.0

> 1.0

Reweighting

0.97

0.91

0.89

0.93

0.98

0.97

0.94

0.96

1.21

1.15

1.11

1.03

1.08

1.07

1.08

1.11

0.97

0.84

0.93

0.83

< 1.0

< 1.0

> 1.0

> 1.0

> 1.0

< 1.0

States in Env

c

States in Env

d

Reweighting

Reweighting

(a) 20 states from environments with different background

(b) 20 states from environments with different robot parameters

T

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 a 'Picking' task in Causal World. This task involves a 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-guided weighted data. Among these features, we singled out s 10 , a feature that represents mass, as the focal feature for environmental variation. Consequently, when the model is given the task of handling an object with a new mass, it becomes critical to carefully examine the correlations 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.

no decorrelation

decorrelation

decorrelation+RFF

(a)

0

1.0

0

300

600

1e6

0

1.0

0

200

400

1e6

saliency-guided decorrelation

no decorrelation

decorrelation

no decorrelation

decorrelation

decorrelation+RFF

MT-cheetah-run

RFF+decorrelation

no decorrelation

decorrelation

620

540

460

500

250

0

MT-cheetah-run

Cheetah-run

reward

(b)

Cheetah-run

reward

In contrast, our SGFD model can effectively reduce the correlations between the mass feature ( s 10 ) 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 in Figure 5 (a), the green dotted boxes indicate several states with "fall" posture. Due to the impact of background noises on the policy, we identified 4 similar states in the environment a and 2 similar states in the environment b . After the reweighting process, both environments a and b had approximately 3 similar states, effectively neutralizing the correlation between background and body posture. A similar outcome was observed in Figure 5 (b), wherein the reweighting process endeavored to balance the number of states in each group across environments with varied body lengths. As a 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 environments 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 to 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 Saliency-Guided Model. We conduct a separate ablation study on the saliency-guided 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 6 (b), the performance significantly deteriorates when the saliency-guided model is removed from SGFD. This outcome is predictable because directly decorrelating all pairs of features is challenging, thereby complicating the model'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 saliencyguided model. Detailed analyses and results of these studies are available in Appendix C.2. Moreover, we report on the learning curve of the classifier, as shown in Appendix C.3.

6 Conclusion, Limitations, and Broader Impact

In this work, we introduce SGFD, a sample reweighting 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, which empower the RL 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 RL tasks.

In terms of limitations, we note that our work, like many others [48, 11, 15], relies on stacking consecutive frames to approximate a fully observable condition. This might not be optimal for all scenarios, particularly when the changed features cannot be directly observed. A 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.

Acknowledgments

This work was supported by the National Key R & D Program of China under Grant Nos. 2021ZD0112501 and 2021ZD0112502; 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.20220402009GH.

References

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.

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.

Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR) , 2021.

O Ahmed, F Träuble, A Goyal, A Neitz, M Wüthrich, Y Bengio, B Schölkopf, and S. Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv peprint arXiv:2010.04296 , 2020.

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.

David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR) , 2022.

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.

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.

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.04226 , 2023.

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.

Jianda Chen and Sinno Jialin Pan. Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. In Proceedings of the International Conference on Learning Representations (ICLR) , 2022.

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.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning (ICML) , 2020.

Peng Cui and Susan Athey. Stable learning establishes some common ground between causal inference and machine learning. Nature Machine Intelligence , 4(2):110-115, 2022.

Mhairi Dunion, Trevor McInroe, Kevin Luck, Josiah Hanna, and Stefano V Albrecht. Temporal disentanglement of representations for improved generalisation in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR) , 2023.

Arthur Gretton, Kenji Fukumizu, Choon Teo, Le Song, Bernhard Schölkopf, and Alex Smola. A kernel statistical test of independence. Advances in neural information processing systems , 20, 2007.

T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning (PMLR) , page 1861-1870, 2018.

Aleksi Hamalainen, Karol Arndt, Ali Ghadirzadeh, and Ville Kyrki. Affordance learning for end-to0end visuomotor robot control. In International Conference on Intelligent Robots and Systems (IROS) , 2019.

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.

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.

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.

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 Neural Information Processing Systems , 35:12619-12632, 2022.

Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint arXiv:1705.06950 , 2017.

Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint arXiv:2111.09794 , 2021.

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.

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.

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.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems , 33:19884-19895, 2020.

Bonnie Li, Vincent François-Lavet, Thang Doan, and Joelle Pineau. Domain adversarial reinforcement learning. arXiv preprint arXiv:2102.07097 , 2021.

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 , 33:3686-3698, 2020.

Bogdan Mazoure, Ahmed M. Ahmed, Patrick MacAlpine, and R Devon Hjelm. Cross-trajectory representation learning for zero-shot generalization in rl. In Proceedings of the International Conference on Learning Representations (ICLR) , 2022.

Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014.

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.

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 in Neural Information Processing Systems , 34: 12686-12699, 2021.

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.

Eric V Strobl, Kun Zhang, and Shyam Visweswaran. Approximate kernel-based conditional independence tests for fast non-parametric causal discovery. Journal of Causal Inference , 7(1), 2019.

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.

Yanchao Sun, Shuang Ma, Ratnesh Madaan, Rogerio Bonatti, Furong Huang, and Ashish Kapoor. Smart: Self-supervised multi-task pretraining with control transformers. In Proceedings of the International Conference on Learning Representations (ICLR) , 2023.

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.00690 , 2018.

Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS) , pages 23-30. IEEE, 2017.

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.

Renzhe Xu, Xingxuan Zhang, Zheyan Shen, Tong Zhang, and Peng Cui. A theoretical analysis on independence-driven importance weighting for covariate-shift generalization. In International Conference on Machine Learning , pages 24803-24829. PMLR, 2022.

Han Yu, Peng Cui, Yue He, Zheyan Shen, Yong Lin, Renzhe Xu, and Xingxuan Zhang. Stable learning via sparse variable independence. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI) , 2023.

Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: Task-aware lipschitz data augmentationfor visual reinforcement learning. arXiv peprint arXiv:2202.09982 , 2022.

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 arXiv:2212.08860 , 2022.

AZhang, R Mcallister, R Calandra, Y Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. In Proceedings of the International Conference on Learning Representations (ICLR) , 2021.

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.

Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstractions for hidden-parameter block mdps. arXiv preprint arXiv:2007.07206 , 2020.

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2021.

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 on pretrained foundation models: A history from bert to chatgpt. arXiv preprint arXiv:2302.09419 , 2023.

Appendix

A Pseudocode of Saliency-Guided Features Decorrelation

In Algorithm 1, we introduce the training process of SGFD. At each step, we apply the policy model π ϕ to interact with training environments and sample a 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 t > 1 e 5 , we update the classifier based on the cross-entropy loss, as 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 (1) and Equation (8) to update state-action value function Q θ and policy model π ϕ based on the weighted data, as described in lines 18-20.

Algorithm 1 Saliency-Guided Features Decorrelation

B Implement Details

Experimental Setting Details. For task-irrelevant variations, we follow the prior work [30] to sample several images from Kinetics [23]. In training, we create four parallel environments with different image backgrounds. In testing, we create a 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 1 and 5. In testing, we exhibit evaluations under two setups: interpolation and extrapolation. In the interpolation setup, task-relevant features changed in a manner interpolated between those in the training environments, whereas in the extrapolation setup, changes exceeded the range of those in the training environments, i.e., the tenth setting with the farthest training distribution. The visualizations are illustrated in Figure 7.

,

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.10GHz. Each run of the experiments spanned about 12-24 hours, depending on the algorithm and the complexity of the task.

Architectures. Following AMBS [11], we use the same encoder architecture, which consists of 4 convolutional layers. The encoder weights are shared between the policy and the state-action value function. Each convolutional layer has a 3 × 3 kernel size and 32 channels. The first layer has a stride of 2, and all other layer has a stride of 1. There is a ReLU activation between each convolutional layer. The convolutional layers are followed by a trunk network with a linear layer, layer normalization, and finally, a tanh activation. As we focus on the sample reweighting, the encoder is updated by referring to the AMBS method.

The policy π ϕ and the state-action value function Q θ are both 2-layer MLPs with a hidden dimension of 1024. We apply ReLU activations after each layer except the last layer. The classifier is implemented with 2-layer MLPs with a hidden dimension of 128. We also apply ReLU activations after each layer except the last layer. Table 3 shows the values of the architectures for the encoder, the policy, the value function, and the classifier.

Hyperparameters. Table 4 shows the hyperparameters used in our SGFD model.

An testing environment

with unseen robot parameters

cheetah

walker

Training environments

with different robot parameters

finger

(a) The visualizations of robot configurations

An testing environment

with unseen robot parameters

Training environments

with different background

(b) The visualizations of background noise

L

(a) Original states

(b) Decorrelation

(c) SGFD

s

1

s

1

s

2

s

2

s

3

s

3

s

10

s

10

…

…

s

9

s

8

s

9

s

8

s

1

s

1

s

2

s

2

s

3

s

3

s

10

s

10

…

…

s

9

s

8

s

9

s

8

s

1

s

1

s

2

s

2

s

3

s

3

s

10

s

10

…

…

s

9

s

8

s

9

s

8

0.0

1.0

0.5

0.0

1.0

0.5

0.0

1.0

0.5

0.71

0.65

0.7

0.92

0.77

0.79

0.64

0.95

0.67

0.72

0.55

0.46

0.51

0.44

0.49

0.53

0.53

0.44

0.180.280.210.26 0.210.280.14 0.14

0.45

0.71

0.62

0.79

0.71

0.78

0.72

0.640.64

0.67

0.52

0.41

0.56

0.59

0.54

0.49

0.54

0.53

0.44

0.180.180.210.089 0.180.140.290.14

0.45

0.71

0.62

0.68

0.69

0.64

0.64

0.91

1

0.640.95

0.9

0.41

0.56

0.51

0.53

0.74

0.49

0.530.53

0.48 0.066 0.210.140.25 0.049 0.11

0.140.14

0.78

0.62

0.75

0.69

0.68

0.69

1

0.91

0.640.64

0.49

0.41

0.46

0.54

0.51

0.5

0.490.540.53

0.720.71

0.710.790.82

0.77

0.110.290.14

0.68

0.66

0.61

0.69

0.85

1

0.690.640.720.79

0.49

0.59

0.460.39

0.51

0.5

0.740.490.49

0.720.750.710.87

0.64

1

0.770.049 0.140.28

0.87

0.76

0.61

0.61

1

0.85

50.680.640.780.77

0.49

0.66

0.52

0.59

0.510.51

0.530.540.44

0.760.71

0.71

0.42

0.64

0.820.250.180.21

0.68

0.89

0.61

0.610.690.690.690.71

0.92

0.78

0.68

0.46

0.59

0.39

0.54

0.510.590.51

0.83

0.71

0.8

1

0.42

0.870.790.140.0890.26

0.68

0.78

0.94

1

0.94

1

0.89

0.61

0.61

0.76

0.61

0.66

0.62

0.75

890

0.62

0.62

0.79

0.65

0.7

0.64

0.4

0.91

0.91

0.680.66

0.460.520.460.46

0.59

0.410.410.410.55

0.560.560.46

0.83

0.83

0.72

1

0.720.710.710.750.710.066 0.180.28

1

0.8

0.710.710.710.210.210.21

0.78

0.680.68

0.87

0.68

0.78

0.71

0.71

0.71

0.4

0.64

0.78

0.490.490.49

60

0.52

0.72

0.830.830.830.760.72

0.72

0.480.180.18

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) on saliency-guided weighted data (SGFD). Among these features, we designate the S 8 , S 9 , and S 10 , which represent the length, width, and height, as the features of interest for environmental variations. Therefore, when the model is required to push an object with a new size, the correlation between size and other features needs to be carefully considered. As shown in Figure 8, due to 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 ( S 8 , S 9 , and S 10 ) and other features.

C.2 Additional Ablation Studies

We performed additional ablation experiments to test the effect of RFF and the saliency-guided model. As depicted in Figure 9 (a) and Figure 9 (b), when RFF is removed, the generalization performance drops significantly. In contrast, the global decorrelation can cause over-adjusting of samples weighting and impair the performance, as demonstrated by the reward curves in Figure 9 (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 saliency-guided 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 1e5 steps. To avoid overfitting, we stopped to update the classifier when the accuracy exceeded 0.9. As shown in Figure 10, the classifier converged quickly and achieved the accuracy of 0.9, which means that the classifier does not take up many computing resources.

no decorrelation

decorrelation

decorrelation+RFF

0

1.0

0

250

500

1e6

0

1.0

0

300

600

1e6

saliency-guided decorrelation

no decorrelation

decorrelation

no decorrelation

decorrelation

decorrelation+RFF

(d) MT-walker-walk

RFF+decorrelation

no decorrelation

decorrelation

430

405

380

560

290

420

(b) MT-walker-walk

(a) walker-run

reward

(c) walker-run

reward

(f) MT-walker-walk

0

1.0

(b) cheetah run

0.9

(c) finger spin

1e6

(g) MT-cheetah-run

(h) MT-finger-spin

200

(i) MT-walker-run

(d) walker run

0

1.0

0.9

1e6

(a) walker walk

0

1.0

0.9

1e6

0

1.0

0.9

1e6

0

1.0

0.9

1e6

0

1.0

0.9

1e6

0

0.9

1e6

1.0

0

0.9

1e6

1.0

(e) finger turn

0

1.0

0.9

1e6

0

1.0

0.9

1e6

(j) MT-finger-turn

C.4 Saliency Maps on the Policy Models

An intuitive approach to explaining visually-based models involves identifying pixels that have a significant impact on the final decision [1]. To determine whether the model focuses on the object or the background during decisions, we visualize the gradient of the action with respect to the input pixels. Experiments are conducted in environments with varying backgrounds and robot parameters, as depicted in Figure 11. Saliency maps of the baseline models reveal that various backgrounds draw considerable attention from the policy model while failing to make decisive contributions to our model. Moreover, SGFD can notice the changing aspects of the robot's body that other baseline models overlook. As a result, the saliency maps demonstrate that the decorrelation benefits the model's generalization in both task-relevant and task-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 5.3. 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 12. As illustrated in Figure 12 (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 A and B had approximately 42 % similar states, effectively neutralizing the correlation between background and body posture. A similar outcome was observed in Figure 12 (b), wherein the reweighting process endeavored to balance the number of states in each group across environments with varied body lengths.

SGFD

AMBS

SGQN

DBC

DR

Observation

C.6 Evaluation on the Generalization of SGFD

Evaluation on 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 1. As illustrated in Figure 13, 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 of SGFD achieves more significance. The results demonstrate that the decorrelation is suitable for RL tasks with different task-irrelevant features, and SGFD exhibits strong generalization capabilities to new values of task-irrelevant features.

Evaluation on Changed Task-Relevant Features. We also report the training curves of SGFD in environments with different task-relevant features, where the end points of the curves correspond to the results in 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 it requires the model to fully understand how changed features affect optimal decisions. Benefiting from decorrelation, our model achieved superior performance compared to 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 understanding of the true associations between changed features and decisions, leading to strong generalization on new values of task-relevant features.

States in Env A

States in Env B

Mean reweighting

1.26

1.23

0.88

0.86

0.81

1.16

0.91

0.82

1.21

1.08

1.33

0.83

States in Env C

States in Env D

(a) 300 states from environments with different background

(b) 300 states from environments with different robot parameters

Mean reweighting

Mean reweighting

Mean reweighting

25.3%

24.7%

50%

11.3%

35.4%

53.3%

40%

44%

16%

31.3%

21.4%

47.3%

人

CL人

0

1.0

0.5

reward

0

500

1000

(b) cheetah-run

(c) finger-spin

0

1.0

0.5

0

300

600

0

1.0

0.5

0

500

1000

1e6

1e6

1e6

1e6

SGFD

AMBS

SGQN

DBC

(d) walker-run

0

1.0

0.5

0

250

500

DR

1e6

(e) finger-turn

0

1.0

0.5

0

500

1000

(a) walker-walk

Extrapolation

0

1.0

0.5

(b) MT-cheetah-run

0

1.0

0.5

Interpolation

reward

600

300

0

400

200

0

200

100

0

(c) MT-finger-spin

0

0.5

1.0

1e6

1e6

200

100

0

0

0.5

1.0

1e6

1e6

reward

600

300

0

0

1.0

0.5

1e6

(g) MT-cheetah-run

0

1.0

0.5

400

200

0

1e6

200

100

0

(h) MT-finger-spin

0

0.5

1.0

200

100

0

(i) MT-walker-run

0

0.5

1.0

1e6

1e6

(d) MT-walker-run

SGFD

AMBS

SGQN

DBC

750

500

0

0

0.5

1.0

1e6

750

500

0

0

0.5

1.0

1e6

DR

(e) MT-finger-turn

(j) MT-finger-turn

(f) MT-walker-walk

(a) MT-walker-walk

—

M

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nDespite 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\nQUESTIONS:\n* 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\nLIMITATIONS:\nWhile the authors have reasonably addressed the methodological limitations of their approach, they assume no potential negative societal impacts stemming from their method.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-30"}

comments: 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.
To 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.

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.
Their experimental results and analysis include only one of task-irrelevant feature cases or task-relevant feature cases.
The proposed method is dependent on the saliency of the learned policy. Additional analysis of policy dependency would be helpful to support the authos' claims.

QUESTIONS:
Additional experiments and analysis of mixture of changes of task-irrelevant feature cases and task-releeant feature cases would be included.
Additional 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-07

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

STRENGTHS:
Code and supplementary material content are good.
The paper is well written with example in intro.
Appendix 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.
Contribution section please re-write and make it crisp to two main claims.
Apart from Pearson 218, other correlation checks could be put into function.

QUESTIONS:
Fig. 2 weighted samples in figure, why those weights and what value it adds?
Not clear why HSIC was discussed in first place, line 142.
Is not Figure 3 too perfect wrt other approaches? Why so? I will also like to see worst results in supplementary to see variance.
Table 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nMentioned in weakness and questions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}