SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To illustrate an example of the Gaussian Equivalence idea, in Figure 1, we consider an MDP which is defined by diffusion through a 2-dimensional (2D) state space (Figure 1(a)). We choose the features  $\psi(s)$  to be a 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 1(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  $\psi(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 a remarkable match of the three curves. The equivalence can be used to predict the speed of TD learning for different features, such as place cells with varying bandwidth as we illustrated in Figure 1 (d)-(f). In Figure 1 (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 error 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 are 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.1. 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
(a) Sparse  $R_{1}$  or Dense  $R_{2}$  Rewards
(b) Code-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 a 2D spatial navigation task can illustrate this effect. (b) The cumulative power distribution  $C(k)$  defined from the spectral decomposition of  $\mathbf{A} = \bar{\mathbf{\Sigma}} - \gamma \bar{\mathbf{\Sigma}}_{+}$ . Concretely we let  $\mathbf{A}\mathbf{u}_k = \lambda_k\mathbf{u}_k$  with  $\lambda_{k}$  ordered by real part and  $w_{TD} = \sum_{k}w_{k}\pmb{u}_{k}$ . In the  $B\to \infty$  limit the task which has rapidly rising  $C(k) = \frac{\sum_{\ell < k}w_{\ell}^{2}}{\sum_{\ell}w_{\ell}^{2}}$  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  $R_{2}$  than for  $R_{1}$ .
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  $\mathbf{A} = \bar{\mathbf{\Sigma}} - \gamma \bar{\mathbf{\Sigma}}_{+}$ , obtaining  $\mathbf{A}\mathbf{u}_k = \lambda_k\mathbf{u}_k$ , noting that eigenvalues  $\lambda_{k}$  can be complex. We then expand the TD solution in this basis  $\mathbf{w}_{TD} = \sum_{k}w_{k}\mathbf{u}_{k}$ .
The theory predicts that, the average learned weights will be  $\langle \pmb{w}_n\rangle = \sum_k|1 - \eta \lambda_k|^ne^{i\theta_kn}w_k\pmb{u}_k$  where  $|\cdot |$  is complex modulus and  $\theta_{k} = \mathrm{Arg}(1 - \eta \lambda_{k})$ . We can therefore order the modes by their convergence timescales  $|1 - \eta \lambda_k|$ . Given this ordering of timescales, we can order the modes  $k$  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  $\pmb{w}_{k}$  in the modes with small timescales. We quantify how well aligned a task is to a given feature representation by computing a cumulative power distribution for the target weights  $C(k) = \frac{\sum_{\ell < k}w_{\ell}^{2}}{\sum_{\ell}w_{\ell}^{2}}$ . If this quantity rises rapidly with  $k$  then the task can be learned from a 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 a sparse or dense reward function (Figure 2(a)). As expected the cumulative power rises more rapidly for the dense reward function  $R_{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
(a) Episodic Batches
(b) Discount Factor
(c) Fixed Learning Rate
(d) Annealed Learning Rate
(e) Optimal Annealing
Figure 3: Finite batch size, discount factor and learning rate all contribute to a stochastic semigroup 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  $\eta$ , the plateau scales as  $\mathcal{O}(\eta \gamma^2 B^{-1})$ . (a) Larger batch sizes  $B$  reduce SGD noise and leads to a lower plateau in the reducible value error for a decoupled power-law feature model. (b) Larger discount factor  $\gamma$  and (c) larger learning rate  $\eta$  lead to higher SGD plateau floor. (d) An annealing strategy  $\eta_n \sim \eta_0 n^{-\chi}$  for  $\chi > 0$  can allow one to avoid the plateau. For slow annealing (small  $\chi$ ), the error scales as  $\mathcal{L}_n \sim \mathcal{O}(n^{-\chi})$ . (e) The value error as a function of the learning rate annealing exponent  $\chi$  defined by  $\eta_n = \eta_0 n^{-\chi}$ . For this task, the optimal exponent balances the scale of the asymptote with the rate of convergence.
The stochastic noise from TD learning has striking qualitative differences from SGD noise in the standard supervised case. In standard supervised learning (such as  $\gamma = 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  $\hat{V}(s)$  is bootstrapped using the model's weights  $\boldsymbol{w}_n$  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  $\eta$  and  $\gamma$ . For small  $\frac{\eta\gamma^2}{B}$ , we deduce that  $M$  satisfies a self-consistent
(a) Geometry of Reward Shaping
(b) Scale-based Shaping
Figure 4: The theory can be used to understand how reward shaping decisions alter temporal difference learning dynamics. (a) A visualization of possible reward shaping potentials  $\phi(s) = \pmb{w}_{\phi} \cdot \pmb{\psi}(s)$  strategies in feature space. Probability density level curves for the features are depicted in blue. Reshaping with  $\pmb{w}_{\phi} = \beta \pmb{w}_{TD}$  for scale factor  $\beta$  merely changes the scale of weights which must be recovered (gold) and does not change timescales of TD dynamics. (b) The value error dynamics for the scale based reward shaping for the features in Figure 3. On the other hand, rotation based reward shaping where  $\pmb{w}_{\phi}$  is not parallel to  $\pmb{w}_{V}$  (red) leads to a potentially helpful mixture of timescales if the new target vector is more aligned with feature dimensions with high variance (purple). In (c), we plot loss curves for rotation angle  $\theta$  between the original mode  $\pmb{w}_{V}$  and the top eigenvector of the feature covariance matrix  $\bar{\Sigma}$ . Dashed black lines are theory.
(c) Rotation-based Shaping
asymptotic scaling of the form  $M = \mathcal{O}\left(\frac{\eta\gamma^2}{B}\right)$  implying an asymptotic value error scaling of  $\mathcal{L} \sim \frac{1}{N} TrM\bar{\Sigma} \sim \mathcal{O}\left(\frac{\eta\gamma^2}{B}\right)$ .
In Figure 3, we demonstrate that our theory predicts the plateaus and their scaling as a function of finite batch size  $B$  (Figure 3(a)), non-zero discount factor  $\gamma > 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  $\eta_{n}$  with iteration  $n$ . In Figure 3(d), we show the effect of annealing the learning rate as a power law  $\eta_{n} = \eta_{0} n^{-\chi}$  for some non-negative exponent  $\chi$ . For  $\chi = 0$  the learning rate is constant and a fixed plateau is reached. For small nonzero  $\chi$ , such as  $\chi = 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  $\mathcal{L}_n \sim \mathcal{O}(n^{-\chi})$ . For large  $\chi$ , the learning rate decreases very quickly and the plateau is never reached. Our approach can be used to find an optimal annealing exponent  $\chi$  and in Figure 3(e), we 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]. In standard supervised learning, the goal is to directly approximate the target objective given a cost function. However, in reinforcement learning, the objective is not to estimate rewards at each state directly but the discounted sum of 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 a theoretical description of the changes in the learning dynamics due to reward shaping which suggests they can be understood through a change of the alignment between the original rewards and the reshaped rewards in the space of the features used to represent the states.
The original ideas around reward shaping were inspired by work in experimental psychology and were closer to what is now studied as curriculum learning [65-67]. Reward shaping as currently used in reinforcement learning directly changes the reward function by adding a potential-based shaping function  $F$  such that  $F(s_{t},a,s_{t + 1}) = \gamma \phi (s_{t + 1}) - \phi (s_{t})$  [64]. In each step of the algorithm we feed
the following reshaped rewards  $\tilde{R}$  to the TD learner
$$
\tilde {R} \left(s _ {t}\right) = \left\{ \begin{array}{l l} R \left(s _ {t}\right) - \gamma \phi \left(s _ {t + 1}\right) & t = 0 \\ R \left(s _ {t}\right) + \phi \left(s _ {t}\right) - \gamma \phi \left(s _ {t + 1}\right) & t > 0 \end{array} . \right. \tag {8}
$$
We note that this transformation simply offsets the target value function by  $\phi(s)$  as the series above telescopes with a cancellation of  $\phi(s_t)$  between the  $t-1$  and  $t$ -th terms [64] (see Appendix C). However, the dynamics of TD learning with these reshaped rewards  $\tilde{R}$  is quite distinct from the dynamics with original rewards  $R$ . Here, we study the case where we can express  $\phi(s)$  as a linear function of our features:  $\phi(s) = \pmb{\psi}(s) \cdot \pmb{w}_{\phi}$ . This leads to a change in the dynamics for  $M_n$  and  $\langle \pmb{w}_n \rangle$  that we describe in the Appendix C.
In Figure 4, we illustrate the possible benefits of reward shaping. We explore two types of reward shaping. First, a scale based reward shaping where  $\boldsymbol{w}_{\phi}$  is parallel to the target TD weights  $\boldsymbol{w}_{TD}$ . This merely changes the overall scale of the weights needed to converge in the dynamics, leading to similar timescales and an identical plateau for TD learning as we show in Figure 4 (b). On the other hand, reward shaping which rotates the fixed point of the TD dynamics into directions of higher feature variance can improve timescales of convergence. In Figure 4 (c), we show an example where we vary the angle  $\theta$  of the shaped-TD fixed point (see also Appendix C).
7 TD Learning Plateaus in More Realistic Settings
In this section, we test if some of the phenomena observed in our theory and experiments also hold in more realistic settings. We perform TD learning with Fourier features to evaluate a pre-trained policy on MountainCar-v0. As expected, we see that the value error plateaus to an error level determined by both the learning rate (Figure 5a) and batch size (Figure 5b) due to semigrainient noise.
(a) Varying Learning Rates
(b) Varying Batch Sizes
Figure 5: Policy evaluation in MountainCar-v0 environment. The policy was learned with tabular  $\epsilon$ -greedy Q-learning (see Appendix F for details). (a) Value error curves for different  $\eta$  when  $B = 1$ . (b) Value error curves for different  $B$  with  $\eta = 0.1$ . Shaded area denotes  $95\%$  confidence interval over 10 seeds.
We show that the plateau obeys the predicted scalings of  $\mathcal{O}(\eta B^{-1})$  in Appendix F.
Our work presents a new approach using concepts from statistical physics to derive average-case learning curve for policy evaluation in TD-learning. However, it is only a first step towards a 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 in the "lazy" regime of training [68, 69], however it 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 (Proposition 3.1). In Figure 1 we evaluate these dynamics for a 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(0). 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( $\lambda$ ) 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 an RL agent is to learn how to act in the work and not merely to 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 a 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]. A 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 a 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 at 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 is 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 provides a new promising direction towards a theory of learning dynamics in reinforcement learning in artificial and biological agents.
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 to Venkatesh Murthy and Naoshige 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.
[71] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[72] Juan Perdomo, Akshay Krishnamurthy, Peter L Bartlett, and Sham Kakade. A sharp characterization 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 reinforcement learning. In International conference on machine learning, pages 1928-1937. PMLR, 2016.
[75] Nishil Patel, Sebastian Lee, Stefano Sarao Mannelli, Sebastian Goldt, and Andrew M Saxe. The rl perceptron: Dynamics of policy learning in high dimensions. In 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):1593-1599, 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 B Baram, Kimberly L Stachenfeld, and Zeb Kurth-Nelson. What is a cognitive map? organizing knowledge for flexible behavior. Neuron, 100(2):490-509, 2018.
[80] Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel J Gershman. The hippocampus as a 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 J Gershman. Flexible modulation of sequence generation in the entorhinal-hippocampal system. Nature neuroscience, 24(6):851–862, 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://proceedings.mlr.press/v119/fedus20a.html.
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 a Markov reward process with transition matrix  $p(s_{t + 1} = s'|s_t = s) = \Pi (s,s')$ . 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) + \gamma \sum_ {s ^ {\prime}} \Pi \left(s, s ^ {\prime}\right) V \left(s ^ {\prime}\right) \tag {A.1}
$$
which gives the vector equation  $\mathbf{V} = (\mathbf{I} - \gamma \mathbf{\Pi})^{-1}\mathbf{R}$  for  $\mathbf{V},\mathbf{R}\in \mathbb{R}^{|S|}$ . Suppose the limiting distribution over states is  $\pmb {p}\in \mathbb{R}^{|\mathcal{S}|}$  which has entries  $p(s) = \frac{1}{T}\sum_{t = 1}^{T}p(s_{t} = s)$ . The fixed point of TD dynamics is
$$
\Psi \operatorname {d i a g} (\boldsymbol {p}) \Psi^ {\top} \boldsymbol {w} _ {T D} = \Psi \operatorname {d i a g} (\boldsymbol {p}) \boldsymbol {R} + \gamma \Psi \operatorname {d i a g} (\boldsymbol {p}) \Pi \Psi^ {\top} \boldsymbol {w} _ {T D}. \tag {A.2}
$$
We now consider the two possible cases for this fixed point condition.
Case 1: Underparameterized Regime First, if the feature dimension  $N$  is smaller than the size of the state space  $|\mathcal{S}|$  and the features are maximal rank, then the TD learning fixed point is
$$
\boldsymbol {w} _ {T D} = \left(\Psi \operatorname {d i a g} (\boldsymbol {p}) \boldsymbol {\Psi} ^ {\top} - \gamma \Psi \operatorname {d i a g} (\boldsymbol {p}) \boldsymbol {\Pi} \boldsymbol {\Psi} ^ {\top}\right) ^ {- 1} \boldsymbol {\Psi} \operatorname {d i a g} (\boldsymbol {p}) \boldsymbol {R} \tag {A.3}
$$
In this case, the value function is not learned perfectly, as can be seen by computing  $\hat{\pmb{V}} = \pmb{\Psi}^{\top}\pmb{w}_{TD}$  and comparing to the ground truth  $\pmb {V} = (\pmb {I} - \gamma \pmb {\Pi})^{-1}\pmb{R}$ . In this case, we would say that TD learning has an irreducible value error due to capturing only a  $N$  dimensional projection of the value function.