SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
Table 1: Generalization to changed task-irrelevant features. For each task, we cost 1e6 steps in training environments and evaluate in the testing environment under background noises.
Tasks	SGFD	TED	AMBS	SGQN	DBC	DR
walker-walk	959.1±26.3	871.5±60.6	926.7±53.2	815.1±53.9	800.9±41.4	712.4±93.7
cheetah-run	599.6±47.2	544.7±22.9	517.7±73.4	332.8±55.1	312.1±20.3	340.0±44.0
finger-spin	965.7±45.9	932.1±71.0	925.1±50.5	943.3±46.2	663.7±68.7	860.8±42.1
walker-run	420.7±39.2	387.8±27.1	398.7±32.0	317.2±34.5	332.4±37.1	231.3±8.9
finger-turn	984.3±11.5	963.9±94.5	966.7±37.0	971.3±26.0	931.2±41.6	947.2±21.7
Total	3929.5	3700.0	3734.9	3379.7	3040.3	3091.7
Table 2: Generalization to changed task-relevant features. For each task, we cost 1e6 steps in training environments and evaluate in the testing environment with different robot parameters. We consider two setups for evaluation: an interpolation setup and an extrapolation setup where the variations in the task-relevant features are interpolations and extrapolations between the variations of the training environment, respectively.
Tasks	SGFD	TED	AMBS	SGQN	DBC	DR
Interpolation	MT-walker-walk	549.4± 42.5	471.9± 18.3	532.5± 81.7	287.1± 34.5	245.7± 47.4	343.8± 70.6
MT-cheetah-run	395.9± 35.2	367.8± 42.2	298.6± 53.5	225.7± 40.9	191.7± 23.5	216.3± 33.1
MT-finger-spin	234.1± 11.9	201.7± 17.9	161.3± 17.3	135.7± 16.9	221.6± 21.5	207.1± 28.2
MT-walker-run	170.3± 05.7	125.6± 13.2	161.4± 06.7	126.3± 18.6	97.2± 19.0	160.3± 16.6
MT-finger-turn	923.7± 26.1	748.9± 44.3	821.8± 52.3	786.6± 65.6	358.5± 83.9	704.7± 70.1
Total	2273.4	1910.6	1975.6	1561.4	1114.7	1632.2
Extrapolation	MT-walker-walk	541.7± 65.4	365.9± 17.7	467.5± 91.7	271.2± 75.4	229.8± 89.9	307.8± 58.9
MT-cheetah-run	392.3± 32.1	311.9± 52.7	270.2± 35.5	167.2± 39.1	174.0± 45.1	196.6± 49.8
MT-finger-spin	231.8± 11.5	199.7± 18.0	160.2± 17.6	135.6± 11.3	221.4± 43.0	197.1± 21.5
MT-walker-run	170.0± 07.2	126.7± 13.2	156.2± 07.5	118.9± 18.2	89.7± 19.7	156.9± 12.7
MT-finger-turn	917.3± 22.6	743.6± 58.3	803.5± 57.4	653.3± 56.6	335.6± 56.5	611.7± 53.6
Total	2253.1	1747.8	1857.6	1346.2	1050.5	1470.1
5.2 Evaluation on the Generalization of SGFD
Evaluation on Changed Task-Irrelevant Features. In this experiment, we assessed the 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
Figure 4: Pearson correlation coefficients among features: a) on raw data, b) on weighted data, c) on saliency-guided weighted data. The features  $s_1$  to  $s_{10}$  represent the state information of the environment and act as conditional inputs for decision-making  $\pi_{\phi}(a|s)$ , such as the coordinates of the object or the angle of the manipulator. The feature  $s_{10}$  is the one that varies across the environments; thus, the correlations between  $s_{10}$  and other features play a crucial role in generalizations. The strength of the correlation is reflected by the color intensity - the darker the color, the stronger the correlation.
(a) 20 states from environments with different background
(b) 20 states from environments with different robot parameters
Figure 5: Visualization of weighted samples obtained from different environments with background noises and robot parameters. States that belong to the same dotted box with a specific color are considered semantically similar. SGFD balances the weights of samples to eliminate the correlations between changed features and other features. As shown in (a), due to the interference of the background noises, the robot fell more often (states in the green dotted box) than it did without the background. SGFD reduces the weighting of states with backgrounds and enhances the weighting of states without backgrounds, effectively neutralizing the correlation between background and body posture. A comparable result was observed in (b), where the reweighting process aimed to equalize the number of states in each group across environments with different body lengths.
setting. These results demonstrate that SGFD can effectively adapt to new values of task-relevant features, showing its robust generalization capability in a variety of scenarios.
5.3 Case Studies on the SGFD
Correlation of the Weighted Samples. To evaluate the decorrelation capability of our model, we implemented 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.
Figure 6: Ablation studies. (a) The performance difference of ablation studies on the RFF. (b) The performance difference of ablation studies on the saliency-guided model. The x-axis denotes the steps in the training environments. The y-axis denotes the cumulative reward recorded in the testing environment with background noises (Cheetah-run) or varied robot parameters (MT-cheetah-run).
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.
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 saliency-guided 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.
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.
[1] Julius Adebayo, Justin Gilmer, Michael Muelly, Ian Goodfellow, Moritz Hardt, and Been Kim. Sanity checks for saliency maps. Advances in neural information processing systems, 31, 2018.
[2] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[3] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
[4] O Ahmed, F Träuble, A Goyal, A Neitz, M Wüthrich, Y Bengio, B Schölkopf, and S. Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv peprint arXiv:2010.04296, 2020.
[5] Akanksha Atrey, Kaleigh Clary, and David Jensen. Exploratory not explanatory: Counterfactual analysis of saliency maps for deep reinforcement learning. In International Conference on Learning Representations, 2020.
[6] David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization in reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
[7] David Bertoin, Adil Zoutine, Mehdi Zoutine, and Emmanuel Rachelson. Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[8] Ioana Bica, Daniel Jarrett, and Mihaela van der Schaar. Invariant causal imitation learning for generalizable policies. In Advances in Neural Information Processing Systems (NIPS), 2021.
[9] Yihan Cao, Siyu Li, Yixin Liu, Zhiling Yan, Yutong Dai, Philip S Yu, and Lichao Sun. A comprehensive survey of ai-generated content (aigc): A history of generative ai from gan to chatgpt. arXiv preprint arXiv:2303.04226, 2023.
[10] Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In International Conference on Robotics and Automation (ICRA), 2019.
[11] Jianda Chen and Sinno Jialin Pan. Learning generalizable representations for reinforcement learning via adaptive meta-learner of behavioral similarities. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
[45] 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.
[46] Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: Task-aware lipschitz data augmentation for visual reinforcement learning. arXiv preprint arXiv:2202.09982, 2022.
[47] Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. arXiv preprint arXiv:2212.08860, 2022.
[48] A Zhang, 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.
[49] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning (PMLR), 2020.
[50] Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstractions for hidden-parameter block mdps. arXiv preprint arXiv:2007.07206, 2020.
[51] Xingxuan Zhang, Peng Cui, Renzhe Xu, Linjun Zhou, Yue He, and Zheyan Shen. Deep stable learning for out-of-distribution generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
[52] Ce Zhou, Qian Li, Chen Li, Jun Yu, Yixin Liu, Guangjing Wang, Kai Zhang, Cheng Ji, Qiben Yan, Lifang He, et al. A comprehensive survey on pretrained foundation models: A history from bert to chatgpt. arXiv preprint arXiv:2302.09419, 2023.
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  $\pi_{\phi}$  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\mathrm{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  $\mathcal{Q}_{\theta}$  and policy model  $\pi_{\phi}$  based on the weighted data, as described in lines 18-20.
Algorithm 1 Saliency-Guided Features Decorrelation
1: Input: Initialize a policy model  $\pi_{\phi}$ , a state-action value function  $\mathcal{Q}_{\theta}$ , a classifier model  $f_{\psi}$  and a set of training environments with environment labels  $\{e^{1}, e^{2}, \ldots, e^{K}\}$ , where  $e^{k}$  is a  $K$ -dimensional one-hot vector and the  $k$ -th component is 1  
2: Output: The learnt policy  $\pi_{\phi^*}$   
3: for  $t = 1$  to  $T$  do  
4: Applying the policy model  $\pi_{\phi}$  to interact with training environments  
5: Sampling a batch of transitions along with environment labels  $(s_i, a_i, r_i, s_i', e_i^k) \sim \mathcal{D}$   
6: for  $iter = 1$  to 10 do  
7: if  $t < 1$ e5 or the accuracy of  $f(s) > 0.9$  then  
8: Break this cycle  
9: else  
10: Sampling a batch of transitions along with environment labels  $(s_i, a_i, r_i, s_i', e_i^k) \sim \mathcal{D}$   
11: Update the classifier model  $f_{\psi}$  by minimizing the following loss  
 $\mathcal{L}(\psi) = -\frac{1}{N} \sum_{i=1}^{N} \sum_{j=1}^{K} e_{i,j}^k \log(f_{\psi}(s_i)_j)$ , where  $e_{i,j}^k$  denotes the  $j$ -th component of vector  $e_i^k$ ,  $f_{\psi}(s_i)_j$  denotes the predicted probability of the  $j$ -th class  
12: end if  
13: end for  
14: Initialize a set of sample weights that can be learned w  
for  $iter = 1$  to 10 do  
16: Update the sample weights by minimizing Equation (7)  
17: end for  
18: Update the state-action value function  $\mathcal{Q}_{\theta}$  by minimizing Equation (1)  
19: Update  $\pi_{\phi}$  by minimizing Equation (8)  
20: end for
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 configurations. 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 configurations, where the values of configurations 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.
Table 3: Architectures of SGFD.
Hyperparameter	Value
Encoder architecture	Convolutional layers	4
Latent representation dimension	50
Image size	(84,84)
Stacked frames	3
kernel size	3 × 3
Channels	3
stride	2 for the first layer, 4 otherwise
Activation	ReLU
Policy and value function	MLP layers	2
Hidden dimension	1024
Activation	ReLU
Classifier architecture	MLP layers	2
Hidden dimension	128
Activation	ReLU
Table 4: Hyperparameters of SGFD.
Hyperparameter	Value
Training for policy	Optimizer	Adam
Learning rate	1e-3
Action repeat	2
Replay buffer capacity	1000000
Batch size	128
Target soft-update rate (τ)	0.01
Actor update frequency	2
Training for sample reweighting	Optimizer	SGD
momentum	0.9
Learning rate	1e-2
weight decay	1e-4
The number of RFFs m	5
For each experiment, we report the mean and standard deviation over 10 random seeds.
Compute. Experiments are carried out on NVIDIA GeForce RTX 3090 GPUs and NVIDIA A10 GPUs. Besides, the CPU type is Intel(R) Xeon(R) Gold 6230 CPU @ 2.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 \times 3$  kernel size and 32 channels. The first layer has a stride of 2, and all other layers have 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  $\pi_{\phi}$  and the state-action value function  $\mathcal{Q}_{\theta}$  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.
Figure 7: The visualization of training environments and testing environments.