O VERCOMING T HE S PECTRAL -B IAS OF N EURAL V ALUE A PPROXIMATION
A BSTRACT
Value approximation using deep neural networks is at the heart of off-policy deep reinforcement learning, and is often the primary module that provides learning signals to the rest of the algorithm.
While multi-layer perceptron networks are universal function approximators, recent works in neural kernel regression suggest the presence of a spectral bias , where ﬁtting high-frequency components of the value function requires exponentially more gradient update steps than the low-frequency ones.
In this work, we re-examine off-policy reinforcement learning through the lens of kernel regression and propose to overcome such bias via a composite neural tangent kernel.
With just a single line-change, our approach, the Fourier feature networks (FFN) produce state-of-the-art performance on challenging continuous control domains with only a fraction of the compute.
Faster convergence and better off-policy stability also make it possible to remove the target network without suffering catastrophic divergences, which further reduces TD (0) ’s estimation bias on a few tasks.
Code and analysis available at https://geyang.github.io/ffn.
I NTRODUCTION
At the heart of reinforcement learning is the question of how to attribute credits or blame to speciﬁc actions that the agent took in the past.
This is referred to as the credit assignment problem (Minsky, 1961).
Correctly assigning credits requires reasoning over temporally-extended sequences of observation and actions, so that trade-offs can be made, for example, to choose a less desirable action at a current step in exchange for higher reward in the future.
Temporal difference methods such as TD ( λ ) and Watkins’ Q-learning (Sutton, 1988; Watkins, 1989) stitch together the immediate rewards local to each state transition, to estimates the discounted sum of rewards over longer horizons.
This is an incremental method for dynamic programming (Watkins & Dayan, 1992) that successively improves the value estimate at each step, which reduces the computation that would otherwise be needed to plan ahead at decision time.
A key tension in scaling TD learning to more challenging domains is that the state space (and action space in continuous control problems) can be very large.
In Neurogammon (Tesauro, 1991), for example, the complexity of the state space is on the scale of O (10 20 ) , which prevents one from storing all state-action pairs in a table.
Such constraints and the need to generalize to unseen board arrangements prompted Tesauro (1991) to replace the look-up table with an artiﬁcial neural network to great effect, achieving master-level play on backgammon.
In general, however, adopting neural networks as the value approximator introduces learning instability that can sometimes lead the iterative learning procedure into a divergent regime where errors are ampliﬁed at each step (Sutton & Barto, 2018; Bertsekas, 1995; Baird, 1995).
A slew of algorithmic ﬁxes have been developed to tackle this problem from different angles, including sampling from a replay buffer (Lin, 1992;
Riedmiller, 2005a); using a delayed copy for the bootstrapping target (Mnih et al., 2013); and using ensembles (Hasselt, 2010; Lee et al., 2021).
Despite of the popularity of these techniques, the neural value approximator itself remains unchanged.
Popular treatments view these networks as universal function approximators that would eventually converge to arbitrary target functions given sufﬁcient number of gradient updates and training data.
In practice, however, state-of-the-art off-policy algorithms interleave optimization with sampling, which greatly limits the amount of optimization that is accessible.
This puts deep reinforcement learning into the same regime as a neural network regularized through early-stopping where the model bias, the amount of optimization available and the sample size interact with each other in intricate ways (Belkin et al., 2018; Nakkiran et al., 2019; Canatar et al., 2021).
Therefore to understand off-policy learning with neural value approximators we ﬁrst need to understand how deep neural networks generalize, and how they converge under the dynamic process of gradient decent.
Then to ﬁnd a solution, we need to ﬁgure out a way to control the generalization and learning bias, so that deliberate bias-variance trade-offs could be made for better convergence.
Recent efforts in deep learning theory and computer graphics offer new tools for both.
Jacot et al.
(2018) uncovers a spectral-bias for shallow, inﬁnitely-wide neural networks to favor low-frequency functions during gradient descent (Kawaguchi & Huang, 2019; Bietti & Mairal, 2019; Ronen et al., 2019).
This result has since been extended to deeper networks at ﬁnite-width of any “reasonable” architecture via the tensor program (Yang, 2019; 2020; Yang & Littwin, 2021; Yang & Salman, 2019).
Our analysis on the toy MDP domain (Section 3) shows that value approximator indeed underﬁts the optimal value function, which tends to be complex due to the recursive nature of unrolling the dynamics.
Learning from recent successes in the graphics community (Sitzmann et al., 2020;
Tancik et al., 2020), we go on to show that we can overcome the spectral-bias by constructing a composite kernel that ﬁrst lifts the input into random Fourier features (Rahimi & Recht, 2008; Yang et al., 2015).
The resulting neural tangent kernel is tunable, hence offering much needed control over how the network interpolates data.
Our main contributions are twofold.
First, we show that the (re-)introduction of the Fourier features (Konidaris et al., 2011) enables faster convergence on higher frequency components during value approximation, thereby improving the sample efﬁciency of off-policy reinforcement learning by reducing the amount of computation needed to reach the same performance.
Second, our improved neural tangent kernel produce localized generalization akin to a Gaussian kernel.
This reduces the cross-talk during gradient updates and makes the learning procedure more stable.
The combined effect of these two improvements further allows us to remove the target network on a few domains, which leads to additional reduction in the value estimation bias with TD (0) .
B ACKGROUND A ND N OTATION
We consider an agent learning to act in a Markov Decision Process (MDP), (cid:104) S, A, R, P, µ, γ (cid:105) where S and A are the state and action spaces, P : S × A (cid:55)→ S is the transition function,
R : S × A (cid:55)→ R is the reward and µ ( s ) is the initial state distribution.
We consider an inﬁnite horizon problem with the discount factor γ .
The goal of the agent is to maximize the expected future discounted
return J = E [ (cid:80) ∞ t =
γ t R ( s t , a t )] by learning a policy π ( a | s ) that maps a state s to a distribution over actions.
The state-action value function (Q-function) is deﬁned as
Q π ( s, a ) =
E [ (cid:80) ∞
[Equation]
=
[Equation]
t R ( s t , a t ) | ( s 0 , a 0 )
=
( s, a )] .
The optimal Q ∗ ( s, a ) is the ﬁxed point of the
Bellman optimality operator B ∗
[Equation]
The ﬁtted Q iteration family of algorithms (Ernst et al., 2003; Riedmiller, 2005a) iteratively ﬁnds the (sub) optimal Q θ ∗ by recursively applying the gradient update
[Equation]
Let X =
Q θ ( s, a ) − B ∗ Q
θ ( s, a ) and apply the chain-rule.
The update in Equation 2 becomes
[Equation]
We can approximate the updated Q in function form through its ﬁrst-order Taylor expansion
[Equation]
where the bilinear form K ( s, a ; s (cid:48) , a (cid:48) )
=
∇ θ Q θ ( s, a ) T ∇ θ Q θ ( s (cid:48) , a
[Equation]
) is the neural tangent kernel (NTK, see Jacot et al.
2018 and Section 4).
This procedure can be interpreted as producing a
[Figure]
V a l u e
Action I
[Figure]
II (a)
[Figure]
MLP@ 0.0 0.2 0.4 0.
0.8 1.0 State [0, 1)
V a l u e
Action I Action II (b) Sup.
MLP@ 0.
0.2 0.4 0.6 0.8 1.0 State [0, 1)
V a l u e
[Figure]
I
[Figure]
II (c)
[Figure] [Figure]
[Figure] [Figure]
V a l u e
Action I Action II (d) Sup
[Figure]
MLP
[Figure]
2k 0.0 0.2 0.4 0.6 0.8 1.
[Figure]
V a l u e
Action I Action II (e) FQI FFN@ sequence { Q 0 , Q 1 , Q 2 , . . .
} using the iterative rule, Q i +1 =
Π B ∗ ( Q i ) starting with Q 0 ( s, a ) =
R ( s, a ) .
Π represents a projection to functions expressible by the class of function approximators, which reduces to the identity in the tabular case.
The i th item in the sequence, Q i , is the projected optimal Q function of a derived MDP with a shorter horizon H
= i .
A M OTIVATING
E XAMPLE
[Figure]
We motivate through a toy MDP adapted from Dong et al.
(2020), and show that neural ﬁtted Q iteration signiﬁcantly underﬁts the optimal value function.
Despite of the simple-looking reward and forward dynamics, the value function is quite complex.
Our spectral analysis in Section 3. indicates that such complexity arises from the recursive application of the Bellman operator.
Toy MDP Distribution Consider a class of toy Markov decision processes M .
The state space is deﬁned on the real-line as S
= R [0 , 1) .
The reward is the identity function.
The action space is a discrete set with two elements { 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 ﬁxed number k = 10 , and uniformly sample the value of each turning point in this dynamic function between 0 and 1 .
The result is a distribution M that we can sample from
M
[Equation]
A Curious Failure of Neural Fitted Q Iteration When we apply standard neural ﬁtted Q iteration (FQI, see Riedmiller 2005b) to this toy problem using a simple four-layer multi-layer perceptron (MLP) with 400 hidden units at each layer, we observe that the learned MLP value approximator, produced in Figure 2a, signiﬁcantly underﬁts in comparison to the optimal value function.
One hypothesis is that Q-learning is to blame, and in fact prior works such as Kumar et al. 2021 have argued that underﬁtting comes from minimizing the TD (Temporal-Difference) error, because bootstrapping resembles self-distillation (Mobahi et al., 2020), which is known to lead to under-parameterization in the learned neural network features.
To show that Q learning is not to blame, we can use the same MLP architecture, but this time directly regress towards the (ground truth) optimal Q function via supervised learning.
Figure 2b shows that the under-ﬁtting still persists under supervised learning, and increasing the depth to 12 layers (see Figure 2c) fails to ﬁx the problem.
This shows an over-parameterized network alone is insufﬁcient to reduce under-ﬁtting.
Finally, if we increase the number of training iterations from 400 epochs to 2000 , the original four-layer MLP attains a good ﬁt.
This shows that solely focusing on the expressivity of the neural network, and making it bigger (Sinha et al., 2020) can be misguided.
The class of function that can be approximated depends on how many gradient updates one is allowed to make, a number capped by the compute budget.
3. S PECTRAL S IGNATURE OF T HE T OY MDP
The Bellman optimality operator imprints a non-trivial spectrum on the resulting Q i as it is applied
[Figure]
We present the evolving spectra marginalized over M in Figure 3.
As the effective horizon increases, the value function accrues more mass in the higher-frequency part of the spectrum, which corresponds to less correlation between the values at near-by states.
In a second experiment, we ﬁx the horizon to 200 while increasing the discount factor from 0 .all
the way up to 0 .
We observe a similar whitening of the spectrum at longer effective recursion depths.
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 “inﬁnity mirror.”
The spectrum of the Bellman operator gets folded into the resulting Q function upon each iterative step.
Although our analysis focuses on the state space, the same effect can be intuitively extrapolated to the joint state-action space.
3. K ERNEL V IEW ON O FF - POLICY D IVERGENCE
Following the formulation of convergence in Dayan (1992); Tsitsiklis (1994);
Jaakkola et al.
Deﬁnition: Consider a complete metric space S with the norm (cid:107) · (cid:107) .
An automorphism f on S is a contraction if ∀ a, b ∼ S , (cid:107) f
( a ) − f ( b ) (cid:107) ≤ γ (cid:107) a − b (cid:107) .
Here γ ∈ [0 , 1) is called the contraction modulus .
When γ = 1 , f is a nonexpansion .
Banach ﬁxed-point theorem.
Let S be non-empty with a contraction mapping f .
Then f admits a unique ﬁxed-point x ∗ ∈ S s.t.
f ( x ∗ )
= x ∗ .
Furthermore, ∀ x 0 ∈ S , x ∗ is the limit of the sequence given by x i +
= f
( x i ) .
a.k.a x ∗ =
lim
i →∞ x i .
Without lost of generality, we can discretize the state and action space S and A .
The NTK becomes the gram matrix K ∈ R | S ×
A |×| S × A | .
Transition data are sampled from a distribution ρ ( s, a ) .
Theorem.
(Achiam et al., 2019) Let indices i , j refer to state-action pairs.
Suppose that K , η and ρ satisfy the conditions:
[Equation]
Then, Equation 4 induces a contraction on Q in the sup norm, with ﬁxed-point Q ∗ and the TD loss optimization converges with enough optimization steps.
For relatively large γ (for instance, γ ∈ (0 . 99 , 0 . 999) ), the above theorem implies that small offdiagonal terms in the NTK matrix are sufﬁcient conditions for convergence.
S PECTRAL -B IAS AND N EURAL K ERNEL R EGRESSION
Consider a simple regression problem where we want to learn a function f ( ξ ; θ ) ∈ R .
ξ is a sample from the dataset.
To understand how the output of the network changes w.r.t small perturbations to the parameters, we can Taylor expand around θ
[Equation]
[Equation]
During stochastic gradient descent using a training sample ˆ ξ with a loss function L , the parameter update is given by the product between the loss derivative L (cid:48) ◦
f
( ˆ ξ ) and the neural tangent kernel
K (Jacot et al., 2018)
[Equation]
where
[Equation]
In the inﬁnite-limit with over-parameterized networks, the function remains close to initialization during training (Chizat et al., 2019).
The learning dynamics in this “lazy” regime, under an L regression loss behaves as a minimum norm least-square solution
[Equation]
where f t is the function under training at time t and f 0 is the neural network at initialization.
Replacing K with the Gramian matrix K between all pairs of the training data, we can re-write Equationin its spectral form K = O Λ O T where each entry in the diagonal matrix Λ is the eigenvalue
λ i >for the basis function O i in the orthogonal matrix O .
Using the identity e A
= O e Λ O T , we can decompose the learning dynamics in Equation 4 into
[Equation]
The key observation is that the convergence rate on the component O i , ηλ i , depends exponentially on its eigenvalue
λ
i .
a.k.a the absolute error
[Equation]
Multiple works (Rahaman et al., 2019; Shah et al., 2020; Yang & Salman, 2019;
Huh et al., 2021) have shown that the NTK spectrum ( λ i ) of a regular ReLU network decays rapidly at higher frequency.
In particular, Bietti & Mairal (2019) provided a bound of Ω( k − d − 1 ) for the k th spherical harmonics 1 .
Such results have been extended to ﬁnite-width networks of arbitrary architecture and depth via the tensor programs (Yang, 2019; 2020).
[Equation]
[Figure]
The Gramian matrix form of the NTK also offers an intuitive way to inspect state-aliasing during gradient descent.
This is because the offdiagonal entries corresponds to the similarity between gradient vectors for different state-action pairs.
The kernel of an multi-layer perceptron is not stationary when the input is not restricted to a hypersphere as in Lee et al.
We can, however, compute K of popular network architectures over R [0 , 1) , shown in Figure 4.
The spectral-bias of the NTK of both
ReLU networks and hyperbolic tangent networks produce large off-diagonal elements, increasing the chance of divergence.
O VERCOMING THE S PECTRAL -B IAS OF N EURAL V ALUE A PPROXIMATION
To correct the spectral-bias of a feed-forward neural network, we construct a composite kernel where a random map z : R d (cid:55)→ R D ﬁrst “ lifts ” the input into a randomized harmonics basis.
This explicit kernel lifting trick was introduced by Rahimi & Recht (2007) and it allowed the authors to ﬁt complex datasets using a linear machine.
The mixing brings high-frequency input signals down to a lower and more acceptable band for the ReLU network.
Data also appears more sparse in the higher-dimensional spectral basis, further simpliﬁes learning.
To emulate arbitrary shift-invariant kernel K ∗ , Rahimi & Recht (2007) offered a procedure that sample directly from the nominal distribution given by K ∗ ’s spectrum F ( K ∗ )
[Equation]
ˆ ξ ) (cid:105)
≈
z ( ξ ) T z (ˆ ξ )
where z ( ξ ) =
(cid:88)
i w i
e
2 πk i and
w i ∼ F
( K ∗ ) .
We choose to initialize FFN by sampling the weights from an isotropic multivariate Gaussian with a tunable cutoff frequency b .
We modify the normalization scheme from Rahimi & Recht (2007) and divide b with the input dimension d , s.t. similar bandwidth values could be applied across a wide variety of reinforcement learning problems with drastically different state and action space dimensions.
We slightly abuse the notion by using b j as the bias parameters, and b (without subscript) for the bandwidth
RFF( x ) j = sin( d (cid:88) i =
w i,j x i + b j ) where w i,j ∼ N (0 , πb/d ) and b j ∼ U
( − π, π ) .
0 d B b=64 b= b=4 b=
[Figure]
Spectrum  0.
b=1 b= b=16 b=
[Figure]
Algorithm Learned Fourier Features (
[Figure]
) class LFF (nn
[Figure]
Linear): def init (self, in, out, [Figure] scale): super()
[Figure]
init (in, out) init
[Figure]
normal (self.weight, std= [Figure] scale/in) init
[Figure]
uniform (self [Figure] bias, − 1.0, 1.0) def forward(self, x):
[Figure]
np
[Figure]
[Figure] [Figure]
To attain better performance, Rahimi & Recht (2008) learns a weighted sum of these random kitchen sink features, whereas Yang et al.
(2015) adapts the sampled spectral mixture itself through gradient descent.
Using the modern deep learning tool chain, we can view the entire network including the random Fourier features as a learnable kernel.
For small neural networks with limited expressivity, we found that enabling gradient updates on the RFF parameters is important for performance.
We refer to this adaptive variant as the learned Fourier features (LFF), and the shallow MLP with an adaptive Fourier feature expansion as the Fourier feature networks (FFN).
Kernel Spectrum, Convergence, and Interpolation
The composite kernel of the FFN has larger eigenvalues on the high-frequency harmonics (see Figure 5), which improves the speed at which the network converges.
We can tune this kernel spectrum by changing the band-limit parameter b .
With a larger bandwidth, the interpolation becomes more local .
On the toy domain used to motivate this work FFN achieves a perfect ﬁt with just 400 optimization epochs whereas the MLP baseline requires at least two thousand gradient steps (see Figure 2e).
The gain is even more prominent on more challenging environments in the results section (Quadruped run, see Figure 9).
Optimization overhead is a key bottleneck in off-policy learning, therefore faster convergence also translates into better sample efﬁciency, and faster wall-clock time.
[Figure]
Improving Off-policy Stability
“cross-talks” between the gradients of near-by states is a key factor in off-policy divergence.
Such “crosstalk” manifest as similarity between gradient vectors, the measure of which is captured by the Gram matrix of the NTK.
The Fourier feature network kernel is both localized (Figure 4c) and tunable (Figure 5), offering direct control over the biasvariance trade-off.
The improved learning stability allows us to remove the target network on a number of domains while retaining substantial performance.
A curious ﬁnd is when the target network is used, the learned value approximation still contains a large constant offset from the optimal value function (yellow curve in Figure 6).
Such offset disappears after we remove the target value network (blue curve).
Visualizing Mountain Car
The Mountain Car environment (Moore, 1990; Sutton, 2000) has a simple, two-dimensional state space that can be directly visualized to show the learned value function.
We compare the value estimate from three baselines In Figure 7: the ground-truth value estimate acquired using tabular value iteration; one obtained from a four-layer ReLU network using ﬁtted Q iteration; and one from the same network, but using a 16 dimensional Fourier features on the input.
All networks use 400 latent neurons and are optimized for 2000 epochs.
[Figure]
S CALING U P T O C OMPLEX C ONTROL T ASKS
We scale the use of FFN to high-dimensional continuous control tasks from the DeepMind control suite (Tassa et al., 2020) using soft actor-critic (SAC, Haarnoja et al. 2018) as the base algorithm.
Our implementation is built off the pytorch codebase from Yarats & Kostrikov (2020) which is the current state-of-the-art on state space DMC.
Both the actor and the critic have three latent layers in this base implementation.
We use FFN for both the actor and the critic by replacing the ﬁrst layer in the MLP with learned Fourier features (LFF) without increasing the total number of parameters.
We provide extensive empirical analysis on eight common DMC domains and additional results with DDPG in Appendix A.9.
Due to space constraints we present the result on four representative domains in Figure 8.
FFN consistently improves over the MLP baseline in complex domains while matching the performance in simpler domains.
We include additional ablations and hyperparameter sensitivity analysis in Appendix A.7.
[Figure]
Faster Convergence via Fourier Feature Networks
The key beneﬁt of the Fourier feature network is that it reduces the computation needed for good approximation.
Off-policy reinforcement learning algorithms tend to be bottlenecked by optimization whereas on-policy algorithms tend to be bottlenecked by simulation time.
This means with a reduced replay ratio
[Figure]
FFN can achieve faster wall-clock time.
Figure 9 shows that we can indeed reduce the update frequency to 1 / 4 in Walkerwalk and 1 / 6 in Quadruped-run while still matching the performance of the MLP baseline.
In a control experiment (see Figure 10), we ﬁnd that increasing the replay ratio causes the MLP to crash
[Figure]
to match the performance of the MLP baseline on both Walker-run and Quadruped-run.
[Figure] [Figure]
on Quadruped-run, illustrating the intricate balance one otherwise needs to maintain.
FFN brings additional improvements on learning stability besides faster convergence.
We additionally visualize the value approximation error in Figure 25.
FFN consistently produces smaller value approximation error compare to the MLP
[Figure]
V a l u e e s ti m a ti
[Figure]
e rr o r
[Figure]
Improved Feature Representation Makes Learning Easier
The random Fourier features sample from a Gaussian distribution in the frequency domain.
This is a more uniform functional distribution compared to the MLP, which exponentially favors low-frequency functions (Bietti & Mairal, 2019).
Let W t and B t refer to the concatenated weight and bias vectors.
In this experiment we inspect the (cid:96) 2 -norm of the changes of weights and biases w.r.t to their initial value, (cid:107)W t −W 0 (cid:107) 2 and (cid:107)B t −B 0 (cid:107) 2 .
We include all eight DMC domains in Appendix A.11 and present the most representative domains in Figure 11.
In seven out of eight domains, the ﬁrst layer in both the MLP and the FFN experience roughly the same amount of change in the weights, but the FFN requires less weight change in the later layers.
This indicates that LFF makes learning easier for the rest of the network.
Quadrupedrun is the only outlier where the LFF experience more weight change than the ﬁrst layer of the MLP.
Our hypothesis is that this is a harder task, and we are either underparameterized in terms of Feature dimension, or there is a distribution-misalignment between the task and the model.
So far we have focused on the weights.
The biases change consistently less in the FFN than the MLP.
Fourier Features Improve Off-Policy Stability Without a target network, gradient updates through the TD objective causes the regression target itself to change, making the learning procedure less stable.
Mnih et al.
(2013) replaces the value network on the r.h.s of the Bellman equation with a copy that is updated at a slower rate, so that the TD objective is more stationary.
In Figurewe remove the target network.
While the MLP baseline completely fails in all environments, the
[Figure]
stability to the extent that learning can happen on some domains even without the target network.
On the contrary, MLP consistently fails without a target value network.
FFN has consistently less variance than the MLP.
[Figure]
improved neural tangent kernel of the FFN sufﬁciently stabilizes Q learning, that only minor performance losses occur with seven out of eight domains.
We offer the complete result in Appendix A.12.
Improving the Convolution Kernel
The ﬁlters in a standard convolution neural network suffer the same spectral bias in the RGB and feature dimension as the state and action space above.
We can in principle extend our spectral-bias ﬁx to convolution networks by replacing the ﬁrst layer with a 1 × 1 convolution with a sine non-linearity.
In this experiment we build upon the DrQv implementation (Yarats et al., 2021) and replace the CNN in both the actor and the critic with this
Fourier-CNN (F-CNN) architecture according to Equation 14.
The results indicate that while FCNN’s performance is within the variance of CNN’s performance, its variance is much lower (see
Figure 14).
Such technique has also been used by Kingma et al.
(2021) to improve image generation with diffusion models.
For more detailed theoretical analysis of the CNN NTK, one can refer to Arora et al.
and
Li et al.
D ISCUSSION
The inspiration of this work comes from our realization that techniques used by the graphics community (Tancik et al., 2020; Sitzmann et al., 2020; Mildenhall et al., 2020) to learn high-ﬁdelity continuous neural representations represent a new way to explicitly control generalization in neural networks.
We refer to Achiam et al.
’s pioneering work for the theoretical setup on off-policy divergence, which predates many of the recent analysis on spectral-bias in neural networks, and the (re-)introduction of the random Fourier features as a solution to correct such learning biases.
Functional regularization through gradient conditioning is an alternative to re-parameterizing the network.
A recent effort could be found in Pich ´ e et al.
A key beneﬁt of reparameterizing the network is speed.
We managed to reduce the wall-time by 32% on the challenging Quadruped environment, by reducing the replay ratio to 1 / 6 of the baseline.
Fourier features for reinforcement learning dates back as early as Konidaris et al.
While this manuscript was under review, a few similar publication and preprints came out, all developed independently.
Li & Pathak (2021) focuses on the smoothing effect that Fourier feature networks have on rejecting excessive noise.
Our ﬁnding is that vanilla feed-forward neural networks are on the biased-side of the bias-variance trade-off.
In other words, the main issue in off-policy learning with neural function approximators is that the network underﬁts real signal, as opposed to being overﬁt to random noises.
The rank of the neural representation describes the portion of the linear space occupied by the largest eigenvalues of the kernel regardless of the spatial frequency of those corresponding eigenfunctions.
Therefore lower rank does not correspond to smoother functions.
We additionally ﬁnd that maintaining a Fourier feature to input feature ratio ( D/d > 40 ) is critical to the expressiveness of the network, which allowed us to scale up to Quadruped without needing to concatenating the raw input as a crutch.
Brellmann et al.
(2022) is a concurrent submission to ours that delightfully includes the random tile encoding scheme from Rahimi & Recht (2007), and on-policy algorithm, PPO.
Our intuition is that policy gradient needs to be considered a life-long learning problem, and locality in the policy network can speed up learning by eliminating the need to repeatedly sample areas the policy has experienced.
We are excited about these efforts from the community, and urge the reader to visit them for diverse treatments and broader perspectives.
A CKNOWLEDGMENTS
The authors would like to thank Leslie Kaelbling for her feedback on the manuscript; Jim Halverson and Dan Roberts at IAIFI for fruitful discussions over neural tangent kernels, and its connection to off-policy divergence; Aviral Kumar at UC Berkeley for bringing into our attention the expressivity issues with off-policy value approximation; and the reviewer and area chair for their kind feedback.
This work is supported in part by the National Science Foundation Institute for Artiﬁcial Intelligence and Fundamental Interactions (IAIFI, https://iaiﬁ.org/) under the Cooperative Agreement
PHY-2019786; the Army Research Ofﬁce under Grant W911NF-21-1-0328; and the MIT-IBM Watson AI Lab and Research Collaboration Agreement No.W1771646.
The authors also acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing high performance computing resources.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the ofﬁcial policies, either expressed or implied, of the Army Research Ofﬁce or the U
S. Government.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
R EPRODUCIBILITY S TATEMENT
We include detailed implementation details in Appendix A.4.
A A PPENDIX A. E XPERIMENTAL D ETAILS ON T HE T OY MDP
Ofﬂine Data Generation We divided the 1-dimensional state space into 1000 discrete bins and used the midpoints of the bins as initial states.
We then took both the actions from these initial states to get corresponding next states.
We used the dataset of these transitions ( 2000 total transitions) as our ofﬂine dataset.
Optimization details We use 4-layer MLP with ReLU Activation, with 400 latent neurons.
We use Adam optimization with a learning rate of 1e-4, and optimize for 400 epochs.
We use gradient descent with a batch size of 200.
For a deeper network, we use a 12-layer MLP, with 400 latent neurons but keep other optimization hyperparameters ﬁxed.
To show that longer training period help MLPs evade the spectral bias, we use a 4-layer MLP, with 400 latent neurons and train it for 2000 epochs.
All the optimization hyperparameters are kept the same.
A. E XPERIMENT D ETAILS ON M OUNTAIN C AR
Environment We use the implementation from the OpenAI gym (Brockman et al., 2016), and discretize the state space into 150 bins.
This is a relatively small number because otherwise the space complexity becomes infeasible large.
We use a four-layer MLP as the Q function, and a bandwidth parameter of 10 .
A. E XPERIMENTAL S ETUP ON D EEP M IND C ONTROL S UITE
DMC domains We list the 8 DMC domains along with their observation space dimension and action space dimension in Table 1.
We also list the optimal bandwidth b used by FFN for each of these envs.
[Table]
-run
Box
[Table] [Table]
Walker-
[Table]
Box
[Table]
[Table] [Table]
[Table] [Table]
-walk
[Table] [Table]
[Table] [Table]
[Table]
A. A RCHITECTURAL D ETAILS
For the state space problems, we parameterize the actor and the critic in the MLP baseline with three latent layers.
For FFN (ours), we replace the ﬁrst layer in the MLP a learned Fourier features layer (LFF).
If d is the input dimension, both MLP and FFN have [40 × d, 1024 , 1024] as hidden dimension for each of their layers.
Note that in the actor, d = observation dimension whereas in the critic, d = observation dimension + actor dimension.
We build upon DrQv2 (Yarats et al., 2021), and derive conv FFN by replacing the ﬁrst convolutional layer with a 1x1 convolution using the sine function as non-linearity.
We use the initialization scheme described in Equation 14 for the weights and biases.
A. S UMMARY OF V ARIOUS F OURIER F EATURES
Using Fourier features to correct the spectral-bias is a general technique that goes beyond a particular parameterization.
Hence we present comparison between
• vanilla MLP is a stack of t linear layers with ReLU activation
[Equation]
• Fourier features network (FFN) (Ours) uses sine activation with random phase-shift, to replace the ﬁrst layer of the network with learned Fourier features
[Equation]
so that the Fourier features network (FFN) is
[Equation]
• RFF (Tancik et al., 2020) that uses sine and cosine pairs concatenated together
[Equation]
• SIREN network (Sitzmann et al., 2020) that stacks learned Fourier layers through our the entire network, using the Sitzmann initialization according to
[Equation]
where the t-layer network
[Equation]
It is critical to note what each layer has a distinct bandwidth scaling parameter.
Hence instead of a single scaling hyperparameter b for the random
matrices,
Siren
has
a set { b i } that each need to be tuned.
We found that it is important to use isotropic distributions for the weights.
This means using Gaussian distributions or orthogonal matrices for the weight matrix.
Uniform distribution is anisotropic.
To stack LFF (which becomes SIREN), the network requires a “fan-out” in the latent dimensions to accommodate the successive increase in the number of features.
SIREN is difﬁcult to tune because the initialization scheme is incorrect, as it depends on the inverse quadratic root of the input dimension.
Correcting the initialization to depend inversely on the input dimension ﬁxes such issue
[Figure]
A. E FFECT OF B ANDWIDTH P ARAMETER b ON T OY MDP
The weight distribution in the random Fourier features offers a way to control the generalization (interpolation) in the neural value approximator.
In Figure 15, we compare the learned Q function using Fourier feature networks initialized with different band-limit parameter b .
With larger b s, the network recovers sharper details on the optimal Q function.
We also provide the result for an FFN with the target network turned off.
[Figure]
A. A BLATIONS AND S ENSITIVITY A NALYSES
Sensitivity to bandwidth b Since FFN introduces a bandwidth hyperparameter b , it is natural to ask how the choice of b affects FFN’s performance.
Figure 16 shows that FFN’s performance does vary with choice of b .
Furthermore, the best performing value of b differs with environment.
This difference arises from the fact that the optimal value function for different environments have different spectral bias and hence requires different bandwidth b for FFN.
[Figure]
We use SAC as the RL algorithm.
We observe that different B values lead to different performances and hence, we must choose the B value carefully for each environment.
MLP Actor + FFN Critic Experiments we have presented so far use FFN for both the actor and the critic.
Our intuition is that the main beneﬁt comes from using FFN in the critic.
Results in
Figure 17 shows that this is indeed the case – using an MLP for the actor does not affect the learning performance, whereas using an MLP for the critic causes the performance to match those of the
MLP baseline.
Sensitivity to (Fourier) Feature-Input Ratio
In addition to the bandwidth b , we need to choose the Fourier feature dimension D for FFN.
We maintained the feature-input ratio Dd to be 40 .
But can we get away with a lower feature-input ratio?
Figure 18 shows that it is important to maintain feature-input ratio to be at least 40 and any reduction in the feature-input ratio hurts the performance of FFN.
[Figure] [Figure]
A. F ULL R ESULTS WITH SAC
We include the full set of results on eight DMC domains in Figure 19 using soft actor-critic (Haarnoja et al., 2018).
A. F ULL R ESULTS WITH DDPG
The beneﬁts FFN brings generalize beyond soft actor-critic to other RL algorithms.
In this section we present results based on deep deterministic policy gradient (DDPG, see Lillicrap et al. 2015).
Figure 20 shows learning curve of the FFN vs the MLP on eight DMC domains.
A. R EDUCING V ALUE A PPROXIMATION E RROR
We include the full result on value approximation error on eight DMC domains in Figure 21.
FFN consistently reduces the value approximation error w.r.t the MLP baseline especially at earlier stages of training when insufﬁcient amount of optimization has occurred.
[Figure] [Figure]
[Figure]
A. M EASURING P ARAMETER C HANGES IN FFN A ND MLP
We can visualize the amount of changes the FFN and the MLP experience in the weight space .
Let
W kt and B kt refer to the ﬂattened weight and bias vectors from layer k in the critic after training over t environment frames which corresponds to the number of optimization steps times a multiple.
Figure 22 shows the evolution of (cid:107)W t − W 0 (cid:107) for all eight domains.
Figure 23 shows the evolution of (cid:107)B t − B 0 (cid:107) .
Both collected with soft actor-critic.
A. R EMOVING THE TARGET NETWORK
We present the full result without the target network in Figure 24.
While MLP completely ﬂat-lines in all environments, FFN matches its original performance with target networks in Finger-turn, and
Quadruped-walk.
It suffers performance loss but still manages to learn in Cheetah-run, Walker-run, and Quadruped-run.
Hopper-hop, Humanoid-run, and Acrobat-swing up are difﬁcult for model-free algorithms, and FFN fails to learn without target network.
We can inspect the value approximation error.
Without a target network, both MLP and FFN do poorly on Hopper-hop, so the result on this domain is uninformative.
In Walker-run, removing the target network reduces the value approximation error even further.
In Quadruped walk, FFN without target network perform worse than FFN with target network, and has slightly higher approximation error.
Both FFN baselines however have lower approximation errors than the MLP.
In Quadrupedrun without the target network, the approximation error diverges to large value, then converges back.
The source of this divergence remains unclear.
[Figure]
[Figure] [Figure]
[Figure] [Figure]