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: This work proposes Random Fourier Networks (RFNs), MLPs with random Fourier featurization within the first layer. RFNs are supposed to learn high-frequency components of the value function more quickly than standard MLPs. The authors present an initialization scheme for the Fourier layer that is better-suited for high-dimensional inputs. RFNs are compared to non-Fourier-feature baselines on high-dimensional environments.