Bellman-consistent Pessimism for Offline Reinforcement Learning

1 Introduction . Using past experiences to learn improved behavior for future interactions is a critical capability for a Reinforcement Learning ( RL ) agent . However , robustly extrapolating knowledge from a historical dataset for sequential decision making is highly challenging , particularly in settings where function approximation is employed to generalize across related observations . In this paper , we provide a systematic treatment of such scenarios with general function approximation , and devise algorithms that can provably leverage an arbitrary historical dataset to discover the policy that obtains the largest guaranteed rewards , amongst all possible scenarios consistent with the dataset . The problem of learning a good policy from historical datasets , typically called batch or offline RL , has a long history [ see e.g. , Precup et al. , 2000 ; Antos et al. , 2008 ; Levine et al. , 2020 , and references therein ] . Many prior works [ e.g. , Precup et al. , 2000 ; Antos et al. , 2008 ; Chen and Jiang , 2019 ] make the so-called coverage assumptions on the dataset , requiring the dataset to contain any possible state , action pair or trajectory with a lower bounded probability . These assumptions are evidently prohibitive in practice , particularly for problems with large state and/or action spaces . Furthermore , the methods developed under these assumptions routinely display unstable behaviors such as lack of convergence or error amplification , when coverage assumptions are violated [ Wang et al. , 2020 , 2021 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Driven by these instabilities , a growing body of recent literature has pursued a so-called best effort style of guarantee instead . The key idea is to replace the stringent assumptions on the dataset with a dataset-dependent performance bound , which gracefully degrades from guaranteeing a near-optimal policy under standard coverage assumptions to offering no improvement over the data collection policy in the most degenerate case . Algorithmically , these works all leverage the principle of pessimistic extrapolation from offline data and aim to maximize the rewards the trained agent would obtain in the worst possible MDP that is consistent with the observed dataset . These methods have been shown to be typically more robust to the violation of coverage assumptions in practice , and their theoretical guarantees often provide non-trivial conclusions in settings where the previous results did not apply . Even though many such best-effort methods have now been developed , very few works provide a comprehensive theory for using generic function approximation , unlike the setting where the dataset satisfies the coverage assumptions [ Antos et al. , 2008 ; Munos , 2003 ; Szepesvári and Munos , 2005 ; Munos and Szepesvári , 2008 ; Farahmand et al. , 2010 ; Chen and Jiang , 2019 ; Xie and Jiang , 2020 ] . For example , [ Kidambi et al. , 2020 ] provides a partial theory under the assumption of an uncertainty quantification oracle , which however is highly nontrivial to obtain for general function approximation . [ Fujimoto et al. , 2019 ; Kumar et al. , 2020 ] develop sound theoretical arguments in the tabular setting , which were only heuristically extended to the function approximation setting . The works that explicitly consider function approximation in their design either use an ad-hoc truncation of Bellman backups [ Liu et al. , 2020 ] or strongly rely on particular parameterizations such as linear function approximation [ Jin et al. , 2021 ] . In particular , [ Liu et al. , 2020 ] additionally requires the ability to approximate stationary distribution of the behavior policy , which is a challenging density estimation problem for complex state spaces and can not be provably performed in the standard linear MDP setting ( see Section 3.1 ) . Our paper takes an important step in this direction . We provide a systematic way to encode pessimism compatible with an arbitrary function approximation class and MDP and give strong theoretical guarantees without requiring any coverage assumptions on the dataset . Our first contribution is an information theoretic algorithm that returns a policy with a small regret to any comparator policy , for which coverage assumptions ( approximately ) hold with respect to the data collection policy . This regret bound is identical to what can be typically obtained when the coverage assumptions hold for all policies [ Antos et al. , 2008 ; Chen and Jiang , 2019 ] . But our algorithm requires neither the coverage assumptions , nor additional assumptions such as reliable density estimation for the data generating distribution used by existing best-effort approaches [ Liu et al. , 2020 ] . We furthermore instantiate these results in the special case of linear parameterization ; under the linear MDP assumption , our sample complexity bound leads to a factor of O ( d ) improvement for a d-dimensional linear MDP , compared with the best known result translated to our discounted setting [ Jin et al. , 2021 ] , when the action set is small in size . In addition to the information theoretic algorithm , we also develop a computationally practical version of our algorithm using a Lagrangian relaxation combined with recent advances in soft policy iteration [ Even-Dar et al. , 2009 ; Geist et al. , 2019 ; Agarwal et al. , 2019 ] . We show that this algorithm can be executed efficiently by querying a ( regularized ) loss minimization oracle over the value function class , although it has slightly worse theoretical guarantees than the information theoretic version . Both our algorithms display an adaptive property in selecting the best possible form of a bias-variance decomposition , where most prior approaches had to commit to a particular point through their choice of hyperparameters ( see the discussion following Theorem 3.1 ) . 2 Preliminaries . Markov Decision Processes We consider dynamical systems modeled as Markov Decision Processes ( MDPs ) . An MDP is specified by ( S , A , P , R , γ , s0 ) , where S is the state space , A is the action space , P : S × A → ∆ ( S ) is the transition function with ∆ ( · ) being the probability simplex , R : S × A → [ 0 , Rmax ] is the reward function , γ ∈ [ 0 , 1 ) is the discount factor , and s0 is a deterministic initial state , which is without loss of generality . We assume the state and the action spaces are finite but can be arbitrarily large . A ( stochastic ) policy π : S → ∆ ( A ) specifies a decision-making strategy , and induces a random trajectory s0 , a0 , r0 , s1 , a1 , r1 , . . . , where at ∼ π ( ·|st ) , rt = R ( st , at ) , st+1 ∼ P ( ·|st , at ) , ∀t ≥ 0 . We denote the expected discounted return of a policy π as J ( π ) : = E [ ∑∞ t=0 γ trt|π ] , and the learning goal is to find the maximizer of this value : π⋆ : = argmaxπ J ( π ) . A related concept is the policy-specific Q-function , Q π : S ×A → R. Qπ ( s , a ) is the discounted return when the trajectory starts with ( s , a ) and all remaining actions are taken according to π. Qπ is the unique fixed point of the ( policy-specific ) Bellman operator T π : RS×A → RS×A , defined as : ∀f , ( T πf ) ( s , a ) = R ( s , a ) + γEs′∼P ( ·|s , a ) [ f ( s′ , π ) ] , where f ( s′ , π ) is a shorthand for Ea′∼π ( ·|s′ ) [ f ( s′ , a′ ) ] . Another important concept is the notion of discounted state-action occupancy , dπ ∈ ∆ ( S × A ) , defined as dπ ( s , a ) : = ( 1 − γ ) E [ ∑∞ t=0 γ t1 [ st = s , at = a ] |π ] , which characterizes the states and actions visited by a policy π. Offline RL In the offline setting , the learner only has access to a pre-collected dataset and can not directly interact with the environment . We assume the standard i.i.d . data generation protocol in our theoretical derivations , that the offline dataset D consists of n i.i.d . ( s , a , r , s′ ) tuples generated as ( s , a ) ∼ µ , r = R ( s , a ) , s′ ∼ P ( ·|s , a ) for some data distribution µ . We will also use Eµ [ · ] for taking expectation with respect to µ . We will frequently use the data-weighted 2-norm ( squared ) ∥f∥22 , µ : = Eµ [ f2 ] , and the definition extends when we replace µ with any other state-action distribution ν . The empirical approximation of ∥f∥22 , µ is ∥f∥22 , D : = 1n ∑ ( s , a , r , s′ ) ∈D f ( s , a ) 2 . Function Approximation Function approximation is crucial to generalizing over large and complex state and action spaces . In this work , we search for a good policy in a policy class Π ⊂ ( S → ∆ ( A ) ) with the help of a value-function class F ⊂ ( S × A → [ 0 , Vmax ] ) to model Qπ , where Vmax = Rmax/ ( 1 − γ ) . Such a combination is commonly found in approximate policy iteration and actor-critic algorithms [ e.g. , Bertsekas and Tsitsiklis , 1996 ; Konda and Tsitsiklis , 2000 ] . For most part of the paper we do not make any structural assumptions on Π and F , making our approach and guarantees applicable to generic function approximators . For simplicity we will assume that these function classes are finite but exponentially large , and use log-cardinality to measure their statistical complexities in the generic results ( Section 3 and Section 4 ) . These guarantees easily extend to continuous function classes where log-cardinalities are replaced by the appropriate notions of covering numbers , which we demonstrate when we instantiate our results in the linear function approximation setting and work with continuous linear classes ( Section 3.1 ) . We now recall two standard expressivity assumptions on F [ e.g. , Antos et al. , 2008 ] . To our knowledge , no existing works on offline RL with insufficient data coverage have provided guarantees under these standard assumptions for general function approximation , and they often require stronger or tweaked assumptions ( see Section 1 ) . Assumption 1 ( Realizability ) . For any π ∈ Π , we have inf f∈F sup admissible ν ∥f − T πf∥22 , ν ≤ εF , where an admissible distribution ν means that ν ∈ { dπ′ : π′ ∈ Π } . Assumption 1 requires that for every π ∈ Π , there exists f ∈ F that well-approximates Qπ . This assumption is often called realizability.1 Technically this is asserted by requiring f to have small Bellman error w.r.t . T π under all possible admissible distributions . As a sufficient condition , we have εF = 0 if Qπ ∈ F , ∀π ∈ Π . Assumption 2 ( Completeness ) . For any π ∈ Π , we have sup f∈F inf f ′∈F ∥f ′ − T πf∥22 , µ ≤ εF , F . Assumption 2 asserts that F is approximately closed under T π .2 Such an assumption is widely used in RL theory and can be only avoided in some rare cases [ Xie and Jiang , 2021 ] , and the hardness of learning with realizability alone has been established in various settings ( e.g. , [ Weisz et al. , 2021 ; Zanette , 2021 ] ) . We also emphasize that we only measure the violation of completeness under µ and do not need to reason about all admissible distributions . 1In the exploratory setting , realizability is usually stated in the form of inff∈F ∥f − T πf∥22 , µ ≤ εF for any π ∈ Π . However , the exploratory setting also usually has data coverage assumptions in the form of supν ∥ν/µ∥∞ ≤ C. Combining them together implies Assumption 1 . 2Sometimes completeness implies realizability , so the latter does not need to be assumed separately [ Chen and Jiang , 2019 ] . However , this often relies on µ being exploratory , which is not the case here . Distribution shift A unique challenge in RL is that the learned policy may induce a state ( and action ) distribution that is different from the data distribution µ , and the issue is particularly salient when we do not impose any coverage assumption on µ . Therefore , it is important to carefully characterize the distribution shift , which we measure using the following definition , which generalizes prior definitions specific to linear function approximation [ Agarwal et al. , 2019 ; Duan et al. , 2020 ] : Definition 1 . We define C ( ν ; µ , F , π ) as follows to measure the distribution shift from an arbitrary distribution ν to the data distribution µ , w.r.t . F and π , C ( ν ; µ , F , π ) : = max f∈F ∥f − T πf∥22 , ν ∥f − T πf∥22 , µ . Intuitively , C ( ν ; µ , F , π ) measures how well Bellman errors under π transfer between the distributions ν and µ . For instance , a small value of C ( dπ ; µ , F , π ) enables accurate policy evaluation for π using data collected under µ . More generally , we observe that C ( ν ; µ , F , π ) ≤ ∥ν/µ∥∞ : = sup s , a ν ( s , a ) µ ( s , a ) , for any π , F . and the RHS is a classical notion of bounded distribution ratio for error transfer ( e.g. , [ Munos and Szepesvári , 2008 ; Chen and Jiang , 2019 ; Xie and Jiang , 2020 ] ) . Moreover , our measure can be tighter than ∥ν/µ∥∞ : Even two distributions ν and µ that are sufficiently disparate might admit a reasonable transfer , so long as this difference is not detected by π and F . To this end , our definition better captures the crucial role of function approximation in generalizing across different states . As an example , in the special case of linear MDPs , full coverage under our definition ( i.e. , boundedness of C for all admissible ν ) can be implied from the standard coverage assumption for linear MDPs that considers the spectrum of the feature covariance matrix under µ ; see Section 3.1 for more details .

Summary: This paper studies the pessimism in offline RL with general function approximation given i.i.d. transition data. An information-theoretic analysis is established for an inefficient method, where instead of point-wise pessimism, a notion of Bellman-consistent pessimism is considered. In this way, the theoretical guarantee is established only under closedness assumptions. An instantiation to linear MDP is considered and improved upon literature. Besides, a practical and efficient policy iteration method is proposed, whose performance is analyzed as well.