Functional Regularization for Reinforcement Learning via Learned Fourier Features

1 Introduction . Most popular deep reinforcement learning ( RL ) approaches estimate either a value or Q-value function under the agent ’ s learned policy . These functions map points in the state or state-action space to expected returns , and provide crucial information that is used for improving the policy . However , optimizing these functions can be difficult , since there are no ground-truth labels to predict . Instead , they are trained through bootstrapping : the networks are updated towards target values calculated with the same networks being optimized . These updates introduce noise that accumulates over repeated iterations of bootstrapping , which can result in highly inaccurate value or Q-value estimates [ 44 , 46 ] . As a result , these RL algorithms may suffer from lower asymptotic performance or sample efficiency . Most prior work has focused on making the estimation of target values more accurate . Some examples include double Q-learning for unbiased target values [ 16 , 47 ] , or reducing the reliance on the bootstrapped Q-values by calculating multi-step returns with TD ( ) [ 41 ] . However , it is impossible to hope that the noise in the target values estimated via bootstrapping will go to zero , because we can not estimate the true expectation over infinite rollouts in practical setups . Hence , we argue that it is equally important to also regularize the function ( in this case , the deep Q-network ) that is fitting these noisy target values . Conventional regularization methods in supervised learning can be associated with drawbacks in RL . Stochastic methods like dropout [ 40 ] introduce more noise into the process , which is tolerable when ground-truth labels are present ( in supervised learning ) but counterproductive when bootstrapping . An alternative approach is early stopping [ 50 ] , which hurts sample efficiency in reinforcement learning because a minimum number of gradient steps is required to propagate value information backwards to early states . Finally , penalty-based methods like L1 [ 45 ] and L2 regularization [ 17 , 22 ] can help in 1Code available at https : //github.com/alexlioralexli/learned-fourier-features 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . RL [ 26 ] , but regularizing the network in weight space does not disentangle noise from reward signal and could make it difficult to learn the true Q-function . This leads us to ask : what is the right way to regularize the RL bootstrapping process ? We suggest that the impact of target value noise can be better reduced by frequency-based functional regularization : direct control over the frequencies that the network tends to learn first . If the target noise consists of higher frequencies than the true target values , discouraging high-frequency learning can help networks efficiently learn the underlying Q-function while fitting minimal amounts of noise . In this work , we propose an architecture that achieves this by encoding the inputs with learned Fourier features , which we abbreviate as LFF . In contrast to using fixed Fourier features [ 34 , 42 ] , we train the Fourier features , which helps them find an appropriate basis even in high dimensional domains . We analyze our architecture using the Neural Tangent Kernel [ 18 ] and theoretically show that tuning the initial variance of the Fourier basis controls the rate at which networks fit different frequencies in the training data . Thus , LFF ’ s initial variance provides a controlled mechanism to improve the stability and performance of RL optimization ( see Figure 1 ) . Tuned to prioritize learning low frequencies , LFF filters out bootstrapping noise while learning the underlying Q-function . We evaluate LFF , which only requires changing a few lines of code , on state-space and image-space DeepMind Control Suite environments [ 43 ] . We find that LFF produces moderate gains in sample efficiency on state-based RL and dramatic gains on image-based RL . In addition , we empirically demonstrate that LFF makes the value function bootstrapping stable even in absence of target networks , and confirm that most of LFF ’ s benefit comes through regularizing the Q-network . Finally , we provide a thorough ablation of our architectural design choices . 2 Preliminaries . The reinforcement learning objective is to solve a Markov Decision Process ( MDP ) , which is defined as a tuple ( S , A , P , R , ) . S and A denote the state and action spaces . P ( s0|s , a ) is the transition function , R ( s , a ) is the reward function , and 2 [ 0 , 1 ] is the discount factor . We aim to find an optimal policy ⇡⇤ ( a|s ) that maximizes the expected sum of discounted rewards . Q-learning finds the optimal policy by first learning a function Q⇤ ( s , a ) such that Q⇤ ( s , a ) = Es0⇠P ( ·|s , a ) [ R ( s , a ) + maxa0 Q⇤ ( s0 , a0 ) ] . The optimal policy is then ⇡⇤ ( a|s ) = argmaxa Q⇤ ( s , a ) . To find Q⇤ , we repeatedly perform Bellman updates to Q , which uses the Q-network itself to bootstrap target values on observed transitions ( s , a , r , s0 ) [ 6 ] . The most basic way to estimate the target values is target = r + maxa0 Q ( s0 , a0 ) , but a popular line of work in RL aims to find more accurate target value estimation methods [ 12 , 16 , 32 , 41 , 47 ] . Once we have these target value estimates , we update our Q-network ’ s parameters ✓ via gradient descent on temporal difference error : ✓ = ⌘r✓ ( Q✓ ( s , a ) target ) 2 . Our focus is on using the LFF architecture to prevent Q-networks from fitting irreducible , high-frequency noise in the target values during these updates . True Q-function Noisy targets LFF , σ = 0.003MLP Figure 2 : Filtering noise with LFF input embedding . The bootstrapped targets in RL are a mix between signal and noise . Right side : we fit different networks to these noisy targets and display their predictions . While the MLP overfits , LFF learns the Q-function and ignores almost all of the noise . 3 Reinforcement Learning with Learned Fourier Features . We present a visualization of the noisy target value problem in Figure 2 . The target value estimates can be noisy due to stochastic transitions , replay buffer sampling , or unintended generalization across state-action pairs due to function approximation [ 48 ] . We simulate this in Figure 2 by adding noise to the optimal Q-function of a small gridworld . MLPs are susceptible to fitting the noise , resulting in inaccurate Q-values that could diverge after repeated bootstrapping . In contrast , our LFF architecture controls how quickly low- and high-frequency signals are learned . Tuned properly , LFF filters out the noise and learns the ground truth Q-function almost perfectly ( Figure 2 , right ) . The problem is that MLPs provide no control over how quickly they learn signals of different frequencies . Prior work in computer vision found this a problem when MLPs blurred desired high frequencies in low-dimensional ( 3-5 dimensions ) graphics regression problems [ 28 , 42 ] . They fixed this blurring problem by transforming the input using a random Fourier feature embedding of the input [ 34 , 42 ] . Specifically , the idea is to map a low-dimensional input x to an embedding ( x ) = sin ( 2⇡Bx ) || cos ( 2⇡Bx ) , where B is a dfourier/2⇥dinput matrix , || denotes concatenation , and sin and cos act elementwise . The embedding ( x ) directly provides a mix of low- and high-frequnecy functions a MLP f✓ can use to learn a desired function . Tancik et al . [ 42 ] use this to improve fidelity in coordinate-based graphics problems . Intuitively , the row vectors bi are responsible for capturing desired frequencies of the data . If they capture only low frequencies , then the MLP will be biased towards only learning the low-frequency signals in the data . Conversely , if the Fourier features capture sufficient high-frequency features , then a MLP can fit high frequency functions by computing simple nonlinear combinations of the features . In these low-dimensional graphics problems , initializing fixed entries Bij ⇠ N ( 0 , 2 ) with large was enough to learn desired high frequency functions ; training B did not improve performance . In the following section , we propose our learned Fourier feature architecture for deep RL , which allows practitioners to tune the range of frequencies that the network should be biased towards learning . We propose several key enhancements that help Fourier features learn in high-dimensional environments . Our work uses learned Fourier features to improve deep RL , in contrast to prior work focused on simple environments with fixed , hand-designed Fourier features and linear function approximation [ 20 , 21 ] . Although we focus on prioritizing low-frequency signals to reduce bootstrap noise , we also present cases in Appendix F.1 where biasing networks towards high-frequency learning with learned Fourier features enables fast convergence and high asymptotic performance in RL . 3.1 Learned Fourier Feature Architecture . Standard MLPs can be written as the repeated , alternating composition of affine transformations Li ( x ) = Wix+ bi and nonlinearity ⌧ , which is usually the ReLU ⌧ ( x ) = max ( 0 , x ) : f✓ ( x ) = Ln ⌧ Ln 1 ⌧ · · · L1 ( x ) ( 1 ) We propose a novel architecture based on Fourier features , shown in Figure 1 . We define a new layer : FB ( x ) = sin ( 2⇡Bx ) || cos ( 2⇡Bx ) ||x ( 2 ) where B is a dfourier/2 ⇥ dinput matrix and || denotes concatenation . dfourier is a hyperparameter that controls the number of Fourier features we can learn ; increasing dfourier increases the degree to which the model relies on the Fourier features . Following Tancik et al . [ 42 ] , we initialize the entries Bij ⇠ N ( 0 , 2 ) , where 2 is a hyperparameter . Contrary to prior work , B is a trainable parameter . Algorithm 1 LFF PyTorch-like pseudocode . class LFF ( ) : def __init__ ( self , input_size , output_size , n_hidden=1 , hidden_dim=256 , sigma=1.0 , f_dim=256 ) : # create B b_shape = ( input_size , f_dim // 2 ) self.B = Parameter ( normal ( zeros ( * b_shape ) , sigma * ones ( * b_shape ) ) ) # create rest of network self.mlp = MLP ( in_dims=f_dim + input_size , out_dims=output_size , n_hidden=n_hidden , hidden_dim=hidden_dim ) def forward ( self , x ) : proj = ( 2 * np.pi ) * matmul ( x , self.B ) ff = cat ( [ sin ( proj ) , cos ( proj ) , x ] , dim=-1 ) return self.mlp.forward ( ff ) normal : sample from Gaussian with specified mean and std dev ; matmul : matrix multiplication ; cat : concatenation . The resulting LFF MLP can be written : f✓ = Ln ⌧ · · · L1 FB ( x ) ( 3 ) We can optimize this the same way we optimize a standard MLP , e.g . regression would be : argmin ✓ , B NX i=1 ( Ln · · · FB ( xi ) yi ) 2 ( 4 ) We propose two key improvements to random Fourier feature input embeddings [ 34 , 42 ] : training B and concatenating the input x to the Fourier features . We hypothesize that these changes help preserve information in highdimensional RL problems , where it is increasingly unlikely that randomly initialized B produces Fourier features well-suited for the task . Training B alleviates this problem by allowing the network to discover them on its own . Appendix G shows that training B does change its values , but its variance remains close to the initial 2 . This indicates that is an important knob that controls the network behavior throughout training . Concatenating x to the Fourier features is another key improvement for high dimensional settings . It preserves all the input information , which has been shown to help in RL [ 38 ] . We further analyze these improvements in Section 6.4 .

Summary: The paper proposes to use input features concatenated with learned Fourier feature (LF2) basis to control the degree of fitting various frequencies in the training data. Theoretical analysis under NTK regime with 2 layer network shows that the initial variance of Fourier basis controls this degree. In value-based RL, the updates using target Q-values can result in noisy optimization and LF2 can reduce dependence on noise by prioritizing low-frequency functions.  Empirically, LF2 leads to small sample-efficiency gains on a subset of DM control environments.