[Paper Review] Temporal-difference learning with nonlinear function approximation: lazy training and mean field regimes

Andréa Agazzi, Jianfeng Lu|arXiv (Cornell University)|May 27, 2019

Reinforcement Learning in Robotics62 references4 citations

TL;DR

This paper analyzes temporal-difference (TD) learning with wide neural networks for value function approximation in Markov Reward Processes, comparing the lazy training regime—where network weights change minimally—and the mean-field regime. It proves exponential convergence to global or local minimizers in the lazy regime and shows all fixed points are global minimizers in the mean-field regime, ensuring optimal approximation.

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 first 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 compare this scaling of the parameters to the mean-field regime, where the approximately linear behavior of the model is lost. Under this alternative scaling we prove that all fixed points of the dynamics in parameter space are global minimizers. We finally give examples of our convergence results in the case of models that diverge if trained with non-lazy TD learning, and in the case of neural networks.

Motivation & Objective

• To analyze the convergence and optimality of TD learning when using wide neural networks for value function approximation.
• To establish theoretical guarantees for TD learning in the lazy training regime, where weight updates are small and the model behaves approximately linearly.
• To compare the lazy regime with the mean-field regime, where nonlinear dynamics dominate, and assess their respective convergence and approximation properties.
• To bridge the gap between theory and practice in deep reinforcement learning by providing convergence results under realistic scaling assumptions.
• To extend existing convergence results from supervised learning to the non-gradient, rotational vector fields typical in TD learning.

Proposed method

• Adapts contraction conditions from linear function approximation to a nonlinear, differential geometric setting for TD learning dynamics.
• Introduces a scaled parameter initialization that leads to the lazy training regime, where neural network parameters evolve slowly during training.
• Applies techniques from mean-field theory to analyze the nonlinear dynamics of wide neural networks in the mean-field scaling regime.
• Proves convergence to global minimizers in the mean-field regime by showing all fixed points of the parameter dynamics are optimal.
• Extends convergence results from gradient-based supervised learning to non-gradient TD learning frameworks, handling rotational vector fields.
• Uses a projected TD error as the objective function and analyzes convergence in both under- and over-parametrized settings.

Experimental results

Research questions

• RQ1 Does TD learning with wide neural networks converge in the lazy training regime, and to what extent?
• RQ2 How does the lazy training regime compare to the mean-field regime in terms of convergence speed and approximation quality?
• RQ3 Are all fixed points of the TD learning dynamics globally optimal in the mean-field regime?
• RQ4 Can the theoretical convergence guarantees of linear function approximation be extended to nonlinear, wide neural networks in reinforcement learning?
• RQ5 What role does the initialization scaling play in ensuring stable and convergent training in nonlinear function approximation?

Key findings

• In the lazy training regime, TD learning converges exponentially fast to local minimizers in the under-parametrized case and to the global minimizer in the over-parametrized case.
• The lazy training regime ensures convergence independently of the reversibility of the underlying Markov process, unlike some prior results.
• In the mean-field regime, all fixed points of the TD learning dynamics are global minimizers, implying perfect approximation of the true value function.
• The mean-field regime exhibits stronger approximation power than the lazy regime, as all stationary points are optimal.
• The theoretical framework extends convergence results from supervised learning to TD learning by generalizing contraction-based analysis to non-gradient vector fields.
• Numerical examples demonstrate the convergence properties and highlight cases where standard TD learning diverges but the lazy regime remains stable.