Learning Generalizable Agents via Saliency-Guided Features Decorrelation

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.

Paper

References (49)

Scroll for more · 37 remaining

Similar papers

Peer review

Reviewer NdCS7/10 · confidence 3/52023-06-30

Summary

This paper introduces an original approach named Stochastic Gradient Feature Decorrelation (SGFD), which aims to amplify the generalization capabilities of reinforcement learning (RL) agents across various environmental variations. These variations can encompass task-irrelevant visual attributes such as backgrounds, as well as task-related factors like physical configurations. The authors propose achieving this through a decorrelation of features, executed via a resampling technique intended to minimize the Frobenius norm of the cross-covariance matrix, derived from Random Fourier features. Recognizing the inherent challenges of complete decorrelation, the authors shift focus towards effectively decorrelating the most variable features. This is facilitated by leveraging the saliencies from an environment classification model, which, under ideal circumstances, makes its decisions based on distinct features that are not common to different environments.

Strengths

The paper is commendably well-written and coherent, effectively explaining complex ideas in an accessible manner. The authors demonstrate a strong theoretical grounding, with well-motivated intuitions supporting their methodology. Their saliency-guided optimization is an interesting approach, backed by a robust ablation study. SGFD successfully tackles generalization issues relating to both task-irrelevant and task-relevant features, demonstrating a broad scope of applicability. The proposed method do present noticeable enhancements in generalization performance, more so in the case of task-relevant features.

Weaknesses

Despite the paper's strengths, there are some areas where it could be improved. Firstly, the methodology requires several environments with variations to train their environment classifier, introducing an element of manual supervision into the learning process, which may not be ideal in all scenarios. Furthermore, the full algorithm can be challenging to comprehend without first referring to Appendix A, suggesting that the main body of the text might benefit from additional clarification. It might also be beneficial to introduce the general objective - namely, the reweighting of the batch sampled from the buffer - earlier in the paper to give readers a clearer understanding of the process. Lastly, while the authors promote their method as an improvement over Soft Actor-Critic (SAC), the achieved results actually utilize the same updates for the encoder as employed by Adaptive Meta-learner of Behavioral Similarities (AMBS). This could lead to misunderstandings, as readers may infer that the saliency-guided resampling alone yields these performances. It would therefore be beneficial for clarity and fairness if the authors explicitly acknowledged this in the experimental section.

Questions

* Could the authors elaborate on the inference procedure employed during testing? This could help elucidate the practical applicability of the methodology. * Does the encoder also benefit from the gradients produced through resampling? Understanding this aspect could contribute to a more comprehensive understanding of the process. * 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

While the authors have reasonably addressed the methodological limitations of their approach, they assume no potential negative societal impacts stemming from their method.

Reviewer SPDk7/10 · confidence 2/52023-07-07

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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

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.

Reviewer jEbK7/10 · confidence 4/52023-07-10

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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The work is limited to the specific datasets tried on in terms of training, otherwise adaptable.

Reviewer NDX85/10 · confidence 4/52023-07-24

Summary

This paper propose SGFD, an novel approach aimed at improving the generalization ability of distinguishing task-irrelevant and task-relevant situations. SGFD leverages two core techniques: Random Fourier Functions (RFF) and the saliency map, to estimate the complex non-linear correlations in high-dimensional images, and identify the changed features to achieve decorrelation. Furthermore, it conducts experiments on two benchmarks to exhibit its improved generalization ability.

Strengths

1. Its proposed method has the capability to address problems in two distinct settings. 2. I think utilizing saliency map to further achieve decorrelation is an interesting idea.

Weaknesses

1. Figure 2 is confusing. This figure is meant to illustrate the core content of the paper, its unclear presentation undermines the overall comprehensibility. For instance, it would be clearer if the authors had depicted only one type of environment. The inclusion of two types of environments creates ambiguity as to whether the training was performed in one or both environments. 2. This paper lacks a clear logical flow in the writing, particularly when it comes to the technical details. 3. 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.

Questions

1. Can you visualize the examples of training enviroments and test environments? I don't know the gap between these two environments and how difficult to generalize to the novel scenes. 2. 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. 3. Why the calculated $w*$ only apply to actor loss? Why not critic? 4. Is it possible to apply SGFD on the environments which contains both task-relevant and task-irrelevant features? If you can address my questions, I would consider increasing the score.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Mentioned in weakness and questions.

Reviewer tXZg7/10 · confidence 4/52023-07-25

Summary

This paper presents a method to weight experiences in reinforcement learning so that the state features they depict are maximally decorrelated. The assumption is that learning with decorrelated features leads to better generalization to unseen task instances because the policy does not confound the role of different features when making decisions. The idea is novel and interesting and the experimental evaluation includes several simulated embodied AI tasks and comparisons to other baselines.

Strengths

- The presented method is novel. The ideas of the paper are well-grounded. - The method seems mathematically sound. The appendix includes all relevant proofs. - The experimental evaluation, including several baselines and different tasks, demonstrates the strengths of the algorithm - The problem addressed in the paper is critical for the robot learning community: improving the generalization capabilities of trained policies.

Weaknesses

- The writing could be improved. The text focuses on describing the steps but would benefit from explaining more the reasoning behind the design decisions. - 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? - 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? - 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.

Questions

- Explain the features per task and where they come from. Discuss the applicability of the method if it requires pre-specifying optimal features. How would the method work with learned latent features? Could you include an evaluation on that? - Can you include completely distracting features to evaluate the robustness of the solution? - 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. - 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

- The limitations section is acceptable, but several significant limitations for the method to be truly impactful are not discussed (see my comments above). Please, include them.

Reviewer NdCS2023-08-11

I appreciate the authors' responses that clarified my questions. After reviewing their explanations, I maintain my opinion that this paper possesses the necessary qualities for acceptance.

Reviewer tXZg2023-08-13

Thanks for the replies, they clarify some misconceptions. I have raised my score.

Reviewer NDX82023-08-15

Reply to Authors

Thanks for your reply. I have several other questions. Regarding Q3, would the calculated weight be effective on any objective? Why must it only be placed on the actor? It feels to me that there is some disconnection between how to obtain $w$ and how to apply it to RL. Moreover, your answer to Q3 seems subjective; many generalization papers (e.g, SVEA, SGQN) have made modifications to the critic objective, and they are also very effective. Regarding Q4, the effect of the finger turn appears very good, but in fact, you have not shown any visualizations related to anything beyond the cheetah task, leaving the reader in the dark about how you modified your environment.

Authorsrebuttal2023-08-16

Reply to the Reviewer NDX8

We thank Reviewer NDX8 for further response to our rebuttal. We want to further clarify some of the questions as below. **Q1. Why must $w$ only be placed on the actor? Would the calculated weight be effective on any objective?** **Why only apply $w$ to the actor:** It is possible to apply $w$ to the critic. We apply $w$ to the actor for the following two reasons. - The core idea of sample reweighting is to eliminate the correlation between features and then facilitate the policy to distinguish the impact of each image feature on decision-making. Since the actor directly controls the training of the policy, incorporating $w$ into the actor's loss is an intuitive and efficient method. Although modifying the critic loss is incorporated into the algorithm in related works, the critic can only indirectly benefit the policy on $w$ in our method, which is demonstrated in experiments (see the next point). - We conduct ablation experiments to test generalization under four settings: no $w$, $w$ only on the actor, $w$ only on the critic, and $w$ on both the actor and the critic. As shown in Table 1, when $w$ acts on the actor, the generalization is significantly improved, while the gain realized by the critic is relatively modest. This is because the actor loss directly controls the learning of the policy so that $w$ applied to the critic alone cannot prevent the policy from establishing false connections with the background noise. In addition, we do not observe significant improvements when adding $w$ to both the actor and the critic, compared to adding $w$ only to the actor. This further demonstrates that actors with $w$ adequately mitigate spurious associations caused by correlations between features. **Table 1:** The generalization of the $w$ acting on different parts of the algorithm in unseen background noises. | | both the actor and critic | only the actor | only the critic | no $w$ | |-------|-------|-------|-------|-------| | walker-walk | 956.3 | $\mathbf{959.1}$ | 915.2 | 906.7 | | cheetah-run | $\mathbf{606.5}$ | 599.6 | 521.6 | 517.7 | | finger-spin | 958.6 | $\mathbf{965.7}$ | 905.5 | 895.5 | **Would the calculated weight be effective on any objective?:** Although the SAC algorithm is the backbone of our method, the form of sample reweighting makes it possible to transfer to other RL algorithms. Based on the experimental results, we propose to apply $w$ to the loss directly associated with the policy. For the RL method without the policy networks, we will further explore the modification of $w$ in future work. **Q2. The visualizations of finger turn.** Due to the page limit of the global rebuttal, we only visualized the configuration of the cheetah. According to the official regulations of NeurIPS 2023, we provided an anonymous link in the AC comment to visualize the configuration of the three robots used in our experiments. After AC's confirmation of anonymity, you can see it in the comments at the top. Specifically, the settings of the Finger and walker are referred to [1], which modifies the Finger's length and walker's foot. For background noise, Finger and cheetah are similar to the walker we showed in the global rebuttal PDF (Figure 3 (a)). The visualization will be updated to the revised version. **Reference** [1] Learning robust state abstractions for hidden-parameter block mdps, ICLR 2021. ----- We hope that our response can address your comments and we would appreciate it if you considered increasing your score. Paper1844 Authors

Reviewer NDX82023-08-18

Reply to Authors

Thanks for your further reply. - SGQN is an algorithm employed in a different setting. Yet It is also proper to apply it on your setting. I wonder in your setting, is it also trained in a noisy environment and subsequently tested in unseen noisy scenarios? - Could you show me the training curve? Is the training performance of SGFD better than other algorithms? Because, I think the level of generalization you have shown appears to be much simpler than DMCGB. If SGQN can achieve high scores in the training scenarios, it is unlikely to experience a significant performance drop.

Authorsrebuttal2023-08-18

Reply to Reviewer NDX8

We thank the reviewer for the response, and the following are our further clarifications. **Q1. I wonder in your setting, is it also trained in a noisy environment and subsequently tested in unseen noisy scenarios?** Testing in the scene with unseen background noise is one of our settings, i.e., task-irrelevant features. The background noise follows the setting of AMBS [1] in the comparison algorithm. However, we want to emphasize that our SGFD can not only handle background noise, but also generalize to scenarios with **unseen configurations of robots** (i.e., task-relevant features), while the latter is not covered by SGQN. Unlike background noise, the differences in robot configurations mean that the optimal action also changes [2]. Therefore, from the causal perspective, we propose feature decorrelation to ensure that the policy recovers the true associations between changed features and optimal behavior, whether the changed features are task-relevant or task-irrelevant. **Q2. Could you show me the training curve? Is the training performance of SGFD better than other algorithms?** Although, according to the rules of Neurips 2023, we cannot visualize the training curve in the comment, Figures 12 and 13 in the appendix show similar results. Figure 12 and Figure 13 visualize the learning curves under two settings of unseen background noises (task-irrelevant) and robot configurations (task-relevant), respectively. The x-axis represents the number of steps in the training environment, and **the y-axis represents the reward obtained in the test environment**. **The training curves in environments with background noises:** For the setting of background noise, the training curves of SGQN are mostly consistent with the curves in Figure 12. Aided by feature decorrelation, our method performs better than SGQN in noisy training environments. **The training curves in environments with different robot configurations:** Since SGQN does not cover task-relevant generalization scenarios, its test performance is worse than its training performance under the robot configuration change scenario. In contrast, our method can generalize well to different robot configurations and outperform SGQN as shown in Figure 13. **The level of generalization and the comparison to DMCGB:** We want to clarify again that SGFD is proposed to expand the types of generalization, i.e., from task-irrelevant features to task-relevant features. It is worth acknowledging that the DMCGB adds more disturbing noise to the test environment, and SGQN proposes an efficient framework. However, we provide an orthogonal approach that attempts to achieve more types of generalization, which is in a different direction than their exploration. We will update this part of the discussion to the revised version. **References** [1] Learning Generalizable Representations for Reinforcement Learning via Adaptive Meta-learner of Behavioral Similarities. ICLR 2022. [2] Temporal disentanglement of representations for improved generalisation in reinforcement learning. ICLR 2023.

Authorsrebuttal2023-08-20

Follow-up to Reviewer NDX8

Dear Reviewer NDX8, We would like to express our sincere gratitude to you for reviewing our paper and providing valuable feedback. Could we kindly know if the responses have addressed your concerns? If there are any further questions, we are happy to clarify. Thank you. Best, Paper1844 Authors

Reviewer NDX82023-08-20

Reply to Authors

Thanks for your response. I will raise my score. Since I still haven't seen the visualization results, my main concern is that whether DMC is a suitable benchmark for your specific two settings. Additionally, from the visualization images you provided, your generalization can be considered as in-distribution generalization. For this type of generalization, I believe that showing the training curve is essential.

Reviewer jEbK2023-08-18

Intermediate comments

Thanks for the clarifications. I am sticking with my initial score while other reviewers are updating.

Authorsrebuttal2023-08-20

Follow-up to Reviewer SPDk

Dear Reviewer SPDk, We value your positive feedbacks and constructive suggestions for our paper and sincerely appreciate your effort in reviewing it. As the end of the discussion is approaching, we kindly request your consideration regarding the possibility of raising the score. We thank the reviewer's effort in the review of our paper. We hope we have effectively addressed all the concerns raised. Should there be any remaining concerns, we stand ready to offer additional clarifications. Thank you again for your dedicated review and invaluable insights. Kind regards, Paper1844 Authors

Reviewer SPDk2023-08-21

Thank you for the responses. I have read the reviews of the other reviewers and the authors' responses. The authors' responses clarified my questions. Also, my concerns about dependency on the classifier part is well described in the appendix. Finally I have raised my score.

Program Chairsdecision2023-09-21

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC