Overcoming The Spectral Bias of Neural Value Approximation

1 INTRODUCTION Approximating value functions with deep neural networks is a critical component of stateof-the-art reinforcement learning ( RL ) algorithms . It has been shown that even for onedimensional continuous state space and a simple ( i.e. , piece-wise linear ) forward dynamics , there exists many MDPs with highly non-linear optimal Q-functions ( Dong et al. , 2020 ) . The authors further suggest that , many real-world problems have similarly complex value functions that are challenging to approximate us- ing a deep neural network . To understand these challenges , we re-created one of the toy MDPs from Dong et al . ( 2020 ) . The forward dynamics and the corresponding optimal Q-function of this MDP are shown in Figure 1 . We use a multi-layer perceptron ( MLP ) ( 4 layers deep with 400 hidden units each layer ) to learn the optimal Q-function of the toy MDP with offline fitted Q iteration ( FQI ) . We use a dataset of transitions collected by first uniformly sampling state-action pairs from the toy MDP and then taking the sampled action on the sampled state to get the next state . Despite training for 400 epochs , we observe the MLP under-fitting the optimal Q-function . It is natural to ask why the under-fitting occurs despite MLPs being universal function approximators . There could be several possible reasons which we now investigate . Prior work of Kumar et al . ( 2021 ) has argued that minimizing TD ( Temporal-Difference ) error in value-based deep RL through bootstrapping resembles self-distillation ( Mobahi et al. , 2020 ) , which is known to induce under-parameterization . To check if the primary cause of under-fitting is bootstrapping , we learn the Q-function using the same MLP architecture with supervised learning using ground truth Q values . Figure 2 shows that the under-fitting still persists after 400 epochs , implying a potential issue with the MLP architecture / training paradigm . To test whether under-fitting is the result of an under-parameterized MLP architecture , we use a 12-layer deep ( with 400 hidden units each layer ) MLP to learn the Q-function with supervised learning . However , as seen in Figure 2 , under-fitting still persists implying that the under-parameterization in the MLP architecture might not be the issue . Finally , to test whether under-fitting is caused by the training paradigm , we train the original MLP architecture for 2000 epochs to learn the Q-function with supervised learning . From Figure 2 , we see that the MLP fit to the Q-function improves , thereby implying that training for a large number of epochs resolves the issue . Given that more gradient updates improve the MLP ’ s fit to the Q-function , could we solve the underfitting issue by making the training period arbitrarily long ? No , given a long training period increases compute requirements in offline FQI and additionally causes data inefficiency in online RL . This is because data collection is interleaved with Q learning in online RL . When MLP fit to the Q-function is bad , the resulting policy might take suboptimal actions with low returns resulting in data that might not be useful for Q learning and would need to be recollected . A potential solution to improve data efficiency in online RL is by taking more gradient steps for learning Q-functions between data collection . However , as pointed in Fu et al . ( 2019 ) , this can cause MLPs to overfit to off-policy samples thereby degrading the overall performance . Therefore , it is imperative to understand the reasons behind the training inefficiency in learning Q-functions so that we can speedup the training . Recent works ( Rahaman et al. , 2019 ; Shah et al. , 2020 ; Yang & Salman , 2019 ; Huh et al. , 2021 ) have suggested presence of spectral bias in MLPs which encourages them to learn the low-frequency components of the target function before the high frequency components . Since the optimal Qfunction of the toy MDP has a large number of high-frequency components in its spectrum ( as seen from Figure 1 ) , the spectral bias in MLPs might explain why they take a large number of epochs to fit the Q function . To evade the spectral bias of MLPs , some recent works ( Huh et al. , 2021 ; Sinha et al. , 2020 ) proposed use of residual connections to increase the rank of MLP features . While this approach improves the performance of RL algorithms on low-dimensional continuous control tasks such as HalfCheetah ( state dim of 17 ) , it fails to provide any dividends in complex high-dimensional continuous control tasks such as Ant ( state dimension of 111 ) and Humanoid ( state dimension of 376 ) ( Sinha et al. , 2020 ) . Another solution , proposed in the graphics literature ( Tancik et al. , 2020 ; Sitzmann et al. , 2020 ; Mildenhall et al. , 2020 ) , is to transform the input with random Fourier features before passing it through a MLP . Taking cues from these works , we propose Random Fourier Network ( RFN ) , as a drop-in replacement of standard MLPs , to improve neural value estimation in MDPs with complex optimal Q-functions . Using RFN ( 4 layers deep with 400 hidden units each layer ) to learn the optimal Q-function of the toy MDP with FQI leads to a near-perfect fit after 400 epochs , as seen in Figure 2 . This shows that RFN is successful in evading the spectral bias of the MLPs . In coming sections , we will first describe RFN in detail and finally demonstrate the advantage of using them with actor-critic algorithms ( RFAC ) on challenging continuous control tasks such as Ant and Humanoid . 2 PRELIMINARIES . A Markov decision process ( MDP , Puterman 2014 ) can be defined by the tuple 〈S , A , R , P , µ , γ〉 where S and A are the state and action spaces , P : S × A 7→ S is the transition function , R : S × A 7→ R is the reward and µ ( s ) is the initial state distribution . We consider an infinite horizon problem where γ is the discount factor . The goal of reinforcement learning is to maximize the expected future discounted return J = E [ ∑∞ t=0 γ tR ( st , at ) ] . We define Qπ ( s , a ) function as expected future discounted return following the policy π when starting from state s and taking action a i.e . Qπ ( s , a ) = E [ ∑∞ t=0 γ tR ( st , at ) | ( s0 , a0 ) = ( s , a ) ] . Furthermore , J = Es µ ( s ) , a π ( a|s ) [ Qπ ( s , a ) ] . The optimal Q∗ ( s , a ) is the fixed point of the Bellman optimality operator B∗ B∗Q ( s , a ) = R ( s , a ) + γEs′∼P ( s′|s , a ) [ max a∗ Q ( s′ , a∗ ) ] ( 1 ) When P isn ’ t known but accessible through a sampled dataset D , we can learn Q∗ ( s , a ) using a parameterized modelQθ ( s , a ) with offline fitted-Q-iteration ( FQI ) ( Riedmiller , 2005 ) that optimizes the TD loss with Adam ( Kingma & Ba , 2015 ) min θ E ( s , a , s′ ) ∼D‖Qθ ( s , a ) −B∗Qθ′ ( s , a ) ‖ ( 2 ) Here , θ′ is a running smoothed average of θ from past iterations of Adam optimizer and Qθ′ is called the target Q-network . We use offline FQI for our offline RL experiments . Additionally , we use DDPG as our base RL algorithm for online RL experiments . 2.1 A KERNEL-SPECTRAL PERSPECTIVE ON THE SIMPLICITY-BIAS . Rahaman et al . ( 2019 ) recently observed that it took less gradient steps for a piece-wise linear network with a ReLU activation to fit to a low-frequency signal than a high-frequency one . This intrinsic bias , called the “ spectral bias ” , is a variant of the “ simplicity bias ” that is observed by Zhang et al . ( 2021 ) and Valle-Perez et al . ( 2018 ) on different occasions . To understand the source of these biases , the community has made two important connections . The first is Conjugate Kernel ( CK ) or the equivalent Gaussian process kernel . CK describes the network parameterization at initialization by computing the gram matrix of the final feature layer ( right before the output ) over the support . The second is the Neural Tangent Kernel , or NTK , which is the gram matrix of the flattened gradient vectors . Its spectrum governs convergence under gradient based updates , in a linear fashion ( Rahaman et al. , 2019 ; Yang & Salman , 2019 ) . When the data is not uniformly distributed , the convergence is linear with respect to the density of the data ( Fan & Wang , 2020 ) . Connecting Effective Rank with Kernel Spectrum Consider a random matrix A sampled from a distribution p ( A ) with the discrete spectrum σ = { σ1 ≥ σ2 ≥ · · · ≥ σQ } , the effective rank ( Roy & Vetterli , 2007 ) is defined as the spectral entropy of the nominal distribution given by the spectrum of the matrix according to the following formula erank ( A ) = expH ( { σi ‖σ‖1 } ) where i ∼ { 1 , . . . Q } . ( 3 ) The ` 1 norm ‖σ‖1 is called the spectral norm , which acts as the partition function of this distribution . H ( σ ) is called the Shannon entropy . In other words , the effective rank is the exponential of the log likelihood-weighted , normalized sum of the kernel spectrum that describes the volume span by the random matrices . 3 OVERCOMING THE SPECTRAL BIAS OF NEURAL VALUE ESTIMATION . We now apply spectral analysis directly to a toy MDP problem . Setup We extend the toy MDP problem from Dong et al . 2020 to consider a general class of toy Markov decision processM with a state space S = R [ 0,1 ) defined on the real-line . The reward is a sine function R ( s ) = sin ( 2πs ) with a simple spectral signature . The action space is a discrete set { 0 , 1 } , each corresponds to a distinct forward dynamics that is randomly sampled from the space of piece-wise linear functions with k “ kinks ” . For all of the examples below , we use a fixed number of kinks k = 10 , and uniformly sample the value of each turning point in this dynamic function between 0 and 1 . The result is a distributionM that we can sample from M ∼M = p ( M ) . ( 4 ) Fitted Q Iteration Consider a sequence of value functions under fitted Q iteration ( Ernst et al. , 2005 ) { Q0 ( s , a ) , Q1 ( s , a ) , Q2 ( s , a ) , · · ·QH ( s , a ) } starting from Q0 ( s , a ) = R ( s , a ) , generated via the repeated application of the Bellman optimality operator BQ ( s , a ) = R ( s , a ) + γmax a∗ Q ( s′ , a∗ ) where s′ ∼ T ( s , a ) . ( 5 ) The Bellman optimality operator is non-linear and imprints a non-trivial spectrum to the resulting Qi as it is applied recursively . We can get an intuitive understanding of the evolution of this process by looking at the spectra of Qi at various horizon H , or by changing the discount factor γ . We present the evolving spectra in figure 3 . As the horizon of the problem increases , the value function attains a larger weight in the higher-frequency part of the spectrum . The same effect also takes place if we increase the discount factor from a small value 0.1 all the way up to 0.99 . We fix the horizon to 200 in this case , and observe a similar increase in the higher frequency components with a longer effective recursion depth . We can quantitatively describe this trend by taking the effective rank of each value function , which is equivalent to computing the Shannon entropy of the spectrum . We present these results in figure 4 . In other words , the complexity of the value function comes from the repeated application of the Bellman optimality operator in a process that is not dissimilar to an “ infinite corridor ” . The spectrum of the Bellman operated gets folded exponentially into each Q iteration . To match this complexity , we need a function approximator class that has a similar long-tail in the higher-frequency spectrum , to avoid an exponential convergence drop-off in these areas .

Summary: The paper discussed the problem of spectral bias in MLPs (i.e., dependence on the higher-frequency part of input feature is under-fitted) within RL domain. The authors first described the problem with toy examples, which clearly explained this problem. Then it was suggested to use random Fourier feature network (RFN) to transform the input to frequency domain and then feed to the network, to resolve this issue.   Experiments demonstrated effectiveness of RFN to learn a complicated Q-function in low-dimensional state space. And it was also shown that RFN contribute to performance gain with DDPG on high-dimensional robotic tasks. The authors also did an “ablation study” (I would rather call it “sensitivity analysis” instead) to empirically show how the performance depends on the choice of hyper-parameters.