Neural tangent kernel analysis advances deep reinforcement learning understanding
Bridging Neural Tangent Kernel Analysis and Reinforcement Learning for Deep Function Approximation
Machine Learning
Summary
Understanding how deep neural networks learn and make decisions is hard, especially in reinforcement learning where the model learns from sequences of actions and rewards. The authors extended a mathematical tool called the neural tangent kernel (NTK), which originally helped explain learning in simple settings, to work with reinforcement learning. This new approach helps track how the model’s estimated value of different states changes over time and suggests a new way to explore environments effectively. They tested their ideas on a grid world and showed it can learn more efficiently and adapt when the data changes.
What this means in practice
- •For reinforcement learning engineers: Improve sample efficiency by using NTK-based models to predict value function evolution and guide exploration in RL tasks with changing data distributions.
- •For autonomous system developers: Design exploration strategies in deterministic environments guided by learned constraints from NTK formulations to better cover feasible states during learning.
Tested on simulated data.
Authors
William St-Arnaud, Margarida Carvalho, Golnoosh Farnadi
Abstract
Despite their widespread success across machine learning, the learning dynamics and generalization properties of neural networks remain poorly understood. The neural tangent kernel (NTK) provides an analytical characterization of the learning dynamics of infinitely wide neural networks and serves as a useful proxy for finite-width networks in supervised learning. In this work, we extend the NTK framework to reinforcement learning (RL) by introducing a temporal-difference (TD) flow formulation that accommodates time-varying mini-batches and TD targets, yielding a closed-form solution for value function evolution in an overparametrized setting. We also devise an optimization-based state-space exploration method that leverages the closed-form solution, coupled with constraints on the environment, to model the set of feasible states. On a deterministic grid environment, we validate that the proposed RL-based NTK formulation is indeed learning efficiently and adapting to the distribution shift in the training data distribution, and we demonstrate that the resulting exploration strategy improves sample efficiency. To our knowledge, this is the first work to extend NTK theory to characterize the learning dynamics of RL training with time-varying mini-batches and TD updates.