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 perceptrons 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 bias to overestimate the value.
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 could be made, to choose a less desirable action at a current step in exchange for higher rewards 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 is otherwise 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 spaces in continuous control problems) can be very large.
In Neurogrammon (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 constrain 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.
Adopting neural networks as the value approximator in general, however, introduces two new issues.
First, the standard perceptron network does not offer explicit ways to specify how it generalizes, making it impossible to control the bias-variance trade-off to better align with the value approximation task at hand (Canatar et al., 2021).
Recent analysis indicates an “expressivity gap” that causes the value approximator to underﬁt (Dong et al., 2020), further indicates that there exists a misalignment between the model and the task, and that we are on the wrong side of the “doubledescent” phase transition (Canatar et al., 2021).
Modern state-of-the-art off-policy algorithms also tend to be bottlenecked by compute, making the convergence dynamics a key candidate for potential improvements in both sample efﬁciency and asymptotic performance.
A second, and older issue around neural function approximators is off-policy divergence (Sutton &
Barto, 2018), where the iterative Q learning procedure tends to “crash” when over-parameterized function approximators are combined with value bootstrapping and off-policy samples (Bertsekas, 1995; Baird, 1995).
A number of contributing causes have been investigated, and corresponding practical techniques have been developed to address such issues.
Riedmiller (2005a) employ experience replay (Lin, 1992) from a buffer of transition tuples to reduce the cross-talk between states during gradient updates.
To make the bootstrapped regression target more stationary and delay the divergence, Mnih et al.
(2015) introduced a target value network that changes at a slower speed.
Additional bias occurs in the TD (0) objective due to the “winner’s curse” (Thaler, 1988), which motivated ensemble based variance-reduction techniques such as double Q-learning (Hasselt, 2010) and bootstrapped Q-ensembles (Lee et al., 2021).
In continuous control tasks, computing the value target further require participation of the actor that is often sub-optimal.
This introduces additional variance in actor-critic methods that motivated delayed actor updates (Fujimoto et al., 2018), entropy regularization (Haarnoja et al., 2018), and trusted region for the policy optimization (Neumann et al., 2009; Abdolmaleki et al., 2018).
None of these instabilities occur, however, under the tabular case.
Nor with many compact parameterization schemes such as decision trees (Ernst et al., 2005), Gaussian kernels (Ormoneit & Sen, 2002; Smart & Kaelbling, 2000), or learned value eigenbasis (Mahadevan & Maggioni, 2007), to name a few.
Recent analysis on neural kernel regression (Neal, 1996; Jacot et al., 2018) uncovers a spectral-bias in multilayer perceptron networks, whose convergence slows down exponentially on high-frequency harmonics (Kawaguchi & Huang, 2019; Bietti & Mairal, 2019; Ronen et al., 2019).
Although originally derived with shallow, inﬁnitely-wide networks, such result has 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).
This spectral bias for the neural tangent kernel to concentrate towards low-frequency bands produce unwanted generalization and cross-talks during gradient descent.
We identify this spectral bias as the key issue, and address both the expressivity gap and off-policy divergence , by explicitly controlling the bias-variance trade-off via an a tunable, and adaptive Gaussian mixture kernel (Rahimi & Recht, 2008; Yang et al., 2015).
Our solution involves changing just a single line of code , yet it achieves state-of-the-art sample efﬁciency and asymptotic performance on challenging continuous control domains with just a fraction of the compute needed by a vanilla perceptron network.
Our main contributions are twofold.
First, we show that use of Fourier features enable faster convergence on high frequency spectral components during value approximation, thereby improving the sample efﬁciency of off-policy reinforcement learning.
Second, we show that the improved neural tangent kernel is more localized with much smaller off-diagonal cross-talk.
Combined with faster convergence, the improved off-policy stability enabled us to remove the target network on a few domains without suffering catastrophic divergence.
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.1 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
[Equation]
0.0 0.2 0.4 0.6 0.8 1.0 State [0, 1)
V a l u e
Action I Action II (a) FQI
+ 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
Action I Action II (c) Sup.
12-layer 0.0 0.2 0.4 0.6 0.8 1.0 State [0, 1)
V a l u e
Action I Action II (d) Sup.
MLP@2k 0.0 0.2 0.4 0.6 0.8 1.0 State [0, 1)
V a l u e
Action I Action II (e) FQI FFN@
A curious failure of neural ﬁtted 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 ground-truth value function.
One hypothesis is that Q learning is to blame, and in fact prior work such as Kumar et al. 2021 has argued that underﬁtting results from minimizing the TD (Temporal-Difference) error, because bootstrapping resembles iteratively applying self-distillation (Mobahi et al., 2020), which is known to lead to under-parameterization in the learned neural network features.
To rule-out this hypothesis, we can use the same MLP architecture, but this time directly regress towards the ground truth 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 used in these experiments to 2000, the original four-layer MLP attains a good ﬁt.
This shows that solely focusing on the expressivity of the neural network by making it bigger (Sinha et al., 2020) can be misguided, as the type of function that can be approximated also depends on the available budget over the number of gradient updates.
B ACKGROUND A ND N OTATIONS
We consider an agent learning to act in a Markov decision process (MDP), parameterized via the tuple (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) ∞ t =
γ t R
[Equation]
s t , a t ) |
[Equation]
s 0 ,
[Equation]
=
( 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 optimal Q ∗ by recursively applying the gradient update
[Equation]
Letting X =
Q θ ( s, a ) − B ∗ Q
θ ( s, a ) and apply the chain-rule.
The update 3 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 θ
[Equation]
s, a ) T ∇ θ Q θ
[Equation]
s
[Equation]
a (cid:48) ) is the neural tangent kernel (NTK, see Jacot et al.
2018 and Section 4).
In the tabular case, this procedure can be simpliﬁed as producing a 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 ) .
The i th item in the sequence, Q i , is the optimal Q function of a derived
MDP with a shorter horizon H
=
i .
Therefore each application of the Bellman optimality operator effectively extends the horizon by one time-step, starting with the 1 -step reward R ( s, a ) .
3. S PECTRAL S IGNATURE OF THE T OY MDP
[Figure] [Figure]
H
[Figure]
H=
[Figure]
H=
[Figure]
H= (a) Value Spectrum vs Horizon (H)
[Figure]
[Figure] [Figure]
Higher frequency component noticeably grows in decibel when the unroll horizon increases.
The Bellman optimality operator imprints a non-trivial spectrum to the resulting Q i as it is applied recursively during ﬁtted Q iteration.
We present the evolving spectra marginalized over M in Figure 3.
As the effective horizon increases, the value function attains a larger weight in the higher-frequency part of the spectrum, which corresponds to less correlation between the value at nearby states.
In a second experiment, we ﬁx the horizon to 200 while increasing the discount factor from 0 .
1 all the way up to 0 .
We observe a similar increase in the higher frequency components, at a longer effective recursion depth.
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 iteration 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
[Equation] [Equation]
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 5 induces a contraction on Q θ in the sup norm, with ﬁxedpoint 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 Lregression 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 work (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 in the input space due to such bias, causing instability.
O VERCOMING THE S PECTRAL -B IAS OF N EURAL V ALUE A PPROXIMATION
Our basic idea is the following: To correct the spectral-bias of the perceptron network, we can constructing 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 kernel.
The mixing brings high-frequency input signals down to a lower and more acceptable band for the perceptron 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
[Equation]
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
0 d B b=64 b= b=4 b=
[Figure]
b=1 b= b=16 b=
[Figure]
Algorithm Learned Fourier Features (
[Figure]
) class LFF [Figure] [Figure] [Figure]
): super()
[Figure] [Figure] [Figure] [Figure]
weight
[Figure]
=
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
[Figure]
np [Figure] [Figure] 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
[Equation]
d
[Equation]
i
[Equation]
w i,j x i + b j ) where w i,j ∼ N (0 , πb/d ) and b j ∼ U
− π, π ) .
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 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).
Improved NTK produce faster convergence The improved composite kernel leads to faster convergence by having a ﬂat kernel spectrum all the way up to the band-limit.
we can compute the kernel spectrum numerically and inspect the improvement in the higher frequencies (see Figure 5).
By tuning the band-limit, we can make explicit bias-variance trade-offs and control the way the network generalizes.
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 (see Figure 2e).
The gain is even more prominent on more challenging environments in the results section (Quadruped run, see Figure 8).
Optimization overhead is a key bottleneck in off-policy learning, therefore faster convergence also translates into better sample efﬁciency, and faster wallclock time.
[Figure]
Improving Off-policy Stability
“cross-talks” between the gradients of near-by states is key factor in offpolicy divergence.
Such “cross-talk” 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 5), offering direct control over the bias-variance trade-off.
The improved learning stability allows us to remove the target network on a number of domains while retaining substantial performance.
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 4-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 uselatent neurons and are optimized for 2000 epochs.
[Figure]
E XPERIMENTS
Setup We scale the use of FFN to high-dimensional continuous control tasks in state DMC domains (Tassa et al., 2020).
We use soft actor critic (SAC) (Haarnoja et al., 2018) as our RL algorithm.
We build off the pytorch SAC codebase (Yarats & Kostrikov, 2020).
We derive the FFN from a MLP by replacing its ﬁrst layer with a LFF layer.
To use FFN with SAC, we replace the 3-layer MLP in the actor and the critic with a 3-layer FFN.
Results We present the result with SAC on 4 DMC domains in Figure 10.
It shows that FFN improve over MLP in complex environments while matching its performance in simpler environments.
We leave the complete results with SAC and DDPG in Appendix A.3.
We also present ablations and hyperparameter sensitivity analysis in Appendix A.4.
As described in section 5, the beneﬁt of FFN in RL comes from its ability to estimate the value function of the task with fewer gradient updates, thereby improving the sample efﬁciency of the RL algorithm.
We empirically verify this observation in Figure 11 where FFN reduces the value estimation error with more sample efﬁciency than a MLP.
We leave the complete results on 8 DMC domains in the Appendix A.3.
6. F ASTER C ONVERGENCE VIA F OURIER F EATURE N ETWORKS
We hypothesize FFN achieves faster convergence in neural value approximation due to its better initialization obtained by sampling from a wide space of frequency and phase.
If this is indeed the case, then FFN’s parameters should undergo less change from the training of an RL agent compared to MLP’s parameters.
Let W t and B t refer to the concatenated weight matrices (in vector form) and the concatenated bias vectors respectively, from different layers of critic of an RL agent after training on t environment frames (steps).
Figure 12 shows evolution of (cid:107)W t − W 0 (cid:107) and (cid:107)B t − B 0 (cid:107) for MLP and FFN during training of SAC agents for Walker-run and Quadruped-run .
We leave the complete results on 8 DMC domains in the Appendix A.3.
Since FFN estimates value function with fewer gradient updates, can we reduce the update frequency during training of an RL agent while still matching the performance of MLP?
Figure 8 shows that we can reduce the update frequency by a factor 4 and 6 in Walker-walk and Quadruped-run using FFN
[Figure] [Figure]
[Figure] [Figure]
1.0 Frames 1e
V a l u e e s ti m a ti on e rr o r
Quadruped-run
MLPFFN
FFN reduces the variance in these regimes by making more appropriate bias-variance trade-off while matching the performance of MLP.
Thus, FFN can help achieve compute efﬁciency in RL.
Conversely, can we improve the performance of MLP by increasing the update frequency?
Figureshows that increasing the update frequency can help MLP get closer to FFN in terms of performance in Walker-walk .
However, this isn’t reliable as it can also lead to overﬁtting in actor and hurt MLP’s performance, as seen in Quadruped-run .
6. F OURIER F EATURES I MPROVE O FF -P OLICY S TABILITY
We now analyze if use of FFN removes the need for a target network.
Target (value) network (Mnih et al., 2013) is a past instantiation of value function which gets updated at a slower rate and is used for calculating targets during bellman update.
Given MLPs often take many gradient updates to ﬁt to a target, the slow changing target network stabilizes the training of the value function.
However, since FFN requires fewer gradient updates to ﬁt to a target, we hypothesize that there will be less of a need for a target network.
Figure 13 shows performance of both MLP and FFN without a target network.
While MLP completely fails in all environments, FFN either matches its performance with target networks or suffers small degradation in most environments except Hopper-hop .
We leave the complete results on 8 DMC domains in the Appendix A.3.
I MPROVING THE C ONVOLUTION K ERNEL
In a standard convolution neural network, the ﬁlter in the ﬁrst layer suffers the same spectral bias in the RGB color-space as the state and action space above.
We can in principle extend our spectralbias ﬁx to convolution networks by replacing the ﬁrst layer with a 1 × 1 convolutional layer and a sinusoidal non-linearity.
In this experiment, we use the same initialization scheme as those described in Equation 14.
Our experiment is conducted with the DrQv2 implementation (Yarats et al., 2021), where we replace the CNN in both the actor and the critic with this Fourier-CNN (F-CNN) architecture.
We present the results in Figure 14.
While F-CNN’s performance is within the variance of
CNN’s performance, its variance is much lower.
Such technique has also been used by Kingma et al.
(2021) to improve image generation with diffusion models.
For more detailed theoretical overview 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
[Figure] [Figure]
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 it.
Function 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 can managed to reduce the walltime by 32% on the challenging Quadruped environment, by reducing the replay ratio to 1 / 6 .
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 perceptron networks are on the biasedside 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 model could 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.
[Figure]
the conﬁdence range of the vanilla CNN.
A CKNOWLEDGMENTS
The authors want to thank 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 Army Research Ofﬁce under Grant W911NF-21-1-0328; the National Science Foundation Institute for Artiﬁcial Intelligence and Fundamental Interactions (IAIFI, https://iaiﬁ.org/) under the Cooperative Agreement PHY-2019786; 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.
A A PPENDIX A. C OMPARING 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
[Equation]
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.
All the above Fourier feature variants are equivalent and should produce similar results when bandwidth is tuned properly.
A. T OY MDP E XPERIMENT D ETAILS
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. R ESULTS ON DMC DOMAINS
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.
Results with SAC and DDPG We scale up use of FFN to 8 DMC domains by replacing the MLP in the actor and the critic of a SAC agent with FFN.
Figure 15 shows learning curve of FFN and
MLP with SAC as the RL algorithm on 8 DMC domains.
To show that FFN also works with other
RL algorithms, we replace the MLP in the actor and the critic of a DDPG agent with FFN.
Figureshow learning curve of FFN and MLP with DDPG as the RL algorithm on 8 DMC domains.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Value approximation error FFN improves over MLP both in sample efﬁciency and the overall performance because of its ability to ﬁt value functions with less gradient updates.
Figure 17 shows that FFN reduces value approximation error with more sample efﬁciency than MLP.
Measuring weight and bias change in FFN and MLP We hypothesize that FFN ﬁts value functions with more sample efﬁciency than MLP because of its better initialization due to sampling from a wide space of frequency and phase.
If this is indeed the case, the weights and biases of FFN should undergo less change than that of an MLP during training of an RL agent.
Let W t refer to the concatenated weight matrices (in vector form) from different layers of critic of an RL agent after training on t environment frames (steps).
Similarly, let B t refer to the concatenated bias vectors from different layers of critic of an RL agent after training on t environment frames.
Figure 18 shows evolution of (cid:107)W t − W 0 (cid:107) for MLP and FFN during training of SAC agents.
Similarly, Figure 19 shows evolution of (cid:107)B t − B 0 (cid:107) for MLP and FFN during training of SAC agents.
Removing the target network Since FFN is able to ﬁt value function with less gradient updates, there is less of a need for a target network.
Figure 20 shows performance of both MLP and FFN without a target network.
While MLP completely fails in all environments, FFN either matches its performance with target networks or suffers small degradation in most environments.
However, there are some environments (Hopper-hop, Humanoid-run, Acrobat-swingup) where FFN completely fails without target network.
Architectural Details for state DMC domains For MLP baseline, we parameterize the actor and the critic with 3 layer MLP.
For our method, we replace the 3 layer MLP in the actor and the critic with a 3 layer FFN.
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 d = observation dimension for actor and d = observation dimension + actor dimension for critic.
[Figure]
[Figure]
V a l u e e s ti m a ti on e rr o r
Quadruped-run
MLPFFN
Architectural Details for image DMC domains For CNN baseline, we parameterize the actor and the critic with CNN model borrowed from DrQv2 (Yarats et al., 2021).
For our method, we derive conv FFN from the CNN model (taken from DrQv2) by replacing the ﬁrst convolutional layer with a 1x1 convolutional layer having sin activation and using the initialization scheme described in
Equation 14.
A. A BLATIONS AND H YPERPARAMETER SENSITIVITY ANALYSIS
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 21 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.
Sensitivity to (Fourier) feature-input ratio In addition to the bandwidth b , we need to choose the
Fourier dimension for FFN.
We maintained the feature-input ratio (i.e. Dd ) to be 40 .
But can we get away with a lower feature-input ratio?
Figure 22 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.
FFN on actor vs FFN on critic In our experiments, we used FFN with both actor and critic.
We analyze how removing FFN from either actor or critic affects the performance of RL agent.
Figure
[Figure] [Figure]
shows that using FFN only critic doesn’t affect the performance of the RL agent.
However, using
FFN with only actor brings down the performance of the RL agent to that using MLP for both actor and critic.
[Figure] [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
[Figure] [Figure]