Loss Dynamics of Temporal Difference Reinforcement Learning
Authors: Blake Bordelon
Paul Masset
Henry Kuo
Cengiz Pehlevan
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement Learning
Statistical Mechanics
Stochastic Gradient Descent
TL;DR: A mean field theory of generalization for temporal difference learning dynamics

Abstract: Reinforcement learning has been successful across several applications in which agents have to learn to act in environments with sparse feedback. However, despite this empirical success there is still a lack of theoretical understanding of how the parameters of reinforcement learning models and the features used to represent states interact to control the dynamics of learning. In this work, we use concepts from statistical physics, to study the typical case learning curves for temporal difference learning of a value function with linear function approximators. Our theory is derived under a Gaussian equivalence hypothesis where averages over the random trajectories are replaced with temporally correlated Gaussian feature averages and we validate our assumptions on small scale Markov Decision Processes. We find that the stochastic semi-gradient noise due to subsampling the space of possible episodes leads to significant plateaus in the value error, unlike in traditional gradient descent dynamics. We study how learning dynamics and plateaus depend on feature structure, learning rate, discount factor, and reward function. We then analyze how strategies like learning rate annealing and reward shaping can favorably alter learning dynamics and plateaus. To conclude, our work introduces new tools to open a new direction towards developing a theory of learning dynamics in reinforcement learning.

Loss Dynamics of Temporal Difference Reinforcement Learning
Blake Bordelon, Paul Masset, Henry Kuo & Cengiz Pehlevan
John Paulson School of Engineering and Applied Sciences,
Center for Brain Science
Kempner Institute for the Study of Natural & Artificial Intelligence, Harvard University
Cambridge MA, 02138
blake_bordelon@g harvard edu, cpehlevan@g harvard, edu
Abstract
Reinforcement learning has been successful across several applications in which agents have to learn to act in environments with sparse feedback. However, despite this empirical success there still a lack of theoretical understanding of how the parameters of reinforcement learning models and the features used to represent states interact to control the dynamics of learning. In this work. use concepts from statistical physics, to study the typical case learning curves for temporal difference learning of value function with linear function approximators. Our theory is derived under a Gaussian equivalence hypothesis where averages over the random trajectories are replaced with temporally correlated Gaussian feature averages and we validate our assumptions on small scale Markov Decision Pro- cesses. We find that the stochastic semi-grade ceni-gradient noise due to subsampling the space of possible episodes leads to significant plateaus in the value error, unlike in traditional gradient descent dynamics. We study how learning dynamics and plateaus depend on feature structure, learning rate, discount factor, and reward function. We then analyze how strategies like learning rate annealing and reward shaping favorably alter learning dynamics and plateaus. To conclude, work introduces new tools to open new direction towards developing a theory learning dynamics in reinforcement learning.
Introduction
Reinforcement learning (RL) is general paradigm which allows agents to learn from experience the relative value of states in their environment and to take actions that maximize long term rewards [1]. RL algorithms have been successfully applied in number of real world scenarios such as strategic games like backgammon and Go, autonomous vehicles, and fine tuning language models [2-7].
Despite these empirical successes, theoretical understanding of the learning dynamics and inductive biases of RL algorithms is currently lacking [8]. A large fraction of the theoretical work has focused on proving convergence and deriving bounds both in the asymptotic [9-14] and asymmottio [15-17] limits, but do not provide full picture of the evolution the learning dynamics.
A desired feature of a candidate theory is to characterize the influence of function approximation RI dynamics and its performance. Early versions of RI operated a tabular setting, similar to dynamic programming [18]. where all the states in the environment could be mapped one-to-one to specific value and policy. In large and complex environments, it is not possible to enumerate all the states in the environment necessitating the use of function approximation for the target value and policy functions. Indeed, the recent success of many RL algorithms relies on deep reinforcement
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

learning architectures that combine an RL architecture with deep neural networks to build effective value estimators and policy networks [19].
One difficulty in analysing these algorithms compared to supervised learning settings is that the distribution of the data received at each time-step not stationary. This non-stationarity arises from two principal sources: First, whether in an episodic or continuous setting, states visited within a learning trajectory are dependent on the recent past. Trajectories might be randomly sampled but points within a trajectory are correlated. Second, when the policy is updated it also changes the distribution of future visited states.
Here, we will focus on the first form of non- non-stationarity when learning a value function in the context of policy evaluation [1] using classical RL algorithm, temporal difference (TD) learning [20]. We develop a theory of learning dynamics for RI. in this setting in a high dimensional asymptotic limit with focus on understanding the role of linear function approximation from a set of nonlinear and static features. In particular, we leverage ideas from recent work in application of statistical physics machine learning theory to perform an average over the possible sequences of features encountered during learning. Our contributions are as follows:
We introduce concepts from statistical physics, including path integral approach to describe dynamics [21-25] and the Gaussian equivalence assumption [26-29], to derive a theory of learning dynamics in TD learning ($3) in online setting. We provide analytical formula for the typical case learning curve for TD learning.
We show that our theory predicts scaling of the learning convergence speed and performance plateaus with parameters of the problem including task-feature alignment [30], learning rate, discount factor or batch size ($4 and $5). Task-feature alignment is a metric that quantifies how features allow fast or slow learning for given task.
We show our theory can be used to understand and guide design principles when choosing meta- parameters. Specifically, we show that we can use our theory to infer optimal schedules of learning rate annealing and the effects of reward shaping ($5 and $6).
2
Problem Setup and Related Works
2.1 Problem Setup
We consider a set of states denoted by s. possibly continuous, and a fixed policy ι which generates a distribution over actions given the state. The state dynamics are defined by a distribution p(r) over trajectories through state space T = {si, $2..... ST]. Note that state transitions do not have to be Markovian, but each trajectory is i.i.d. sampled from P(T). We consider trajectories of length T. Each state is represented by an N-dimensional feature vector (b()) E R-V, so that trajectory generates a collection of feature vectors {vb(>,))?=1 The rewards are generated by a reward function R(s) which depends on the state. (In general, the features and rewards can depend on action as well: transition dynamics are still fixed the policy is fixed, but variance over rewards at given state may need to be modeled, see Appendix B.5).
At any time, we are interested in characterizing the value function associated with a state, which measures the expected discounted sum of future rewards when starting in state Sg
Vsso) R(so) +
Bs R(s:) = R(so) + 71E,, 1s0 V V(1)..
We use linear function approximation to learn the value function ((() = (b)) w. Similar to kernel learning [31], the features vb should be high dimensional so that they can express large set of possible value functions.
We study TD learning dynamics given this setup. At each step of the TD iteration, we sample a batch of B independent trajectories from the distribution and compute the TD update
Wn+1 Wn +
1/ TB
(()))((((())
(!(()= = (
(2)
2

---

We operate in online batch regime as the trajectories in each batch are resampled at each iteration. This is distinct from an offline setting where the batches would be resampled from a finite-sized buffer [1]. Convergence considerations for infinite-batch online TD learning width different types of features ab are outlined in Appendix A. The specific form for the TD-error (t) depends on the precise variant of TD learning that is used. Here, we will focus on TD(O) but our approach can be extended to other TD learning rules and definitions of the return function. We see that the iterates Wn, will form a stochastic process as each sequence of states in episode (t)] are drawn randomly from P(T). In general, we allow the learning rate The to depend on iteration, an important point we will revisit later. The distribution of features { {ab(s) (t))} over random trajectories is in general quite complicated, depending on the details of the state transitions and the nonlinear feature maps, which motivates the following question:
Question: How can the stochastic dynamies of temporal difference learning be characterized for complicated trajectory distributions p(T) and feature maps {s()?
To address this question, in this work, we provide an analysis of TD learning that explicitly models the statistics of stochastic cemi-gradient updates to Wn, Our framework is based on a Gaussian equiv- alence ansatz for TD learning and high dimensional mean field theory which predicts the statistics of TD errors (t) and the weight iterates wn- The theory reveals a rich set of phenomena including plateaus unique to SGD noise in TD learning which can be ameliorated with learning rate annealing.
2.2 Related Works
The dynamics of TT learning have been notoriously difficult to analyse. Unlike supervised learning settings, sampled states are correlated across a trajectory and the algorithms involve bootstrapping: using estimates of the value function for future states in the temporal difference update [1]. Some prior works study the least-square TD learning rule, which solves, at each step n of the algorithm, a linear system for the instantaneous best fit n samples [32-34]. Alternatively, many works focus on the on- line SGD version of TD learning, where incremental updates are made to the parameters at each step, using fresh samples. This is the setting of our work. The focus of this literature has initially been to prove convergence and bounds on asymptotic behavior [1 1-14, 35]. More recently, progress has been made in deriving bounds in the non-asymptotic regime. Initial work assumed that data samples were ii.dd [15-17,36 36] and recent work has extended those approaches to Markovian noise [15, 37-39]. The majority of these proofs use the ODE-like method for stochastic approximation [11, 40], which cor- responds to limit of the stochastic semi-gu aomi-gradient dynamics where the effects of mini batch noise are neglected. This is also known as the "mean-path dynamics of TD learning and will correspond to the infinite batch limit of our theory. Furthermore, many of these methods require the use of iterative aver- aging of the learned value function, whereas we study the final iterate convergence. The approach we take here differs from many of these results as our goal is not to provide bounds on worst-case behavior but instead to provide a full description of the dynamics of the typical case scenario during learning.
Our approach also highlights the importance of the structure of the representations in controlling the dynamics of learning. This had been long been recognized in reinforcement learning and previous works proposed to improve feature representations to improve algorithmic performance [41-43] This line of work has shown the importance of the relative smoothness of the representations and target functions in the ODE limit of TD dynamics [43, 44]. Similarly, several methods have been proposed empirically learn a better shaping function [45, 46). In policy learning has also been recognized that using a gradient aligned to the statistics of the tasks, such as the natural gradient [47] can greatly speed up convergence [48]. Our work does not explore such feature learning per se but could be used as a diagnostic tool to analyse how representations impact learning speed.
We adopt the perspective of statistical physics, by working with a simplified feature distribution which captures the learning dynamics and solving the theory in a high-dimensional limit [49-51]. We derive TD reinforcement learning curves from a mean field theory formalism which exact for infinite dimensional features and batch size. Similar calculations for supervised learning on Gaussian data have been shown to provide an accurate description of high dimensional dynamies [52-54]. Further, even when data is not actually Gaussian, several algorithms, such as kernel or mandom-feature regression, exhibit universality in their loss behavior, enabling analysis of the learning curve with a simpler Gaussian proxy [26-28, 30, 55). We exploit this idea in the TD learning setting to some success. We note that Gaussian equivalence or universality is not a panacea, and in many cases the Gaussian proxy can fail to capture important machine learning phenomena [27, 56, 57].
3

---

3 Theoretical Results for Online TD Learning
3.1 Computation of Learning Curves
We develop dynamical mean field theory (DMFT) formalism can be utilized to compute the learning curves. We provide the full derivation of the DMFT in Appendix B. This computation consists of tracking the moment generating function for the iterates w,, over the trajectories of randomly sampled features (4)( (t) In an appropriate high dimensional asymptotic limit, the results of our theory can be summarized as the following proposition.
Proposition 3.1. Let N, B + oo with B/N = =1) and episode length = O(1), ). Let the ground truth reward function be R(s) wi (s) and value function V(s) WTD *b(5) in the basis of our features. Define matrices
1 - E+ = ΣE(L-T)) A= - -
Σ
(3)
and assume that the features are such that matrix A of extensive rank N. Then the typical value estimation error L.. (((())-Vu(9) 2)
after steps has the form
Cn N
TrEMa;
(4)
Mat1 (1 - MnA)M -I- -
a2T2 [Q(((t))((t))
(5)
1 Qn(1,11) N (wwn w.n) EE(,,'") wm)) + ? (wwn -w,)² w.aTTI(t, [(,, + 1)w,)
N
+ 2 N (w,, El + 1,t')(WR wn)) + 22 (w, E(t 1,t' 1)w,),
(6)
where a B/N and at episodic times t, t' and iteration n. The average over weights () denotes Gaussian average whose moments are related M,, The correlation function Qn(t,t') depends on M,, and the average weights (w,)); we provide its full formula in Appendix B.3, equation (B.17).
is the correlation of randomly sampled TD-errors
Proof. The full derivation is in Appendix B. At high level, we track the moment gen- erating function of the iterates wn over random draws of features (t), Z((ia)] = (() exp (i En In Wn) x f exp ( N S(g, (ix where S is O(1) action and q are a set of order parameters of the theory which include the following overlaps (t,t')) = w ('t,t')w, and (2) B E ((()) In this high dimension N, B - do limit with B/N = 0(1) and episode length T = O1), the order parameters can be obtained from saddle point integration, which requires solving as = 0. This procedure results in a deterministic learning curve given in By
equations (4))().(6) even though the realization of sampled states are disordered. The TD-error variables (t) become mean zero Gaussians and the {w,, also follow Gaussian distribution with mean and variance determined by the order parameters.
Before we explore the predictions of this theory, we first make a few remarks about this result. Remark 1. Though the theory is technically derived for large batch size B, we will show that provides accurate description of the loss trajectory even for batches as small as B 1. An alternative formulation in terms of recursive averaging reveals transparently which approximations lead to the same result the mean field theory (Appendix B.6).
Remark 2. The case where the reward function and/or the value function are inexpressible by the features 2a can also be handled within this framework. In this case, the unlearnable components of the value function act as additional noise which limits performance [29]. These can also be handled by our theory, see Appendix A.
Remark 3.
The limit where 7 = 0 recovers known results in online supervised learning with stochastic gradient methods [29, 58, 59]. In this limit, the dynamics will converge to zero loss provided the model features are sufficiently rich to represent the true value function.

---

Remark 4. The TD learnr with perfect coverage (infinite batch size) at each step will converge to the ground truth WTD ( ~ ERR (see Appendix A). Remark 5. Mn is equivalently defined M, ((w WTD)(W WTD) (T (wn) which measures deviation from the fixed point of gradient flow (vanishing learning rate) dynamics wTD over random sets of sampled episodes (Appendix B).
3.2 Gaussian Approximation
The theory presented in Section 3.1 relies on approximation of the feature distribution Gaussian. Similar approximations have been successfully utilized in high dimensional regression problems even when the true features are non- Gaussian [26-29]. We note that exact, non-asymptotic theory for non-Gaussian features can be provided which closes under knowledge of the fourth cumulants of the features we show in Appendix though this theory especially cumbersome to analyze or evaluate compared to the theory of Section 3.1. Concretely, Proposition 3. relies on the following.
Gaussian Feature Assumption. The learning curves for a TD learner with high dimensional features (sb(8))) over random T are well approximated by the learning curves of a TD learner trained with Gaussian features DG NN, E pm )) with matching mean and correlations
p(() (1b(
E(L, = (ab(s)) "mppss)
where averages are taken over sequences of states {s(t)} ~ p(T).
Example Trajectories
Value Errar
20-1
20.2
16-1
107
15F
Steps
(a) 2D Exploration
(b) Place Cell/RBB Features (8())
Equivalence Phenomenon
1.0
3.25
1.0
Value Error
Value Error
0.6.
10-4
30-1
0.4
BW 025
BW 0.25
597 0.5
2A 0.5
517 0.75
3W 0.75
8W 1.0
3W 1.0
16P
10
15
101
12
Steps
Steps
(d) Place Cell Bandwidths
(e) Large batch
30
Small batch
Figure 1: An illustration of our theory for TD learning. (a) A diffusion process 2D grid world generates many possible trajectories through state space. Each colored line is different trajectory. Reward function shown red, with darker red indicating higher reward. (b) When combined with nonlinear place cell feature representation, the state transitions generate distribution over observed features (*)(s:)) (c) The value eтror associated with TD learning for bump reward function on the true features generated from a single set of MDP trajectories (blue) is compared to training on sampled Gaussian vectors (1/.) with matching within-episode covariance structure. These single runs of TD learning on either set of features are consistent with the typical case theory (black dashed). (d) The structure the features alters learning dynamics. We consider, for simplicity, altering the bandwidth (BW) of the place cell features. (e) Varying place cell BW changes the dynamics for both large batch (B- = 30) and (f) small batch (B 3) TD learning. There is an optimal BW for a given step size. Small batch stochastic semi-gradien noise is more severe.
One interpretation of this ansatz is that the dependence of the learning curve on higher order cumulants of the features is negligible in high dimensional feature spaces under the square loss. This
5

---

approximation has been shown to provide an accurate description on realistic supervised learning settings with non-Gaussian data with the square loss in prior works [26, 27, 29. 30, 55, 58]. As shown in these works, for standard supervised learning, even highly non- non-Gaussian features (st)) } have least squares learning curves which are only sensitive to the first two cumulants of the distribution. We do not aim to provide rigorous proof of this ansatz for TD learning but instead compute the learning curve implied by this assumption and compare to experiments on simple Markov Decision Processes (MDPs). The benefit of this hypothesis in the RL setting is that it abstracts away details of transitions in the state space and instead deals with the correlations of sampled features through time.
To illustrate an example of the Gaussian Equivalence idea, in Figure 1, we consider an MDP which is defined by diffusion through 2-dimensional (2D) state space (Figure ((a)). We choose the features (s)) to be collection of localized 2D Radial Basis Function (RBF) bumps which tile the 2D space, similarly to the "place cell"" neurons found in the mammalian hippocampus [60, 61] (Figure ((b).. The feature map is parameterized by the bandwidth of individual "place cells". In Figure 1(c), we show the value error learning curve as a function of the number of steps n (blue) and compare the value estimation error of the MDP with a Gaussian distribution for (t) with matching first and second moments (orange). Lastly, we plot the theoretical prediction of our theory (described in Section 3), which is computed under the Gaussian equivalence ansatz (black dashed). We see remarkable match of the three curves. The equivalence can be used to predict the speed of TD learning for different features, such place cells with varying bandwidth (e) and (f), we plot the loss trajectories for a single run of TD for each feature set. We observe that bandwidth affects both the learning dynamics and the asymptotic eπror with an optimal bandwidth at any step. One of our goals will be to elucidate the role of feature quality in learning dynamics. While the large batch dynamics
illustrated in Figure I (d)-(f). In Figure I
approximately self-averaging, as shown by the fact that single runs of TD learning coincide with our theoretical typical case theory curves, there is significant semi-gradient variance in the value error at small batch sizes. While we expect Gaussian equivalence to hold for high dimensional features, in low dimensions non-Gaussian effects can significantly alter the learning curves as we show in Appendix D.I. However, for high dimensional features, the equivalence holds for many other feature distributions such as polynomial and fourier features (Appendix E).
4 Spectral Perspective on Hard Reward Functions
(s()
R2(5)
0.8
0.6
0.4
Value
20
6.2
Steps
(a) Sparse R1 or Dense R2 Rewards
(b) Code-TTaee Task Alignment
(c) B 20 Learning Curves
Figure 2: Reward functions and dynamics which lead to value functions with high spectral alignment to the features can be learned more quickly than those that do not. (a) A sparse and dense reward function in 2D spatial navigation task can illustrate this effect. (b) The cumulative power distribution C(k) defined from the spectral decomposition of A Σ E+. Concretely we let Au = λ uy with D ordered by real part and WTD= = Σx Ug. In the B > do limit the task which has rapidly rising C(k) Eist w²
will converge more quickly than the task with slowly rising C(k). (c) Indeed, for large batch regime (B- 20) the value error decreases more rapidly for R2 than for Ri.
Σg wi
Our theory can provide some insights into the structure of tasks which can be learned easily and which require more sampled trajectories to estimate based on spectral decompositions of the feature covariances. We note that similar spectral arguments have been given in the ODE-limit [44] and are intimately related to the source conditions used in recent work to identify power-law rates in the large batch regime [39].
To build our argument, we diagonalize the matrix A Σ --7E4, obtaining Au = u noting that cigenvalues X can be complex. We then expand the TD solution in this basis WTD = En WkMk

---

The theory predicts that, the average learned weights will be (Wn) where 1- is complex modulus and 0g Arg(1 g7a)). We can therefore order the modes by their convergence timescales 1 - d Given this ordering of timescales, we can order the modes from those with smallest to largest timescales. Given this ordering, we see that tasks can be learned efficiently are those with most of the norm of wy in the modes with small timescales. We quantify how well aligned a task is to given feature representation by computing a cumulative power distribution for the target weights (kk)
Ex dd """"""" ua
Erer wi Ee wi If this quantity rises rapidly with k then the task
can be learned from small number of samples [30].
We consider again, the setting of Figure 1, the 2D exploration MDP but now contrast two different reward functions. In Figure 2 we show that this spectral decomposition can account for the gaps in loss for a place cell code in learning sparse or dense reward function (Figure 2(a)). As expected the cumulative power rises more rapidly for the dense reward function (2(s) (Figure 2(b)). As a consequence, the value error converges to zero more rapidly than for the sparse rewards.
5 Stochastic Semi-Gradient Learning Plateaus and Annealing Strategies
30"
10
10"
4 0.00
Value Errar
0.02
10-2
10-
Value
0-0.1
10
20
0.15
30-2
0.5
10
10
p-0.
16-
15
10
107
107
101
steps
steps
steps
(a) Episodic Batches
(b) Discount Factor
(c) Fixed Learning Rate
10a
50
Value Error
10-1
Value Error
13-2
22-1
2: 5.00
21 0:22
10-3
31 8.75
1.25
13-3
102
107
steps
(d) Annealed Learning Rate
(e) Optimal Annealing
Figure 3: Finite batch size, discount factor and learning rate all contribute to a stochastic semi- gradient plateau in the TD dynamics. The features are generated from a synthetic power law covariance with exponential temporal autocorrelation (see Appendix G). Dashed black lines are theory. In general, for fixed learning rate 7), the plateau scales as D(n²BB-) (a) Larger batch sizes B reduce SGD noise and leads to lower plateau the reducible value error for decoupled power-law feature model.
(b) Larger discount factor and (c) larger learning rate 1) lead to higher SGD plateau floor. (d) An annealing strategy Th ~ T)o X for X > 0 can allow one to avoid the plateau. For slow annealing (small x). the error scales as Cn function of the learning rate annealing exponent x defined by 7h2 = ron -x, For this task, the optimal exponent balances the scale of the asymptote with the rate of convergence.
O(n x). (e) The value error
The stochastic noise from TD learning has striking qualitative differences from SGD noise in the standard supervised case. In standard supervised learning (such as y = 0 version of this theory), the stochastic gradient noise does not prevent the model from fitting the target function with zero error provided the features are sufficiently rich to represent the target function. However, this is not the case in TD learning, where the predicted value ((ss) is bootstrapped using the model's 's weights w,, at each iteration n. This leads to asymptotic plateaus in learning curves. Our theory can predict these plateaus and their scaling whose proof is given in Appendix B.7.
Proposition 5.1. Our theoretical learning curves exhibit a fixed point for the value error dynamics for finite B and non-zero 7 and 7. For small n² B, we deduce that M satisfies self-consistent
7

---

Value Enoo
Value Enor
(1 Dr
Wyr (os(e) P)sin(0)
20%
0.3
16-1
104
361
Steps
Steps
(a) Geometry Reward Shaping
(b) Scale-based Shaping
(c) Rotation- bassod Shaping
Figure The theory can be used to understand how reward shaping decisions alter temporal difference learning dynamics. (a) A visualization of possible reward shaping potentials Ø(s) Wa (b(s) strategies in feature space. Probability density level curves for the features are depicted in blue. Reshaping with w2 = BWTD for scale factor 3 merely changes the scale of weights which must be recovered (gold) and does not change timescales of dynamics. (b) The value error dynamies for the scale based reward shaping for the features in Figure 3. On the other hand, rotation based reward shaping where wy is not parallel to wy (red) leads to potentially helpful mixture of timescales the new target vector is more aligned with feature dimensions with high variance (purple). In (c), plot loss curves for rotation angle & between the original mode wy and the top eigenvector of the feature covariance matrix Σ. Dashed black lines are theory.
10T" impliying asymptotic value error scaling of asymptotic scaling of the form M O (mn ~~ TTME m
In Figure 3, we demonstrate that our theory predicts the plateaus and their scaling as function of finite batch size B (Figure 3(a)), non-zero discount factor y >> 0 (Figure 3(b)) and non-negligible learning rate (Figure 3(c)).
A strategy used in the literature to increase rates of convergence and improve asymptotic behavior is adaptation of the learning learning through an annealing schedule [1, 16, 62, 63]. To overcome this plateau in the loss, we consider annealing the learning rate T/2 with iteration n. In Figure 3(d), we show the effect of annealing the learning rate as a power law Tm = 'Jun" X for some non-negative exponent X. For x 0 the learning rate is constant and fixed plateau is reached. For small nonzero X. such as X 0.2, the value error is, after an initial transient, always near its instantaneous fixed point plateau so the loss scales linearly with the learning rate, giving the asymptotic rate En ~ O(n ). For large X. the learning rate decreases very quickly and the plateau is never reached. Our approach can be used to find an optimal annealing exponent X and in Figure 3(c), show that the optimal annealing exponent balances these effects and is well predicted by our theory.
6 Reward Shaping
Another strategy to improve the learning dynamics in reinforcement learning algorithms is reward shaping [64]. standard supervised learning, the goal is to directly approximate the target objective given cost function. However, in reinforcement learning, the objective is not to estimate rewards at each state directly but the discounted sum future rewards, the value function. Importantly, many different reward schedules can lead to identical value functions. Reward shaping exploits this symmetry to speed up learning by altering the structure of TD updates and SGD noise. Here, we provide theoretical description of the changes in the learning dynamics due to reward shaping which suggests they can be understood through change of the alignment between the original rewards and the reshaped rewards in the space of the features used represent the states.
The original ideas around reward shaping were inspired by work experimental psychology and were closer to what is now studied as curriculum learning [65-67]. Reward shaping as currently used reinforcement learning directly changes the reward function by adding potential-based shaping function F' such that F(st, , St+1) 70((++1) ((s.) [64]. In each step of the algorithm we feed

---

the following reshaped rewards R to the TD learner
[[(s)) 70($(+1)
R(st) [(ss) - ((s)) - (s++1) t>
We note that this transformation simply offsets the target value function by ((s) as the series above telescopes with a cancellation of o(s,) between the t - and t-th terms [64] (see Appendix C). However, the dynamics of TD learning with these reshaped rewards is quite distinct from the dynamics with original rewards R. Here, we study the case where we can express ((s) linear function of our features: ((s) = (b(8) We. This leads to a change in the dynamics for M,. and (w,, that we describe in the Appendix C.
In Figure 4, we illustrate the possible benefits reward shaping. We explore two types of reward shaping. First, scale based reward shaping where w is parallel to the target TD weights WTD- This merely changes the overall scale of the weights needed to converge in the dynamics, leading to similar timescales and identical plateau for TD learning
show in Figure 4 (b). On the other hand, reward shaping which rotates the fixed point the TD dynamics into directions of higher feature variance can improve timescales of convergence. In Figure (c), show example where we vary the angle 0 of the shaped-TD fixed point (see also Appendix C).
7 TD Learning Plateaus in More Realistic Settings
In this section, test if some of the phenomena observed in theory and experiments also hold in more realistic settings. We perform TD learning with Fourier features to evaluate pre-trained policy on MountainCar-vol vO. As expected, we see that the value error plateaus to error level determined by both the learning rate (Figure 5a) and batch size (Figure 5b) due to semigradient noise.
18'
30/
B = 0=1
==2
Value Error
14
Value Error
300
B =
B
n= 0.01
20
7= 0.02
n= 0.05
0.1
140
10
0.2
Batch
Batch
(a) Varying Learning Rates
(b) Varying Batch Sizes
Figure 5: Policy evaluation in MountainCar-wi environment. The policy was learned with tabular e-greedy Q-learning (see Appendix F for details). (a) Value error curves for different 7) when B 1. (b) Value error curves for different B with = 0.1. Shaded area denotes 95% confidence interval over 10 seeds.
We show that the plateaus obey the predicted scalings of On)B'l) in Appendix F.
Discussion
Our work presents new approach using concepts from statistical physics to derive average-case learning curve for policy evaluation in TD-learnin. However, it is only a first step towards new theory of learning dynamics in reinforcement learning.
One major limitation of the present work is that it concerns linear function approximation where the features representing states/actions are fixed throughout learning. This limit can apply to neural networks the "lazy" regime of training [68, 69], however cannot account for neural networks that adapt their internal representations to the structure of the reward function. This differs from the setting of most practical algorithms, including in deep reinforcement learning, that specifically adapt their representations.

---

Our theory provides a description of learning dynamics through a set of iterative equations (Propo- sition 3.1). In Figure I we evaluate these dynamics for simple MDP but although the predicted dynamics present an excellent fit to the empirical simulations, the iterative equations can be difficult to interpret and computationally expensive to evaluate in a larger network and more realistic tasks. Nevertheless, our equations can be used to derive some scaling between key parameters of the algorithm for example by studying their fixed points as in Proposition 5.1.
Here, we considered the simplest form of temporal difference learning, batched online TD( In future work, it will be important to further characterize the behavior for online TD(0) with batch size B = 1 and to expand our approach to TD(A) and other return distributions. Similarly, expanding our theory to the offline setting, in which the buffer of resampled trajectories would be of finite size, could provide an understanding of how the interactions between parameters govern convergence and divergence [1. 70-72].
Another limitation of our work is that we only considered the setting of policy evaluation with a fixed policy. The goal of RL agent is to learn how to act in the work and not merely represent the value is its states. Unlike in supervised learning, the changes in the value function affect the policy but in many of RL algorithms, for example in actor-critic architecture, there is separation of the policy evaluation (critic) and the policy learning (actor) [73, 74]. Such algorithms estimate the value associated with state/action pairs under a given policy and then use this information to make beneficial updates to the policy, usually with the value and policy functions approximated by separate neural networks. In this paper, we only treated the first part of this process. Recently, a related approach has been used to analyse the dynamics of policy learning in an "RL. perceptron" setup [75]. full theory of reinforcement learning combining policy evaluation and policy learning remains difficult due to the interaction between the two processes, but combining these approaches would be fruitful. One promising direction is in settings where the timescales of the two processes are different [76], such as when policy learning occurring at much slower rate which is often the case in practice.
Beyond developing a theory of learning dynamics in reinforcement learning, the approach could be used in neuroscience to understand how neural representation of space or value can shape the learning dynamics the behavioral level. Ideas from reinforcement learning have been extremely influential to understand phenomena observed in neuroscience and have been mapped directly onto specific brain circuits [77-79]. The place cells of the hippocampus [60] exhibit localized tuning as the example in Figure 1 and together with grid cells in enthorinal cortex are thought to be crucial for navigation in spatial and cognitive spaces and their tuning shaped by experience [61, 79-81]. Our theory specifically link the structure of representations, policy and reward to learning rates, which can all be experimentally measured simultaneously and could shed on light on how the spectral properties of representations govern learning and navigation [79, 82]. similarly to how the mean field theories we have used here can explain learning of sensory features [83]. Future work could straightforwardly extend this DMFT formalism to deal with replay of sampled experiences during TD learning [84] at the cost of tracking correlations of weight updates across iterations of the algorithm [52].
To summarize, our work provide new promising direction towards theory of learning dynamics in reinforcement learning in artificial and biological agents.
Acknowledgments and Disclosure of Funding
BB is supported by a Google PhD Fellowship. CP and BB were supported by NSF grant DMS- 2134157. CP is further supported by NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. PM was supported by NIH grant 5R01DC017311 Uchida. HK was supported by the Harvard College Research Program. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. We thank Jacob Zavatone- Veth for useful discussions and comments on this manuscript.
Venkatesh Murthy and Naoshige
10

---

References
[1] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018.
[2] Gerald Tesauro. Temporal difference learning and td ad-gammon. Communications of the ACM, 38(3):58 -9, 1995.
[3] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
[4] David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam,, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
[5] Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature, 602(7896): 223-228, 2022.
[6] Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions on Intelligent Transportation Systems, 2366):499--4926 2021.
[7] Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. 2019.
arXiv preprint
[8] Matteo Hessel, Hado van Hasselt, Joseph Modayil, and David Silver. On inductive biases in deep reinforcement learning. arXiv preprint arXiv: 00290, 2019.
[9] Peter Dayan. The convergence of () for general. Machine learning, 8(3):311-32,, 1992.
[10] Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning. 8:279-292, 1992.
[11] JN Tsitsiklis and B Vanroy. An analysis of temporal-difference learning with function approxi- mation. IEEE Transactions on Automatic Control, 42(5):674-690, 1997.
[12] Geoffrey J Gordon. Reinforcement learning with function approximation converges to region. Advances in neural information processing systems, 13, 2000.
[13] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Machine learning, 49:209-232. 2002.
[14] Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems, 21, 2008
[15] Jalaj Bhandari, Daniel Russo, and Raghav Singal. A finite time analysis of temporal difference learning with linear function approximation, 2018. URL /taps: org/ wtrgg/bbs/800 02450.
[16] Gal Dalal, Balázs Szörényi, Gugan Thoppe, and Shie Mannor. Finite sample analyses for td (0) with function approximation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[17] Chandrashekar Lakshminarayanan and Csaba Szepesvari. Linear stochastic approximation: How far does constant step-size and iterate averaging go? In Amos Storkey and Fernando Perez-Cruz, editors, Proceedings of the Twenty-Fin -Fist International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 1347 1355. PMLR, 09-11 Apr 2018. URL https:////// htttps:/ppooceeeeeee http///pooeeenneee///e////e mlm, html.
[18] Richard E Bellman. Dynamic programmiing Princeton university press, 2010.

---

[19] Kai Arulkumaran, Mare Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. Deep reinforcement learning: brief survey. IEEE Signal Processing Magazine, 34(6):26-38, 2017.
[20] Richard Stuart Sutton. Temporal credit assignment in reinforcement learning. University of Massachusetts Amherst, 1984.
[21] Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A, 8(1):423, 1973.
[22] Crisanti and H Sompolinsky. Path integral approach to random neural networks. Physical Review E, 98(6):062120 2018.
[23] Moritz Helias and David Dahmen. Statistical field theory for neural networks, volume 970. Springer, 2020.
[24] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. arXiv preprint P2XXi:2205.00655: 2022.
[25] Blake Bordelon and Cengiz Pehlevan. The influence of learning rule representation dynamics in wide neural networks. arXiv preprint arXiv:22/0.02/57/ 2022.
[26] Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learning, pages 1024 1034. PMLR, 2020.
[27] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mezard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with teacher-student model. Advances in Neural Information Processing Systems, 34:18137-18151, 2021.
[28] Hong Hu and Yue M Lu. Universality laws for high-dimensionnl learning with random features. IEEE Transactions on Information Theory, 2022.
[29] Blake Bordelon and Cengiz Pehlevan. Learning curves for SGD on structured features. In International Conference on Learning Representations, 2022. URL https://openrevier: 1d=WPI2vbA1111.
[30] Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization kernel regression and infinitely wide neural networks. Nature communications, 12(1):2914, 2021.
[31] Bernhard Schölkopf, Alexander J Smola, Francis Bach, al. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
[32] Manel Tagorti and Bruno Scherrer. On the rate of convergence and error bounds for 1std (A). In International Conference on Machine Learning, pages 1521-1529. PMLR, 2015.
[33] Yangchen Pan, Adam White, and Martha White. Accelerated gradient temporal difference learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
[34] Alborz Geramifard, Michael Bowling, Martin Zinkevich, and Richard S Sutton. ilstd: Eligibility traces and convergence analysis. Advances Neural Information Processing Systems, 19, 2006.
[35] Fernando Pineda. Mean-field theory for batched td (A). Neural computation, 977:1403- 1997.
[36] Gandhary Patil, LA Prashanth, Dheeraj Nagaraj, and Doina Precup. Finite time analysis of tem- poral difference learning with linear function approximation: Tail averaging and regularisation. In International Conference on Artificial Intelligence and Statistics, pages 5438-5448, PMLR, 2023.
[37] Rayadurgam Srikant and Lei Ying. Finite-time error bounds for linear stochastic approximation andtd learning. In Conference Learning Theory, pages 2803-2830. PMLR, 2019.
12

---

[38] LA Prashanth, Nathaniel Korda, and Rémi Munos. Concentration bounds for temporal difference learning with linear function approximation: the case of batch data and uniform sampling. Machine Learning, 1 110:59-668, 2021.
[39] Eloise Berthier, Ziad Kobeissi, and Francis Bach. A non-asymptotic analysis of non non-parametric temporal-difference learning. In Koyejo, Mohamed, A. Agar- wal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Informa- tion Processing Systems, volume 35, pages 7599-7613. Curran Associates, Inc., 2022. URI. https: //proceee///cc///ssse neurips. ctt.cc//ppa.o http/ppper______ faleerrffsss//eeeeeeeeeereereerr wellfflllepppepeee patp//ppp/////////eeeee 424464722710006666777777777788 Pappeerr77aaapapap-aaaape- Ceree-Cooeeeeeeee pdf.
[40] Vivek Borkar and Sean P Meyn. The ode method for convergence of stochastic approximation and reinforcement learning. SIAM Journal on Control and Optimization, 38(2):447-46, 2000. 2000.
[41] Ishai Menache, Shie Mannor, and Nahum Shimkin. Basis function adaptation in temporal difference reinforcement learning. Annals of Operations Research, 134(1):215-235 2005.
[42] Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learning representation and control in markov decision processes. Journal of Machine Learning Research, 8(10), 2007.
[43] Mare Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. A geometric perspective optimal representations for reinforcement learning. Advances in neural information processing systems, 32, 2019.
[44] Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dynamics and generalization in deep reinforcement learning. In Kamalika Chaudhuri, Ste- fanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceed- ings of Machine Learning Research, pages 14560-14581, PMLR, 17-23 Jul 2022. URL mlr
https:/ https://proceedinng
html.
[45] Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: A new approach of reward shaping. Advances in Neural Information Processing Systems, $3:15931-15941 2020.
[46] Haosheng Zou, Tongzheng Ren, Dong Yan, Hang Su, and Jun Zhu. Learning task-distribution reward shaping with meta-learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 11210-1122 1218, 2021.
[47] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2): 251-276, 1998.
[48] Sham M Kakade. natural policy gradient. Advances in neural information processing systems, 14, 2001.
[49] Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical review A, 45(8):6056, 1992.
[50] Andreas Engel and Christian Van den Broeck. Statistical mechanics of learning. Cambridge University Press, 2001.
[51] Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam $ Schoenholz, Jascha Sohl- Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, 11:501-528, 2020.
[52] Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zeleorová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540- 9055009550000.0 9550, 2020.
[53] Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zde- borova, Rigorous dynamical mean field theory for stochastic gradient descent methods, arXiv preprint 220.0:2211000.11,, 2022.
13

---

[54] Michael Celentano, Chen Cheng, and Andrea Montanari. The high- dith-ddimmnnnnl asymptoties of first order methods with random data. arXiv preprint arX1v:2112.07572. 2021.
[55] James B Simon, Madeline Dickens, Dhruva Karkada, and Michael Deweese. The eigen- learning framework: A conservation law perspective on kernel ridge regression and wide neural networks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openrevial .et. www/fooumm/ee
[56] Maria Refinetti, Alessandro Ingrosso, and Sebastian Goldt. Neural networks trained with sgd learn distributions of increasing complexity. In International Conference on Machine Learning. PMLR, 2023.
pages
[57] Alessandro Ingrosso and Sebastian Goldt. Data-driven emergence of convolutional structure in neural networks. Proceedings of the National Academy of Sciences, 119(40):c2211854119
2022.
[58] Aditya Vardhan Varre, Loucas Fillaud-Vivien, and Nicolas Flammarion. Last iterate convergence of sgd for least-squares in the interpolation regime. Advances in Neural Information Processing Systems, 34:21581-21591, 2021.
[59] Maksim Velikanov, Denis Kuznedelev, and Dmitry Yarotsky. A view of mini-batch SGD via generating functions: conditions of convergence, phase transitions, benefit from negative momenta. In The Eleventh International Conference on Learning Representations, 2023. URL https://openrevaal
net/f orum? cerrum?id - bzaPGEllsjE.
[60] John O' Keefe. Place units in the hippocampus of the freely moving rat. Experimental neurology. 51(1):78-109, 1976.
[61] Edvard Moser, Emilio Kropff, and May-Britt Moser. Place cells, grid cells, and the brain's spatial representation system. Annu. Rev. Neurosci., 31:69-89 89, 2008.
[62] Robert Jacobs. Increased rates of convergence through learning rate adaptation. Neural networks, 1(4):295 -007 1988.
[63] William Dabney and Andrew Barto. Adaptive step-size for online temporal difference learning, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 26, pages 872-878, 2012.
[64] Andrew Ng, Daishi Harada, and Stuart J Russell. Policy invariance under reward transforma- tions: Theory and application reward shaping. In Proceedings of the Sixteenth International Conference on Machine Learning, pages 278-287, 1999.
[65] Burrhus Frederic Skinner. Science and human behavior. Number 92904, Simon and Schuster, 1965.
[66] Vijaykumar Gullapalli and Andrew G Barto. Shaping as a method for accelerating reinforcement learning. In Proceedings of the 1992 IEEE international symposium on intelligent control, pages 554-559. IEEF, 1992.
[67] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41-48, 2009.
[68] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[69] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemnl:: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchh, and R. Gamnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 8571-8580, Curran Associates, Inc., 2018. URL https://prrceeedigg pdf.
neurips. http:////wppoooo......aaaaaaaea.
[70] Hado Van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad. arXiv preprint arXiv: 18:e:81220000 02648, 2018.
14

---

[71] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.016443 2020.
[72] Juan Perdomo, Akshay Krishnamurthy, Peter Bartlett, and Sham Kakade. A sharp characteri- zation of linear estimators for offline policy evaluation. Journal of machine learning research, 2023.
[73] Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
[74] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning. In International conference on machine learning, pages 1928-1937. PMLR, 2016.
[75] Nishil Patel, Sehastian Lee, Stefano Sarao Mannelli, Sehastian Goldt, and Andrew M Saxe. The perceptron: Dynamics of policy learning in high dimensions. ICLR 2023 Workshop on Physics for Machine Learning, 2023.
[76] Vijay R Konda and John N Tsitsiklis. Convergence rate of linear two-time-scale stochastic approximation. Annals of Applied Probability, pages 796-819, 2004.
[77] Wolfram Schultz, Peter Dayan, and P Read Montague. A neural substrate of prediction and reward. Science, 275(5306):1533-15555 1997.
[78] Kenji Doya. Modulators of decision making. Nature neuroscience, 11(4):410 416, 2008.
[79] Timothy EJ Behrens, Timothy H Muller, James CR Whittington, Shirley Mark, Alon Baram, Kimberly L. Stachenfeld, and Zeb Kurth-Nelson. What is a cognitive map? organizing knowl- edge for flexible behavior. Neuron, 100(2):490-50,, 2018.
[80] Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel Gershman. The hippocampus as predictive map. Nature neuroscience, 20(11):1643-1653 2017.
[81] Marielena Sosa and Lisa M Giocomo. Navigating for reward. Nature Reviews Neuroscience, 22 (8):472-487, 2021.
[82] Daniel C McNamee, Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel Gershman. Flexible modulation of sequence generation in the entorhinal hippocampal system. Nature neuroscience, 24(6):851-800 2021.
[83] Blake Bordelon and Cengiz Pehlevan. Population codes enable learning from few examples by shaping inductive bias. Elife, 11:e78606, 2022.
[84] William Fedus, Prajit Ramachandran, Rishabh Agarwal, Yoshua Bengio, Hugo Larochelle, Mark Rowland, and Will Dabney. Revisiting fundamentals of experience replay. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3061-3071. PMLR, 13-18 Jul 2020. URL https:://///. http:///oooeeeeeeee mlr www.s/res/eeeee /vaeas/vvvi v119//// w0w.///f22aa....... html.
15

---

Appendix
A General Convergence Considerations for MDPs in Finite State Space
In this section, we will discuss the infinite batch limit and compare the value function obtained with TD to the ground truth value function. We will, for simplicity, consider in this section Markov reward process with transition matrix p(8++1 = = 8) = [x 8'). The general theory described in the main text does not only apply to MDPs, but the convergence analysis for MDPs is much more straightforward so we describe it here. In this case, the ground truth value function satisfies
V(s) = R(s) + Σ [Ts 8V)(())
(A.1)
which gives the vector equation V (I) R for V, R E R $1, Suppose the limiting
distribution over states is p psls which has entries p(s) Σ p(s: s). The fixed point of TD dynamics is
dpiag(p)
v WTD (p)aag()) (diag(p) I
WTD-
(A.2)
We now consider the two possible cases for this fixed point condition.
Case I: Underparameterized Regime First, if the feature dimension N is smaller than the size of the state space [S] and the features maximal rank, then the TD learning fixed point is
WTD = (Iriaa p)
y Viia (Pdiagg))IT II
Ediag (p) R
(A.3)
In this case, the value function is not learned perfectly, as can be seen by computing and comparing to the ground truth V = (II R. In this case, we would say that TD learning has an irreducible value error due capturing only a N dimensional projection of the value function.
Vp WTD
Case Overparameterized Regime Alternatively. if the feature dimension exceeds the total number of states, then the fixed point equation for TD is underspecified, However, throughout TD learning WTD span span( (1)(s) so we can instead consider the decompostion wy = E. (s)) where a RIS| satisfies
(I = diag(p)
(A.4)
where K E [S|Ss(Ss is the kernel computed with features K (ss, (0))))=b((1)) The solution to the above equation is unique and the learned value function V v TWTD = KK- -1 R (I (I- 1 R V. Therefore, in the pere-aatmtttrizee limit, the irreducible value error for TD learning is zero. This limit was considered dynamically in the infinite batch (vanishing SGD noise) setting by [44].
B Derivation of Learning Curves
In this section, we now consider the dynamics of TD learning when B random episodes are sampled time. In this calculation, the finite batch of episodes leads non-negligible SGD effects which can cause undesirable plateaus in TD dynamics.
B.1 Field Theory Derivation
In this section we use a Gaussian field theory formalism to compute the learning curve in the high dimensional asymptotic limit N, B + oo with B/N a. The episode length T treated as O(1). While this paper focuses on the online setting, where fresh trajectories (TØ² ) are sampled at each iteration n, this model can be straightforwardly extended the case where a fixed number of experience trajectories {r"} are replayed repeatedly during TD learning. We leave the experience
S1

---

replay dynamic mean field theory calculation for future work. The starting point of our analysis is tracking the moment generating function for the iterate dynamics
ZHi ((..).(t)))) exp
Jn
wn
(B.1)
To compute this object over random draws of training trajectories, express the joint average over Wn; {t)] into conditional averages over wx, (t) ((b))))))))))) (t)). To simplify the computation, in
this section, will compute the learning curve for mean zero features pu(s) 0 and
Z -E (WEC))
/ dIw, duw, Wn-1
Wn
Σ
exp
E ==0
In Wn
VBT
/ II d) 6 (M(() (WR wr)
(t)())
?
wn
2ha w/( (1
(B.2)
VN
VN
Expressing the Dirac-delta function as Fourier integral 6(2) J constraints. Under the Gaussian equivalence ansatz, we can easily average over Gaussian zp to obtain
2n exp (1ēz) for each of our
Z DADADWDU exp (-
172 2BT
ΣΣ A0 Tr'
exp
wn (w++1- - wn)
1
Σ
(wn w., w()) E(t,t) [wwn - w., w..))((()(]
cxp
2N
Σ
Án(tt 1) - 1) E(t, t2
exp
2N
nigtt!'
Σ
Án(t - 1) E(t, t')((()) w.))
exp
N
nutt
7
Σ
()(1)WOR win + - 1)20m "(,,(),,
exp
NBT
exp
()())A(()) +i > fn W..
(B.3)
THE
where adopted the shorthand DA Ia,,, (t) for the measure for the collection of variables (A)) Likewise should interpret Dwo II, dw,, To analyze the high dimensional limit of the above moment generating function, we introduce order parameters for the theory
B AH()))((), C.(t,)) = N u, Et, ((t,t")un L')wn
1
1
(t,tt))
B
==1
1
(t,t'') N
WR (tt,t)) Dn(t,t)= =- いい E(t, ywn, (x,tt'') - E(t, WR -
(B.4)
S2

---

For each of these order parameters, we enforce the definition of the order parameter using the Fourier representation of Dirac-delta function
1 B aQuit:t)0 (Qqulitt)
Σ ΣAHOAHO)
B (dQ((,())Q(((1)) exp
B (0,( - Σ A.0AN0)00(11)
4mi
2
(B.5)
Repeating this procedure for order parameters y y (Q, (Q, C, C, CR, disregarding irrelevant prefactors, we have the following formula for the moment generating function
CR D, D, DR, D and
Z x / Dq ( Slg]
(B.6)
where the action has the form
S >>
2 Σ
D.(0-))D,(tt))) + N
Bw + 2x In 23
1/2 Σ (n((, n T E(1, (')enn + Σ
2 =
exp
wn (Wa+1 w.n
27
n
Σ
Cn E(t, (')on 1 -00,)))0000 f)on
exp
(t, E(t,t')o,, -
(t, E E(t,
exp
natt
ZA = DADDexp exp
Σ
Σ
A A (1)
n2..
Σ (t) (t)
+ +(()))))
exp
N
ntE
1
Σ A (n)) (C_(1,)) (00+,))))))())))++++++)))))))
((((0)) (001))
exp
2
_An(t)A (t' (1())))))
exp
Σ Ennt 1)An(t - 1)cu(t,t))
exp
mp VaT
D... An A!)
exp
nt.
(B.7)
The function has the interpretation of an effective partition function conditional on order parameters To study the N oo limit, use the steepest descent method and analyze the saddle point
S3

---

0$
0. These saddle point equations give
0s
aQult, (A) =0
(Q)(t,tt)
0S
= aQQ(t,t) 1?
T (w, Ett, =0
(Q)
05
Cn(t,t) 1
N (w, ((t,t)),))= =0
(t,
0S
Cl(tt) (An (t) - - 1) - 1)) =
(t,
08
CR(,)) 1 (w,, =0
(CR(1,) 0s
N
Ca((,)) (1) + (I. - 11)
(t,,(,
0S
2i
-2Dx(1,))
(1) E(1, ()mm)))))) =0
(t,(t,tt)
N
0s
2i
(t,tt
(
E(t, t))wn) =
aD (t)
N
OS
(t,
2ami
A 1)A,(() - (t)
(t, (t)
VaT
(D,(t, -2D
as
20mi
(1))
(B.8)
VaT
The brackets () denote averaging over the stochastic processes defined by moment generating functions Za, Zw. After these saddle point equations are solved the order parameters 9 are treated as transformation is employed. For example.
non-random and a Hubbard-Stratonon
exp (-2
ihn
42 Σ CC
Qnlt, (n((,
i.,
Euy exp
E 72
u
(B.9)
T
where the average is over " ~ N Hubbard fields E. and uA (t), we can perform the integrals over to, and Á,(() which collapse to Dirac-Delta functions. The resulting identities of the delta functions define the following stochastic processes on wn and A
After introducing these
Wn+1 = Wn + " + Σ DR ( ()) + Σ Dn(t,,))((,
6e'
2
(t = (t) VOT EDD"(t Dx(t,t) - 7D, (t', (f))+))))+))))
(B.10)
Using similar trick, we can show that for any observable depending on or (t) that
i O(w,,((w))) = tun O(wm))
OtA(t) 0(14,(111)
(B.11)
Since wn is independent. This can be used to conclude
(t, = 0, (t,ttt') =
(B.12)
which implies that (1) = ( (t). Consequently the response functions have trivial structure
a/5 18(t - t') 70(t - 1- (t,tt)')=
Van 8(t - t').
Dr((t) T
(B.13)
T
S4

---

We therefore obtain stochastic process of the form
9V
.lva ((1,))wn
Wn+11 = Wn + +u +
(1,, - -1 wn
T
T
un ~ 0,
N(1,,
7/2 T
Enitt
1((.1)) ()(,,t)))((())))
( N(0,Qnn)
1
Qu(t,t) (An
N un (t, (R. c(t,t)
1 Cn(t,t) N (w, Ett, ((t,t')wn) (t,tt) 1 N (w)'E(tt()n,)
These are the final equations defining the stochastic evolution of w,, and A (t).
B.2 Simplifying the Saddle Point Equations
Using the above saddle point equations, we see that the variables (t) and {w,] will be Gaussian random variables. It thus suffices to track their mean and covariance. The (t), variables have zero mean and covariance given by the Qn (t, t') function. The {w,} variables have the following mean evolution
(wn-1) = (wm) + +1v0 Ewnn - [ --+ (wn))
= (ww) + + Σ - - [WTD- - (wn))
(B.14)
where WTD = [ - - Ewn is the fixed point of the TD dynamics. We next compute M, = (ww, WTD) (w, - wrn) which admits the recursion
" Mn+1 (I ạ E --E+) M, (I - - -++]) + (en(()))(()))
T
(B.15)
To obtain our formulas which hold for finite batch size, we rescale the learning rate by va giving the following evolution
(Wn+1) -(w,) 17 [ -++] (WID (ww))
Ti T [Q((,t))(tt (B.16)
Ma+1 (I - [Σ - 71+]) M (I -n [ [--)
After this rescaling, picks up an additive term on each step on the order of O(1'a a- which vanishes in the infinite batch limit B/N + cc. The error for value learning be obtained from M,, with C, 1 &TM,, Σ. Lastly, note that express the formula for Qn (t,t,tt E') entirely in terms of M,, and (w,). This
that the mean evolution for wn is independent of but that the variance
S5

---

gives the lengthy expression
Qu((,tt) ((wn w., T ('ttt'))
wn)) + 2
((wn wn) "Et, t' 1)w,))
N
7 (w), El + 1,tt)(wn - wmn) + 72 (w,,i(t [t + 1, + 1)wa)
+
N
N
1
TrM, (t,t)) +
(WTD - (w,)) [E( +) + E(t, ,)] (wR WTD)
N
N
1
+
(WR WTD) (t,,t') (WR-wLD)
N
Y TrMa E'(( 1) 1(++1,()
N
+ Y
(WTD (w.)) [E(1 (1 + + Σ(1 1, 1,) WTD
N
+ Y
(WR WTD) [t 1' +)+ + E(t + 1, t,t')] (wn)
N
2
27
+
Tr + +1, + 1)
((w,) - WTD) E(t + 1, t'+ + L)TTD
N
N
7²
WTD (t + 1, - 1) 1)RD
(B.17)
N
B.3 Final Result
Below we state in compact form the full final result for our TD learning curves. The below equations give the evolution of the first and second moments of w, obtained from the mean-field density of the previous section. Concretely, these moments obey dynamics
(wa+1) (w.) 1 [ E wwv (ur,,)
n²2 aTT2
Ma+1 [I - 11- + .7: E+] Mn [I T/ + E
Qu((,,t)= N (wwn w.) E'(,tt) wn,) + 22
1
? ((wn w.) Elt, + 1)
N
7 " El + 1, ( wn)) + (w, E( 1, 1)wr,) + N
(B.18)
These equations can be solved iteratively for wn Mn. Qn. Finite dimensional versions of this result can be obtained by replacing with B/N as written in the main text. The value estimation error is
Ln Tr
(B.
B.4 Non-Zero Mean Feature
We can also simply modify the DMFT equations if the mean feature is nonvanishing pus) # O. In this case, when averaging over all possible trajectories through state space, there is a mean feature vector at each episodic time utt). The above equations exact for non-zero mean features if E(t, t') is regarded as the (non-centered) correlation matrix (tt))))
B.5
Action Dependent Rewards
B.5.1
Expected Q-Learning Reduces to Previous Model
In the case where consider using features that depend on both states and actions (b(s, a) then we can use expected value learning to identify the expected value of a state-action pair under policy π
as R(s, a) + (V(s', al) s
(B.20)
This V function quantifies the expected reward associated with taking action a when state s and subsequently following policy π. This problem is structurally identical to the state dependent case
S6

---

by recognizing that state action pairs (s, a) act new states 3. As before, the policy defines the probability distribution over transitions on We can thus use Equation (4) to calculate the learning curve for this problem.
B.5.2 Action Dependence Generates Target Noise in State Dependent Value Learning
In the case where the rewards depend on both state and action Rss, a) but features only depend on state (s), we need a slight modification of our theory which models the reward at each state as mean value (over actions) plus a noise. For each state, we decompose the reward function into mean and fluctuation
R(s, a) = R(s) + c(s,a), R(s) (a)n(())) (a)a) R(s, a)
(B.21)
The function R(s) can again be decomposed into the basis of features (s)) However, we need to consider the correlation structure of a))
E, E(e,, ae (e) ar)
B., Ea, c(st, ar) [Bupllsa, (Bar c(sy ae))]
= Va a).
(B.22)
The above average vanishes for * t' since <(se, agd is zero mean over als. We introduce the notation of Vara R(se, a). Thus, we effectively have model where our TD errors obey
(1) R(st) e(s,, 4.) >V.)) - (ss)
(B.23)
The addition of this term leads simple modification of our Q(t,)) function
T (t(t,t') N ((WR w.. E(ttt') ww)) + Y ((wR w.. Elt,t' + 1)w,)
N
2² (w, E 1, 1)
+ 7 N
n (1 + 1, 1' wn)) +
+ 0.
(B.24)
This change to the Qu (1, correlation function alters the dynamics of Mn. Lastly, our population risk for the value estimation takes the form
1 Ln = N TrM, + Σ of
(B.25)
where 1 E. of exactly quantifies the variance in rewards unexplained by state-dependent features.
B.6 Tracking Iterate Moments with Direct Recurrence Relation
In this section give direct calculation of the first two moments of the collection of randomly sampled features ( (t) and show which terms are disregarded in the proportional limit examined in the main text.
Letting A Σ 7, we note that the average evolution of w has the form (Wn+1) = (Σ +E+) (WTD- - (ww))
(B.26)
Thus, we disregarded fluctuations in wn, due to SGD, the model will converge to the correct fixed point. Next, we look at M, ((w,, WTD) (w, WTD)). Under the Gaussian equivalence ansatz, we have
Mn+1 M, - -AM,, - A, AT
Σ (M
T male'
1/2 AM,AT
" TBB
-1 -AT
Σ (A.0AA
B
= (I-
(
T
+ T22
Σ (A
(B.27)
a
S7

---

The mean field theory derived from saddle point integration consists of the first two terms in the final expression. Therefore mean field theory disregards the last term which computes cross time correlations of RPEs with features, effectively making the approximation
112 B
~ O.
(B.28)
After making this approximation, we recover the learning curve obtained in the previous Section B.3. We show in our experiments that dropping this term does not significantly alter the learning curves.
B.7 Scaling of Asymptotic Fixed Points
To identify fixed points the value error dynamics, we can seek non-vanishing fixed points for the weight errrr covariance M ((w wrD)(w w)D)) We note that (w) WTD D asymptotically, Again, letting A = - y +,, we obtain the following fixed point condition for M under these assumptions
AM+MA nAMA BT E Q(t,
Q(t,t') = TrME 8) -TrM (t, E(++ + "TTM +
+ + 1 E4 (t, E) - WTD + + E(t + 1,t' + 1)WTD
+72WrD E E4 [E(t, t' 1) + E(t+ 1,t') WTD-
(B.29)
Where we used the formula for Qn, t²) from Appendix B.6, evaluated (w) = WTD and used the fact that
WTD >Σ Σ., WTD. The solution M 0 is a valid fixed point for M in the 7 + and B - do limits because the constant terms on right-hand side vanish. Similarly, 0 (which corresponds to the standard supervised learning case), the right hand side is linear in M, allowing M to be a valid fixed point.
However, for finite B and non non-zero 17 and there exists solution to the above fixed point equation. For small n² B
we can deduce that M must satisfy self-consistent asymptotic scaling of the form
777
M =
(B.30)
B
implying asymptotic value scaling of L ~ TrME in Figure where experiments obey the expected behavior.
O
y2" (B]
These scalings examined
C Reward Shaping
In this section, consider the role of reward shaping the dynamics of TD learning. As discussed in the main text, we consider potential based shaping with potential function decomposable in the features d(s) w. (s)). We first describe the change the average weight evolution (w,) and then describe the dynamics the correlations. In potential based shaping, the TD errors take the form
A(I) R(s(())) + ((x(1) 70(((t + 1)) + (+(( + (((1))
(C.1)
Computing from the DMFT equations the evolution of (w,,) we have
(w,+1) (w,, + ( + ws (w,)) mn - ws)
-w., A (WD ur (w.,))
(C.2)
We see that including the reward shaping function & offsets the fixed point of the algorithm to be WTD + wg. This occurs precisely because the potential-based reward shaping generates an additive correction to the target value function by dbs) [64]. When we predict value at evaluation, we use the reshifted value (s) (()). The natural quantity to track at the level of the mean field equations is the adapted version of M,
M, = (wn, WTD w (wn, WTD -wo) T)
(C.3)
S8

---

This correlation matrix has dynamics
o
Mn+1 (A (I DA)"
Σ
(C.4)
BT
and the TD-error correlations (1, (') have the form
Qu(t,t) = ((wR + ww wn) "ttt + we wnn)
+ ((w - -w,) [I(1, (() + E((', (WR + w0 wwn)
+ 72 ((w, w " ((( 1, + 1)(w,, wwp)
(C.5)
The value estimation error is again Ln = TrM,, E. We see that the two primary ways that reward shaping alters the loss dynamics is
A change in the initial condition for M, to be Mo = (WTD + wo)(WDD +
A change in the TD error covariance term (t,
Both effects can generate significant changes in the dynamics and plateaus of the model.
D Non-Gaussian Features
The full non-asymptotic theory (no assumptions on N, B large) of TD dynamics with linear function approximation closes under the fourth moments of the features. this setting do not incorporate explicit factors of N in the defintion the value estimator (t) w <b(t). As before, we track the update to the M matrix
IB 1) AM,A
MR+1 Mn nAM, AM
B
7
+ (A
(D.1)
B
To calculate the last term, we introduce the following tensor of fourth moments
Kiyza (11, t2, t3. t4) (t)) (ta))
(D.2)
In the Gaussian case, this reduces to expression involving only the correlations. For example, if the features are zero mean, then we can use Wick's thcorem obtain the decomposition
4.Gauss (t,, 12. ta, (4) = Eig(ta 12) Ena 14) + (tttt, 44) Ex Fejhl
14)
(D.3)
Now, using the fact that (1) (WR w.) <b(() ,,o, 4/(1 1), find
(An(t)A, (+))
= (1)(0))))) vb(1) (wn w.)) - T (())
(v(((()) (b)) (WR wn) w sb((t 1)) (tx(())))) sbt 1)T "w (WR ww) (b(())) +~ (M:(1))))))))) (b(t + +) T WT (b('' +1))
(D.4)
S9

---

Putting this all together, we find the following recurrence for M, in the non- non-Gaussian case
Mn+1 = Mn nAM, - "M, (
A
B
172
+
Tralt,t, ( - wwn)(w ww,) wn)()
BT
T
(t, t,,' ()) - wmn
f,t!
Σ
(Tr 1,() ((WR wn)wn,
BT
2212
Σ E
Tra(t, 1' 1 1,1 + (1 (w.,,,))
(D.5)
BT
where all traces are taken against the last two time and feature indices. The averages over the weights close in terms of the average (w,) and the covariance Mn, This equation exact for any feature distribution, but requires significant computational resources to evaluate and is less illuminating than the mean field limit analyzed in the previous sections.
D.1 Breakdown of Gaussian Theory in Low Dimension
In this section, we discuss the breakdown of Gaussian theory at low dimension N. In Figure D.1 we provide an example where the non-Gaussian distributions exhibit noticeably different learning curves than the Gaussian approximate theory (dashed black) and Gaussian samples with matching covariance (dashed color). We the features
10F
Non-Catssian
10k
Non-Gaussian
N 1 Gaassan
1 Gauasinn
N 10 Non-Caussian Gaussian
Non-Caussian
20
Value Error
10
Value
10
10-1
10-4
107
101
16a
100
101
162
steps
steps
(a) B = 5
(b) N = 100
Figure D.1: The Gaussian theory can break down for non-Gaussian features in low dimension N. Illustration of the possible gap between Gaussian and non-Gaussian performance in the power law features of Figure and defined in Appendix G. (a) The learning curves for batchsize B = 5 and varying dimension N. As N increases the gap between the non-Gaussian experiment (solid) and the Gaussian theory (black dashed) decreases.
D.2 A Simple Solveable Example
We next examine a very simple case where we can exactly characterize the gap between the non- Gaussian and Gaussian distributions. In this section, we examine the special case of and look at features which are independent p(ıb) = II. p(tx) (form a factor distribution). In this case, we obtain the following exact learning curve
f. (1 1) + PN 1+
== -WD) - (1:222
(D.6)
B
S10

---

which holds for any N, B. We note that in the limit where N, B + oo with B/N = a, we see that the dependence on x. disappears and we arrive the universal behavior
-)= 17
L (1-n)²+
N,
(D.7)
For example, we can consider vectors on the hypercube where the E (+1) with equal probability for k (1,-.., .N for the non-Gaussian distribution and compare to the Gaussian with identical covariance x/ ~ (0,
(N+1) 1))
Gaussian &b
fn
(D.8)
1)
(1- B
Hypercube sb
The reason for the discrepancy between the Gaussian and Bernouli/Hypercubee loss curves is exactly the negative kurtosis of the hypercube features
KGanss
Bemoulli
3Σ
Bemoolll = (W))) 3(2)2
(D.9)
An example of this result for low and high dimensions N with B = aN with a = 0.1 is provided in Figure D.2. In low dimension (N 10) the Bernoulii/Hypereube feature have noticeably different dynamics than the Gaussian features. In the proportional limit N, B - oo with a B/N these learning curves are identical and all have the form Ln (1 7) + 7
N=10
N=500
10
Bemoulii
10-1
Bernoulli
Gouss
Gauss theory
10
theory
10-1
10-7
10-1
10-13
10-16_
10
10
10-14
36
10-13
10-22
10
100
(a) Low Dimension
(b) High Dimension
Figure D.2: A simple explicitly solveable case shows how non-Gaussian corrections appear at finite size but disappear the proportional limit where N, B + co with a B/N O())
E
Tests on Other Feature Distributions
In this section, include additional tests of our theory on alternative features with the same random walk policy as in Figure
F Plateau Scaling in MountainCar-vo Environment
We verified the results of the theory on the environment MountainCar-vo) First, we train policy with tabular e-greedy Q-Learning (e = 0.1, 7 = 0.99, r) = 0.01) to Iearn policy л. The position and velocity are discretized into 42 and 28 states, respectively. The Icarned policy is not optimal but consistently reaches goal within 350 timesteps. Therefore, each episode is set to have a length 350 timesteps. Next, we take and evaluate with TD learning.
Since MountainCar-v0 is continuous environment, there no closed solution to the ground truth of the value function. To estimate the ground truth value function, ran TD learning with small learning rate for 10M batches (n 0.01, B 1, 0.99) obtain Va se ViOM-
$11

---

Polynomial Features
Fourier Features
10a
10a
scale 0.2
scale
Value Error
10-1
Value
4x 10
3x-10-1
2x
degroe degroe
10a
10
10a
107
10P
10'
10a
107
Iteration
Iteration
(a) Polynomial features
(b) Fourier Features
Figure E.I: this Figure, we simulate the same random walk policy on the 2D grid world but use other types of features beyond RBF place cells. Learning dynamics still accurately described by our theory. (c) The polynomial basis over states with random powers is constructed as (), 82) = $1 C1,1 61 where C,,1, C22 chosen at random from {0, 1,.,, k} where k is the degree. (f) Fourier where is the scale/bandwidth.
2
features with spectral power density 1+0 +
G Numerical methods and additional details
The code generate the Figures is provided in the Supplementary Material as Jupyter Notebook at the following Github repository https: /stps://githh thub. cww/ Pehleean SaaaanGrrupTT--- RE-TTT-RR----E ALL-dynaicss, Here, we briefly highlight some of the parameter choices.
For Figures and 4 we use diagonally decoupled, but temporally correlated power law features with (') = deg k-12 exp (- ((1/Tx) with тą K+I and w = 1-1.1 for A E [N] with N = 300. This type of feature structure is especially easy to evaluate the theoretical learning curves for. Unless otherwise stated, these figures used 0.9 and batch size B 10.
For the 2D MDP grid world, we defined a discrete set of states on 17 x 17 grid. The agent starts in the middle position and follows a random diffusion policy where each possible movement (up, down, left, right) is taken with equal probability. The features were generated as bell-shaped place cells (shown). We computed (t, t') for the theory by sampling 5000 random draws of length T 50. The Gaussian learning curve is obtained with TD learning with &aG ~ (0,, Σ).
Numerical experiments were performed on NVIDIA SMX4-A 100-80GB GPU using JAX to vectorize repetitive aspects of the experiments. With the exception the MountainCar-vl simulations, the numerical experiments (both preliminary experiments and those presented in the paper) took around 1 hour of compute time.
S12

---

0
-20
10
10
15
15
-0
20
20
-3
35
35
-80
40
40
(a) V estimated with tabular e-greedy (b) 1/8 TD Converged Value Function Q-Learning
Convergence Value Error
0.0.1
Convergence Error
0.02
14.00
n= 0.05
0.1
602-
3009
n= 0.2
Linear fit
B 2
B 4
Linear fit
2505 sais b.
Learning Rate
1/B
(c) Scaling of value error with learning rate
(d) Scaling of value error with batch size
Figure F.1: Simulation in MountainCar-vo environment. (a) Value function learned by Tabular Q-Learning that approximates the value function of optimal policy. (b) An example value function of policy (V*) learned by TD learning. Notice that the value function does not equate to that in (a)
due the policy и not reaching all states the environment. (c-d) Linear scaling of convergence value error with the learning rate and the inverse of batch size. Target value function is the same across both experiments. Each dot represents a different seed. A total of 10 seeds were used. (c) Convergence value errors were computed by averaging the 100k batches before batch 10M. (d) Convergence value errors were computed by averaging the 100k batches before batch IM.
S13


Reviews:
[{"review_id": "lTaP4tcfuV", "reviewer": "Reviewer_5FZT", "summary": "This paper introduces the concept of statistical physics to analyze reinforcement learning models. The authors propose a theory of learning dynamics for RL, with an emphasis on the role of linear function approximation. They investigate how strategies such as learning rate annealing and reward shaping can positively modify learning dynamics and surmount plateaus. The paper concludes by revealing new tools for developing a theory of learning dynamics in RL, thereby setting the stage for future research in this domain.", "questions": "Is there a particular reason for choosing batch sizes that are not powers of 2, such as N=3, N=30, N=20, for the experiments? As far as I am aware, batch sizes that are powers of 2 are most commonly used for practical considerations.\n\n\n", "limitations": "No concerns about negative societal impact.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper introduces a novel method using statistical physics to formulate a theory of learning dynamics in RL.\n2. The authors derive an analytical formula for the typical learning curve, demonstrating how their theory can predict the scaling of both learning convergence speed and performance plateaus based on problem parameters. This predictive ability is a significant advantage of the paper.\n3. The paper outlines how this theory can assist in understanding and guiding design principles when selecting meta-parameters in RL algorithms. Such insights can enable practitioners to gain a deeper comprehension of how these factors influence the RL learning process.", "weaknesses": "1. The paper does have limitations regarding practical implications. Although it presents a theoretical framework and makes predictions about learning dynamics, the methods utilized, such as linear approximation, deviate significantly from current reinforcement learning (RL) methodologies.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and useful questions. We respond to the weaknesses and questions below.\n\n**Responses to Weaknesses**\n\nThe reviewer is correct that our theory is limited to linear function approximation. This allows us to make strong predictions, but does not capture interesting aspects of deep learning such as feature learning. We address this in our limitations section but will expand on it further. \n\nHowever, we think that even in the linear model, some interesting statistical and dynamical phenomena such as SGD effects arise that are worthy of study (Paquette et al 2022 arXiv:2205.07069, Simon et al arXiv:2303.15438 2023, Mignacco et al 2020 arXiv:2006.06098).\n\nFurther, if the Gaussian equivalence/universality ideas hold in RL, it could allow for future analyses of two layer neural neural networks which learn features as they optimize the value function (such as in Goldt et al arXiv:2006.14709 2020).\n\n**Response to Questions**\n\nThere is no particular reason we chose the batch sizes we did. We can add a plot with powers of two if the reviewer thinks this is important. It would look very similar with plateaus evenly spaced on the log-scaled y-axis (see also scaling as a function of batch size on the MountainCar simulation in Figure 2 of rebuttal PDF).\n\n"}, {"review_id": "AeIuDgYXRG", "reviewer": "Reviewer_KZaZ", "summary": "This paper provides a theoretical model that predicts the dynamics of TD learning. The theory assumes that the distribution of feature vectors is, in some sense, equivalent to a Gaussian distribution and predicts the value estimate at each iteration. The theory reveals a rich set of phenomena such as plateaus in TD learning.", "questions": "Is T a constant or a random variable? If T is a constant you are considering the finite-horizon setting and the value function is horizon-dependent. If T is a random variable you need to specify how it is defined.\n\n\"the algorithms bootstraps its current predictions to estimate future states\" Not sure what you mean here. I think you mean algorithm, not algorithms. And I think you mean bootstrapping from the estimated next state's value.\n\nline 134: what are \"the learning curves\"? \"high dimensional features\" how many dimensions are needed? What do you mean when you say a set of learning curves is equivalent to another set of learning curves?\n\neqn 3: what is the meaning of <x>_{y} in your notation?\n\nLine 138: \"higher order cumulants of the features is negligible in high dimensional feature spaces under the square loss\". Again, what do you mean by square loss. What do you mean by \"higher\" order \"cumulants\" of the features?\n\nFigure 1: what is this diffusion process?\n\nProposition 3.1: what is the meaning of <>_s in the definition of L_n?\n\nLine 191: this remark seems to be much more useful. Why not present your result without assuming that B -> infty? Could you briefly explain under what conditions the theory holds when B is not infinity?\n\nLine 195: this remark also seems to be much more useful. Again, why not present your result without assuming that the value function is representable?\n\nLine 203: how is this remark different from remark 2?\n\nLine 205: what is w (without _n)?\n\nLine 209: \"which can be learned easily and which require more sampled trajectories\" Why is it \"easily\" when it requires more sampled trajectories?\n\nLine 216: \"The theory predicts that,\" I do not see which part of your theory gives this result. Furthermore, how is this result different from the rate of convergence in a linear dynamical system?\n\nProposition 5.1: what is a fixed point of a learning curve?\n\nHow is eqn B.24 derived?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The complete dynamics of TD learning have not been previously studied. And the paper provides a unique perspective on this problem.\n\nThe dynamics of TD learning according to the theoretical prediction are close to those in experiments, verifying the usefulness of the theory in these experiments.\n\nThe theory is also consistent with some existing tricks to improve TD learning such as reward shaping and step-size annealing.", "weaknesses": "The paper does not mention the limitations of its main assumption: the Gaussian Equivalence Conjecture. It verifies it using experiments in some small MDPs but these experiments do not tell when this assumption holds.\n\nThe paper does not explain clearly its theoretical results, making it difficult to understand these results.\n\nThe writing of the paper needs to be improved. Specifically, many sentences and notations remain confusing to readers. See Questions for details.\n\ntypos:\n\nline 73: the features -> the vector of features\nline 78: the offline setting\nline 79: width -> with\nline 245: two \"learning\"s", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for their careful reading of our paper and their detailed questions. We tried our best to clarify and address each of the weaknesses and questions raised below. \n    \n**Response to Weaknesses**\n\n1. We will add an acknowledgement that our paper relies on an assumption (Gaussian equivalence) which can not be validated or proven in general. Since the original submission, we have been able to derive more general error estimates (see Global response) and in some simple cases establish that Gaussian equivalence will hold in high dimension (learning curves don't depend on higher order cumulants, see Figure 1 of the rebuttal PDF). Based on this new result, we expect Gaussian equivalence under a variety of conditions including: (a) if the fourth moments are close to those of the corresponding Gaussian or (b) the features are close to independent (in some basis) and $N,B$ are both large. However, a more general set of conditions under which universality should hold would require a more careful proof so we will acknowledge this limitation explicitly in the discussion. \n2. We thank the reviewer for pointing out areas where we can improve our writing and explanation of our results. We are adressing the points raised by the reviewer and listed below in Responses to Questions\n    \n### Responses to Questions\n\nWe thank the reviewer for such detailed reading and questions. Below are our answers which we will include in the revised manuscript:\n\n1. In our model $T$ is a constant that does not scale with the size of the features $N$ or the batch size $B$. We will state this more clearly in the problem setup.\n2. We agree that we wrote this in a confusing way. We will fix this sentence so it reads *\"the algorithm bootraps from the current estimate of the next state's value\"*\n3. line 134: by learning curve, we mean $\\mathcal{L}_n$ the value estimation error as a function of iteration $n$. Our new theoretical work and simulations in Figure 1 of the rebuttal PDF, show that for high dimensional features, the contibution of the higher order cumulants is neglibible.\n4. The notation $\\left< x \\right>_y$ is borrowed from physics and denotes an average of the function $x$ which depends on the random variable $y$.\n5. Line 138: by the square loss, we mean that the evaluation metric is a square error between predicted and true value $(\\hat{V} - V)^2$. By cumulant of features we mean fourth, fifth, etc. cumulants of the random variable $\\psi_i(s_t)$. For example, the second cumulant is the covariance $\\left< \\psi_i(s_t) \\psi_j(s_{t'}) \\right> - \\left< \\psi_i(s_t) \\right> \\left<\\psi_j(s_{t'}) \\right>$. The fourth cumulant for mean zero random variables $\\{ \\psi_i \\}$ is $\\left< \\psi_i \\psi_j \\psi_k \\psi_l \\right> - \\left< \\psi_i \\psi_j \\right>\\left< \\psi_k \\psi_l \\right> - \\left< \\psi_i \\psi_k \\right>\\left< \\psi_j \\psi_l \\right> -\\left< \\psi_i \\psi_l \\right>\\left< \\psi_j \\psi_k \\right>$. Generally, the cumulants are the coefficients in the Taylor series for the function $m(c) = \\ln \\mathbb{E}_{\\psi} \\exp\\left( c \\cdot \\psi \\right)$ around $c = 0$. Our new theoretical work shows we can handle arbitrary fourth cumulants $\\kappa$ and our simulations show (Figure 1 of rebuttal PDF) that as assumed, the contribution of higher order cumulants becomes negligible as dimension increases\n6. The diffusion process is an unbiased (isotropic) random walk in the 2D grid world state space.\n7. $\\left< \\right>_s$ denotes average over states. \n8. Line 191: We expect our result can hold at small $B$ if either the features are close to Gaussian (small fourth cumulant $\\kappa$) or if the learning rate is small so that the SGD noise is small. In general, we expect the Gaussian equivalence to only kick in for large $N,B$, but we empirically observe that it can work well at small $N,B$ as well.\n9. Line 195: We present the main text result with a representable value function to simplify the main text expression. We provide the full formula in the Appendix. \n10. Thank you for pointing out the redundany in line 203. We will combine these two remarks. \n11. Line 205: nice catch! It should be $w_n$. We will change this.\n12. Line 209: This was a badly worded sentence. What we mean to say is that our theory can distinguish between easy and hard reward functions. We will reword the sentence appropriately. \n13. Line 216: the theory predicts this expression since the equation for $M$ in the zero SGD noise limit is $M_{n+1} = (I-\\eta A) M_n (I-\\eta A)^\\top$. Diagonalizing $A$ in its right eigenvectors, we can arrive at the spectral equation in line 216. We will add a detailed derivation in the appendix. \n14. In proposition 5.1, the fixed point of the learning curve or the dynamics would correspond to a value where $L_{n+1}=L_{n}$ and $M_{n+1} = M_n$. We show that there can exist non-zero solutions to these equations that have a scale of $\\mathcal{O}(\\eta \\gamma^2 /B)$.\n15. There are many ways to derive equation B.24. Perhaps the simplest is to expand $M$ and $Q$ in a power series in $\\eta$ so that $M = M_0 + \\eta M_1 + \\eta^2 M_2 + ...$ and $Q = Q_0 + \\eta Q_1 + ...$. Plugging this into the fixed point condition in equation B.23 and solving order by order yields $M_0 = 0$ and $A M_1 + M_1 A = \\frac{1}{T^2 B} \\sum_{tt'} Q_0(t,t') \\Sigma(t,t')$ \nwhere $Q_0$ is independent of $M_1$. \nWe see that the leading order scaling (in $\\eta$) is thus $M \\sim \\mathcal{O}(\\eta)$. This procedure can be repeated by expanding $M, Q$ in power series in $\\gamma$ and $B^{-1}$ and extracting the leading order scalings, namely that $M\\sim \\mathcal{O}(\\gamma^2)$ and $M \\sim \\mathcal{O}(B^{-1})$. Lastly, one can verify that $M \\sim \\mathcal{O}(\\eta \\gamma^2 /B )$ is self-consistent. We will add a more detailed set of arguments for this scaling in the Appendix.\n\nWe hope these questions helped clarify our results. We will aim to improve the writing around these topics in the text and provide additional explanation of our notation and assumptions. \n"}, {"review_id": "7fr3X2MSmK", "reviewer": "Reviewer_u1tm", "summary": "This work utilizes concepts from statistical physics to analyze the learning dynamics of TD(0) (policy evaluation) under the assumption of Gaussian equivalence, online batch update and linear function approximation. Using the theory, it demonstrates how learning rate annealing and reward shaping can improve learning in various aspects, based on multiple synthetic experiments.", "questions": "Q1: Any comments on why the Gaussian equivalence assumption would hold for general MDPs?\n\nQ2: How are the current convergence results different from the typical results from the LSTD literature?", "limitations": "The paper did not sufficiently explain its connection to the literature.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "The paper studies the learning dynamics of TD(0) from a novel perspective, which could provide readers with new insights and understandings", "weaknesses": "1. The whole framework relies on an important hypothesis that is not sufficiently verified in the TD(0) context. The Gaussian equivalence conjecture is only justified using a toy MDP (Fig.1) without explaining how accurate this hypothesis is for general MDPs. The provided example is using RBF features and RBF reward, which could be favorable to the given theory.\n\n2. The connection to least-square TD methods is missing. In the same regime as the current paper, LSTD has been extensively studied for its convergence and learning dynamics (see, for example, Tagoti and Scherrer (2015) and Pan et al., (2017)). It does not rely on the Gaussian assumption and should be compared with the convergence result in the current paper.\n\nRef:\n- Tagorti, M. and Scherrer, B., 2015, June. On the Rate of Convergence and Error Bounds for LSTD (λ). In *International Conference on Machine Learning* (pp. 1521-1529). PMLR.\n- Pan, Y., White, A. and White, M., 2017, February. Accelerated gradient temporal difference learning. In *Proceedings of the AAAI Conference on Artificial Intelligence* (Vol. 31, No. 1).\n\n3. It is unclear what Sec.4 is trying to demonstrate. The matrix A naturally appears in LSTD and has nothing special w.r.t. the current Gaussian equivalence setup.\n\nMinor\n- L136: correlation here means autocorrelation instead of Pearson product-moment correlation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for their good questions and for pointing us to the LSTD literature. Below we address the weaknesses and questions. \n### Response to Weaknesses \n\nWe thank the reviewer for their comment and have added both derivations and simulations to show the generality of our framework. We analyzed non-RBF features including polynomial and fourier features (see Figure 1c-d of the rebuttal PDF). We also verify the predicted scaling of the fixed point of the dynamics with hyperparameters ($\\eta, B$, equation B.24) for policy evaulation using TD(0) in a Mountain car environment (Figure 2 of rebuttal PDF)\n\nWe also extended our theory to account for the non-asymptotic and non-Gaussian cases. Further, we also investigated more generally the question of Gaussian equivalence and looked at some simple exactly solveable models (at any finite $N,B$) where the joint limit is well described by Gaussian learning curve (is insensitive to higher moments in features, see Figure 1a-b of rebuttal PDF). There is a rich literature trying to establish universality in these high dimensional learning settings (see for instance Goldt et al, 2020, arXiv:2006.14709; Hu and Lu, 2020,arXiv:2009.07669; Gerace et al., 2022, arXiv:2205.13303).\n    \nSecond, we will add a detailed comparison with the LSTD literature to our related works section. In summary, we are merely studying online stochastic temporal difference learning where randomly sampled episodes (state sequences $\\{ s_\\mu(t) \\}$) are used to estimate the gradient at each incremental step (more like SGD dynamics). This means at each iteration, the weights $w_n$ are updated incrementally with fresh samples of data. This is in contrast to LSTD algorithms where after observing $n$ samples, the best linear fit to the weights is performed (finding weights $w$ so that $\\sum_{t=1}^n [\\psi_t - \\gamma \\psi_{t+1}] \\psi_t \\cdot w = \\sum_{t=1}^n \\psi_t R_t$). In the papers cited by the reviewer, the authors are computing bounds on the error of the linear solution as a function of $n$. This generates a different kind of convergence rate (for example $O(n^{-1/2})$ in the first paper cited) than what we observe in the online/SGD setting, which are very dependent on feature structure and SGD noise. Further, the style of the analysis carried out in these papers (often worst case bounds) is different than what we pursue in the present work (average/typical case analysis in high dimension). Therefore, we believe our work presents a novel and complementary approach to study dynamics of convergence in reinforcement learning, specifically by studying the online setting of TD learning. \n    \nSection 4 is trying to demonstrate what kinds of dynamics are observed when the SGD noise is negligible (The setting used in most LSTD papers and the approach taken by Lyle et al arXiv:2206.02126). This gives an ordering of reward/value functions which are easier or harder to learn. We agree that the $A$ matrix will also appear as the limiting value of the matrix in the LSTD algorithm and will comment on this similarity. This is a consequence of convergence of empirical covariances $\\frac{1}{B} \\sum_{\\mu} \\psi_\\mu \\psi_\\mu^\\top \\to \\Sigma$ to the population average of the feature covariance. \n\nWe will specify that the correlations are not the pearson correlations in line 136.\n\n### Response to Questions\n    \n1. We have worked out a general solution to the learning curves which closes at the level of the fourth cumulants (see global response). We also looked at some special cases where we can see why Gaussian equivalence should hold in high dimension. First, consider the case where each features $\\psi_i$ is statistically independent. The variable $\\hat{V}  = \\psi \\cdot w$ should obey a central limit theorem and behave as a Gaussian random variable when $N$ is large. Similarly, the random variable $\\frac{1}{\\sqrt B}\\sum_{\\mu=1}^B \\Delta_\\mu \\psi_{\\mu i}$ is also a sum of a large number of independent variables. We would thus expect the algorithm to depend only on the mean and variance of these variables. A concrete example where $T=1$ is provided in Figure 1 of the rebuttal PDF. We provide an exact solution and show that dependencies on higher cumulants vanish as $N,B\\to\\infty$. \n2. The current convergence results are for online learning with SGD rather than for LSTD which can be considered as solving a full least squares problem at each iteration $n$. In the first case, the model does not perfectly fit the examples it has seen at finite $n$, whereas in LSTD, the model is fitting the observed samples as well as possible by solving the linear system \n    $A w = b \\ , \\ A = \\sum_{t} (\\phi_t - \\gamma \\phi_{t+1}) \\phi_t^\\top \\ , \\ b=\\sum_t \\phi_t R_t$\nBoth algorithms are interesting as methods to learn a value function from samples, but the dynamics are both conceptually and practically distinct. We will clarify this in the paper. \n\n"}, {"review_id": "E5mFsYpmIV", "reviewer": "Reviewer_kcB9", "summary": "This paper looks to introduce a new theory for the learning dynamics in the online batch policy evaluation setting for temporal difference learning. The paper introduces the Gaussian Equivalence Conjecture, which postulates that the learning curves in TD can be modeled by Gaussian features (with per-time-step mean and covariance over features and all trajectories). Along with this modeling assumption, the work does not prove this conjecture, but instead assumes this conjecture to be true to show empirical results in modeling policy evaluation learning curves analytically on a simple gridworld MDP with “place cell” features.\n\nUsing this assumption and the example MDP, the work first shows results fitting their model with the learning curve of online policy evaluation, which is interesting in itself. Using this modeling assumption, the authors provide insight into the “hardness” of certain policy evaluation tasks and how it relates to their TD learning model with an example comparing sparse vs dense reward functions in policy evaluation. Beyond this, the work also shows examples (on the same domain) of their theory of learning dynamics predicting the effects of batch size, discount factor, learning rate and learning rate annealing. Finally, the work also uses this model to explain and predict what happens to learning dynamics in the case of reward shaping. The work uses this proposed model to reshape rewards to improve timescales of convergence by rotating the reward-scaling weights to be more aligned with features of high variance.", "questions": "Throughout reading this work, I only have a few section-by-section questions:\n\n******3.1******\n\nCould you give more specific details with regards to this environment? Maybe even in the appendix? I don’t exactly understand what you mean when you say ‘The feature map is parameterized by the bandwidth of individual “place cells”’.\n\n******3.2******\n\nFor Remark 2 and 5, could you elaborate a bit on “fully explainable by features”? Is this referring to the case of partial observability? or is it just the most general formulation of R as a function of (s, a, s’)?\n\nThe notation for the proof for Proposition 3.1 is a bit confusing. I would suggest either simplifying by removing the notation that’s not defined here in the main paper and moving this to the appendix, or moving more of the appendix here.", "limitations": "Please see the general comment made in weaknesses to address limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Overall, I think the work is of great interest to the reinforcement learning community with a good amount of both theoretical and empirical results to back their proposed model of learning dynamics. The paper poses and answers well-founded questions, and is also well written. I only have a few (albeit important) questions I’ll pose below. ", "weaknesses": "My biggest issue with this paper in its current state is its set up. Currently, the proposed model is just a conjecture that seems to fit well with this one set of features (that are Gaussian it seems!) in this one gridworld environment. All the results seem to hinge on this conjecture being true, as evinced by the claim “We do not aim to provide a rigorous proof of this conjecture for TD learning but instead compute the learning curve implied by this assumption and compare to experiments on simple Markov Decision Processes”. Different forms of feature construction have been proposed throughout reinforcement learning history, including tile coding, polynomials, fourier basis etc. I’m still not convinced that this Gaussian model will carry over to learning dynamics with different features and different environments. While this is essentially my only concern with this paper (it’s very well written!), I think it is quite an important concern to address. Due to this concern, I have decided to only give this paper a weak accept. If more feature construction methods and more environments were included in these results, I would definitely be open to raising this score!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their support and useful questions and suggestions.  \n    \n### Response to Weaknesses \nWe thank the reviewer for their comment. These concerns were shared by others and we have added new theoretical results and simulations to show the generality of our approach. \nSpecifcially, we have extended our theory to include a general formulation (see global response) in terms of fourth order moments of the features and show in Figure 1a-b of the rebuttal PDF that the gaussian model is a good approximation. \nIn Figure 1c-d of the rebuttal PDF, we show that our theory also predicts the learning curves for agents with polynomial and fourier feature.\nIn Figure 2 of the rebuttal PDF, we show that the scaling in terms of hyperparameters ($\\eta, B$) predicted by our theory (equation B.24) is verified in a Mountain car environment\n    \n### Response to Questions \n    \n1. About Section 3.1: The 'place cell' terminology comes from neuroscience and refers to basis functions were each feature is a 2-D gaussian shape that is localized at a single location in the 2-D space (neurons with such receptive fields are found in the hippocampus). These are also called RBF features. In the simulations, we parametrize the feature maps by the width of the individual gaussian bumps. The environment is a 2-grid in which the agent performs an unbiased random walk. We will give an extended description of the environment and simulation details in the Appendix. \n2. About Section 3.2: Note that we will merge this remarks into a single one. Here, we do not investigate the case of a partially observable MDP. Instead, 'fully explainable' refers to the case where the true value (or reward) function can be described with zero error within the space spanned by the features. If the true value (or reward) function has components that lie outside of the space spanned by the features, these components will never be explainable by a TD algorithms using this feature space. This is analagous to a target function outside the hypothesis class in supervised learning. In the Appendix, we provide the full formula taking into account the unexplainable case. \n3. About the proof of Proposition 3.1: We will simplify our description of the proof outline in the main text. We will convey the intuition and methods used and we will link to the relevant equations in the Appendix instead.\n\n"}, {"review_id": "43YbIDeQNv", "reviewer": "Reviewer_mCei", "summary": "In summary, they study how learning dynamics and plateaus depend on feature structure, learning rate, discount factor, and reward function in the case of batch TD(0) for policy evaluation.\n\nThis paper applies concepts from statistical physics to study typical case learning curves for TD learning in linear FAs with nonlinear but fixed features (specifically, their study is targeted at policy evaluation with batch TD(0)). They conjecture a Gaussian feature equivalence: roughly, having a high-dim feature vector, the learning curves of a TD learner are conjectured to be equivalent to that of one with Gaussian features having matching mean and correlations. Further, this view allows them to use a Spectral Perspective to find a proxy for what would constitute a harder reward function to optimize (showing alignment between learning rate progress vs. the measure; shaped reward being simpler than sparse reward).   \n\n ", "questions": "1. Are the results of the Discount Factor analysis (Fig. 3b) for a dense reward scenario? Wouldn’t this be a very different case with sparse rewards due to the issues studied by van Seijen et al. (2019)?\n\n2. Isn’t annealing the learning rate a condition (even for tabular TD(0)) to guarantee asymptotic convergence to the optimal V with probability 1? If yes, what is the comparison of Fig. 5c vs 5d really showing us in the context of this paper?  \n\n\n- van Seijen *et al.* [NeurIPS 2019], \"Using a Logarithmic Mapping to Enable Lower Discount Factors in Reinforcement Learning\".\n", "limitations": "Limitations have been discussed sufficiently well.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- A step in an important direction: towards better understanding the learning dynamics of TD methods\n\n- Proposing a creative tool for the analysis of TD methods    \n\n- Good exposition of related works\n\n- Paper is well-motivated and clearly written \n\n- Claims are carefully stated and supported: no overstated claims \n\n- Experiments are pretty interesting and carefully designed", "weaknesses": "- Considering tasks have a finite horizon (fixed T), it has not been made clear if the experiments are run using features that have a good representation of the remaining time in them or not. Of course, this could be my misunderstanding (to be clarified by authors in the Questions section).\n\n- Not sure whether the comparison between fixed and annealing learning rates make too much sense in the context of this study (to be clarified by authors in the Questions section).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the supportive review and good questions. We attempt to address the questions below.\n    \n### Response to Questions\n\n1. About the discount factor analysis: The scenario we study is slightly different than the one studied by van Seijen et al. (2019). In our simualtions, we fix the reward function and compute the corresponding value function for each discount factor. We then compute the empirical and theoretical learning curves associated with each discount factor (which have different target value functions). We are therefore not studying the 'regularizing' effect of the discount factor, where there is a discrepancy between the metric used for performance and the discount factor used in hte learning rule, as investigated in van Seijen et al. (2019) or Amit. R., Meir, R. and Ciosek, K., 2020, November. Discount factor as a regularizer in reinforcement learning. ICML. Using our theory to understand this process would be an interesting avenue for future studies. In the experiments, the reward function was a sparse localized bump near the corner of the grid world.\n2. About learning rate annealing: Learning rate annealing can be a condition to guarantee asymptotic convergence but here we examine a slightly different issue. The usual condition is that the annealing schedule for the learning rate $\\eta_n$ at step n should satify: $\\sum_{n=1}^{\\infty} \\eta_n =\\infty$ and $\\sum_{n=1}^{\\infty} \\eta_n^2 \\lt  \\infty$. But many possible annealing schedules satify these conditions and although they will eventually converge they will do so at different speeds. In Figure 3d-e, we show that our theory allows us to estimate the effect of the annealing schedule on the full learning dynamics and help choose a schedule that converges faster. In Figure 3c, there is no annealing and the dynamics reach a plateau which our theory can predict. In the additonal simulations in Figure 2 of the rebuttal PDF, we show that the predicted scaling of these plateaus is verified in a MountainCar environment. More generally, this point highlight how our approach differs from some theoretical work in RL which analyses bounds at convergence but does not provide a description of the learning dynamics."}]

Author rebuttal:
We thank the reviewers for their comments and suggestions for additional theoretical justifications and experiments. Based on the reviews, we have added a more in depth analysis of non-Gaussian features at arbitrary dimension and can show that the learning curves close under fourth moments of the features and that the predicted plateau of $\mathcal{O}(\frac{\eta \gamma^2}{B})$ is preserved. We have added more experiments testing the range of validity of our theory's predictions and figures which we will include in the revised manuscript (see the attached PDF) which show:
1. Simulations showing why the Gaussian equivalence ansatz is reasonable in the high dimension limit based on a toy model where $T=1$ (SGD) and each feature is independent. We solve the model for arbitrary fourth cumulant $\kappa$ at any dimension $N$ and find that the learning curve scales as $\left[(1-\eta)^2 + \frac{\eta^2(N + 1+\kappa)}{B} \right]^n$, which loses dependence on $\kappa$ in the large $N,B$ limit. This is intuitive since $\hat{V} = \frac{1}{\sqrt N} \sum_i \psi_i w_i$ should behave as a Gaussian random variable during training even if $\psi_i$ is non-Gaussian.
2. More experiments showing the accuracy of our theory on the random walk MDP with polynomial features of varying degrees and fourier features with different lengthscales/bandwidths. This provides more evidence that our theory (and the Gaussian approximation) is reasonable for many types of features. 
3. We have run policy evaluation using TD(0) with a pre-trained policy on Mountain-Car ( we perform policy evaluation with Fourier features) and show that the loss curves exhibit asymptotes that scale linearly with $\eta$ and inversely with $B$. Since the feature space dimension $N$ and episode length $T$ are both very large, we do not attempt to estimate the $\Sigma(t,t')$ matrix. However, our theory's prediction of the asymptote allows us to accurately capture the scaling of the plateau with various hyperparameters.  

We hope that these simulations give more evidence that our setting and assumptions are reasonable. Given more time, we can add additional tests of other types of features.


In addition, we also will make several changes to the paper
1. In the new version of our paper's Appendix we will also provide a derivation of a theoretical learning curve which is non-asymptotic (does not require $N,B$ large) and for any distribution. The SGD noise term in the resulting theory depends on the fourth moments of features $\kappa_{ijkl}(t_1,t_2,t_3,t_4) = \left< \psi_i(t_1)\psi_j(t_2)\psi_k(t_3)\psi_l(t_4) \right>$. The loss is still $\mathcal{L}_n = \frac{1}{N} \text{Tr} M_n \bar{\Sigma}$ but with the dynamics

\begin{align}
M_{n+1} &= M_n - \eta A M_n - \eta M_n A^\top + \frac{\eta^2 (B - 1)}{B} A M_n A^\top \nonumber
    \\
    &+ \frac{\eta^2}{B T^2} \sum_{t,t'} \text{Tr} \kappa(t,t',t,t') \left< (w_R -w_n)(w_R-w_n)^\top \right> \nonumber
    \\
    &+ \frac{\gamma \eta^2}{B  T^2}\sum_{t,t'} \text{Tr} \\ \kappa(t,t',t,t'+1) \left< w_n (w_R -w_n)^\top \right> \nonumber
    \\
    &+ \frac{\gamma \eta^2}{B T^2}\sum_{t,t'} \text{Tr} \\  \kappa(t,t',t+1,t') \left<  (w_R -w_n) w_n^\top \right> \nonumber
    \\
    &+ \frac{\gamma^2 \eta^2}{B T^2} \sum_{t,t'} \text{Tr} \\ \kappa(t,t',t+1,t'+1) \left<  w_n w_n^\top \right>
\end{align}
The trace is taken over the last two of the feature indices of $\kappa$. All weight averages can still be expressed in terms of $M_n$ like in the current Appendix. This prediction still recovers a potential fixed point in the dynamics at scale $\mathcal{O}(\eta\gamma^2/B)$, so this prediction is actually universal. Computing this theory is even less tractable than our original result due to the number of entries associated with the $\kappa$ object $\sim N^4 T^4$.

2. We also will provide more room to explain and introduce our notation (for instance that $\left< \right>$ denotes averaging) and will more carefully introduce the mathematical objects which appear in our main text. 
3. We will provide a more extensive comparison of the algorithm considered in our work (online TD learning) to other algorithms in the literature including TD Least Squares (TDLS). 

We hope that the reviewers will take these updates into consideration when re-evaluating our work. 

Meta-review: The paper uses statistical physics methods to analyse temporal difference learning in simple settings.

Strengths:

-Derives analytical results for typical learning curves in RL

-Predictions agree very well with experiments

-Tackles an important question and provides complementary insights to existing theoretical analyses in RL

-The results are clearly presented

Weaknesses:

-Relies on unproven assumptions (as with many statistical physics methods)

-Can only treat simple settings at present, and the broader applicability remains unclear
Decision: Accept (poster)