Understanding and Preventing Capacity Loss in Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning has achieved remarkable successes in a variety of environments ( Mnih et al. , 2015 ; Moravčı́k et al. , 2017 ; Silver et al. , 2017 ; Abreu et al. , 2019 ) , but the precise reasons for its successes largely remain mysterious . Existing algorithms are highly sensitive to hyperparameters and seemingly innocuous design choices , to the extent that even minor variations to state-of-the-art methods can fail to make learning progress on tasks originally solved with ease . These instabilities are particularly pronounced in sparse-reward environments , where even different random seeds of the same algorithm can attain dramatically different performance outcomes . In these settings , even if the agent has experienced the rewards necessary to learn a high-scoring policy , it will often fail to translate those rewards into successful policy improvements . This presents a stark contrast to supervised learning , where existing approaches are reasonably robust to small hyperparameter changes , random seeds , and GPU parallelisation libraries . Moreover , naive applications of supervised learning methods such as momentum-based optimization or data augmentation to the RL problem often require additional modification to yield performance improvements ( Bengio et al. , 2020 ; Raileanu et al. , 2020 ) . We hypothesize that the non-stationary prediction problems agents face in RL may be a driving force in the challenges described above . RL agents must solve a sequence of similar prediction problems as they iteratively improve their prediction accuracy and their policy ( Dabney et al. , 2021 ) , and solving each subproblem in this sequence is necessary to progress to the next subproblem . As shown in work on active learning , warm-starting a network by fitting similar data to that used in the downstream task can hurt the network ’ s final accuracy , even if this data is drawn from the same distribution ( Ash & Adams , 2020 ) . This suggests that the slowly-shifting input and target distributions RL agents face may be particularly ill-suited to function approximation by deep neural networks . Indeed , several prior works studying the effect of re-initializing network parameters in reinforcement learning have found this enables agents to break through plateaus and improve generalization performance ( Igl et al. , 2021 ; Fedus et al. , 2020 ) . The principal thesis of this paper is that over the course of training , deep RL agents lose some of their capacity to quickly fit new prediction tasks , and in extreme cases this capacity loss prevents the agent entirely from making learning progress . We present a rigorous empirical analysis of this phenomenon which considers both the ability of networks to learn new target functions via gradientbased optimization methods , and their ability to linearly disentangle states ’ feature representations . We confirm that agents ’ ability to fit new target functions declines over the course of training in several environments from the Atari suite ( Bellemare et al. , 2013 ) and non-stationary reward prediction tasks . We further find that the ability of representations to linearly distinguish different states , a proxy for their ability to represent certain functions , quickly diminishes in sparse-reward games , leading to representation collapse , where the feature outputs for every state in the environment inhabit a low-dimensional – or possibly even zero – subspace . Finally , we show evidence that representation collapse is a key factor in agents ’ failure to make learning progress in sparse-reward environments . We then propose a simple regularization technique , Initial Feature Regularization ( InFeR ) , to prevent representation collapse : regress a set of feature projection heads to their values at initialization . This method improves performance on a number of sparse-reward environments and also increases the measures of capacity we are interested in , preventing representation collapse in sparse-reward environments and improving target-fitting capacity in dense-reward environments . One striking take-away from our results is that agents trained on so-called ‘ hard exploration ’ games such as Montezuma ’ s Revenge can attain significant improvements over existing competitive baselines without using smart exploration algorithms . This suggests that the poor performance of deep RL agents in sparse-reward environments is not solely due to inadequate exploration , but rather also in part due to poor representation learning . Essentially , agents which are ‘ too good ’ at predicting the zero value function lose their ability to fit the non-zero targets necessary for policy improvement . We believe this has significant implications for how the community views the interplay of exploration and representation learning in sparse reward environments . 2 BACKGROUND . We consider the reinforcement learning problem wherein an agent seeks to maximize expected return in an MDP M = ( X , A , R , P , ) , where X denotes the state space , A the action space , R the reward function , P the transition probability function , and the discount factor . We will be primarily interested in value-based RL , where the objective is to learn the value function Q⇡ : X ⇥ A ! R associated with some ( possibly stochastic ) policy ⇡ : X ! P ( A ) , defined as Q⇡ ( x , a ) = E⇡ , P [ P1 k=0 kR ( xk , ak ) |x0 = x , a0 = a ] . In particular , we are interested in learning the value function associated with the optimal policy ⇡⇤ which maximizes the expected discounted sum of rewards from any state . In Q-Learning ( Watkins & Dayan , 1992 ) , the agent performs updates to minimize the distance between a predicted action-value function Q and the bootstrap target defined as T Q ( x , a ) = E [ R ( x0 , a0 ) + max a02A Q ( x1 , a 0 ) |x0 = x , a0 = a ] . ( 1 ) In most practical settings , updates are performed with respect to sampled transitions rather than on the entire state space . The target can be computed for a sampled transition ( xt , at , rt , xt+1 ) as T̂ Q ( xt , at ) = rt + maxa Q ( xt+1 , a ) . When a deep neural network is used as a function approximator ( the deep RL setting ) , Q is defined to be the output of a neural network with parameters ✓ , and updates are performed by gradient descent on sampled transitions ⌧ = ( xt , at , rt , xt+1 ) . A number of tricks are often used to improve stability : the sample-based objective is minimized following stochastic gradient descent based on minibatches sampled from a replay buffer of stored transitions , and a separate set of parameters ✓̄ is used to compute the targets Q✓̄ ( xt+1 , at+1 ) which is typically updated more slowly than the network ’ s online parameters . This yields the following loss function ( where ⌧ denotes the sampled transition ) : ` TD ( Q✓ , ⌧ ) = ( Rt+1 + max a0 Q✓̄ ( Xt+1 , a 0 ) Q✓ ( Xt , At ) ) 2 . ( 2 ) In this work we will be interested in how common variations on this basic learning objective shape agents ’ learning dynamics , in particular the dynamics of the learned representation , or features . We will refer to the outputs of the final hidden layer of the network ( i.e . the penultimate layer ) as its features , denoted ✓ ( x ) . Our choice of the penultimate layer is motivated by prior literature studying representations in RL ( Ghosh & Bellemare , 2020 ; Kumar et al. , 2021 ) , although many works studying representation learning consider the outputs of earlier layers as well . In general , the features of a neural network are defined to be the outputs of whatever layer is used to compute additional representation learning objectives . 3 CAPACITY LOSS . In this section we demonstrate conditions under which neural networks progressively lose their ability to quickly fit new targets when trained on sequences of prediction tasks including , but not limited to , those found in value-based RL . We find that the loss of this notion of capacity loss , which we refer to as target-fitting capacity , is particularly pronounced in sparse prediction tasks , where many of the target values the agent seeks to predict are zero . To study the effect of extreme capacity loss on performance in greater depth , we present a special case of the target-fitting capacity measure which is efficient to compute and has the intuitive interpretation of measuring the ability of the representation to linearly disentangle states . We find evidence that agents which have greater capacity according to this metric tend to achieve better performance in challenging environments from the Atari suite where agents fail to match human performance , and that those suffering from representation collapse according to this metric fail to make any learning progress at all . 3.1 TARGET-FITTING CAPACITY . Neural networks trained with temporal difference learning objectives must solve a sequence of target prediction problems . Ideally , the network should be able to fit later targets as quickly and accurately as it is able to fit early targets . However , prior work suggests that using ‘ warm starts ’ , i.e . weights from a network trained on a different dataset or target function , can slow down training on new data points ( Ash & Adams , 2020 ) . We are therefore interested in studying whether this notion of interference may affect the ability of reinforcement learning agents to make learning progress . We begin by formally defining a notion of capacity loss which measures whether the network can quickly fit new targets , for example after a policy improvement step increases the expected reward obtained in a given state or after the target network is updated . Definition 1 ( Target-fitting capacity ) . Let PX 2 P ( X ) be some distribution over inputs X and PF a distribution over a family of functions F with domain X . Let N = ( g✓ , ✓0 ) represent the pairing of a neural network architecture with some initial parameters ✓0 , and O correspond to an optimization algorithm for supervised learning . We measure the capacity of N under the optimizer O to fit the data-generating distribution D = ( PX , PF ) as follows : C ( N , O , D ) = Ef⇠PF [ Ex⇠PX [ ( g✓0 ( x ) f ( x ) ) 2 ] ] where ✓0 = O ( ✓0 , PX , f ) . ( 3 ) This definition of capacity takes into account the interaction between the optimization algorithm and the initial parameters , and can be adapted to a broad range of function classes such as the network ’ s TD targets on the current dataset , the reward function on the environment , or randomlygenerated functions . The choice of function class is crucial to the notion of capacity measured . Using bootstrap targets from the network ’ s current parameters can identify networks which can not fit their current target function , but may reveal more about the complexity of the network output than the network ’ s capacity to fit novel targets . A network which can only output the zero function , for example , will attain very low Bellman error in a sparse-reward environment , but may not be able to fit more interesting value functions . In our evaluations , we will therefore consider a random class of functions which is independent of the current network parameters ; the effect of this choice is discussed further in Appendix B . To evaluate target-fitting capacity in deep RL agents , we sample target functions by randomly initializing neural networks with new parameters , and use the outputs of these networks as targets for regression . We then load initial parameters from an agent checkpoint at some time t , and regress on random network outputs over a fixed set of inputs sampled from the earliest available checkpoint ’ s replay buffer . This ensures that the same regression problem is being evaluated at every checkpoint . We then evaluate the mean squared error after training for fifty thousand steps . We observe that as training progresses agents ’ checkpoints on average get modestly worse at fitting these random targets in most environments ; due to space limitations we only show two representative environments where this phenomenon occurs in Figure 1 , and defer the full evaluation to Appendix C. While the resulting curves are somewhat noisy due to the randomness in the sampled target function , we do see consistent results across both the double DQN agent and the agent trained with an additional auxiliary loss . We now turn our attention to a less computationally expensive experimental setting where it is easier to obtain statistically robust results . We are particularly interested in understanding why capacity loss occurs . Two possible causes are immediate : the effect of bootstrapping , and the effect of sequential training . The effect of bootstrapping on capacity has been studied in other contexts ( Mobahi et al. , 2020 ; Kumar et al. , 2021 ) . We aim to isolate the effect of sequential prediction tasks on capacity loss . To minimize the potential for confounding factors to influence our results , we construct our toy iterative prediction problems on the MNIST data set , which consists of images of handwritten digits and corresponding labels , and manually construct a sequence of targets which the network must fit over the course of training . We first consider labels computed by a randomly initialized neural network f✓ : we transform input-label pairs ( x , y ) from the canonical MNIST dataset to ( x , f✓ ( x ) ) , where f✓ ( x ) is the network output . To generate a new task , we simply reinitialize the network ; our evaluations consist of 10 iterations of label-generation followed by a training period during which we run a gradient-based optimizer on the network starting from the parameters we obtained in the previous iteration . We further consider a ‘ sparse-reward ’ version of MNIST : for each of 10 iterations i , we use the label ŷi = [ y < i ] , where y is the true label of the image . For example , at the first iteration , all images are assigned label zero . At the second iteration , the images of the digit zero are assigned label one , while all other inputs retain the zero label . This continues until all inputs are assigned label one at the final iteration . We follow the same training procedure in both cases : optimizing the network for a fixed number of steps on one set of labels , then generating new labels and running the optimization algorithm again from the parameters obtained in the previous phase . In Figure 2 we see that most networks trained in these two experiments exhibit decreasing ability to fit later target functions under a fixed optimization budget . This effect is strongest in small networks with ReLU activations , suggesting that this capacity loss may be driven by saturated units and that this phenomenon will be easiest to detect in settings where the network architecture is not highly over-parameterized relative to the prediction task . The sparse reward setting is particularly intriguing : we do not expect to see a monotone increase in error as the later label functions correspond to ‘ easier ’ learning problems ( i.e . predicting the majority class will already yield reasonably low prediction error ) , but we do see that for equal difficulty , the network obtains greater error on the later target set than the earlier one , and this effect is significantly more pronounced than in the random labels tasks . This suggests that sparse reward signals can be particularly damaging to the ability of networks to fit new target functions .

Summary: Summary: - Over course of training, deep RL agents experience "capacity loss", where networks are unable to quickly fit new functions. This problem is further exacerbated by non-stationary predictions (such as bootstrapping) over the course of training. - To prove this problem, the author runs two experiments (Atari and MNIST), and show that training error against a randomly-initialized network increases over time. They also define effective dimension (rank of feature space) and show that effective dimension is tightly correlated with performance. - The authors introduce INFER, which regularizes Q-value loss with error w.r. to randomly initialized network, and evaluate this on Atatri-57, showing most benefits in sparse reward environments (Montezuma's revenge)