SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Blake Bordelon, Paul Masset, Henry Kuo & Cengiz Pehlevan
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.
Reinforcement learning (RL) is a 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 a 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, a 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 non-asymptotic [15-17] limits, but do not provide a full picture of the evolution of the learning dynamics.
A desired feature of a candidate theory is to characterize the influence of function approximation to RL dynamics and its performance. Early versions of RL operated in a tabular setting, similar to dynamic programming [18], where all the states in the environment could be mapped one-to-one to a 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
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 is 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-stationarity when learning a value function in the context of policy evaluation [1] using a classical RL algorithm, temporal difference (TD) learning [20]. We develop a theory of learning dynamics for RL in this setting in a high dimensional asymptotic limit with a 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 to 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 a 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 an online setting. We provide an 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 ( $\S 4$  and  $\S 5$ ). Task-feature alignment is a metric that quantifies how features allow fast or slow learning for a given task.
- We show our theory can be used to understand and guide design principles when choosing metaparameters. Specifically, we show that we can use our theory to infer optimal schedules of learning rate annealing and the effects of reward shaping ( $\S 5$  and  $\S 6$ ).
2 Problem Setup and Related Works
We consider a set of states denoted by  $s$ , possibly continuous, and a fixed policy  $\pi$  which generates a distribution over actions given the state. The state dynamics are defined by a distribution  $p(\tau)$  over trajectories through state space  $\tau = \{s_1, s_2, \dots, s_T\}$ . Note that state transitions do not have to be Markovian, but each trajectory is i.i.d. sampled from  $p(\tau)$ . We consider trajectories of length  $T$ . Each state is represented by an  $N$ -dimensional feature vector  $\psi(s) \in \mathbb{R}^N$ , so that trajectory generates a collection of feature vectors  $\{\psi(s_t)\}_{t=1}^T$ . 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 as the policy is fixed, but variance over rewards at a 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  $s_0$
$$
V \left(s _ {0}\right) = R \left(s _ {0}\right) + \sum_ {t \geq 1} \mathbb {E} _ {s _ {t} \mid s _ {0}} \gamma^ {t} R \left(s _ {t}\right) = R \left(s _ {0}\right) + \gamma \mathbb {E} _ {s _ {1} \mid s _ {0}} V \left(s _ {1}\right). \tag {1}
$$
We use linear function approximation to learn the value function  $\hat{V}(s) = \psi(s) \cdot w$ . Similar to kernel learning [31], the features  $\psi$  should be high dimensional so that they can express a 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
$$
\begin{array}{l} \boldsymbol {w} _ {n + 1} = \boldsymbol {w} _ {n} + \frac {\eta_ {n}}{T B} \sum_ {\mu = 1} ^ {B} \sum_ {t = 1} ^ {T} \Delta_ {n} ^ {\mu} (t) \boldsymbol {\psi} \left(s _ {n} ^ {\mu} (t)\right), \\ \Delta_ {n} ^ {\mu} (t) \equiv R \left(s _ {n} ^ {\mu} (t)\right) + \gamma \hat {V} \left(s _ {n} ^ {\mu} (t + 1)\right) - \hat {V} \left(s _ {n} ^ {\mu} (t)\right). \tag {2} \\ \end{array}
$$
We operate in a 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  $\psi$  are outlined in Appendix A. The specific form for the TD-error  $\Delta_n^\mu (t)$  depends on the precise variant of TD learning that is used. Here, we will focus on TD(0) but our approach can be extended to other TD learning rules and definitions of the return function. We see that the iterates  $w_{n}$  will form a stochastic process as each sequence of states in an episode  $\{s_n^\mu (t)\}$  are drawn randomly from  $p(\tau)$ . In general, we allow the learning rate  $\eta_{n}$  to depend on iteration, an important point we will revisit later. The distribution of features  $\{\psi (s_n^\mu (t))\}$  over random trajectories  $\tau$  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 dynamics of temporal difference learning be characterized for complicated trajectory distributions  $p(\tau)$  and feature maps  $\psi(s)$ ?
To address this question, in this work, we provide an analysis of TD learning that explicitly models the statistics of stochastic semi-gradient updates to  $\pmb{w}_{n}$ . Our framework is based on a Gaussian equivalence ansatz for TD learning and high dimensional mean field theory which predicts the statistics of TD errors  $\Delta_{n}^{\mu}(t)$  and the weight iterates  $\pmb{w}_{n}$ . 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.
The dynamics of TD 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 to  $n$  samples [32-34]. Alternatively, many works focus on the online 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 [11-14, 35]. More recently, progress has been made in deriving bounds in the non-asymptotic regime. Initial work assumed that data samples were i.i.d. [15-17, 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 corresponds to a limit of the stochastic semi-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 averaging 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 to empirically learn a better shaping function [45, 46]. In policy learning it 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 is 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 dynamics [52-54]. Further, even when data is not actually Gaussian, several algorithms, such as kernel or random-features 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 Theoretical Results for Online TD Learning
3.1 Computation of Learning Curves
We develop a 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  $\pmb{w}_n$  over the trajectories of randomly sampled features  $\{\psi_{\mu}^{n}(t)\}_{t=1}^{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 \to \infty$  with  $B / N = \mathcal{O}(1)$  and episode length  $T = \mathcal{O}(1)$ . Let the ground truth reward function be  $R(s) = \pmb{w}_R \cdot \pmb{\psi}(s)$  and value function  $V(s) = \pmb{w}_{TD} \cdot \pmb{\psi}(s)$  in the basis of our features. Define matrices
$$
\bar {\boldsymbol {\Sigma}} \equiv \frac {1}{T} \sum_ {t} \boldsymbol {\Sigma} (t, t), \quad \bar {\boldsymbol {\Sigma}} _ {+} \equiv \frac {1}{T} \sum_ {t} \boldsymbol {\Sigma} (t, t + 1), \quad \boldsymbol {A} \equiv \bar {\boldsymbol {\Sigma}} - \gamma \bar {\boldsymbol {\Sigma}} _ {+}, \tag {3}
$$
and assume that the features are such that matrix  $\mathbf{A}$  is of extensive rank in  $N$ . Then the typical value estimation error  $\mathcal{L}_n = \left\langle \left(V(s) - \hat{V}_n(s)\right)^2\right\rangle_s$  after  $n$  steps has the form
$$
\mathcal {L} _ {n} = \frac {1}{N} \operatorname {T r} \bar {\boldsymbol {\Sigma}} M _ {n}, \tag {4}
$$
$$
\boldsymbol {M} _ {n + 1} = (\boldsymbol {I} - \eta \boldsymbol {A}) \boldsymbol {M} _ {n} (\boldsymbol {I} - \eta \boldsymbol {A}) ^ {\top} + \frac {\eta^ {2}}{\alpha^ {2} T ^ {2}} \sum_ {t t ^ {\prime}} Q _ {n} (t, t ^ {\prime}) \boldsymbol {\Sigma} (t, t ^ {\prime}) \tag {5}
$$
$$
\begin{array}{l} Q _ {n} (t, t ^ {\prime}) = \frac {1}{N} \left\langle \left(\boldsymbol {w} _ {R} - \boldsymbol {w} _ {n}\right) ^ {\top} \boldsymbol {\Sigma} (t, t ^ {\prime}) \left(\boldsymbol {w} _ {R} - \boldsymbol {w} _ {n}\right) \right\rangle + \frac {\gamma}{N} \left\langle \left(\boldsymbol {w} _ {R} - \boldsymbol {w} _ {n}\right) ^ {\top} \boldsymbol {\Sigma} (t, t ^ {\prime} + 1) \boldsymbol {w} _ {n} \right\rangle \\ + \frac {\gamma}{N} \left\langle \boldsymbol {w} _ {n} ^ {\top} \boldsymbol {\Sigma} \left(t + 1, t ^ {\prime}\right) \left(\boldsymbol {w} _ {R} - \boldsymbol {w} _ {n}\right) \right\rangle + \frac {\gamma^ {2}}{N} \left\langle \boldsymbol {w} _ {n} ^ {\top} \boldsymbol {\Sigma} \left(t + 1, t ^ {\prime} + 1\right) \boldsymbol {w} _ {n} \right\rangle , \tag {6} \\ \end{array}
$$
where  $\alpha = B / N$  and  $Q_{n}(t,t^{\prime}) = \langle \Delta_{n}(t)\Delta_{n}(t^{\prime})\rangle$  is the correlation of randomly sampled TD-errors at episodic times  $t,t^{\prime}$  and iteration  $n$ . The average over weights  $\langle \rangle$  denotes a Gaussian average whose moments are related to  $M_{n}$ . The correlation function  $Q_{n}(t,t^{\prime})$  depends on  $M_{n}$  and the average weights  $\langle \pmb{w}_n\rangle$ ; we provide its full formula in Appendix B.3, equation (B.17).
Proof. The full derivation is in Appendix B. At a high level, we track the moment generating function of the iterates  $\pmb{w}_n$  over random draws of features  $\{\pmb{\psi}_n^\mu(t)\}$ ,  $Z[\{\pmb{j}_n\}] = \mathbb{E}_{\{\pmb{\psi}_n^\mu(t)\}} \exp(i \sum_n \pmb{j}_n \cdot \pmb{w}_n) \propto \int \mathcal{D}q \exp\left(\frac{N}{2} S[q, \{\pmb{j}_n\}]\right)$  where  $S$  is a  $\mathcal{O}(1)$  action and  $q$  are a set of order parameters of the theory which include the following overlaps  $C_n(t, t') = \frac{1}{N} \pmb{w}_n^\top \pmb{\Sigma}(t, t') \pmb{w}_n$  and  $Q_n(t, t') = \frac{1}{B} \sum_{\mu=1}^{B} \Delta_n^\mu(t) \Delta_n^\mu(t')$ . In this high dimension  $N$ ,  $B \to \infty$  limit with  $B/N = \mathcal{O}(1)$  and episode length  $T = \mathcal{O}(1)$ , the order parameters can be obtained from saddle point integration, which requires solving  $\frac{\partial S}{\partial q} = 0$ . This procedure results in a deterministic learning curve given in equations (4),(5),(6) even though the realization of sampled states are disordered. The TD-error variables  $\Delta_n(t)$  become mean zero Gaussians and the  $\{\pmb{w}_n\}$  also follow a 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 it provides an 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 as 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  $\psi$  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  $\gamma = 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 learner with perfect coverage (infinite batch size) at each step will converge to the ground truth  $\pmb{w}_{TD} = \left(\bar{\pmb{\Sigma}} -\gamma \bar{\pmb{\Sigma}}_{+}\right)^{-1}\bar{\pmb{\Sigma}}\pmb{w}_{R}$  (see Appendix A).
Remark 5.  $M_{n}$  is equivalently defined as  $\begin{array}{r}M_{n} = \big{\langle}(\pmb {w} - \pmb{w}_{TD})(\pmb {w} - \pmb{w}_{TD})^{\top}\big{\rangle}_{\{\tau_{n^{\prime}}^{\mu}\}_{n^{\prime} <   n}} \end{array}$  , which measures deviation from the fixed point of gradient flow (vanishing learning rate) dynamics  $\pmb{w}_{TD}$  over random sets of sampled episodes (Appendix B).
3.2 Gaussian Approximation
The theory presented in Section 3.1 relies on an approximation of the feature distribution as 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 an exact, non-asymptotic theory for non-Gaussian features can be provided which closes under knowledge of the fourth cumulants of the features as we show in Appendix D, though this theory is especially cumbersome to analyze or evaluate compared to the theory of Section 3.1. Concretely, Proposition 3.1 relies on the following.
Gaussian Feature Assumption. The learning curves for a TD learner with high dimensional features  $\{\psi(s_t)\}_{t=1}^T$  over random  $\tau$  are well approximated by the learning curves of a TD learner trained with Gaussian features  $\psi_G \sim \mathcal{N}(\mu, \Sigma + \mu \mu^\top)$  with matching mean and correlations
$$
\boldsymbol {\mu} (t) = \left\langle \psi \left(s _ {t}\right) \right\rangle_ {\tau \sim p (\tau)}, \quad \boldsymbol {\Sigma} (t, t ^ {\prime}) = \left\langle \psi \left(s _ {t}\right) \psi \left(s _ {t ^ {\prime}}\right) ^ {\top} \right\rangle_ {\tau \sim p (\tau)}. \tag {7}
$$
where averages are taken over sequences of states  $\{s(t)\} \sim p(\tau)$
(a) 2D Exploration
(b) Place Cell/RBF Features  $\psi (s_t)$
(d) Place Cell Bandwidths
(e) Large batch  $B = 30$
(c) Equivalence Phenomenon
(f) Small batch  $B = 3$
Figure 1: An illustration of our theory for TD learning. (a) A diffusion process in a 2D grid world generates many possible trajectories through state space. Each colored line is a different trajectory. Reward function is shown in red, with darker red indicating higher reward. (b) When combined with nonlinear place cell feature representation, the state transitions generate a distribution over observed features  $\{\psi(s_t)\}$ . (c) The value error associated with TD learning for a bump reward function on the true features generated from a single set of MDP trajectories (blue) is compared to training on sampled Gaussian vectors  $\{\psi_t\}$  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 of 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-gradient 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
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-Gaussian features  $\{\psi (s_t)\}$  have least squares learning curves which are only sensitive to the first two cumulants of the distribution. We do not aim to provide a 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.