Look where you look! Saliency-guided Q-networks for visual RL tasks
Anonymous Author(s)
AffiliationAddress email
Abstract
Deep reinforcement learning policies, despite their outstanding efficiency in simulated visual control tasks, have shown disappointing ability to generalize acrossdisturbances in the input training images.
Changes in image statistics or distracting background elements are pitfalls that prevent generalization and real-worldapplicability of such control policies.
We elaborate on the intuition that a goodvisual policy should be able to identify which pixels are important for its decision,and preserve this identification of important sources of information across images.
This implies that training of a policy with small generalization gap should focuson such important pixels and ignore the others.
This leads to the introduction ofsaliency-guided Q-networks (SGQN), a generic method for visual reinforcementlearning, that is compatible with any value function learning method.
SGQN vastlyimproves the generalization capability of Soft Actor-Critic agents and outperformsexisting state-of-the-art methods on the Deepmind Control Generalization benchmark, setting a new reference in terms of training efficiency, generalization gap,and policy interpretability.
Introduction
Deploying reinforcement learning (RL) [Sutton and Barto, 2018] algorithms in real-life situationsrequires overcoming a number of still open challenges.
Among these is the ability for the trainedcontrol policies to focus their attention on causal state features and ignore confounding factors[Machado et al., 2018, Henderson et al., 2018].
In visual RL tasks, this implies for instance being ableto ignore the background and other distracting factors, even when they might be somehow correlatedwith progress within the task at hand.
Despite a very active trend of research on the topic of closingthe generalization gap for RL agents [Cobbe et al., 2019, 2020, Song et al., 2019, Hansen et al., 2021,
Hansen and Wang, 2021], current algorithms are still rather brittle when it comes to filtering out suchdistracting factors, which hinders their applicability to real-life scenarios.
In the present work we propose a novel method which encourages the agent to identify efficientlycrucial input pixels, and strengthen the policy’s dependency on those pixels.
In plain words, weencourage the agent to pay attention to, and be self-aware of where it looks in input images, inorder to make its decision policy more focused on important areas and less sensitive to ambiguousor distracting pixels.
This intention is expressed within the generic method of saliency-guided
Q-networks (SGQN), which can be applied on any approximate value iteration based, deep RLalgorithm.
SGQN relies on two core mechanisms.
First, it regularizes the value function learningprocess with a consistency term that encourages the value function to depend in priority on pixelsthat are identified as decisive.
The second mechanism pushes the agent to be self-aware of whichpixels are responsible for making decisions, and encode this information within the extracted features.
This second mechanism translates to a self-supervised learning objective, where the agent trains topredict its own Q-value’s saliency maps.
In turn, this improves the regularization of the value function learning phase, which provide better labels for the self-supervised learning phase, overall resulting ina virtuous improvement circle.
SGQN is a simple, generic method, that permits many variants in the way the two core mechanismsare implemented.
In the present paper, we demonstrate that applying SGQN to soft actor-criticagents [Haarnoja et al., 2018] dramatically enhances their quality on the DMControl generalization benchmark [Hansen and Wang, 2021], a standard evaluation benchmark for generalization incontinuous actions RL.
SGQN already improves the training efficiency of such agents in domainswithout distractions.
But most importantly, it sets a new state-of-the-art in terms of generalizationperformance, in particular in especially difficult benchmarks where previous methods suffered fromconfounding factors.
As a side benefit, it also provides explanations of its own decisions at run time,under the form of interpretable attribution maps, with no overhead cost and no need to compute adhoc saliency maps, which is another desirable property in the pursuit of deployable RL.
Section 2 of this paper introduces the necessary background and state-of-the-art in closing thegeneralization gap in RL, as well as attribution methods, leading to the key intuitions underpinning
SGQN.
Section 3 introduces the method itself and implements it within soft actor-critic agents.
Section 4 evaluates SGQN’s training efficiency, generalization capabilities, and policy interpretability.
It also discusses the different design choices made along the way and some foreseeable limitations.
Section 5 summarizes and concludes this paper.
Background and related work
Reinforcement learning (RL).
RL [Sutton and Barto, 2018] considers the problem of learning adecision making policy for an agent interacting over multiple time steps with a dynamic environment.
At each time step, the agent and environment are described through a state s ∈ S , and an actiona ∈ A is performed; then the system transitions to a new state s ′ according to probability P ( s ′ | s, a ) ,while receiving reward R ( s, a ) .
The tuple M = ( S , A , P , R ) forms a Markov Decision Process(MDP)
[Puterman, 2014], which is often complemented with the knowledge of an initial statedistribution p 0 ( s ) .
A decision making policy parameterized by θ is a function π θ ( a | s ) mapping statesto distributions over actions.
Training a reinforcement learning agent consists in finding the policythat maximizes the discounted expected return J ( π θ ) =
E [ (cid:80) ∞ t =0 γ t R ( s t , a t )] .
Poor generalization in RL.
Despite the recent progress of (deep) RL algorithms in solving complextasks, a number of studies have pointed out their poor generalization capabilities.
Using a grid-worldmaze environment, Zhang et al.
[2018c] demonstrate the ability of deep RL agents to memorize anon-trivial number of training levels with completely random rewards.
Using attribution methods,
Song et al.
[2019] highlight what they define as observational overfitting i.e., the propensity of
RL agents to base their decision on background uninformative elements observed during training,instead of the semantic pieces of information one could intuitively expect such as object positions orrelations.
Zhang et al.
[2018a] measure the generalization error in continuous control environmentsby training and testing agents on different sets of seeds.
Zhao et al.
[2019] define generalization in
RL as robustness to a distribution of environments, and samples environments from this distributionto learn a robust policy.
Overall, these works illustrate the lack of generalization abilities of vanilladeep RL algorithms, either to states that were not encountered during training, or to variations in thetransition dynamics.
In the present work, we aim to shape the policy learning process, so that it relieson meaningful features that permit such generalization and robustness.
Evaluating generalization in RL.
Under the impetus of these works and the need for benchmarkswith separate training and testing environments [Whiteson et al., 2011, Machado et al., 2018, Henderson et al., 2018], original benchmarks for evaluating the generalization capacities of RL agents havebeen designed.
Packer et al.
[2018] propose a modified version from control problems in OpenAI
Gym [Brockman et al., 2016] and Roboschool [Schulman et al., 2017] that lets the user change thesystem dynamics.
Machado et al.
[2018] propose a modified version of the ALE environments [Bellemare et al., 2013] allowing one to change modes and difficulties.
Without modifying the underlyingtransition model, Zhang et al.
[2018b], Grigsby and Qi [2020], Stone et al.
[2021] add distractingelements (e.g., addition of real images or videos in the background, change of colors) to the ALEenvironments and the Deepmind control suite.
Even if the modifications to the original environmentsdo not alter the semantic information, they already appear to be challenging for agents prone toobservational overfitting.
Cobbe et al.
[2019, 2020], Juliani et al.
[2019] use procedural content generation to design highly randomized sets of environments with different level layouts, game assets,and objects locations, letting the user study robustness to several independent variation factors.
Onemay note that the diversification of learning environments is in itself a first practical way to inducegeneralization
[Tobin et al., 2017, Cobbe et al., 2019, 2020] and also permits curriculum-basedlearning [Jiang et al., 2021].
Nevertheless, when the diversity of training scenarios is lacking, threesets of methods are generally employed, as detailed below.
Regularization.
Farebrother et al.
[2018], Cobbe et al.
[2019, 2020] demonstrated the beneficialeffects of popular regularization methods from the supervised learning literature.
Igl et al.
[2019]mitigate the adverse effect that classical regularization may have on the gradient quality with selectivenoise injection and combine it with an information bottleneck regularization.
Inspired by mixup[Zhang et al., 2018d], Wang et al.
[2020] use mixtures of observations to stimulate linearity in thepolicy’s outputs in-between states.
Data augmentation.
Laskin et al.
[2020a] evidence the benefits of training RL agents with augmenteddata (RAD).
Yarats et al.
[2020] average both the value function and its target over multiple imagetransformations (DrQ).
Hansen et al.
[2021] only apply data augmentation in Q -value estimationwithout augmenting Q -targets used for boostrapping.
Raileanu et al.
[2021] combine the previousmethod with UCB [Auer, 2002] to pick the most promising augmentation, and apply it to PPO[Schulman et al., 2017].
Yuan et al.
[2022] propose a task-aware Lipschitz data augmentation method(TLDA) to augment task irrelevennt pixels.
Fan et al.
[2021] use weak data augmentation to train anexpert without hindering its performance and distill its policy to a student trained with substantialdata augmentation.
Besides augmentations of raw inputs, other augmentations operate directly withinthe agents’ network.
Lee et al.
[2020] introduces a random convolutions layer at the earliest levelof the agent’s network to modify the texture of the visual observations.
Zhou et al.
[2020] adaptmixup [Zhang et al., 2018d] with style statistics encoded in early instance normalization layers toincrease data diversity.
Bertoin and Rachelson [2022] apply channel-consistent local permutationsof the feature map to induce robustness to spatial spurious correlations.
Finally, data augmentationcan also be used in an auxiliary loss to promote invariance to distributional shift in representations.
Hansen and Wang [2021] propose a soft-data augmentation method (SODA) by adding an auxiliaryself-supervised learning phase to SAC [Haarnoja et al., 2018], similar to BYOL [Grill et al., 2020].
Representation learning.
Higgins et al.
[2017b] demonstrate zero-shot adaptation to unseen configurations in testing environments, using a β -VAE [Higgins et al., 2017a] to learn disentangledrepresentations.
Fan and Li [2021] jointly maximize the mutual information between sequencesof observations to remove the task-irrelevant information.
Fu et al.
[2021] learn a disentangledworld model that separates reward-corrlated features from background.
Wang et al.
[2021] extract,using visual attention, the observation foreground to provide background invariant inputs to thepolicy learner.
Raileanu and Fergus [2021] separate the actor from the critic in the agent’s networkarchitecture and add an adversarial auxiliary objective on the actor’s representations to remove theinformation needed to estimate the value function that is not irrelevant to a general policy.
Zhang et al.[2020] train an encoder to project states so that their distances match with the bisimulation distancesin state space.
Other recent works use behavioral similarities combined with contrastive learning[Agarwal et al., 2020] or clustering [Mazoure et al., 2022] to map behaviorally similar observationsto similar representations.
Attributing decisions to inputs.
Although not directly aiming at generalization, a related topic isthat of attribution , where one wishes to identify which parts of an input are responsible for majorchanges in the output of a function.
Intrinsically, computing attributions boils down to computing(some transformation of) the gradient of the function’s output with respect to the input’s components.
Computational graphs of differentiable functions, such as neural networks, are particularly suited tocomputing attributions by using the back propagation algorithm
[Simonyan et al., 2014, Springenberget al., 2015, Smilkov et al., 2017, Selvaraju et al., 2017, Chattopadhay et al., 2018].
When thesemethods are applied to images, one obtains a map which is known as a saliency map or attributionmap .
Such attribution maps indicate which input pixels are determining for a policy’s output (or the
Q-value of action a ) and thus permit interpretation of the function itself, rather than its point-wisedecision alone.
Mousavi et al.
[2016], Greydanus et al.
[2018], Atrey et al.
[2020] have used saliencymaps to analyze and explain the behavior of RL agents.
Rosynski et al.
[2020] indicates in particularthat guided backpropagation [Springenberg et al., 2015] provides good visualizations of RL policiesacross a span of environments.
Most existing works, however, only exploit saliency maps as toolsfor interpretation.
Closely related to our contribution is that of Ismail et al.
[2021], who incorporate attributions into their training process in supervised learning.
They construct a binary mask from thesaliency map by keeping only pixels above a certain quantile.
Their procedure iteratively uses thismask to remove features with small and potentially noisy gradients while maximizing the similarityof model outputs for both masked and unmasked inputs.
This contribution.
The rationale of the method we introduce in the next section is to encouragethe agent to generalize to new states, based on which pixels are identified as important decisionfactors.
For this purpose, we perform pixel-level masking on the input image, depending on thecomputed attribution, and regularize the value function learning process with the difference in Qvalues.
This way, we encourage the value function to focus specifically on the pixels with highattribution.
Leveraging data augmentation, self-supervised learning methods have demonstrated theirability to induce features that are insensitive to lighting, background, and high-frequency noise.
Themethod we propose does not directly use data augmentation in the value function or policy updatephase.
Instead, it is introduced during an auxiliary phase where the augmented state’s encoding is usedto predict the attribution mask of the original state.
This way, the encoder is encouraged to preserveinformation that is useful for predicting which pixels were important in the agent’s decision-making.
A parallel with SODA can thus be made by considering that the projector used in the BYOL objectiveis here replaced by a surrogate of the derivative of the value function, allowing to refine the quality ofthe projection and to learn which pixels and visual features are consistently important across states topredict the Q-values.
The value function regularization and the self-supervised learning objective aremutually beneficial: the former outputs sharper saliency maps from the value function, that serve asbetter labels for the auxiliary self-supervised learning task, which in turn induces better features andbetter attributions.
In short, we encourage the agent to pay attention to where it is looking, with theintention that this triggers more efficient learning and more interpretable output.
Saliency-guided Q-networks
We propose a generic saliency-guided Q-networks (SGQN) method, for visual deep reinforcementlearning.
In a nutshell, SGQN considers the application of the binarized attribution map as a maskover the input state and regularizes the value function learning objective with a consistency termbetween the Q-values of the masked and the original state images.
It also defines an auxiliaryself-supervised learning task that aims to match a prediction of the attribution map on an augmentedimage, with the attribution map of the original image.
Such an auxiliary task orients the gradientdescent towards features that are shared across states, as illustrated by the work of Hansen and
Wang [2021].
SGQN can be combined with any value function learning objective, any attributionmap computation technique, any image augmentation method for self-supervised learning, and issuited for both discrete and continuous actions.
We first present SGQN as a generic enhancement ofapproximate value iteration methods.
Then we derive a specific version built on SAC [Haarnoja et al.
,2018] and on the guided backpropagation
algorithm
[Springenberg et al., 2015].
A vast number of deep RL algorithms belong to the family of approximate value iteration methods.
Such methods build a sequence of ( Q n ) n ∈ N and ( π n ) n ∈ N functions that aim to asymptotically tendto Q ∗ and π ∗ .
Q n +1 is defined as a minimizer of L Q =
∥ Q − T π n Q n ∥ , where T π n is the Bellmanevaluation operator with respect to π n .
Then π n +1 is defined by applying a greediness operator G to
Q n +1 and the process is iterated.
Geist et al.
[2019] showed how one could introduce regularizationwithin the expression of L Q , yielding the class of regularized MDPs.
The classical DQN algorithm[Mnih et al., 2015] approximates the solution to L Q by taking a number of gradient steps with respectto a target network Q n and uses an arg max greediness operator.
When actions are continuous,actor-critic methods introduce a surrogate model of the Q n +1 -greedy policy, under the form ofan actor network π n +1 =
G ( Q n +1 ) obtained by gradient ascent.
In what follows, we denote by
L Q ( θ ) the loss minimized by a generic learning procedure for Q θ , based on the Bellman operator,independently of whether it is regularized, uses double critics, etc.
Similarly, we note L π ( θ ) the lossminimized to yield a greedy policy π θ , when applicable.
We denote by M ( Q, s, a ) an attribution map for Q ( s, a ) , in the space of images S .
Vanilla grad
[Simonyan et al., 2014] for instance will compute such a map under the form M ( Q, s, a ) =
∂Q ( s, a ) /∂s ,while guided backpropagation [Springenberg et al., 2015] will mask out negative gradients, yielding a different attribution map.
We note M ρ ( Q, s, a ) the binarized value attribution map where
M ρ ( Q, s, a ) j = 1 if attribution pixel M ( Q, s, a ) j belongs to the ρ -quantile of highest values for
M ( Q, s, a ) , and 0 otherwise.
The proposed method is built on a classical Q -network architecture.
The value function is dividedinto 2 parts: an encoder f θ : S → Z and a Q -function Q θ : Z × A → R built on top of this encoder.
We add a decoder function M θ after the features encoder f θ , such that M θ ( f θ ( s ) , a ) aims to predictthe attribution map of Q θ ( f θ ( s ) , a ) .
Many algorithms require defining double critics [Fujimoto et al.
,2018] or target networks f ψ and Q ψ which are often updated with an exponential moving average ofθ [Polyak and Juditsky, 1992].
We omit them here for clarity, although their introduction in SGQN isstraightforward.
When needed, a policy head π θ : Z → A is built on top of the encoder f θ to definethe actor network.
The backbone architecture and training process are summarized in Figure 1.
The SGQN training procedure involves two alternating phases: the critic update phase and theself-supervised learning auxiliary phase.
The critic update follows the classical critic loss L Q ( θ ) ,enhanced with a consistency regularization objective L C ( θ ) =
E s,a [[ Q θ ( f θ ( s ) , a ) − Q θ ( s ⊙ M ρ ( Q θ ◦f θ , s, a ) , a )] 2 ] , where ⊙ denotes the Hadamard product.
This loss function encourages the Q-network
Q θ ◦ f θ to make its decision based in priority on the salient pixels in M ( Q, s, a ) , hence promotingconsistency between the masked and original images.
The new critic objective function is thus definedas L Q ( θ ) + λL C ( θ ) .
During the self-supervised learning phase, given a generic image augmentationfunction τ , the parameters of f θ are updated so that ( f θ ( τ ( s )) , a ) contains enough information toaccurately reconstruct the attribution mask M ρ ( Q θ , s, a ) .
This defines a self-supervised learningobjective function L SL ( θ ) =
E s,a [ BCE ( M θ ( f θ ( τ ( s ) , a ) , M ρ ( Q θ , s, a )] , where BCE is the binarycross entropy loss, which could be replaced by any other measure of discrepancy between attributionmaps.
The interplay between these two phases acts as a virtuous circle.
The consistency regularizationloss, similar to that of Ismail et al.
[2021], pushes the network to focus its decision on a selectedset of pixels (hence relying on the assumption that initial saliency maps are reasonably good).
Thisenhances the contrast between pixels in the gradient image, and thus yields sharp saliency maps, evenbefore binarization.
These maps serve as a target labels during the self-supervised learning phase;since they are less noisy than without the consistency loss, they provide a stronger incentive to encodethe information of which pixels are important, within f θ .
In turn, as exemplified by Hansen and
Wang [2021] and Grill et al.
[2020], the features obtained by the self-supervised learning proceduretend to be insensitive to background, noise, or exogenous conditions, and provide features that areshared across observations.
These features benefit from the better labels (less irrelevant pixels inthe attribution map).
Finally, this helps provide good pixel attributions that will be used in theminimization of the consistency loss during the critic phase.
Appendix G proposes an extendeddiscussion on this virtuous circle.
Note that the thresholding operation performed to obtain M ρ is not strictly necessary, either in theconsistency loss or in the self-supervised learning one.
Instead of a hard thresholding, one couldturn to a normalization of the attribution map, such as a softmax for instance.
Such a soft-attributionimage remains fully compatible with SGQN.
The choice to keep the thresholded ρ -quantile maskis motivated by the arguments of Ismail et al.
[2021] who extensively study such variations andconclude to the benefits of this binarized mask.
One could also remark that SGQN does not requireto use the target network during the self-supervised learning phase, which contrasts with the choicesof SODA or BYOL and makes the method somewhat more versatile.
[Figure]
Algorithm 1 presents the pseudo-code of combining SGQN with SAC, yielding an SG-SAC algorithm.
Note that, for the sake of simplicity, we write θ for the full set of network parameters, which are thusshared by the encoder, the Q-value head, the policy head, and the attribution reconstruction head.
Parameters: frequency of auxiliary updates N SL , attribution quantile value ρ , learning rate α , data augmentation function τ .
for each interaction time step do
[Equation]
// Sample a transition
B ←
[Equation]
// Add transition to replay buffer
[Equation]
// Sample a mini-batch of transitions
[Equation]
Critic update phase θ ← θ − α ∇ θ L π ( θ ) //
Actor update
Every N SL steps: θ ← θ − α ∇ θ L SL ( θ ) //
Self-supervised learning
Note: L Q and L π are as defined by Haarnoja et al.
[2018], temperature update, double critics and target network updates are omitted here for clarity.
Experimental results and discussion
This section evaluates SGQN’s training efficiency, generalization capabilities, and policy interpretability.
It also discusses the different design choices made along the way.
We compare our approach withcurrent state-of-the-art methods for generalization in continuous actions RL (RAD [Laskin et al.
,2020b], DrQ [Yarats et al., 2020], SODA [Hansen and Wang, 2021], SVEA [Hansen et al., 2021])on five environments from the DMControl Generalization Benchmark (DMControl-GB)
[Hansenand Wang, 2021].
The DMControl-GB presents a variety of vision-based continuous control tasksbased on the Deepmind control suite.
Agents are trained in a fixed background environment andevaluated under two challenging distribution shifts, consisting in replacing the training backgroundwith natural videos.
Figure 2 illustrates the effects of both domain shifts.
All the compared methodsherein are variants of SAC, for which we use the same architecture for all agents.
These methodsall use data augmentation in one of their stages.
Following the experimental protocol of the competing approaches, we used a random overlay augmentation
[Hansen and Wang, 2021] (consistingin blending together original observations with random images from the Places365 dataset [Zhouet al., 2017]) for all the methods except for RAD and DrQ, for which we used random crops andrandom shifts respectively, as it is reported as producing the best results [Laskin et al., 2020b, Yaratset al., 2020].
We trained all agents for 500 000 steps using the vanilla training environment with novisual variation.
Appendix A, B, and E discuss all the hyperparameters, network architectures, andimplementation choices used for this benchmark.
Appendix C includes extra experimental results on
DMControl-GB and Appendix D provides additional results on a vision-based robotic environment.
SGQN improves value iteration in the training domain.
We first compare the performance inthe training domain, with no visual distractions, of SAC and SGQN, on five environments from the
DMControl-GB (Figure 3).
The SGQN agents outperform, by a considerable margin, the SAC agentsboth in terms of asymptotic performance and sample efficiency on 4 out of the 5 environments.
Inaddition to obtaining better results, the variance of the agents trained with SGQN is also significantlylower than that of the agents trained with SAC, demonstrating that the enhancements employed in
SGQN have a beneficial effect on the stability of the training, regardless of the ability to generalizeacross domains.
[Table]
SGQN improves generalization.
We assess the zero-shot generalization ability of SGQN on thevideo easy and video hard benchmarks from the DMControl-GB.
The easy version only replaces thebackground of the training image with a distracting image, while the hard version also replaces theground and the shadows (Figure 2).
We report the average sum of rewards after 500 000 training stepsfor the video easy benchmark in the top part of Table 1.
Agents trained with SGQN outperform agentstrained with other state-of-the-art methods on all tasks but one (where it is on par with other agents),thus demonstrating the generalization capabilities induced by the method.
By removing the groundand shadows, the video hard benchmark (bottom part of Table 1) causes a larger, more confusing,and more challenging distributional shift.
All the competitors of SGQN experience a radical decreasein their generalization performance.
SGQN is significantly less impacted and outperforms all itscompetitors with an average margin over the second-best of 57% on all environments, and a gainrange of 14 to 144% across environments.
Figure 4 reports the evolution of each agent’s score on thevideo hard environments, along training.
SODA and SVEA’s scores drop drastically when the groundand shadows are removed.
SGQN is less sensitive to this change, notably through the consistencyloss, which encourages agents to make decisions based on the subsets of pixels they deem mostinteresting.
Overall, the interplay of the two phases of SGQN seems to be key to a major leap forwardin terms of generalization gap in difficult environments, setting a new reference state-of-the-art.
SGQN yields sharp saliency maps.
We use guided backpropagation to visually compare the SGQNagents’ ability to discriminate the essential information with that of other agents.
Figure 5 shows anexample of the binarized attribution maps for all agents in a video hard state.
While the other agentsseem to be disturbed by background elements and retain some dependency on background pixels intheir decision, the attributions of the agent trained with SGQN are precisely located on the importantinformation, hence suggesting better generalization potential.
[Figure]
SGQN is interpretable by design.
In all our experiments, we trained the SGQN agents using guidedbackpropagation.
We emphasize that any other attribution method could be used instead.
Some ofthese methods are expensive and require several forward (or backward) passes within the network(e.g., RISE [Petsiuk et al., 2018], or the work of Fel et al.
[2021]) to yield attribution maps whichexplain the agent’s decision.
In contrast, SGQN’s auxiliary phase trains a predictor to estimate themost important pixels according to the chosen attribution method.
Therefore, this predictor is asurrogate of the attribution method used during training.
It allows identifying the essential imagefeatures that condition the agent’s decision in the same forward pass as the prediction of the actionitself, without incurring the cost of one or several costly additional backward or forward passes.
Figure 6 illustrates the proximity between the actual saliency map and the attribution surrogate model
M θ .
This makes SGQN both self-aware (of its own saliency maps) and intrinsically interpretablefrom a human perspective, with no computational overhead at evaluation time.
[Figure]
Ablation study.
SGQN relies on two enhancements of vanilla approximate value iteration: theauxiliary self-supervised learning phase and the consistency regularization term in the critic’s loss.
We perform an ablation study to assess their individual contribution and Table 2 reports the averagesum of rewards in the training domain and for zero-shot generalization in all environments (Appendix
F takes a different perspective and compares SGQN with the combination of SVEA and SODA).
Individually, each of these features greatly improves both training and zero-shot generalizationperformance on all environments.
The auxiliary self-supervised learning phase provides the mostsignificant performance gains over vanilla SAC.
The average training performance on all environmentsof agents trained with this auxiliary objective is more than 87% higher than that of agents trainedwith vanilla SAC.
The same applies to the performance in zero-shot generalization, which increasesby more than 141% on video easy and by more than 182% on video hard environments.
One can notethat agents trained with our self-supervised learning objective obtain performance of the same orderof magnitude as SODA agents (Table 1).
Recall that SODA relies on the BYOL [Grill et al. , 2020]self-supervised feature learning procedure, whose target labels differ notably from those proposedherein.
The reported performance, compared to that of SODA suggests that attribution maps constitutea good labeling function that could be considered in the more general context of self-supervisedlearning.
To a slightly lesser extent, adding the consistency loss to SAC also significantly improvesits performance.
The average score obtained improves by more than 49% on the training domain andby 86% and 166% respectively on the video easy and video hard domains.
Similarly to the resultsobtained by Ismail et al. [2021] in supervised learning, the regularization of the critic’s loss with theconsistency term sharpens the attribution maps obtained (Figure 7).
In SGQN, these fine-grainedattributions provide higher quality labels to the auxiliary self-supervised learning phase, thus yieldingsignificant performance improvements in training ( +94% on average, range up to +242%), videoeasy generalization (+183% on average, range up to +399%), and video hard generalization (+440%on average, range up to +1977%).
[Table]
[Figure]
Conclusion
The ability to filter out confounding variables is a long-standing goal in machine learning.
For visualreinforcement learning, it is a pre-requisite for real-world deployment of learned policies, since wewant to avoid at all costs situations where an agent makes the wrong decision due to distracting visualfactors.
In this work, we introduced saliency-guided Q-networks (SGQN), a generic method for visualreinforcement learning, that is compatible with any value function learning method.
SGQN relies onthe positive interaction of two core mecanisms of self-supervised learning and attribution consistencythat jointly encourage the RL agent to be self-aware of the decisive factors that condition its valuefunction.
We implement an SGQN agent based on the soft actor-critic algorithm, and evaluate iton the DMControl generalization benchmark.
This agent displays a dramatically more efficientlearning curve than vanilla SAC on the various environments it is trained on.
Most importantly, thepolicy it learns closes the generalization gap on environments that include confusing and distractingvisual features, setting a new reference in terms of generalization performance.
Since they rely onself-awareness of important pixels, SGQN agents are also very interpretable, in the sense that theyprovide a prediction of their own saliency maps, with no computational overhead.
The introduction of SGQN is an exciting milestone in RL generalization, and we wish to concludethis contribution by highlighting some limitations and perspectives for research that we believe arebeneficial to the community.
Attribution maps seem to be an efficient proxy for encouraging causalrelationships within policies, but they are strongly grounded in an anthropomorphic point of view ofwhat a visual policy should be.
Extending their definition to more abstract notions of attribution is stilla challenge and begs for important contributions, both theoretical and algorithmic.
Similarly, suchattribution maps appear to be relevant self-supervised learning targets, in order to learn good featuresfor RL agents.
Exploring whether this still holds for different tasks than RL is an open question.
It islikely that one could design variations of SGQN that perform better than SG-SAC.
The extensionof SGQN agents to discrete action agents (e.g. DQN), or policy gradient methods (e.g. PPO) is apromising perspective in itself.
Although generic and grounded in sound algorithmic mechanisms,the losses introduced by SGQN lack a formal connection to some measure of the generalization gap.
Such a connection could provide insights to better self-aware, explainable agents with improvedgeneralization capabilities.
Finally, using SGQN as a building brick, among all those required tobridge the gap between simulation and real-world applications, is an exciting perspective.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e g., type of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly [Table]
[Table]
[N/A]
Look where you look! Saliency-guided Q-networks for
visual RL tasks
Appendix
A Network architecture and training hyperparameters
All the methods evaluated in Section 4 are variants of SAC.
We based our implementation of
SAC, RAD, SODA, and SVEA on the ones proposed by Nicklas Hansen, available at https://github.com/nicklashansen/dmcontrol-generalization-benchmark .
To ensure a faircomparison, we used the same network architecture and hyperparameters for all agents.
Networks architectures.
The shared encoder f θ is composed of a stack of 11 convolutional layers,each with 32 filters of 3 × 3 kernels, no padding, stride of 2 for the first one and stride of 1 for allothers, yielding a final feature map of dimension 32 × 21 × 21 (inputs have dimension 84 × 84 × 3 ).
The policy head π θ and the value function head Q θ consist of independent fully connected networks,each composed of a linear projection of dimension 100 with layer normalization, followed by 3 linearlayers with 1024 hidden units.
SGQN also uses a predictor head M θ responsible for reconstructingthe attribution mask
M ρ ( Q θ , s, a ) .
The M θ network follows the 100-unit layer of Q θ and has 6 layers.
It is composed of a first linear layer projecting the 100-dimensional embedding to 32 × 21 × 21features, then followed by two convolutional+upsampling blocks, each having 64 filters of 3 × 3kernels (padding of 1 to preserve the feature map size) and an upsampling factor of 2, and finallya last convolutional layer with 9 filters of 3 × 3 kernels (padding of 1).
Figure 8 illustrates SGQNagents’ architecture.
Hyperparameters.
SGQN introduces three additional hyperparameters to SAC: the quantile valueρ used to binarize the attribution masks, the data augmentation function τ used during the selfsupervised learning updates, and the frequency of these auxiliary updates N SL .
Table 3 summarizesthe hyperparameters used in all experiments.
[Figure]
B Reproducibility
All the experiments from Section 4 were run on a desktop machine (Intel i9, 10th generation processor,64GB RAM) with a single NVIDIA RTX 3090 GPU.
All scores were calculated on an average of5 repetitions.
Details about all experiments are reported in Table 4.
Besides this information, weprovide the full source code of our implementation and experiments, along with the data files of theexperimental results we obtained.
After the review phase, this code will be available at [anonymousurl].
Algorithm
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
C Additional experimental results on DMControl-GB
For the sake of simplicity and consistency, all experiments in Section 4 have been conducted withthe same set of hyperparameters (except for the quantile value ρ , recalled in Table 4).
A particularattention needs to be paid to the number of times an action is repeated in the environments of the
DMControl-GB, since it has an important influence on the scores reached by the agents.
Hansenet al.
[2021] indicate they run their experiments with an action repetition covering 4 time steps forall environments, except for Cartpole (8 time steps) and Finger spin (2 time steps).
To avoid suchheterogeneity across environments, we chose to report scores in Tables 1 and 2 with a constant actionrepetition parameter of 4 time steps for all environments.
We repeated the experiments of Section4 by setting the value of this parameter to 8 and 2 for the Cartpole and Finger spin environments
[Table]
respectively as suggested by Hansen et al.
Table 5 reports the results obtained.
This finetuning benefits almost all agents and, more particularly SGQN, which increases the gap with itscompetitors both in video easy and video hard domains.
We report all the corresponding trainingcurves and testing scores in Figures 9, 10, and 11.
We also report a comparison of the attribution mapsfor all agents on all environments in video hard environments in Figure 12, illustrating how SGQNconsistently relies on pixels that really belong to the system to control and discards confoundingfactors. Observation
D Vision-based robotic manipulation experiments
To demonstrate the genericity of our method, and following the recommendation of the reviewers, wealso consider two goal-reaching robotic manipulation tasks from the vision-based robotic manipulation simulator introduced in [Jangir et al., 2022]: Reach, a task where a robot has to reach for a goalmarked by a red disc placed on a table, and Peg in box, a task where a robot has to insert a peg tied toits arm into a box. We modified the original simulator to include three testing environments for bothtasks, similar to the training ones but with different colors and textures for the background and thetable as illustrated in Figure 13. Note that no fine-tuning of hyperparameters (learning rates, quantilethreshold, etc.) was performed whatsoever.
[Figure]
We trained SGQN agents for 250 000 steps with a quantile threshold ρ = 0 .
95 and compared theirgeneralization scores with those of agents trained with SAC [Haarnoja et al., 2018], SODA [Hansenand Wang, 2021] and SVEA [Hansen et al., 2021].
We used random convolutions [Lee et al., 2020] asimage augmentations for all the methods except for SAC, since it produces color and texture increasesthat better match our test environments than the structured distractions induced by a random overlay.
The results are reported in Table 6.
For the Reach task, the agents trained with SAC, SODA, and
SVEA fail to maintain their performance when evaluated on the test environments.
The agents trainedwith SGQN maintain performance almost identical to those obtained during training on two out ofthree testing environments and outperform the other agents on average by more than 124%.
Figure14 shows examples of the saliency maps obtained for the Reach training and testing environments.
In the Peg in box task, the generalization scores are degraded for all the agents.
Nevertheless, intwo out of the three testing environments (Test 2 and 3), the agents trained with SGQN seem tobe the least affected.
The third environment (Test 1) seems to feature textures and colors that areparticularly difficult for generalization and would require further investigation.
Overall, the SGQNagents’ generalization scores over the three environments are still 40% better on average than thoseof their competitors.
[Table]
[Figure]
In our experiments, the value of the mask threshold parameter ρ was selected after a quick visualsearch, as illustrated in Figure 15.
During our experiments on DMControl-GB, we set ρ to 0.95 onall environments but Cartpole and Ball in Cup (for which the ratio of foreground/background pixelsis smaller than on the others environments).
To assess SGQN’s sensitivity to ρ , we also performedan experiment on these two environments with ρ set to 0.95.
Results, reported in Table 7, show that agents trained with ρ = 0 .
95 perform slightly worse than the ones trained with ρ = 0 .
98 thusindicating that this parameter has an impact on performance and is worth fine-tuning.
However, it isimportant to note that on the most difficult ( video hard ) environments, the scores obtained by SGQNremain notably above those of other methods, whichever the value of ρ .
Visual search is a (loose)measure of the amount of information actually present in the image and necessary for predicting thevalue.
It is likely that one could exhibit worst case environments for which all pixels are necessary topredict the value.
In such environments, SGQN might perform poorly.
However, we argue that suchenvironments are not representative of most realistic visual RL tasks, either real-world or simulated,where pixel information is very redundant (which is a cause for overfitting).
It is likely that in other(maybe more difficult) environments (e.g. with more complex important objects moving across thescreen), the necessary threshold for ρ could be lowered. Benchmark Environment ρ = 0 .ρ = 0 .
Easy
Ball in cup
[Table] [Table] [Table] Hard
Ball
[Table]
cup
[Table]
[Table]
[Table]
[Figure]
F Comparison with SVEA + SODA
The losses of SGQN can in some way be related to those present in SODA
[Hansen and Wang, 2021]and SVEA [Hansen et al., 2021].
The auxiliary self-supervised learning objective of SGQN is similarto replacing the projector used by SODA (and originally introduced in BYOL [Grill et al., 2020])by
∂Q ( s, a ) /∂s .
Besides, during the critic update, SVEA performs a new estimate of Q from anaugmented state τ ( s ) .
The consistency loss used in SGQN can be expressed in a similar fashion byconsidering that the augmentation τ consists in the application of a mask resulting in a reductionof the information contained in s .
From this perspective, SGQN can be seen as the combinationof particular cases of SVEA and SODA.
Thus we compared the performance of SGQN with thecombination of SVEA and SODA.
Table 8 reports the corresponding results.
Except on the Walkerwalk video hard testing levels, the combination of SVEA and SODA diminishes the scores obtainedby using each method independently.
Agents trained with both SVEA and SODA also obtain worseperformance than the agents trained with SGQN, thus showing the advantages brought by each lossintroduced in SGQN.
[Table]
G Discussion on initial saliency maps
Since SGQN uses thresholded saliency maps to mask out input images from the very first steps ofthe algorithm, and since there is no reason for these initial maps to point towards important pixelsfor the true value function, it is legitimate to wonder how the initial saliency maps actually affectthe optimization path of SGQN.
We propose the following discussion which aims at explainingwhy the mechanism of SGQN is robust to the Q θ network initialization and to which pixels arebeing selected by the initial saliency maps.
Initial saliency maps are likely to appear random sincethey represent the gradients of random functions close to zero (as per the classical initialization ofneural networks).
Consequently, when thresholded, these saliency maps are likely to yield pixelsthat are uniformly spread across the image.
In turn, the masking operation initially acts as a randomsubsampling operation.
Since many close pixels in the input image hold redundant information,the application of M ρ to input images is likely to preserve enough information to correctly predictthe value function (minimize L Q ).
It is important to note that many other subsampling masks areequally likely to preserve the information necessary to correctly predict the value function.
So theinitialization of the Q θ network does not prevent learning the true Q -function based on s ⊙ M ρ .
Adirect consequence is that many Q θ functions, that differ only by which subsampled pixels they relyon, can actually fit the true Q -function, but few will generalize to unseen states, which is the veryissue of observational overfitting.
Even if the original maps indicating which pixels are informativeare wrong, i.e. if too many confounding pixels are retained in M ρ , then, L Q , the first part of thecritic’s loss, still encourages that Q θ be a solution to the Bellman equation.
In that sense, L C (thesecond part of the critic’s loss) acts as a regularizer: it allows discriminating between functions thatwould otherwise be equivalent approximate solutions to the Bellman equation.
The self-supervisedlearning procedure, in turn, yields features f θ that are useful at predicting one’s own saliency map.
Asnoted by Grill et al.
[2020] or Hansen and Wang [2021], such features are good representations of theinput state.
Finally, as the Q -function becomes more accurate and as the f θ features become better atpredicting M ρ , the saliency maps become sharper.
Overall, the interplay between the two phases of
SGQN is mostly a virtuous circle, that starts from the observation that the initial M ρ preserve enoughinformation in s to permit correct learning of Q θ .