An NTK Theory Approach to UCB Estimation in Semi-gradient TD Learning

Efficient exploration is a central challenge in reinforcement learning ({RL}), especially under sparse rewards, large state-action spaces, or adversarial environment dynamics. A principled approach to exploration is to choose actions that maximize the upper bound return. Under certain conditions, this can be formulated in the {Bayesian} framework as exploration based on the posterior variance. In practice, the predictive variance of a neural network is often used as a surrogate for the posterior variance by techniques such as ensembles, randomized priors, and random network distillation. However, for wide neural networks, neural tangent kernel ({NTK}) theory proves there is a discrepancy between the predictive variance of a wide neural network trained by gradient descent and its posterior variance. This is a known issue in supervised learning and there exist methods to reconcile this discrepancy. In this work, we show that, in contrast to the supervised setting, reconciliation is not possible for semi-gradient temporal difference ({TD}) learning, which is a commonly used algorithm in deep {RL}. Instead, we derive the upper confidence bound ({UCB}) of the value function in semi-gradient {TD} learning directly and show that this {UCB} can be estimated by a neural network. To aid our analysis, we extend {NTK} theory to regularized semi-gradient {TD} learning.