Implicit Under-Parameterization Inhibits Data-Efficient Deep Reinforcement Learning

1 INTRODUCTION . Many pervasive deep reinforcement learning ( RL ) algorithms estimate value functions using bootstrapping , that is , by sequentially fitting value functions to target value estimates generated from the value function learned in the previous iteration . Despite high-profile achievements ( Silver et al. , 2017 ) , these algorithms are highly unreliable due to poorly understood optimization issues . Although a number of hypotheses have been proposed to explain these issues ( Achiam et al. , 2019 ; Bengio et al. , 2020 ; Fu et al. , 2019 ; Igl et al. , 2020 ; Liu et al. , 2018 ; Kumar et al. , 2020a ) , a complete understanding remains elusive . We identify an “ implicit under-parameterization ” phenomenon that emerges when value networks are trained using gradient descent combined with bootstrapping . This phenomenon manifests as an excessive aliasing of features learned by the value network across states , which is exacerbated with more gradient updates . While the supervised deep learning literature suggests that some feature aliasing is desirable for generalization ( e.g. , Gunasekar et al. , 2017 ; Arora et al. , 2019 ) , implicit under-parameterization exhibits more pronounced aliasing than in supervised learning . This over-aliasing causes an otherwise expressive value network to implicitly behave as an under-parameterized network , often resulting in poor performance . Implicit under-parameterization becomes aggravated when the rate of data re-use is increased , restricting the sample efficiency of deep RL methods . In online RL , increasing the number of gradient steps in between data collection steps for data-efficient RL ( Fu et al. , 2019 ; Fedus et al. , 2020b ) causes the problem to emerge more frequently . In the extreme case when no additional data is collected , referred to as offline RL ( Lange et al. , 2012 ; Agarwal et al. , 2020 ; Levine et al. , 2020 ) , implicit under-parameterization manifests consistently , limiting the viability of offline methods . We demonstrate the existence of implicit under-parameterization in common value-based deep RL methods , including Q-learning ( Mnih et al. , 2015 ; Hessel et al. , 2018 ) and actor-critic ( Haarnoja et al. , 2018 ) , as well as neural fitted-Q iteration ( Riedmiller , 2005 ; Ernst et al. , 2005 ) . To isolate the issue , we study the effective rank of the features in the penultimate layer of the value network ( Section 3 ) . We observe that after an initial learning period , the rank of the learned features drops steeply . As the rank decreases , the ability of the features to fit subsequent target values and the optimal value function generally deteriorates and results in a sharp decrease in performance ( Section 3.1 ) . ∗Equal Contribution . Correspondence to Aviral Kumar < aviralk @ berkeley.edu > and Rishabh Agarwal < rishabhagarwal @ google.com > . To better understand the emergence of implicit under-parameterization , we formally study the dynamics of Q-learning under two distinct models of neural net behavior ( Section 4 ) : kernel regression ( Jacot et al. , 2018 ; Mobahi et al. , 2020 ) and deep linear networks ( Arora et al. , 2018 ) . We corroborate the existence of this phenomenon in both models , and show that implicit underparameterization stems from a pathological interaction between bootstrapping and the implicit regularization of gradient descent . Since value networks are trained to regress towards targets generated by a previous version of the same model , this leads to a sequence of value networks of potentially decreasing expressivity , which can result in degenerate behavior and a drop in performance . The main contribution of this work is the identification of implicit under-parameterization in deep RL methods that use bootstrapping . Empirically , we demonstrate a collapse in the rank of the learned features during training , and show it typically corresponds to a drop in performance in the Atari ( Bellemare et al. , 2013 ) and continuous control Gym ( Brockman et al. , 2016 ) benchmarks in both the offline and data-efficient online RL settings . We verify the emergence of this phenomenon theoretically and characterize settings where implicit under-parameterization can emerge . We then show that mitigating this phenomenon via a simple penalty on the singular values of the learned features improves performance of value-based RL methods in the offline setting on Atari . 2 PRELIMINARIES . The goal in RL is to maximize long-term discounted reward in a Markov decision process ( MDP ) , defined as a tuple ( S , A , R , P , γ ) ( Puterman , 1994 ) , with state space S , action space A , a reward function R ( s , a ) , transition dynamics P ( s′|s , a ) and a discount factor γ ∈ [ 0 , 1 ) . The Q-function Qπ ( s , a ) for a policy π ( a|s ) , is the expected long-term discounted reward obtained by executing action a at state s and following π ( a|s ) thereafter , Qπ ( s , a ) : = E [ ∑∞ t=0 γ tR ( st , at ) ] . Qπ ( s , a ) is the fixed point of the Bellman operator T π , ∀s , a : T πQ ( s , a ) : = R ( s , a ) + γEs′∼P ( ·|s , a ) , a′∼π ( ·|s′ ) [ Q ( s′ , a′ ) ] , which can be written in vector form as : Qπ = R + γPπQπ . The optimal Q-function , Q∗ ( s , a ) , is the fixed point of the Bellman optimality operator T : T Q ( s , a ) : = R ( s , a ) + γEs′∼P ( ·|s , a ) [ maxa′ Q ( s′ , a′ ) ] . Practical Q-learning methods ( e.g. , Mnih et al. , 2015 ; Hessel et al. , 2018 ; Haarnoja et al. , 2018 ) convert the Bellman equation into an bootstrapping-based objective for training a Q-network , Qθ , via gradient descent . This objective , known as mean-squared temporal difference ( TD ) error , is given by : L ( θ ) = ∑ s , a ( R ( s , a ) + γQ̄θ ( s ′ , a′ ) −Q ( s , a ) ) 2 , where Q̄θ is a delayed copy of the Q-function , typically referred to as the target network . These methods train Q-networks via gradient descent and slowly update the target network via Polyak averaging on its parameters . We refer the output of the penultimate layer of the deep Q-network as the learned feature matrix Φ , such that Q ( s , a ) = wTΦ ( s , a ) , where w ∈ Rd and Φ ∈ R|S||A|×d . Algorithm 1 Fitted Q-Iteration ( FQI ) 1 : Initialize Q-network Qθ , buffer µ . 2 : for fitting iteration k in { 1 , . . . , N } do 3 : Compute Qθ ( s , a ) and target values yk ( s , a ) = r + γmaxa′ Qk−1 ( s ′ , a′ ) on { ( s , a ) } ∼ µ for training 4 : Minimize TD error for Qθ via t = 1 , · · · , T gradient descent updates , minθ ( Qθ ( s , a ) − yk ) 2 5 : end for For simplicity of analysis , we abstract deep Q-learning methods into a generic fitted Q-iteration ( FQI ) framework ( Ernst et al. , 2005 ) . We refer to FQI with neural nets as neural FQI ( Riedmiller , 2005 ) . In the k-th fitting iteration , FQI trains the Q-function , Qk , to match the target values , yk = R+γPπQk−1 generated using previous Q-function , Qk−1 ( Algorithm 1 ) . Practical methods can be instantiated as variants of FQI , with different target update styles , different optimizers , etc . 3 IMPLICIT UNDER-PARAMETERIZATION IN DEEP Q-LEARNING . In this section , we empirically demonstrate the existence of implicit under-parameterization in deep RL methods that use bootstrapping . We characterize implicit under-parameterization in terms of the effective rank ( Yang et al. , 2019 ) of the features learned by a Q-network . The effective rank of the feature matrix Φ , for a threshold δ ( we choose δ = 0.01 ) , denoted as srankδ ( Φ ) , is given by srankδ ( Φ ) = min { k : ∑k i=1 σi ( Φ ) ∑d i=1 σi ( Φ ) ≥ 1− δ } , where { σi ( Φ ) } are the singular values of Φ in decreasing order , i.e. , σ1 ≥ · · · ≥ σd ≥ 0 . Intuitively , srankδ ( Φ ) represents the number of “ effective ” unique components of the feature matrix Φ that form the basis for linearly approximating the Qvalues . When the network maps different states to orthogonal feature vectors , then srankδ ( Φ ) has high values close to d. When the network “ aliases ” state-action pairs by mapping them to a smaller subspace , Φ has only a few active singular directions , and srankδ ( Φ ) takes on a small value . Definition 1 . Implicit under-parameterization refers to a reduction in the effective rank of the features , srankδ ( Φ ) , that occurs implicitly as a by-product of learning deep neural network Q-functions . While rank decrease also occurs in supervised learning , it is usually beneficial for obtaining generalizable solutions ( Gunasekar et al. , 2017 ; Arora et al. , 2019 ) . However , we will show that in deep Q-learning , an interaction between bootstrapping and gradient descent can lead to more aggressive rank reduction ( or rank collapse ) , which can hurt performance . Experimental setup . To study implicit under-parameterization empirically , we compute srankδ ( Φ ) on a minibatch of state-action pairs sampled i.i.d . from the training data ( i.e. , the dataset in the offline setting , and the replay buffer in the online setting ) . We investigate offline and online RL settings on benchmarks including Atari games ( Bellemare et al. , 2013 ) and Gym environments ( Brockman et al. , 2016 ) . We also utilize gridworlds described by Fu et al . ( 2019 ) to compare the learned Q-function against the oracle solution computed using tabular value iteration . We evaluate DQN ( Mnih et al. , 2015 ) on gridworld and Atari and SAC ( Haarnoja et al. , 2018 ) on Gym domains . Data-efficient offline RL . In offline RL , our goal is to learn effective policies by performing Qlearning on a fixed dataset of transitions . We investigate the presence of rank collapse when deep Q-learning is used with broad state coverage offline datasets from Agarwal et al . ( 2020 ) . In the top row of Figure 2 , we show that after an initial learning period , srankδ ( Φ ) decreases in all domains ( Atari , Gym and the gridworld ) . The final value of srankδ ( Φ ) is often quite small – e.g. , in Atari , only 20-100 singular components are active for 512-dimensional features , implying significant underutilization of network capacity . Since under-parameterization is implicitly induced by the learning process , even high-capacity value networks behave as low-capacity networks as more training is performed with a bootstrapped objective ( e.g. , mean squared TD error ) . On the gridworld environment , regressing toQ∗ using supervised regression results in a much higher srankδ ( Φ ) ( black dashed line in Figure 2 ( left ) ) than when using neural FQI . On Atari , even when a 4x larger offline dataset with much broader coverage is used ( blue line in Figure 2 ) , rank collapse still persists , indicating that implicit under-parameterization is not due to limited offline dataset size . Figure 2 ( 2nd row ) illustrates that policy performance generally deteriorates as srank ( Φ ) drops , and eventually collapses simultaneously with the rank collapse . While we do not claim that implicit under-parameterization is the only issue in deep Q-learning , the results in Figure 2 show that the emergence of this under-parameterization is strongly associated with poor performance . To prevent confounding from the distribution mismatch between the learned policy and the offline dataset , which often affects the performance of Q-learning methods , we also study CQL ( Kumar et al. , 2020b ) , an offline RL algorithm designed to handle distribution mismatch . We find a similar degradation in effective rank and performance for CQL ( Figure A.3 ) , implying that underparameterization does not stem from distribution mismatch and arises even when the resulting policy is within the behavior distribution ( though the policy may not be exactly pick actions observed in the dataset ) . We provide more evidence in Atari and Gym domains in Appendix A.1 . Data-efficient online RL . Deep Q-learning methods typically use very few gradient updates ( n ) per environment step ( e.g. , DQN takes 1 update every 4 steps on Atari , n = 0.25 ) . Improving the sample efficiency of these methods requires increasing n to utilize the replay data more effectively . However , we find that using larger values of n results in higher levels of rank collapse as well as performance degradation . In the top row of Figure 3 , we show that larger values of n lead to a more aggressive drop in srankδ ( Φ ) ( red vs. blue/orange lines ) , and that rank continues to decrease with more training . Furthermore , the bottom row illustrates that larger values of n result in worse performance , corroborating Fu et al . ( 2019 ) ; Fedus et al . ( 2020b ) . We find similar results with the Rainbow algorithm ( Hessel et al. , 2018 ) ( Appendix A.2 ) . As in the offline setting , directly regressing to Q∗ via supervised learning does not cause rank collapse ( black line in Figure 3 ) .

Summary: This paper discusses a phenomenon wherein the feature vectors of the learned value function in reinforcement learning (RL) lose their diversity as training progresses. The paper analyzes the rank of the final hidden layer in the model parameterizing the value function and shows experimentally that for offline-RL and online-RL setups on Atari and Gym benchmarks, this rank collapse occurs with a drop in the average return. The paper further develops two models for understanding this phenomenon, (i) where the value function is modeled using the neural tangent kernel, and (ii) where the value function is modeled using a deep linear network. The paper argues that bootstrapping results in reduction of the rank of the feature matrix as training progresses for these models. A regularization term that equalizes the singular values of the feature matrix is used to mitigate this rank collapse and experimental results on Atari benchmarks are shown with this regularizer.