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 image augmentations for bothoff-policy and on-policy algorithms by augmenting image observations processed by the agent duringtraining.
Yarats et al.
[2020] average both the value function and its target over multiple imagetransformations.
Raileanu et al.
[2021] combine the previous method with UCB [Auer, 2002] to pickthe most promising augmentation, and apply it to PPO [Schulman et al., 2017].
Fan et al.
[2021] useweak data augmentation to train an expert without hindering its performance and distill its policyto a student trained with substantial data augmentation.
Besides augmentations of raw inputs, otheraugmentations operate directly within the agents’ network.
Lee et al.
[2020] introduces a randomconvolutions layer at the earliest level of the agent’s network to modify the texture of the visualobservations.
Zhou et al.
[2020] adapt mixup [Zhang et al., 2018d] with style statistics encoded inearly instance normalization layers to increase data diversity.
Bertoin and Rachelson [2022] applychannel-consistent local permutations of the feature map to induce robustness to spatial spuriouscorrelations.
Finally, data augmentation can also be used in an auxiliary loss to promote invariance todistributional shift in representations.
Hansen and Wang [2021] propose a soft-data augmentationmethod (SODA) by adding an auxiliary self-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.
Wang et al.
[2021] extract, using visualattention, the observation foreground to provide background invariant inputs to the policy learner.
Raileanu and Fergus [2021] separate the actor from the critic in the agent’s network architectureand add an adversarial auxiliary objective on the actor’s representations to remove the informationneeded 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 distances in statespace.
Other recent works use behavioral similarities combined with contrastive learning [Agarwalet al., 2020] or clustering [Mazoure et al., 2022] to map behaviorally similar observations to similarrepresentations.
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 incorporateattributions 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 this mask 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.
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
[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 approachwith current state-of-the-art methods for generalization in continuous actions RL (RAD [Laskinet al., 2020b], SODA
[Hansen and Wang, 2021], SVEA [Hansen et al., 2021]) on five environmentsfrom the DMControl Generalization Benchmark (DMControl-GB)
[Hansen and Wang, 2021].
The DMControl-GB presents a variety of vision-based continuous control tasks based on the Deepmindcontrol suite.
Agents are trained in a fixed background environment and evaluated under two challenging distribution shifts, consisting in replacing the training background with natural videos.
Figure2 illustrates the effects of both domain shifts.
All the compared methods herein are variants of SAC,for which we use the same architecture for all agents.
These methods all use data augmentation in oneof their stages.
Following the experimental protocol of the competing approaches, we used a randomoverlay augmentation [Hansen and Wang, 2021] with the Places365 dataset [Zhou et al., 2017] forall the methods except for RAD, for which we used a random crop, as it is reported as producingthe best results [Laskin et al., 2020b].
We trained all agents for 500 000 steps using the vanillatraining environment with no visual variation.
Appendix A and B summarize all the hyperparameters,network architectures, and implementation choices used for this benchmark.
Appendix C includeextra experimental results.
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.
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 the [Table] ground and the shadows (Figure 2).
We report the average sum of rewards after 500 000 trainingsteps for the video easy benchmark in the top part of Table 1.
Agents trained with SGQN outperformagents trained with other state-of-the-art methods on all tasks but one (where it is on par with otheragents), thus demonstrating the generalization capabilities induced by the method.
By removing theground and shadows, the video hard benchmark causes a larger, and more confusing distributionalshift and is, therefore, more challenging. The bottom part of Table 1 reports the performance of allmethods on this benchmark.
All the competitors of SGQN experience a radical decrease in theirgeneralization performance.
SGQN is significantly less impacted and outperforms all its competitorswith an average margin over the second-best of 57% on all environments, and a gain range of 14 to144% across environments.
Figure 4 reports the evolution of each agent’s score on the video hardenvironments, along training.
If SODA and SVEA are performing well when the ground and shadowinformation is still present, removing it leads to a drastic drop in their performance.
SGQN is lesssensitive to this change, notably through the consistency loss, which encourages agents to makedecisions based on the subsets of pixels they deem most interesting.
Overall, the interplay of the twophases of SGQN seems to be key to a major leap forward in terms of generalization gap in difficultenvironments, 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.
Weperform an ablation study to assess their individual contribution and Table 2 reports the average sumof rewards in the training domain and for zero-shot generalization in all environments.
Individually,each of these features greatly improves both training and zero-shot generalization performance on allenvironments.
The auxiliary self-supervised learning phase provides the most significant performancegains over vanilla SAC.
The average training performance on all environments of agents trained withthis auxiliary objective is more than 87% higher than that of agents trained with vanilla SAC.
Thesame applies to the performance in zero-shot generalization, which increases by more than 141% onvideo easy and by more than 182% on video hard environments.
One can note that agents trainedwith our self-supervised learning objective also outperform SODA agents (Table 1) on 4 out of 5video easy environments and on 3 out of 5 video hard environments.
Recall that SODA relies onthe BYOL [Grill et al., 2020] self-supervised feature learning procedure, whose target labels differnotably from those proposed herein.
The reported performance, compared to that of SODA suggeststhat attribution maps constitute a good labeling function that could be considered in the more generalcontext of self-supervised learning.
To a slightly lesser extent, adding the consistency loss to SACalso significantly improves its performance.
The average score obtained improves by more than 49%on the training domain and by 86% and 166% respectively on the video easy and video hard domains.
Similarly to the results obtained by Ismail et al.
[2021] in supervised learning, the regularizationof the critic’s loss with the consistency term sharpens the attribution maps obtained (Figure 7).
In SGQN, these fine-grained attributions provide higher quality labels to the auxiliary self-supervisedlearning phase, thus yielding significant performance improvements in training ( +94% on average,range up to +242%), video easy generalization (+183% on average, range up to +399%), and videohard 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 it onthe DMControl generalization benchmark.
This agent displays a dramatically more efficient behaviorthan vanilla SAC on the various environments it is trained on.
Most importantly, the policy it learnscloses the generalization gap on environments that include confusing and distracting visual features,setting a new reference in terms of generalization performance.
Since they rely on self-awareness ofimportant pixels, SGQN agents are also very interpretable, in the sense that they provide a predictionof 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.