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: This paper proposes a neural network based contextual bandit algorithm in the offline setting where a dataset of contexts and rewards are given by a logging policy. The goal of the proposed algorithm is to learn an optimal policy from the offline dataset. The proposed algorithm NeuralLCB is similarly structured as the NeuralUCB algorithm (Zhou et a., 2019) in the online setting. The difference is that it uses a lower confidence bound for estimating the reward function instead of an upper confidence bound, and that the optimization procedure for learning the neural network representation is based on the loss on one data point instead of the whole historical data. The authors established an upper bound of the optimal gap of the learned policy and evaluated the algorithm on both simulation and UCI datasets.