Title: TEMPORAL-DIFFERENCE LEARNING FOR NONLINEAR VALUE FUNCTION APPROXIMATION IN THE LAZY TRAINING REGIME

Abstract: We discuss the approximation of the value function for infinite-horizon discounted Markov Reward Processes (MRP) with nonlinear functions trained with the Temporal-Difference (TD) learning algorithm. We consider this problem under a certain scaling of the approximating function, leading to a regime called lazy training. In this regime the parameters of the model vary only slightly during the learning process, a feature that has recently been observed in the training of neural networks, where the scaling we study arises naturally, implicit in the initialization of their parameters. Both in the under-and over-parametrized frameworks, we prove exponential convergence to local, respectively global minimizers of the above algorithm in the lazy training regime. We then give examples of such convergence results in the case of models that diverge if trained with non-lazy TD learning, and in the case of neural networks.

1. INTRODUCTION:

In recent years, deep reinforcement learning has pushed the boundaries of Artificial Intelligence to an unprecedented level, achieving what was expected to be possible only in a decade and outperforming human intelligence in a number of highly complex tasks. Paramount examples of this potential have appeared over the past few years, with such algorithms mastering games and tasks of increasing complexity, from playing Atari to learning to walk and beating world grandmasters at the game of Go (Haarnoja et al., 2018; Mnih et al.; 2013; Silver et al., 2016; 2017; 2018) . Such impressive success would be impossible without using neural networks to approximate value functions and / or policy functions in reinforcement learning algorithms. While neural networks, in particular deep neural networks, provide a powerful and versatile tool to approximate high dimensional functions (Barron, 1993; Cybenko, 1989; Hornik, 1991) , their intrinsic nonlinearity might also lead to trouble in training, in particular in the context of reinforcement learning. For example, it is well known that nonlinear approximation of the value function might cause divergence in classical temporal-difference learning due to instability (Tsitsiklis & Van Roy, 1997) . Several algorithms have been proposed in the literature to address the issue of non-convergence (Bhatnagar et al., 2009; Maei & Sutton, 2010; Riedmiller, 2005; Sutton et al., 2009a; b; Szepesvári, 2010) , while practical deep reinforcement learning often employs and prefers basic algorithms such as temporal-difference (Sutton, 1988) and Q-learning (Watkins, 1989) due to their simplicity. It is thus crucial to understand the convergence of such algorithms and to bridge the gap between theory and practice. The theoretical understanding of deep reinforcement learning is of course rather challenging, as even for supervised learning, which can be viewed as a special case of reinforcement learning, deep neural networks are still far from being understood despite the huge amount of research focus in recent years. On the other hand, recent progress has led to an emerging theory for neural network learning at least in the regime of over-parametrization, including recent works on mean-field point of view of training dynamics (Chizat & Bach, 2018a; Mei et al., 2018; Rotskoff et al., 2019; Rotskoff & Vanden-Eijnden, 2018; Wei et al., 2018) and also the linearized training dynamics in the over-parametrized regime (Allen-Zhu et al., 2018a; b; Chizat & Bach, 2018b; Du et al., 2018a; b; Ghorbani et al., 2019a; Jacot et al., 2018; Lee et al., 2019; Oymak & Soltanolkotabi, 2019; Zou et al., 2018) . The main goal of this work is to analyze the dynamics of a prototypical reinforcement learning algorithm -temporal/difference (TD) learning -based on the recent progress in deep supervised learning. In particular, we will focus on the lazy training regime, inspired by the recent work Chizat & Bach (2018b) , and analyze TD learning in both over-parametrized and under-parametrized regimes with scaled value function approximations. Related Works. This work is closely related to the recent paper Chizat & Bach (2018b) , addressing the problem of lazy training in the supervised learning framework when models are trained through (stochastic) gradient descent. In particular, that paper introduced the scaling that we consider in this work as an explanation, e.g., of the small relative displacement of the weights of over-and under-parametrized neural networks for supervised learning. That work, however, leverages the gradient structure of the underlying vector field, which we lack in the present framework when the underlying policy is not reversible (Ollivier, 2018) . The linear stability analysis is also considered in the recent work Achiam et al. (2019) based on the neural tangent kernel (Jacot et al., 2018) for off-policy deep Q-learning. The groundbreaking paper Tsitsiklis & Van Roy (1997) proves convergence of TD learning for linear value function approximation, unifying the manifold interpretations of this convergence phenomenon that preceded it by highlighting that convergence of the algorithm is to be understood in the norm induced by the invariant measure of the underlying Markov process. Furthermore, the paper gives an illuminating counterexample for the extension of the linear result to the general, nonlinear setting. Our result shows that divergence does not occur in the lazy training regime. Concurrent work (Brandfonbrener & Bruna, 2019) has shown convergence and non-divergence of TD learning in the over-parameterized, respectively the under-parametrized regime, provided that the environment is sufficiently reversible. We note that working in the lazy training regime allows to ensure convergence independently on the reversibility of the environment and quantify the error of the fitted model in the under-parametrized regime. Finally, another concurrent work (Cai et al., 2019) analyzes global convergence of a modified TD algorithm for two-layer neural networks with ReLu nonlinearity when the width of the hidden layer diverges. In contrast, in the present paper we focus on the original TD(λ) learning algorithm for general approximators. Contributions. This paper proves that on-policy TD learning for policy evaluation (on-policy policyevaluation for short), a widely used algorithm for value function approximation in reinforcement learning, is convergent (asymptotically with probability one), in the lazy training regime, when the model is a nonlinear function of its parameters. More specifically, we prove convergence of this algorithm in both the under-and over-parametrized regime to local and global minima, respectively, of a natural, weighted error function (the projected TD error), and illustrate such convergence properties through numerical examples.

To obtain the result summarized above, we adapt the contraction conditions developed in the framework of linear function approximations to a nonlinear, differential geometric setting. Furthermore, we extend some existing results on the convergence in the lazy training regime of nonlinear models trained by gradient descent in the supervised learning framework to the world of reinforcement learning. This requires a generalization of the techniques developed in the gradient flow setting to non-gradient (i.e., rotational) vector fields such as the ones encountered in the TD learning framework.

2. MARKOV DECISION PROCESSES:

We denote a Markov Reward Process (MRP) by the 4-tuple (S, P, r, γ), where S is the state space, P = P (s, s ) s,s ∈S a transition kernel, r(s, s ) s,s ∈S is the real-valued, bounded immediate reward function and γ ∈ (0, 1) is a discount factor. In this context, the value function V : S → R + maps each state to the infinite-horizon, expected discounted reward obtained by following the Markov process defined by P . We assume that this Markov process satisfies the following assumption: Assumption 1. The Markov process with transition kernel P is ergodic and its stationary measure µ has full support in S. Furthermore we assume that S is compact.

In this note we are interested in learning the value (or cost-to-go) function V * (x) of a given MRP (S, P, r, γ), which is given by

EQUATION

where E s [ • ] denotes the expectation of the stochastic process s t starting at s 0 = s. More specifically we would like to estimate this function through a set of predictors V w (s) in a Hilbert space F parametrized by a vector w ∈ W := R p . We make the following assumption on such predictors:

Assumption 2. The parametric model V : R p → F mapping w → V w ( • ) is differentiable with Lipschitz continuous derivative DV : w → DV w (where DV w is a linear map from R p → F) with Lipschitz constant L DV defined WRT the operator norm.

A popular algorithm to solve this problem is given by value function approximation with TD(λ) updates (Sutton & Barto, 2018) . Starting from an initial condition w(0) ∈ W, for any λ ∈ [0, 1), this learning algorithm updates the parameters w of the predictor by the following rule: w(t + 1) := w(t) + β t δ(t)z λ (t) ,

(2) for a fixed sequence of time steps {β t } to be specified later, where the temporal-difference error δ(t) and eligibility vector z λ (t) are given by

δ(t) := r(s t , s t+1 ) + γV w(t) (s t+1 ) -V w(t) (s t ) z λ (t) := t τ =0 (γλ) t-τ ∇ w V w(t) (s τ ) . (3)

This work focuses on the asymptotic regime of small constant step-sizes β t → 0. In this adiabatic limit, the stochastic component of the dynamics is averaged out before the parameters of the model can undergo a significant change. This allows to consider the TD update as a deterministic dynamical system emerging from the averaging of the underlying stochastic algorithm. We focus on analysis of this deterministic system to highlight the aspect of nonlinear function approximation. The averaged, deterministic dynamics is given by the set of ODEs d dt

EQUATION

where E µ denotes the expectation with respect to the invariant measure of the underlying dynamics.

In the case of finite state space (|S| = d) we can represent V w as a vector in R d , while in general it is a function S → R, which we will restrict to the space L 2 (S, µ), namely square integrable function with respect to the measure µ.

To streamline our analysis of the TD algorithm, we define the TD operator T λ : L 2 (S, µ) → L 2 (S, µ):

T λ V (s) := (1 -λ) ∞ m=0 λ m E s m t=0 γ t r(s t , s t+1 ) + γ m+1 V (s m+1 ) .

Note that when λ = 0 the above operator acquires the simple form T 0 V := r + γP V for r(s) := E s [r(s, s )]. Then, denoting throughout by DV w the Fréchet derivative of V at w, it can be shown (Tsitsiklis & Van Roy, 1997, Lemma 8) (and is immediately verified in the special case λ = 0) that the continuous dynamics (4) for general λ < 1 can be written as

EQUATION

where we define throughout the inner product induced by the invariant measure µ (acting componentwise in expressions such as the one above) as

EQUATION

and denote by • µ the corresponding norm. Note that in the case |S| = d, denoting by Γ the d-dimensional diagonal matrix whose entries are the (positive) values of the invariant measure µ(s), one has a, b µ = a Γb. The extension of convergence results for the limiting, average dynamics we consider in this paper to convergence with probability one of the underlying, stochastic algorithm can be obtained through standard stochastic approximation arguments (Borkar & Meyn, 2000; Borkar, 2009) . More details on this straightforward extension are given in Remark 3.4 in Section 3 and in the appendix.

In this work, we are interested in a certain scaling of the TD learning algorithm with function approximation. More specifically, we consider the rescaled update

d dt w(t) = 1 α T λ (αV w(t) ) -αV w(t) , DV w(t) µ (7)

for large values of the scaling parameter α > 1 . One of the reasons why this scaling of the model is of practical interest is because it arises naturally when training neural networks, implicit in some widely applied choices of initial conditions, as we explain in Section 4.2. Furthermore, as we shall see below, under some mild assumptions for large values of α the parameters w of the model vary only slightly during training, inducing what is called the "lazy training" regime. A visual representation of the geometric effect of this scaling in the case where p < d < ∞ is given in Fig. 1 . (e.g., in (11) ) in the under-parametrized setting. The space of parameters (left) is mapped to the space of predictors (right) by the parametric model V . The scaling V → αV changes the manifold F w that the parameter space is mapped to (different surfaces on the right). In particular, this scaling "widens" the reach in the space of functions of the predictors within a ball of small radius in W, but at the same time it "flattens" that space (locally in W) bringing it closer to the tangential plane to the initial model V w(0) . Choosing V w(0) = 0 as in the picture above leaves the initial point of the dynamics (in predictor space) invariant under such transformation.

T V (w(0)) F w V (w(0)) αV (W) w(0) W F

3. MAIN RESULTS:



3.1. OVER-PARAMETRIZED REGIME:

In the over-parametrized setting we assume that DV w(0) is surjective, i.e., its singular values are uniformly bounded away from 0. This is only possible in the finite state space setting and is automatically the case if the number of parameters p is larger than the size of the state space S.

Admittedly, in applications such as AlphaGo (Silver et al., 2016; 2017) , it is unrealistic to overparametrize, but we start with this regime as it parallels the study of over-parametrized supervised learning for global convergence of the training loss. Analysis of the under-parametrized regime will be discussed in the next subsection. In order to state our first result, we introduce the scalar product in F defined by a, b 0 = a, g w(0) b where g w := (DV w • DV w ) -1 , and denote by • 0 the norm it induces. Note that g w is the metric tensor associated to the pushforward metric induced by the parametric model V : R p → F. We note that if DV w(0) has singular values that are uniformly bounded away from 0, the norms • µ , • 0 are equivalent, i.e., there exists κ > 0 such that κ -1 f 0 < f µ < κ f 0 for all f ∈ F . Theorem 3.1 (Over-parametrized case). Assume that σ min > 0, where σ min is the smallest singular value of DV w(0) . Assume further that w(0) is such that

EQUATION

Recall that V * is the exact value function given by (1). Moreover, if V w(0) 0 ≤ Cα -1 for a constant C > 0, then sup t>0 w(t) -w(0) = O(α -1 ).

Similarly to the proof in Chizat & Bach (2018b) , we first show that DV w and V w do not change much assuming that w stays in a small ball of radius . Then, combining this result with the Lipschitz continuous character of DV in w, one shows that w does indeed stay in the desired ball of radius . A similar computation can be done in our case. To bypass the absence of a strongly convex cost functional in our framework, which was crucial in the analysis of Chizat & Bach (2018b) , we adopt a strategy based on the use of a local Lyapunov function

EQUATION

where V * is the sought for value function (1). The theorem is based on some preparatory lemmas, proofs of which can be found in appendix. The first one states that for large values of the scaling parameter α the pushforward metric g w varies in a negligible way during training. Throughout, we denote by 1 the identity map in the corresponding space and by B µ (v), B 0 (v) and B (v) the balls with radius around v in • µ , • 0 and • 2 respectively. Lemma 3.2 (Perturbation of the metric). Let G 0 be a compact subset of a linear space G. For v(0) ∈ G 0 , let g v be a continuous, self-adjoint linear operator that is positive definite in a neighborhood of v(0) when restricted on G. Then for all ε > 0 there exists δ > 0 such that, for all

EQUATION

for a linear operator gv with gv < ε . More specifically, let σ min be the smallest singular value of DV w(0) . Then if ≤ (1 -γ)σ 2 min /(48L DV ), (10) holds with gV (w) < 1-γ 4 for all w ∈ B (w(0)). We also recall from Tsitsiklis & Van Roy (1997) the following contraction property of the TD operator in the • µ norm. For the convenience of readers, we recall the proof in the appendix.

Lemma 3.3. (Tsitsiklis & Van Roy, 1997, Lemmas 1, 3, 7) Under Assumption 1, for any V, Ṽ ∈ F we have that

T λ V -T λ Ṽ µ ≤ γ λ V -Ṽ µ for γ λ := γ 1-λ 1-γλ ≤ γ < 1 .

In particular there exists a unique fixed point of T λ , V * ∈ F given by (1).

The proof of Theorem 3.1 relies on the above lemma to establish decay of the local Lyapunov function U as long as w stays within a ball. The nonlinear effects become negligible when α is sufficiently large. The control of U in turn gives the bound of the change of w, which closes the argument. The details are given in the supplementary materials. Remark 3.4. Our results can be extended to show stability and convergence in the stochastic approximation setting, similarly to Bhatnagar et al. (2009) ; Tsitsiklis & Van Roy (1997), under the additional assumption that the step size {β t } satisfies the Robbins-Monro condition (Robbins & Monro, 1951) . For example, one can apply (Borkar & Meyn, 2000, Thms. 2.2, 2.4) guaranteeing almost sure convergence and exponential contraction of the expected error with probability one over the initial condition provided that the limiting vector field (in our case ( 7)) has a unique fixed point and is Lipschitz continuous. Lipschitz continuity is an immediate consequence of the linearity of T λ and the boundedness of closed balls in F together with the Lipschitz continuity of the models Assumption 2. The existence of a fixed point (1) in F of the limiting vector field is trivial while its uniqueness is shown in the proof of Theorem 3.1 in the appendix.

3.2. UNDER-PARAMETRIZED REGIME:

We now proceed to state and prove a convergence theorem in the under-parametrized case. The underlying assumption in this section is that the size of state space is larger than the number of parameters, which in turn bounds the rank r of DV w(0) from above: r < p < d (where possibly d = ∞). In this regime, in general, there is no hope that TD will converge to the true value function V * . In fact, the image of the operator T λ might not even lie in the space F w of approximating functions. However, the derivative DV w(t) in the TD update acts as a projection (WRT the product

• , • µ ) onto the tangent space of F w at V w(t) (more specifically, DV w(t) projects the image of T λ onto W, which is then mapped back to T V (w(t)) F w by DV w(t) ). We denote throughout by Π and Π 0 the projection operator under (6) onto T V (w(t)) F w and T V (w(0)) F w respectively. What one can hope for is that the TD algorithm converges to a locally "optimal" approximation Ṽ * of V * on the manifold F w , which is close to the best approximator Π 0 V * of V * on the linear tangent space T V (w(0)) F w .

Theorem 3.5 (Under-parametrized case). Assume that r := rank(DV w ) is constant in a neighborhood of w(0) and V w(0) = 0. Then there exists α 0 > 0 such that for any α > α 0 the dynamics (7) (and the corresponding approximation V w ) converge exponentially fast to a locally (in W) attractive fixed point Ṽ * , for which

Π(T λ Ṽ * -Ṽ * ) µ = 0 and Ṽ * -V * µ < 1-λγ 1-γ Π 0 V * -V * µ + O(α -1

). Note that for random initialization the constant rank assumption is generically satisfied. Indeed, the maximal rank property holds generically in W and thus WP1 at w(0) when the model parameters are initialized randomly. Furthermore, by the lower semicontinuity of the rank function the Jacobian DV will have maximal rank in an open subset of W. The main difference of the proof of the above result WRT the one in the over-parametrized regime is that DV w • DV w does not have full rank anymore. This implies on one hand that the norms • µ and • 0 are not equivalent in F, even though we still have • 0 ≤ κ • µ for a κ > 0, provided that Assumption 1 holds. On the other hand, as mentioned above, this implies that the model V w evolves on a submanifold F w of F, and that T λ does not, in general, map onto the tangential plane T V (w) F w of F w at V w . The action of T λ is then projected back onto T V (w) F w by the operator DV w(t) . The nonlinear structure of the space F w slightly complicates the proof WRT the over-parametrized case, and we apply standard differential geometric tools to map the problem back to a linear space.

W 0 F 0 W 0 F 0 V φ ψ πr

Proof. We apply the rank theorem (Boutaib, 2015; Lee, 2003) ( (Abraham et al., 2012) for the ∞-dimensional setting) to show that there exist sets W 0 , W 0 ⊆ R p , F 0 , F 0 ⊆ F and diffeomorphic maps φ : W 0 → W 0 , ψ : F 0 → F 0 where ψ•V •φ -1 = π r , φ(w(0)) = 0, ψ(V w(0) ) = 0 and, for an appropriate choice of bases, π r maps the coordinates of W 0 to the first r coordinates of F 0 , i.e., (x 1 , . . . , x p ) → (x 1 , . . . , x r , 0, 0, . . . ), where r is the rank of the operator DV w(0) . We denote by Π r the hyperplane in F spanned by the first r vectors of the basis. We recall that by Abraham et al. (2012) ; Boutaib (2015) ; Lee (2003) the maps, ψ, φ, π r are continuous with Lipschitz derivatives Dψ, Dφ, Dπ r respectively. We consider the trajectory of V w(t) := π r • φ(w(t)) = ψ(V w(t) ). Denoting by D• the Fréchet derivative at the corresponding point of the dynamics and noting that DV = Dψ -1 Dπ r Dφ we have

EQUATION

so V remains in Π r . As a consequence of the above we can naturally define a metric (the pushforward metric) on F 0 by the tensor ḡv = (Dπ r DφDφ Dπ r ) -1 . In fact, by choosing the metric tensor to be constant on F 0 , i.e., equal to ḡ0 for all v ∈ F 0 , we equip the linear space F 0 with a scalar product • , • 0 . This, in turn, directly induces a norm • 0 on the same space. We now proceed to use such simple metric structure to establish the existence and uniqueness of a fixed point of (11) in F 0 for α large enough. The result of our theorem follows from (Simpson-Porco & Bullo, 2014, Proposition 4.1), which establishes uniqueness and exponential contraction at rate > 0 of a dynamical system evolving under the flow of a vector field X given by the RHS of (11) in a forward invariant set F 0 provided that for every geodesic γ(s) in F 0 (12) holds. Therefore, the proof of convergence is concluded by applying Lemma 3.6 and Lemma 3.7, whose proofs can be found in supplementary materials. The proof of the optimality of the fixed point is postponed as Lemma A.1 in the appendix.

Lemma 3.6. There exists δ > 0 and α 0 > 0 such that the ball B 0 δ (0) ⊆ F 0 is forward invariant and forward complete with respect to the dynamics of ( 7) for all α > α 0 . Lemma 3.7. There exists > 0, δ > 0 and α 0 > 0 such that for all α > α 0 and all geodesics γ(s) contained in the ball B 0 δ (0) ⊆ F 0 , the function γ (s), X(γ(s)) 0 -s γ (0), γ (0) 0 ,

is strictly decreasing in s.

Remark 3.8. The proof of Theorem 3.5 can be straightforwardly generalized to the case where the initial condition V 0 is not identically 0 but within B µ (α) (0) for (α) going to 0 with α → ∞. This generalization, however, requires the map V to be uniformly Lipschitz smooth for w ∈ W 0 . Among other things, this extension allows to explicitly cover the training of randomly initialized, single layer neural networks.

4. NUMERICAL EXAMPLES 4.1 A DIVERGENT NONLINEAR APPROXIMATOR:

We illustrate the convergence properties of TD learning in the lazy training regime in the underparametrized case by applying it to the classical framework of (Tsitsiklis & Van Roy, 1997, Section X). This reference gives an example of a family of nonlinear function approximators that diverge when trained with the TD method. The intuition behind this counterexample is that one can construct a manifold of approximating functions F w in the form of a spiral, with the same orientation as the rotation of the vector field induced by the TD update in the space of functions. By choosing the windings of the spiral to be dense enough, the projection of the TD vector field follows the spiral in the outward direction, leading to a divergence of the algorithm, as displayed schematically in Fig. 2a . More specifically, consistently with Tsitsiklis & Van Roy (1997), we parametrize the manifold F w as V ϑ := e εϑ (a cos( λϑ) -b sin( λϑ)) -V * for a = (10, -7, -3), b = (2.3094, -9.815, 7.5056), ε = 0.01, λ = 0.866. We choose the discount γ = 0.9 and a step-size of β t ≡ 2 × 10 -3 , while the underlying Markov chain is defined by the transition matrix P ij = (δ j,mod(i,3)+1 +δ i,j )/2, where δ i,j is the Kronecker delta function and equals 1 if i = j and 0 else. We note that the step-size does not affect the convergence properties of the algorithm, as argued in Tsitsiklis & Van Roy (1997) , where the immediate reward was set to r = (0, 0, 0). Note that, as realizing the conditions of Theorem 3.5 would start the simulation at the solution V * = (0, 0, 0), we shift both the solution and the manifold of approximating functions by the same vector in the embedding space, leaving the new solution V * = -V 0 = -a at the center of the spiral, i.e., realized at ϑ = -∞. This corresponds to choosing an average reward r = (-6.85, 8.35, -1.5 ). We note that by the affine nature of the TD update, this change in r results in a global shift of the TD vector field in F and does not affect the update of ϑ. In particular, this means that the TD update remains divergent for every initial condition different than the solution V * . We run the TD update in the off-centered situation both for values of α = 1 (the classical, divergent regime) and α = 100. As explained in the previous sections, this scaling of the approximating function makes the TD update convergent, as displayed in Fig. 3a . Indeed, under this scaling the solution converges to a local minimum of the dynamics. The intuition behind the convergence of the algorithm is outlined in Fig. 2 : when α is large we are in an almost linear regime where the TD update converges.

4.2. SINGLE LAYER NEURAL NETWORKS:

We show that the regime of study arises naturally in one hidden layer neural networks for a certain family of initialization. We consider the example of ReLu activation, i.e., when the model is given by

EQUATION

) for s ∈ R m and N distinct (m + 2)-dimensional vectors w i = (a i , (b i ) 1 , . . . , (b i ) m , c i ) i∈(1,...,N ) .

Typical initialization of the weights of the above model is of the form a i iid

∼ N (0, 1/ √ N ), (b i ) j iid ∼ N (0, 1/ √ m

) for all j and c i iid ∼ N (0, 1). However, by the linearity of ( 13) in a i , by the rescaling property of normal distribution this is equivalent to writing

EQUATION

for an N -dependent α(N ) = √ N (diverging in N ), a i iid ∼ N (0, 1), (b i ) j iid ∼ N (0, 1/ √ d) and c i iid

∼ N (0, 1) 1 . Therefore, this common choice of initial conditions implicitly starts the training of the above model in the lazy regime (Ghorbani et al., 2019b) . We train the model ( 14) by TD learning (7) with fixed step-size β t ≡ 10 -3 both in the over-and under-parametrized regime. To do so, we draw an objective function V * randomly with distribution V * (s) iid ∼ N (0, 1) for all s ∈ S on a grid 1 A heuristic justification that the scaling the parameters of the neural network by α(N )/N = 1/ √ N leads to lazy training while the scaling N -1 is natural for the model Vw and does not lead to the lazy regime can be found in Chizat & Bach (2018b) . This natural scaling is studied in depth in Chizat & Bach (2018a) of d equally spaced points on the interval [-1, 1]. We then compute the corresponding average reward by solving the TD equation: r = (1 -γP )V * , and train the model ( 7) for λ = 0, γ = 0.9 (when not specified otherwise) with transition matrix P ij = (δ j,mod(i,d)+1 + δ i,j )/2. To respect the conditions of Theorem 3.5, we initialize half of the parameters of the neural network as explained above, while the other half is obtained by replicating the values of b i , c i and inverting the one of a i → -a i . This "doubling trick" introduced in Chizat & Bach (2018b) produces a neural network with V w(0) ≡ 0 and randomly initialized weights with the desired distribution. We consider situations where N = 10, d = 50 (under-parametrized, taking α = 100) and N = 100, d = 30 (over-parametrized, with α = 500), and plot the convergence to local, respectively global minima in Fig. 3b .

5. DISCUSSION AND CONCLUSION:

In this work we have proven the convergence properties of the TD learning algorithm with nonlinear value function approximation in the lazy training regime. In this regime, the algorithm behaves essentially like a linear approximator spanning the tangential space of the approximating manifold (in function space) at initialization. As such, the training converges exponentially fast with probability one to the global minimum or a local fixed point depending on the codimension of the approximating manifold in the search space. This guarantees convergence with little parametric displacement. This phenomenon can be intuitively understood as an effect of the linearized regime in which the neural networks are trained which reduces them, in the limit, to a randomized kernel method (more precisely a Neural Tangent Kernel (Jacot et al., 2018) ). In this sense, convergence of lazy models may come at the expense of their expressivity. Recent works (Chizat & Bach, 2018b; Ghorbani et al., 2019b) discuss the approximating power of lazy neural networks in the supervised setting, highlighting their limits WRT their non-lazy counterparts and naturally comparing them with random feature models (Yehudai & Shamir, 2019) , but an exhaustive study of the expressivity of these models, in particular in the context of reinforcement learning is still lacking. Nonetheless, the results proven in this work emphasize the interest of this regime in the framework of deep reinforcement learning, where models often suffer from divergent behavior especially during early stages of training.

Future directions of research include the extension of these results to more complex, nonlinear reinforcement learning algorithms such as Q-learning, and the development of more refined, nonasymptotic versions of the above theorems. Furthermore, a more thorough exploration of the relationship between the limiting results in Chizat & Bach (2018a) and the ones presented here and in Chizat & Bach (2018b) while transposing those to the framework of reinforcement learning would be important for the understanding of the limiting dynamics of neural networks in this domain.

A SUPPLEMENTARY PROOFS:

To simplify the notation in the forthcoming analysis, we slightly abuse the notation used when the state space is finite-dimensional extending it, when necessary, to the infinite-dimensional setting. This naturally generalizes matrix multiplication to the action of linear operators. In particular the action of Γ, which we recall in the finite-dimensional setting is a diagonal matrix with entries µ(s), is to be intended as (a Γb) ij = S a i (s)b j (s) µ(ds) .

Furthermore, we introduce the following decomposition of the TD operator:

T λ V = rλ + γP λ V , where rλ (s) := (1-λ) ∞ m=0 λ m E s m t=0 γ t r(s t , s t+1 ) , P λ V (s) := (1-λ) ∞ m=0 (λγ) m E s [V (s m+1 )] ,

or, in vector notation

rλ := (1 -λ) ∞ m=0 λ m m t=0 γ t P t r , P λ V (s) := (1 -λ) ∞ m=0

(λγ) m P m+1 V .

In the proofs below, we will use the above, simplified notation to obtain contraction estimates on the dynamical system (4). These estimates will leverage the fact that P λ is nonexpansive and γ < 1, and from this notation contraction rates in terms of γ will arise naturally. However, by Lemma 3.3, we know that the contraction rate of T λ is γ λ . Rewriting the proofs with γ → γ λ will show the stronger contraction.

A.1 OVER-PARAMETRIZED REGIME Lemma 3.2 (Perturbation of the metric). Let G 0 be a compact subset of a linear space G. For v(0) ∈ G 0 , let g v be a continuous, self-adjoint linear operator that is positive definite in a neighborhood of v(0) when restricted on G. Then for all ε > 0 there exists δ > 0 such that, for all

v ∈ B δ (v(0)) ⊆ G 0 g v(0) = (1 + gv )g v ,

for a linear operator gv : F → F with gv < ε . More specifically, let σ min be the smallest singular value of DV w(0) . Then if ≤ (1 -γ)σ 2 min /(48L DV ), (10) holds with gV (w) < 1-γ 4 for all w ∈ B (w(0)).

Proof of Lemma 3.2. Let B w = DV w(0) DV w(0) -DV w DV w . We carry out the proof for the case σ min < 1 (else the result holds with σ min = 1 in ), in which case we have for all w ∈ B (w(0)) that

B w ≤ 2L DV w(0) -w + (L DV w(0) -w ) 2 ≤ 3L DV w(0) -w .

Then we can write

g w(0) = (DV w(0) DV w(0) ) -1 = (DV w DV w + B w ) -1 = (g -1 w (1 + g w B w )) -1 = (1 + g w B w ) -1 g w = ∞ n=0 (-1) n (g w B w ) n g w = g w + ∞ n=1 (-1) n (g w B w ) n g w .

Furthermore, by the assumptions on the regularity of V and on the initial condition w(0) we have that g w 4/σ 2 min 1, provided that w ∈ B (w(0)) for as in Lemma 3.2. Therefore, the perturbation gw :=

∞ n=1 (-1) n (g w B w ) n satisfies gw = ∞ n=1 (-1) n (g w B w ) n ≤ ∞ n=1 g w B w n ≤ ∞ n=1 3L DV σ 2 min /4 w(0) -w n ≤ 1 -γ 4 .

The same proof applies in the general case with different, implicit constants.

Lemma 3.3. (Tsitsiklis & Van Roy, 1997, Lemmas 1, 3, 7) Under Assumption 1, for any V, Ṽ ∈ F we have that

T λ V -T λ Ṽ µ ≤ γ λ V -Ṽ µ for γ λ := γ 1 -λ 1 -γλ ≤ γ < 1 . (A.1)

In particular there exists a unique fixed point of T λ , V * ∈ F given by (1).

Proof of Lemma 3.3. We first prove that P V µ ≤ V µ . This follows by Jensen inequality and by the invariance of µ:

P V 2 µ = V P ΓP V = S µ(ds)( S P (s, ds )V (s )) 2 ≤ S 2 µ(ds)P (s, ds )V (s ) 2 = S µ(ds)V (s) 2 = V 2 µ . (A.2)

Then, writing

T λ V (s) = (1 -λ) ∞ m=0 λ m E s m t=0 γ t r(s t , s t+1 ) + γ m+1 V (s m+1 ) = (1 -λ) ∞ m=0 λ m m t=0 γ t E s [r(s t )] + γ m+1 E s [V (s m+1 )] = (1 -λ) ∞ m=0 λ m m t=0 γ t P t r(s) + (γP ) m+1 V (s) ,

where s t is the process on S induced by P with initial condition s 0 , we have contraction of the operator T λ in L 2 (S, µ) by

T λ (V -Ṽ ) µ = (1 -λ) ∞ m=0 λ m (γP ) m+1 V (s) -Ṽ (s) µ ≤ (1 -λ) ∞ m=0 λ m γ m+1 V (s) -Ṽ (s) µ = γ(1 -λ) 1 -γλ V (s) -Ṽ (s) µ ,

where in the inequality above we have used (A.2). This proves that T λ is a contraction in F, and as such it must have a unique fixed point. That this fixed point corresponds to (1) is immediately checked by direct computation.

Theorem 3.1 (Over-parametrized case). Assume that σ min > 0, where σ min is the smallest singular value of DV w(0) . Assume further that w(0

) is such that V w(0) 0 < M := (1 -γ) 2 σ 2 min /(192κ 2 L DV DV w(0) ), then for α > α 0 := V * 0 /M we have for all t ≥ 0 that V * -αV w(t) 2 0 ≤ V * -αV w(0) 2 0 e -1-γ 2κ 2 t . (A.3)

Recall that V * is the exact value function given by (1). Moreover, if V w(0) 0 ≤ Cα -1 for a constant C > 0, then sup t>0 w(t) -w(0) = O(α -1 ).

Proof of Theorem 3.1. By setting := (1 -γ)σ 2 min /(48L DV ) and by the assumed Lipschitz smoothness of V , DV w • DV w σ 2 min /4 as long as w ∈ B (w(0)). We would like to check a local exponential contraction condition, i.e., that for all w(t) ∈ B (w(0)) we have for any ε > 0 by using the linearity of the Hessian and bounding its operator norm of ψ -1 on a compact space in F 0 while choosing α large enough and δ small enough, since γ(s) ∈ B 0 δ (0). Note that if Dψ -1 is not differentiable, the above computation is to be understood in the sense of distributions. We now focus on the second term of (A.23). In this case we incorporate the operator Γ in the inner product and write this term as Getting back to the last two terms in (A.12), we immediately see from Lemma 3.2 that gγ(s) is a small, Lipschitz continuous perturbation. Hence, the product γ (s), gγ(s) X (γ(s))

d

can be bounded from above similarly to (A.13), while the second order derivative in the third term of (A. This directly gives (A.21) by choosing large enough.

The next lemma estimates the distance between the fixed point Ṽ * of the dynamics (7) and V * given by ( 1), showing that it is close, for large values of α to the best linear model in the tangent space of F w at V w(0) , given by Π 0 V * . We recall that the projection operator Π 0 onto the linear space spanned by the columns of DV is given by (Tsitsiklis & Van Roy, 1997, Eq. ( 1))

Π 0 W := arg min

{DV w(0) ∆w : ∆w∈R p } DV w(0) ∆w -W µ = DV w(0) (DV w(0) ΓDV w(0) ) -1 DV w(0) ΓW , for all W ∈ F where, if necessary, we interpret (DV w(0) ΓDV w(0) ) -1 as a pseudo-inverse.

Lemma A.1. Let Ṽ * be the fixed point of ( 7) and V * be the global fixed point of the TD operator, given by (1). Then under the assumptions of Theorem 3.5 there exists constants α 0 > 0 and C * > 0 (independent of α 0 ), such that

Ṽ * -V * µ < 1 -λγ 1 -γ Π 0 V * -V * µ + C * α -1 , (A.27)

where Π 0 is the projection operator onto T V (w(0)) F w .

:

We recall that by the construction of the mappsings ψ, φ, πr and by our assumption in Theorem

3.5 the metric tensor ḡt has full rank on F0 and being the latter set compact its eigenvalues are uniformly bounded from below. At the same time, we can equip F0 with the metric induced by Γ by restricting it to its first r elements, which are uniformly bounded from below. Hence, the two metrics are equivalent on this space for some equivalence constant κ.

Figures and Tables:

Figure 1: Schematic representation of the effect of the linear scaling of the approximating function(e.g., in (11)) in the under-parametrized setting. The space of parameters (left) is mapped to the space of predictors (right) by the parametric model V . The scaling V → αV changes the manifold F w that the parameter space is mapped to (different surfaces on the right). In particular, this scaling "widens" the reach in the space of functions of the predictors within a ball of small radius in W, but at the same time it "flattens" that space (locally in W) bringing it closer to the tangential plane to the initial model V w(0) . Choosing V w(0) = 0 as in the picture above leaves the initial point of the dynamics (in predictor space) invariant under such transformation.

Figure 2: Schematic representation of the manifold F w for the example in Section 4.1 before (a) and after (b) scaling of α. The underlying vector field represents the TD error δ(V ) from (3), whose projection on T ϑ F w gives the dynamics of the TD update in F w . In (a) this projection points "outwards" along the spiral, while (b) it has a fixed point close to 0. The scaling yields an effective "linearization" of the manifold around 0. The red point marks the global fixed point of the vector field.

Figure 3: Results of the training of nonlinear value function approximation with TD learning for the examples described in Section 4.1 (a) and Section 4.2 (b). In (a), we plot the µ-norm of the projected TD error Π(T λ V -V ). This quantity measures the increments of the model parameters during training and vanishes at a local minimum of the TD dynamics. We see that the algorithm diverges for α = 1 (blue curve), but converges to a local minimum for α = 100. In (b, above) we plot the MSE of single layer neural network during training in the over-parametrized regime (N = 100, d = 30, α = 500 ) for different choices of γ (0.8, 0.83, 0.85, 0.87, 0.9), showing exponential convergence (at different rates) to the global minimum claimed in Theorem 3.1. In (b, below) we again plot the norm of the the projected TD error for a neural network in the under-parametrized regime (N = 10, d = 50, α = 100) for different initial conditions, showing that the dynamics converge to a local fixed point.

∆v, DT λ Dψ -1 γ(s) ∆v µ -Dψ -1 γ(s) ∆v 2 µ .Now, by the contraction property of T λ onto the tangential space T ψ -1γ(s)F in the norm • µ we can writeDψ -1 γ(s) ∆v, DT λ Dψ -1 γ(s) ∆v µ ≤ Dψ -1 γ(s) ∆v µ P λ Dψ -1 γ(s) ∆v µ ≤ γ Dψ -1 γ(s) ∆v 2 µ , so that Dψ -1 γ(s) ∆v, DT λ Dψ -1 γ(s) ∆v µ -Dψ -1 γ(s) ∆v 2 µ ≤ (γ -1) Dψ -1 γ(s) ∆v 2 µ . (A.25)Denoting by σ Dψ -1 max , σ Dψ -1 min the largest and smallest, respectively, singular values of the map Dψ -1 in B 0 δ (0) (which are bounded away from 0 upon possibly making this set smaller), by nondegeneracy of Dψ -1 and by the equivalence of the • µ and • 0 norms on F 0 we have thatκ -1 σ Dψ -1 min ∆v 0 ≤ ∆v µ σ Dψ -1 min ≤ Dψ -1 γ(s) ∆v µ ≤ ∆v µ σ Dψ -1 max ≤ κ ∆v 0 σ Dψ -1 max .

22) can be dealt with analogously to what is done in (A.24), giving terms ε ∆v 2 0 and ε (3) ∆v 2 0 respectively, both going to 0 as δ → 0. Therefore, combining the above with (A.24), (A.25) and (A.26) we have

Peer review (iclr2020):
The paper discusses the policy evaluation problem using temporal-difference (TD) learning with nonlinear function approximation. The authors show that in the “lazy training” regime both over- and under-parametrized approximators converge exponentially fast, the former to the global minimum of the projected TD error and the latter to a local minimizer of the same error surface. Simply put, the lazy regime refers to the approximator behaving as if it had been linearized around its initialization point. This can happen if the approximation is rescaled, but can also occur as a side-effect of its initialization. The authors present simple numerical examples illustrating their claims. Although I did not carefully check the math, this seems like a solid contribution on the technical side. My main concern about the paper is that it falls short in providing intuition and contextualizing its technical content. Regarding the presentation, I believe it is possible to have a less dry prose without sacrificing mathematical rigor. If some of the technical material is moved to the appendix --like auxiliary results, discussion on proof techniques, etc--, the additional space could be used to discuss the implications of the theoretical results in more accessible terms. For example, a subject that ought to be discussed more clearly is the nature of the approximation induced by the lazy training regime. As far as I understand, this regime can be thought of as a sort of regularization that severely limits the capacity of the approximator. Although the authors mention in the conclusion that “...convergence of lazy models may come at the expense of their expressivity”, after reading the paper I do not have a clear sense of how expressive such models actually are. In their experiments, Chizat et al. (2018) observed that the performance of commonly used neural networks degrades when trained in the lazy regime --to a point that they consider it unlikely that the successes of deep learning can be credited to this regime of training. It seems to me that this subject should be more explicitly discussed in a paper that sets out to provide theoretical support for deep reinforcement learning. Still regarding the behavior of lazy approximators, my intuitive understanding is that they work as a linear model using random features. If this interpretation is correct, this makes the theoretical results a bit less surprising. They are still interesting, though, for they can be seen as relying on a “smoother” version of the linearity assumption often made in the related literature. Maybe this is also something worth discussing? Still on this subject, it seems to me that one potential disadvantage of lazy models with respect to their linear counterparts is that it is less clear how to enforce the lazy regime in practice. In Section 4.2 the authors discuss how this can naturally happen as a side-effect of the initialization, but it is unclear how applicable the particular strategy used to illustrate this phenomenon, with the “doubling trick”, is in practice. This is another example of a less technical discussion that would make the paper a stronger contribution.