Towards Deeper Deep Reinforcement Learning with Spectral Normalization

1 Introduction . In computer vision and natural language processing ( NLP ) , competitive models are growing increasingly large , and researchers now train billion-parameter models [ 13 , 36 ] . The earliest neural networks were often shallow [ 40 ] with performance dropping for excessively deep models [ 29 ] . However , ever since the introduction of batch normalization [ 31 ] and residual connections [ 29 ] , performance has improved more or less monotonically with model scale [ 64 ] . As a result , competitive models in computer vision and NLP are growing ever larger [ 10 , 56 ] , and further architectural innovations are continuously researched [ 14 ] . In stark contrast with this trend , state-of-the-art ( SOTA ) reinforcement learning ( RL ) agents often rely on small feedforward networks [ 35 , 37 , 39 ] and performance gains typically originate from algorithmic innovations such as novel loss functions [ 58 , 62 , 68 ] rather than increasing model capacity . Indeed , a recent large-scale study has shown that large networks can harm performance in RL [ 2 ] . It is natural to suspect that high-capacity models might overfit in the low-sample regime common in RL evaluation . Imagenet contains over a million unique images whereas RL is often evaluated in contexts with fewer environment samples [ 37 , 39 ] . However , to date , this overfitting hypothesis remains largely untested . To address this , we study the effects of using larger modern architectures in RL . By modern architectures we mean networks with high capacity , facilitated by normalization layers [ 6 , 31 ] and skip connections [ 29 ] . We thus depart from the trend in RL of treating networks as black-box function ∗Correspondence to : Johan Bjorck < njb225 @ cornell.edu > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . approximators . To limit the scope and compute requirements we focus on continuous control from pixels [ 65 ] with two actor-critic agents based on the Soft Actor-Critic ( SAC ) [ 25 , 35 ] and Deep Deterministic Policy Gradients ( DDPG ) [ 43 , 71 ] algorithms . Actor-critic methods form the basis of many SOTA algorithms for continuous control [ 35 , 38 , 39 ] . As can be expected , we demonstrate that naïvely adopting modern architectures leads to poor performance , supporting the idea that RL does not benefit from larger models . However , we show that the issue is not necessarily overfitting , but instead , that training becomes unstable with deeper modern networks . We hypothesize that taking the gradient of the actor through the critic network creates exploding gradients [ 53 ] for deeper networks . We connect this setup with generative adversarial networks ( GANs ) [ 22 ] , and propose to use a simple smoothing technique from the GAN literature to stabilize training : spectral normalization [ 47 ] . We demonstrate that this simple strategy allows the training of larger modern networks in RL without instability . With these fixes , we can improve upon state-of-the-art RL agents on competitive continuous control benchmarks , demonstrating that improvements in network architecture can dramatically affect performance in RL . We also provide performance experiments showing that such scaling can be relatively cheap in terms of memory and compute time in RL from pixels . Our work suggests that model scaling is complementary to algorithmic innovations and that this simple strategy should not be overlooked . We summarize our contributions as follows : • We verify empirically that large modern networks fail for two competitive actor-critic agents . We demonstrate dramatic instabilities during training , which casts doubt on overfitting being responsible . • We argue that taking the gradients through the critic is the cause of this instability . To combat this problem we propose to adopt spectral normalization [ 47 ] from the GAN literature . • We demonstrate that this simple smoothing method enables the use of large modern networks and leads to significant improvements for SOTA methods on hard continuous control tasks . We further provide evidence that this strategy is computationally cheap in RL from pixels . 2 Background . 2.1 Reinforcement Learning . Reinforcement learning tasks are often formulated as Markov decision processes ( MDPs ) , which can be defined by a tuple ( S , A , P , r ) [ 63 ] . For continuous control tasks the action space A and the state space S are continuous and can also be bounded . Each dimension of the action space might for example correspond to one joint of a humanoid walker . At time step t , the agent is in a state st∈S and takes an action a∈A to arrive at a new state st+1 ∈ S. The transition between states given an action is random with transition probability P : S × S ×A → [ 0 , ∞ ) . The agent receives a reward rt from the reward distribution r at each timestep t. The goal is typically to find a policy π : S → A that maximizes expected discounted reward E [ ∑ t γ trt ] , where 0 < γ < 1 is a pre-determined constant . In practice , it is common to measure performance by the cumulative rewards ∑ t rt . 2.2 Actor-Critic Agents . Actor-critic methods have roots in Q-learning [ 69 ] and form the backbone of many recent stateof-the-art algorithms [ 35 , 38 , 39 ] . We focus on two popular and performant actor critic methods : SAC [ 25 ] and DDPG [ 43 ] . Given some features φ ( s ) that can be obtained from any state s ∈ S , the actor-network maps each state-feature φ ( s ) to a distribution πθ ( s ) over actions . The action distribution is known as the policy . For each dimension , the actor-network outputs an independent normal distribution , where the mean µθ ( s ) and possibly the standard deviation σθ ( s ) come from the actor-network . The action distribution is then obtained as πθ ( s ) = tanh ( µθ ( s ) + σθ ( s ) ) , ∼ N ( 0,1 ) . ( 1 ) The tanh non-linearity is applied elementwise , which bounds the action space to [ −1 , 1 ] n for n = dim ( A ) . Sampling allows us to sample actions aθ from the policy πθ ( s ) . DDPG typically adds the exploration noise outside the tanh non-linearity . To promote exploration , SAC instead adds the entropy of the policy distribution H ( πθ ( st ) ) times a parameter α to the rewards . The critic network outputs a Q-value Qψ ( a , s ) for each state s and action a . In practice , one can simply concatenate the feature φ ( s ) with the action a and feed the result to the critic . The critic network is trained to minimize the soft Bellman residual : min ψ E ( Qψ ( a , st ) − [ rt + γE [ Q̂ ( at+1 , st+1 ) + αH ( π ) ] ] ) 2 ( 2 ) Here , rt is the obtained reward and E [ Q̂ ( at+1 , st+1 ) ] is the Q-value estimated by the target critic – a network whose weights are the exponentially averaged weights of the critic . One can also use multi-step targets for the rewards [ 62 ] . The loss ( 2 ) is computed by sampling transitions from a replay buffer [ 48 ] . Since the Q-values that the critic outputs measure expected discounted rewards , the actor should simply take actions that maximize the output of the critic network . Thus , to obtain a gradient step for the actor , both SAC and DDPG sample an action aθ from πθ ( st ) and then take derivatives of Qψ ( aθ , s ) with respect to θ . That is ∇θ = ∂Qψ ( aθ , s ) ∂θ ( 3 ) As the critic is a differentiable neural network , the derivatives are taken through the Q-values Qψ ( aθ , s ) . In practice , the features φ ( s ) can be obtained from a convolutional network trained by backpropagating from ( 2 ) . For further details , see Haarnoja et al . [ 25 ] , Lillicrap et al . [ 43 ] . 3 Motivating Experiments . 3.1 Experimental Setup . For experimental evaluation , we focus on continuous control from pixels , a popular setting relevant to real-world applications [ 33 , 34 ] . Specifically we evaluate on the DeepMind control suite [ 65 ] , which has been used in [ 26 , 27 , 35 , 39 ] . We use the 15 tasks considered in Kostrikov et al . [ 35 ] and evaluate after 500,000 samples , which has become a common benchmark [ 35 , 39 , 41 ] . If learning crashes before 500,000 steps ( this occasionally happens for modern networks ) , we report the performance before the crash . We will primarily focus on the image augmentation based SAC agent of Kostrikov et al . [ 35 ] , known as DRQ , and adopt their hyperparameters ( listed in Appendix B ) . This agent reaches state-of-the-art performance without any nonstandard bells-and-whistles , and has an excellent open-source codebase . Unless specifically mentioned , all figures and tables refer to this agent . Its network consists of a common convolutional encoder that processes the image into a fixed-length feature vector . Such feature vectors are then processed by the critic and actor-network , which are just feedforward networks . We will refer to these feedforward networks as the heads . The convolutional encoder consists of four convolutional layers followed by a linear layer and layer normalization [ 6 ] and the heads consist of MLPs with two hidden layers that concatenate the actions and image features ; see Appendix C for details . In Table 1 we also consider the augmentation-based DDPG agent of Yarats et al . [ 71 ] , using the default hyperparameters proposed therein . Both agents use the same network design . In all experiments , we only vary the head architecture . We evaluation over ten seeds . 3.2 Testing Modern networks . Two crucial architectural innovations which have enabled deep supervised models are residual connections [ 29 ] and normalization [ 31 ] . We test if adopting such modern techniques allows us to successfully scale networks . Specifically , we use the feedforward network ( and not the attention module ) found in a Transformer block [ 67 ] and will refer to this architecture as modern . Each block consists of two linear transformations w1 , w2 with a ReLu between and layer normalization [ 6 ] that is added to a residual branch . I.e . the output of a layer is x+w2Relu ( w1norm ( x ) ) . Compared to the original architecture , normalization is applied before the feedforward blocks instead of after , which is now favored in practice [ 50 , 70 ] . We use four transformer blocks for the critic and two for the actor . Results from these experiments are shown in Figure 1 , with one standard error given . We see that these modifications often decrease performance . Thus , actor-critic agents have seemingly not reached the level where model capacity improves performance monotonically , and tricks used in supervised learning do not seem to suffice . This is perhaps not very surprising as state-of-the-art RL agents typically use simple feedforward networks .

Summary: The paper hypothesizes that the poor performance of modern NN architectures with the SAC algorithm stems from large gradient updates caused by propagating gradients for actor through critic. The authors mitigate the issue by borrowing a spectral normalization technique from GANs, which was designed for a similar problem. Experimentally, they verify that the approach improves the performance and allows usage of modern architectures, with gains comparable to algorithmic advances of RL algorithms themselves.