Offline Neural Contextual Bandits: Pessimism, Optimization and Generalization

1 INTRODUCTION . We consider the problem of offline policy learning ( OPL ) ( Lange et al. , 2012 ; Levine et al. , 2020 ) where a learner infers an optimal policy given only access to a fixed dataset collected a priori by unknown behaviour policies , without any active exploration . There has been growing interest in this problem recently , as it reflects a practical paradigm where logged experiences are abundant but an interaction with the environment is often limited , with important applications in practical settings such as healthcare ( Gottesman et al. , 2019 ; Nie et al. , 2021 ) , recommendation systems ( Strehl et al. , 2010 ; Thomas et al. , 2017 ) , and econometrics ( Kitagawa & Tetenov , 2018 ; Athey & Wager , 2021 ) . Despite the importance of OPL , theoretical and algorithmic progress on this problem has been rather limited . Specifically , most existing works are restricted to a strong parametric assumption of environments such as tabular representation ( Yin & Wang , 2020 ; Buckman et al. , 2020 ; Yin et al. , 2021 ; Yin & Wang , 2021 ; Rashidinejad et al. , 2021 ; Xiao et al. , 2021 ) and more generally as linear models ( Duan & Wang , 2020 ; Jin et al. , 2020 ; Tran-The et al. , 2021 ) . However , while the linearity assumption does not hold for many problems in practice , no work has provided a theoretical guarantee and a practical algorithm for OPL with neural network function approximation . In OPL with neural network function approximation , three fundamental challenges arise : Distributional shift . As OPL is provided with only a fixed dataset without any active exploration , there is often a mismatch between the distribution of the data generated by a target policy and that of the offline data . This distributional mismatch can cause erroneous value overestimation and render many standard online policy learning methods unsuccessful ( Fujimoto et al. , 2019 ) . To guarantee an efficient learning under distributional shift , common analyses rely on a sort of uniform data coverage assumptions ( Munos & Szepesvári , 2008 ; Chen & Jiang , 2019 ; Brandfonbrener et al. , 2021 ; Nguyen-Tang et al. , 2021 ) that require the offline policy to be already sufficiently explorative over the entire state space and action space . To mitigate this strong assumption , a pessimism principle that constructs a lower confidence bound of the reward functions for decision-making ( Rashidinejad et al. , 2021 ) can reduce the requirement to a single-policy concentration condition that requires the coverage of the offline policy only on the target policy . However , Rashidinejad et al . ( 2021 ) only uses this condition for tabular representation and it is unclear whether complex environments such as ones that require neural network function approximation can benefit from this condition . Moreover , the single-policy concentration condition still requires the offline policy to be stationary ( e.g. , the actions in the offline data are independent and depend only on current state ) . However , this might not hold for many practical scenarios , e.g. , when the offline data was collected by an active learner ( e.g. , by an Q-learning algorithm ) . Thus , it remains unclear what is a minimal structural assumption on the distributional shift that allows a provably efficient OPL algorithm . Optimization . Solving OPL often involves in fitting a model into the offline data via optimization . Unlike in simple function approximations such as tabular representation and linear models where closed-form solutions are available , OPL with neural network function approximation poses an additional challenge that involves a non-convex , non-analytical optimization problem . However , existing works of OPL with function approximation ignore such optimization problems by assuming free access to an optimization oracle that can return a global minimizer ( Brandfonbrener et al. , 2021 ; Duan et al. , 2021 ; Nguyen-Tang et al. , 2021 ) or an empirical risk minimizer with a pointwise convergence bound at an exponential rate ( Hu et al. , 2021 ) . This is largely not the case in practice , especially in OPL with neural network function approximation where a model is trained by gradient-based methods such as stochastic gradient descents ( SGD ) . Thus , to understand OPL in more practical settings , it is crucial to consider optimization in design and analysis . To our knowledge , such optimization problem has not been studied in the context of OPL with neural network function approximation . Generalization . In OPL , generalization is the ability to generalize beyond the states ( or contexts as in the specific case of stochastic contextual bandits ) observed in the offline data . In other words , an offline policy learner with good generalization should obtain high rewards not only in the observed states but also in the entire ( unknown ) state distribution . The challenge of generalization in OPL is that as we learn from the fixed offline data , the learned policy has highly correlated structures where we can not directly use the standard concentration inequalities ( e.g . Hoeffding ’ s inequality , Bernstein inequality ) to derive a generalization bound . The typical approaches to overcome this difficulty are data splitting and uniform convergence . While data splitting splits the offline data into disjoint folds to break the correlated structures ( Yin & Wang , 2020 ) , uniform convergence establishes generalization uniformly over a class of policies learnable by a certain model ( Yin et al. , 2021 ) . However , in the setting where the offline data itself can have correlated structures ( e.g. , an offline action can depend on the previous offline data ) and the model used is sufficiently large that renders a uniform convergence bound vacuous , neither the data splitting technique in ( Yin & Wang , 2020 ) nor the uniform convergence argument in ( Yin et al. , 2021 ) yield a good generalization . Thus , it is highly non-trivial to guarantee a strong generalization in OPL with neural network function approximation from highly correlated offline data . In this paper , we consider the problem of OPL with neural network function approximation on the axes of distributional shift , optimization and generalization via studying the setting of stochastic contextual bandits with overparameterized neural networks . Specifically , we make three contributions toward enabling OPL in more practical settings : • First , we propose an algorithm that uses a neural network to model any bounded reward function without assuming any functional form ( e.g. , linear models ) and uses a pessimistic formulation to deal with distributional shifts . Notably , unlike any standard offline learning methods , our algorithm learns from the offline data in an online-like manner , allowing us to leverage the advantages of online learning into offline setting . • Second , our theoretical contribution lies in making the generalization bound of OPL more realistic by taking into account the optimization aspects and requiring only a milder condition for distributional shifts . In particular , our algorithm uses stochastic gradient descent and updates the network completely online , instead of retraining it from scratch for every iteration . Moreover , the distributional shift condition in our analysis , unlike in the existing works , does not require the offline policy to be uniformly explorative or stationary . Specifically , we prove that , under mild conditions with practical considerations above , our algorithm learns the optimal policy with an expected error of Õ ( κd̃1/2n−1/2 ) , where n is the number of offline samples , κ measures the distributional shift , and d̃ is an effective dimension of the neural network that is much smaller than the network capacity ( e.g. , the network ’ s VC dimension and Rademacher complexity ) . • Third , we evaluate our algorithm in a number of synthetic and real-world OPL benchmark problems , verifying its empirical effectiveness against the representative methods of OPL . Notation . We use lower case , bold lower case , and bold upper case to represent scalars , vectors and matrices , respectively . For a vector v = [ v1 , . . . , vd ] T ∈ Rd and p > 1 , denote ∥v∥p = ( ∑d i=1 v p i ) 1/p and let [ v ] j be the jth element of v. For a matrix A = ( Ai , j ) m×n , denote ∥A∥F =√∑ i , j A 2 i , j , ∥A∥p = maxv : ∥v∥p=1 ∥Av∥p , ∥A∥∞ = maxi , j |Ai , j | and let vec ( A ) ∈ Rmn be the vectorized representation of A . For a square matrix A , a vector v , and a matrix X , denote ∥v∥A = √ vTAv and ∥X∥A = ∥ vec ( X ) ∥A . For a collection of matrices W = ( W1 , . . . , WL ) and a square matrix A , denote ∥W ∥F = √∑L l=1 ∥Wl∥2F , and ∥W ∥A = √∑L l=1 ∥Wl∥2A . For a collection of matrices W ( 0 ) = ( W ( 0 ) 1 , . . . , W ( 0 ) L ) , denote B ( W ( 0 ) , R ) = { W = ( W1 , . . . , WL ) : ∥Wl −W ( 0 ) l ∥F ≤ R } . Denote [ n ] = { 1 , 2 , . . . , n } , and a ∨ b = max { a , b } . We write Õ ( · ) to hide logarithmic factors in the standard Big-Oh notation , and write m ≥ Θ ( f ( · ) ) to indicate that there is an absolute constant C > 0 that is independent of any problem parameters ( · ) such that m ≥ Cf ( · ) . 2 BACKGROUND . In this section , we provide essential background on offline stochastic contextual bandits and overparameterized neural networks . 2.1 STOCHASTIC CONTEXTUAL BANDITS . We consider a stochastic K-armed contextual bandit where at each round t , an online learner observes a full context xt : = { xt , a ∈ Rd : a ∈ [ K ] } sampled from a context distribution ρ , takes an action at ∈ [ K ] , and receives a reward rt ∼ P ( ·|xt , at ) . A policy π maps a full context ( and possibly other past information ) to a distribution over the action space [ K ] . For each full context x : = { xa ∈ Rd : a ∈ [ K ] } , we define vπ ( x ) = Ea∼π ( ·|x ) , r∼P ( ·|xa ) [ r ] and v∗ ( x ) = maxπ vπ ( x ) , which is attainable due to the finite action space . In the offline contextual bandit setting , the goal is to learn an optimal policy only from an offline data Dn = { ( xt , at , rt ) } nt=1 collected a priori by a behaviour policy µ . The goodness of a learned policy π̂ is measured by the ( expected ) sub-optimality the policy achieves in the entire ( unknown ) context distribution ρ : SubOpt ( π̂ ) : = Ex∼ρ [ SubOpt ( π̂ ; x ) ] , where SubOpt ( π̂ ; x ) : = v∗ ( x ) − vπ̂ ( x ) . In this work , we make the following assumption about reward generation : For each t , rt = h ( xt , at ) +ξt , where h : Rd → [ 0 , 1 ] is an unknown reward function , and ξt is a R-subgaussian noise conditioned on ( Dt−1 , xt , at ) where we denote Dt = { ( xτ , aτ , rτ ) } 1≤τ≤t , ∀t . The R-subgaussian noise assumption is standard in stochastic bandit literature ( Abbasi-Yadkori et al. , 2011 ; Zhou et al. , 2020 ; Xiao et al. , 2021 ) and is satisfied e.g . for any bounded noise .

Summary: The paper studies the problem of offline contextual bandits, where policy learning can only leverage a fixed dataset collected a priori by behavior policies. Using a pessimism principle, the authors propose a new algorithm called NeuralLCB with overparameterized neural networks and provide theoretical regret guarantees based on the analysis framework of the neural tangent kernel. Experiments on both synthetic and real-world data are conducted, which confirms the theoretical results.