Can Temporal-Difference and Q-Learning Learn Representation? A Mean-Field Theory

TL;DR

This paper uses mean-field theory to analyze how overparameterized neural networks employed in temporal-difference and Q-learning evolve their feature representations, showing convergence to optimal features and global minimization of Bellman error.

Contribution

It provides a novel mean-field theoretical framework demonstrating convergence and optimality of feature representations in neural network-based reinforcement learning.

Findings

Neural networks minimize Bellman error at a sublinear rate.

Feature representations converge to the optimal one.

Analysis extends to soft Q-learning and policy gradient methods.

Abstract

Temporal-difference and Q-learning play a key role in deep reinforcement learning, where they are empowered by expressive nonlinear function approximators such as neural networks. At the core of their empirical successes is the learned feature representation, which embeds rich observations, e.g., images and texts, into the latent space that encodes semantic structures. Meanwhile, the evolution of such a feature representation is crucial to the convergence of temporal-difference and Q-learning. In particular, temporal-difference learning converges when the function approximator is linear in a feature representation, which is fixed throughout learning, and possibly diverges otherwise. We aim to answer the following questions: When the function approximator is a neural network, how does the associated feature representation evolve? If it converges, does it converge to the optimal one? We prove that, utilizing an overparameterized two-layer neural network, temporal-difference and Q-learning globally minimize the mean-squared projected Bellman error at a sublinear rate. Moreover, the associated feature representation converges to the optimal one, generalizing the previous analysis of Cai et al. (2019) in the neural tangent kernel regime, where the associated feature representation stabilizes at the initial one. The key to our analysis is a mean-field perspective, which connects the evolution of a finite-dimensional parameter to its limiting counterpart over an infinite-dimensional Wasserstein space. Our analysis generalizes to soft Q-learning, which is further connected to policy gradient.