Improved Variance-Aware Confidence Sets for Linear Bandits and Linear Mixture MDP

• For linear bandits , we obtain an rOppolypdq b 1 ` řK k “ 1 σ 2 kq data-dependent regret bound , where d is the feature dimension , K is the number of rounds , and σ2k is the unknown variance of the reward at the k-th round . This is the first regret bound that only scales with the variance and the dimension but no explicit polynomial dependency on K. When variances are small , this bound can be significantly smaller than the rΘ ` d ? K ˘ worst-case regret bound . • For linear mixture MDPs , we obtain an rOppolypd , logHq ? Kq regret bound , where d is the number of base models , K is the number of episodes , and H is the planning horizon . This is the first regret bound that only scales logarithmically with H in the reinforcement learning with linear function approximation setting , thus exponentially improving existing results , and resolving an open problem in [ Zhou et al. , 2020a ] . We develop three technical ideas that may be of independent interest : 1 ) applications of the peeling technique to both the input norm and the variance magnitude , 2 ) a recursion-based estimator for the variance , and 3 ) a new convex potential lemma that generalizes the seminal elliptical potential lemma . 1 Introduction . In sequential decision-making problems such as bandits and reinforcement learning ( RL ) , the agent chooses an action based on the current state , with the goal to maximize the total reward . When the state-action space is large , function approximation is often used for generalization . One of the most fundamental and widely used methods is linear function approximation . For ( infinite-actioned ) linear bandits , the minimax-optimal regret bound is rΘpd ? Kq [ Dani et al. , 2008 , Abbasi-Yadkori et al. , 2011 ] , where d is the feature dimension and K is the number of total rounds played by the agent.2 However , oftentimes the worst-case analysis is overly pessimistic , and ˚Equal contribution . 2We follow the reinforcement learning convention to use K to denote the total number of rounds / episodes . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . it is possible to obtain data-dependent bound that is substantially smaller than rOpd ? Kq in benign scenarios . One direction to study is the variance magnitude . As a motivating example , in linear bandits , if there is no noise ( variance is 0 ) , one only needs to pay at most d regret to identify the best action because d samples are sufficient to recover the underlying linear coefficients ( in general position ) . This constant-type regret bound is much smaller than the ? K-type regret bound in the worst case where the variance magnitude is a lower bounded constant . Therefore , a natural question is : Can we design an algorithm that adapts to the variance magnitude , and its regret degrades gracefully from the benign noiseless constant-type bound to the worst-case ? K-type bound ? In RL , exploiting the variance information is also important . For tabular RL , one needs to utilize the variance information , e.g. , Bernstein-type exploration bonus to achieve the minimax optimal regret [ Azar et al. , 2017 , Zanette and Brunskill , 2019 , Zhang et al. , 2020c , a , Menard et al. , 2021 , Dann et al. , 2019 ] . For example , the recently proposed MVP algorithm [ Zhang et al. , 2020a ] , enjoys an rOppolylogpHq ˆ p ? SAK ` S2Aqq regret bound , where S is the number of states , A is the number of actions , H is the planning horizon , and K is the total number of episodes . 34 Notably , this regret bound only scales logarithmically with H . On the other hand , without using the variance information , e.g. , using Hoeffding-type bonus instead of Bernstein-type bonus , algorithms would suffer a regret that scales polynomially with H [ Azar et al. , 2017 ] . Going beyond tabular RL , a recent line of work studied RL with linear function approximation with different assumptions [ Yang and Wang , 2019 , Modi et al. , 2020 , Jin et al. , 2020 , Ayoub et al. , 2020 , Zhou et al. , 2020a , Modi et al. , 2020 ] . Our paper studies the linear mixture Markov Decision Process ( MDP ) setting [ Modi et al. , 2020 , Ayoub et al. , 2020 , Zhou et al. , 2020a ] , where the transition probability can be represented by a linear function of some features or base models . This model-based assumption is motivated by problems in robotics and queuing systems . We refer readers to Ayoub et al . [ 2020 ] for more discussions . For this linear mixture MDP setting , previous works can obtain regret bounds in the form rOppolypd , Hq ? Kq , where d is the number of base models . While these bounds do not scale with SA , they scale polynomially with H , because the algorithms in previous works do not use the variance information . In practice , H is often large , and even a polynomial dependency on H may not be acceptable . Therefore , a natural question is Can we design an algorithm that exploits the variance information to obtain an rOppolypd , logHq ? Kq regret bound for linear mixture MDP ? 1.1 Our Contributions . In this paper , we develop new , variance-aware confidence sets for linear bandits and linear mixture MDP and answer the above two questions affirmatively . Linear Bandits . For linear bandits , we obtain an rOppolypdq b 1 ` řK k “ 1 σ 2 kq regret bound , where σ2k is the unknown variance at the k-th round . To our knowledge , this is the first bound that solely depends on the variance and the feature dimension , and has no explicit polynomial dependency on K. When the variance is very small so that σ2k ! 1 , this bound is substantially smaller than the worst-case rΘpd ? Kq bound . Furthermore , this regret bound naturally interpolates between the worst-case ? K-type bound and the noiseless-case constant-type bound . Linear Mixture MDP . For linear mixture MDP , we obtain the desired rOppolypd , logHq ? Kq regret bound . This is the first regret bound in RL with function approximation that 1 ) does not scale with the size of the state-action space , and 2 ) only scales logarithmically with the planning horizon 3 rOp¨q hides logarithmic factors . Sometimes we write out polylogH explicitly to emphasize the logarithmic dependency on H . 4This bound holds for setting where the transition is homogeneous and the total reward is bounded by 1 . We focus on this setting in this paper . See Section 2 and 3 for more discussions . H . Therefore , we exponentially improve existing results on RL with linear function approximation in term of the H dependency , and resolve an open problem in [ Zhou et al. , 2020a ] . More importantly , our result conveys the positive conceptual message for RL : it is possible to simultaneously overcome the two central challenges in RL , large state-action space and long planning horizon . 1.2 Main Difficulties and Technical Innovations . We first describe limitations of existing works why they can not achieve the desired regret bounds described above . Limitations of Existing Variance-Aware Confidence Sets Faury et al . [ 2020 ] , Zhou et al . [ 2020a ] applied Bernstein-style inequalities to construct a confidence sets of the least square estimator for linear bandits . However , their methods can not be applied directly to obtain the desired data-dependent regret bound . Abeille et al . [ 2021 ] also designed an variance-dependent confidence set for logistic bandits . However in their problem the rewards are Bernoulli and the variance is a function of the mean . We give a simple example to illustrate their limitations . Consider the case where the variance is always σ2 ! 1 . Let px1 , y1q , . . . , pxk´1 , yk´1q be the samples collected before the k-th round . Their confidence set at the k-th round is Θk “ tθ|||θ ´ θ̂k||Λk´1 ď Cpσ ? d ` 1 ` λ1 { 2qu ( See In Equation ( 4.3 ) of Zhou et al . [ 2020a ] and Theorem 1 of Faury et al . [ 2020 ] ) . where Λk´1 “ řk´1 τ “ 1 xτx J τ ` λI is the un-normalized covariance matrix , θ̂k “ Λ´1k´1 řk´1 τ “ 1 yτxτ is the estimated linear coefficients by least squares , λ is a regularization parameter and C is a constant . Consider the case d “ 1 and xk “ a 1 { K for k “ 1 , . . . , K. Their regret bound is roughly K ÿ k “ 1 pσ ? d ` 1 ` λ1 { 2q } xk } Λ´1k ě p1 ` λ 1 { 2q K ÿ i “ 1 } xk } Λ´1k ě p1 ` λ 1 { 2q c K 1 ` λ ě ? K , which is much larger than our bound , O ` ? Kσ2 ` 1 ˘ when σ is very small . For more detailed discussion , please refer to Appendix B . Below we describe our main techniques . Elimination with Peeling . Instead of using least squares and upper-confidence-bound ( UCB ) , we use an elimination approach . More precisely , for the underlying linear coefficients θ˚ P Rd , we build a confidence interval for pθ˚qJ µ for every µ in an -net of the d-dimensional unit ball , and we eliminate θ P Rd if θJµ fails to fall in the confidence interval of pθ˚qJµ for some µ . To build the confidence intervals , we use 1 ) an empirical Bernstein inequality ( cf . Theorem 4 ) and 2 ) the peeling technique to both the input norm and the variance magnitude . As will be clear in the proof ( cf . Section D ) , this peeling step is crucial to obtain a tight regret bound for the example above . The new confidence region provides a tighter estimation for θ˚ , which helps address the drawback in least squares . Generalization of the Elliptical Potential Lemma . Since we use the peeling technique which comes with a clipping operation , we can not use the seminal elliptic potential lemma Dani et al . [ 2008 ] any more . Instead , we propose a more general lemma below , which provides a bound of potential for a general class of convex functions though with a worse dependency on d than the bound in the elliptical potential lemma . We believe this lemma can be applied to other problems as well . Lemma 1 ( Generalized Quadratic Potential Lemma ) . Let fpxq ě 0 be a convex function over R such that fpxqx2 ď fpyq y2 ď 1 and fpxq ě fpyq if x 2 ě y2 ą 0 . Let Bp1q denote the d-dimensional unit ball . Fix ` P p0 , 1s . For any x1 , x2 , . . . , xt P Bp1q and µ1 , µ2 , . . . , µt P Bp1q , we have that t ÿ i “ 1 min # fpxiµiq ři´1 j “ 1 fpxjµiq ` ` 2 , 1 + ď Opd4 logpdt { ` qq . Note that by choosing fpxq “ x2 and µi “ xiΛ ´1 i } xiΛ´1i } with Λi “ ři´1 j “ 1 xjx J j ` ` I , Lemma 1 reduces to the classical elliptic potential lemma [ Dani et al. , 2008 ] . Our proof consists of two major parts . We first establish a symmetric version of Equation ( ? ? ) using rearrangement inequality , and then bound the number of times the energy for some µ ( i.e. , ři j “ 1 fpxjµq ` l2 ) doubles . The full proof is deferred to Appendix C. For linear mixture MDP , we propose another technique to further reduce the dependency on d. Recursion-based Variance Estimation . In linear bandits , generally it is not possible to estimate the variance because the variance at each round can arbitrarily different . On the other hand , for linear mixture MDP , the variance is a quadratic function of the underlying coefficient θ˚ . Furthermore , the higher moments are polynomial functions of θ˚ . Utilizing this rich structure and leveraging the recursion idea in previous analyses on tabular RL [ Lattimore and Hutter , 2012 , Li et al. , 2020 , Zhang et al. , 2020a ] , we explicitly estimate the variance and higher moments to further reduce the regret . See Section 5 for more explanations .

Summary: This paper proposes a new design of the confidence set in the linear bandit problem and linear mixture RL problem. The new confidence set leads to a different regret bound, that has worse dependence on the dimension but only logarithmic dependence on the horizon length. This new design of confidence sets is the main contribution, which composes the convex potential lemma and the peeling argument.