SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sili Huang
Yanchao Sun $^{2}$  Yi Chang $^{6*}$
Jifeng Hu $^{3}$   
Lichao Sun $^{7}$
Siyuan Guo $^{4}$  Bo Yang $^{8*}$
Hechang Chen*5
$^{1,8}$ Key Laboratory of Symbolic Computation and Knowledge Engineering of Ministry of Education 1,3,4,5,6 School of Artificial Intelligence, Jilin University, China
$^{2}$ Department of Computer Science, University of Maryland, College Park, USA
$^{7}$ Lehigh University, Bethlehem, Pennsylvania, USA
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 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
*Corresponding Author. Bo Yang, Hechang Chen and Yi Chang.
(a) Testing tasks with changed task-irrelevant features
(b) Training tasks
Figure 1: Visualizing tasks across different variations.  $a^*$  denotes the optimal decision, which should be invariant to task-irrelevant features but different to changed task-relevant features. Note that 'changed feature' refers to a feature whose value has changed across tasks.
(c) Testing tasks with changed task-relevant features
[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.
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 decision-making. 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.
Markov Decision Process. We define an environment as a Markov Decision Process (MDP) [34], represented by a tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma)$ , where  $\mathcal{S}$  denotes the high-dimensional state space,  $\mathcal{A}$  is the action space,  $\mathcal{P}(s'|s,a)$  is the state-transition function,  $\mathcal{R}:S\times \mathcal{A}\to \mathbb{R}$  is the reward function, and  $\gamma \in [0,1)$  is the discount factor. A decision-making policy, parameterized by  $\phi$ , is a function  $\pi_{\phi}(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  $\mathbb{E}_{\pi_{\phi}}[\sum_{t=0}^{\infty}\gamma^{t}\mathcal{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:  $\mathbb{E}[\sum_t\gamma^t\mathcal{R}(s_t,a_t) + \alpha \mathcal{H}(\pi (\cdot |s_t))]$ . To optimize this objective, SAC learns two models: a state-action value function  $\mathcal{Q}_{\theta}(s,a)$  and a stochastic policy  $\pi_{\phi}(a|s)$ , where  $\theta$  and  $\phi$  are the parameters for  $\mathcal{Q}_{\theta}(s,a)$  and  $\pi_{\phi}(a|s)$ , respectively. The parameters  $\theta$  of the state-value function are trained by minimizing the soft Bellman residual, given as:
$$
\mathcal {J} _ {\mathcal {Q}} (\theta) = \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \mathcal {D}} [ \frac {1}{2} (\mathcal {Q} _ {\theta} (s _ {t}, a _ {t}) - (r (s _ {t}, a _ {t}) + \gamma \mathbb {E} _ {s _ {t + 1} \sim \mathcal {P} (\cdot | s _ {t}, a _ {t})} [ V _ {\bar {\theta}} (s _ {t + 1}) ]) ^ {2} ], \tag {1}
$$
where  $\mathcal{D}$  is the replay buffer storing the data,  $V_{\bar{\theta}}$  denotes the value function, and  $\bar{\theta}$  are the target parameters of  $\theta$ . The stochastic policy  $\pi_{\phi}(a|s)$  is trained by maximizing the following objective
$$
\mathcal {J} _ {\pi} (\phi) = \mathbb {E} _ {s _ {t} \sim \mathcal {D}} [ \mathbb {E} _ {a _ {t} \sim \pi_ {\phi}} [ \alpha \log (\pi_ {\phi} (a _ {t} | s _ {t})) - \mathcal {Q} _ {\theta} (s _ {t}, a _ {t}) ] ], \tag {2}
$$
where  $\alpha$  is a temperature coefficient.
Figure 2: The architecture of our SGFD. 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.
Notations. For existing visual RL methods [10, 18, 43], the state  $s$  is typically compressed into a compact vector representation  $\mathbf{z}$ . Therefore, we use  $\mathbf{z}$  as the default in our calculations unless stated otherwise. For clarity, we define  $\mathbf{Z}_j$  as the  $j$ -th feature in the  $d$ -dimensional vector, and  $\mathbf{z}_{i,j}$  as the value of  $\mathbf{Z}_j$  in sample  $i$ . For instance, if  $\mathbf{Z}_j$  corresponds to the angle of a robot arm,  $\mathbf{z}_{i,j}$  would represent the encoded value of the angle in image  $i$ .
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  $(\mathbf{Z}_i,\mathbf{Z}_j)$  in the RL states, we gather a batch of  $n$  samples from the replay buffer  $\mathcal{D}$ , denoted as  $(\mathbf{z}_{1,i},\mathbf{z}_{1,j}),(z_{2,i},z_{2,j})\ldots (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:
$$
\hat {\sum} _ {\mathbf {Z} _ {i} \mathbf {Z} _ {j}} = \frac {1}{n - 1} \sum_ {k = 1} ^ {n} \left[ \left(\mathbf {u} \left(\mathbf {z} _ {k, i}\right) - \mathbb {E} \left[ \mathbf {u} \left(\mathbf {Z} _ {i}\right) \right]\right) ^ {T} \cdot \left(\mathbf {v} \left(\mathbf {z} _ {k, j}\right) - \mathbb {E} \left[ \mathbf {v} \left(\mathbf {Z} _ {j}\right) \right]\right) \right], \tag {3}
$$
where  $\mathbb{E}[\mathbf{u}(\mathbf{Z}_i)] = \frac{1}{n}\sum_{k=1}^{n}\mathbf{u}(\mathbf{z}_{k,i})$ ,  $\mathbf{u}$  and  $\mathbf{v}$  are vectors, and each component is a function sampled from the RFF space, formulated as follows:
$$
\begin{array}{l} \mathbf {u} (\mathbf {z} _ {:, i}) = \left(u _ {1} (\mathbf {z} _ {:, i}), u _ {2} (\mathbf {z} _ {:, i}), \dots , u _ {M} (\mathbf {z} _ {:, i})\right), u _ {m} (\mathbf {z} _ {:, i}) \in \mathcal {H} _ {\mathrm {R F F}}, \forall m, \\ \mathbf {v} (\mathbf {z} _ {:, j}) = \left(v _ {1} \left(\mathbf {z} _ {:, j}\right), v _ {2} \left(\mathbf {z} _ {:, j}\right), \dots , v _ {M} \left(\mathbf {z} _ {:, j}\right)\right), v _ {m} \left(\mathbf {z} _ {:, j}\right) \in \mathcal {H} _ {\mathrm {R F F}}, \forall m, \\ \end{array}
$$
where  $\mathcal{H}_{\mathrm{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  $\mathcal{H}_{\mathrm{RFF}} = \{h : x \to \sqrt{2} \cos(\omega x + \psi) | \omega \sim \mathcal{N}(0, 1), \psi \sim \mathrm{Uniform}(0, 2\pi)\}$ , where  $\omega$  is sampled from a standard normal distribution and  $\psi$  is sampled from a uniform distribution. The independence between  $\mathbf{Z}_i$  and  $\mathbf{Z}_j$  can be calculated by the Frobenius norm of the cross-covariance matrix  $||\hat{\sum}_{\mathbf{Z}_i \mathbf{Z}_j}||_F^2$ . If  $||\hat{\sum}_{\mathbf{Z}_i \mathbf{Z}_j}||_F^2$  is zero, the two features  $\mathbf{Z}_i$  and  $\mathbf{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  $||\hat{\sum}_{\mathbf{Z}_i\mathbf{Z}_j}||_F^2$ . To this end, we introduce a set of learnable sample weights, denoted by  $\mathbf{w}\in \mathbb{R}$ , subject to the constraint that  $\sum_{k = 1}^{n}w_{k} = n$ . Based on the weighted samples, the independence can be defined as:
$$
\hat {\sum} _ {\mathbf {Z} _ {i} \mathbf {Z} _ {j}; \mathbf {w}} = \frac {1}{n - 1} \sum_ {k = 1} ^ {n} \left[ \left(w _ {k} \mathbf {u} \left(\mathbf {z} _ {k, i}\right) - \mathbb {E} \left[ \mathbf {w u} \left(\mathbf {Z} _ {i}\right) \right]\right) ^ {T} \cdot \left(w _ {k} \mathbf {v} \left(\mathbf {z} _ {k, j}\right) - \mathbb {E} \left[ \mathbf {w v} \left(\mathbf {Z} _ {j}\right) \right]\right) \right], \tag {5}
$$
where  $\mathbb{E}[\mathbf{w}\mathbf{u}(\mathbf{Z}_i)] = \frac{1}{n}\sum_{k = 1}^{n}w_k\mathbf{u}(\mathbf{z}_{k,i})$ . The correlation among all features in RL states can be eliminated by solving the following optimization problem:
$$
\mathbf {w} ^ {*} = \arg \min  _ {\mathbf {w}} \sum_ {1 \leq i <   j \leq d} | | \hat {\sum} _ {\mathbf {Z} _ {i} \mathbf {Z} _ {j}; \mathbf {w}} | | _ {F} ^ {2}, \tag {6}
$$
where  $d$  is the dimension of  $\mathbf{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) = \partial f(\mathbf{Z}) / \partial \mathbf{Z}_{\mathrm{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}} / \sum_{1\leq j\leq d}e^{M_{j}}$ . We incorporate this term into the Equation (6), revising it as:
$$
\mathbf {w} ^ {*} = \arg \min  _ {\mathbf {w}} \sum_ {1 \leq i <   j \leq d} p \left(\mathbf {Z} _ {i}\right) p \left(\mathbf {Z} _ {j}\right) \| \hat {\sum} _ {\mathbf {Z} _ {i} \mathbf {Z} _ {j}; \mathbf {w}} \| _ {F} ^ {2}. \tag {7}
$$
Based on the weighted samples, we can reformulate the Equation (2) as:
$$
\mathcal {J} _ {\pi} (\phi) = \sum_ {s _ {t} \sim \mathcal {D}} w _ {t} \mathbb {E} _ {a _ {t} \sim \pi_ {\phi}} [ \alpha \log \left(\pi_ {\phi} \left(a _ {t} \mid s _ {t}\right)\right) - \mathcal {Q} _ {\theta} \left(s _ {t}, a _ {t}\right) ]. \tag {8}
$$
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
TED
6-10: task-relevant variations
AMBS
SGQN
1: Walker-walk 6: MT-Walker-walk
2: Cheetah-run 7: MT-Cheetah-run
3: Finger-spin 8: MT-Finger-spin
4: Walker-run 9: MT-Walker-run
5: Finger-turn 10: MT-Finger-turn
1-5: task-irrelevant variations
Figure 3: The normalized performance of 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-turn. Note that the results related to task-relevant features are the average of the two settings presented in Table 2.
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.