Sample-Efficient Reinforcement Learning Is Feasible for Linearly Realizable MDPs with Limited Revisiting

1 Introduction . Emerging reinforcement learning ( RL ) applications necessitate the design of sample-efficient solutions in order to accommodate the explosive growth of problem dimensionality . Given that the state/action space could be unprecedentedly enormous , it is often infeasible to request a sample size exceeding the fundamental limit set forth by the ambient dimension in the tabular setting ( which enumerates all combinations of state-action pairs ) . As a result , the quest for sample efficiency can not be achieved in general without exploiting proper low-complexity structures underlying the problem of interest . 1.1 Linear function approximation . Among the studies of low-complexity models for RL , linear function approximation has attracted a flurry of recent activity , mainly due to the promise of dramatic dimension reduction in conjunction with its mathematical tractability ( see , e.g. , Wen and Van Roy ( 2017 ) ; Yang and Wang ( 2019 ) ; Jin et al . ( 2020 ) ; Du et al . ( 2020a ) ) . Two families of linear function approximation merit particular attention , which we single out below . Here and throughout , we concentrate on a finite-horizon Markov decision process ( MDP ) , and denote by S , A and H its state space , action space , and horizon , respectively . • Model-based linear representation . Yang and Wang ( 2019 ) ; Jin et al . ( 2020 ) ; Yang and Wang ( 2020 ) studied a scenario when both the probability transition kernel and reward function of the MDP can be linearly parameterized . This type of MDPs is commonly referred to as linear MDPs . Letting Ph ( · | s , a ) be the transition probability from state s when action a is executed at step h , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the linear MDP model postulates the existence of a set of predetermined d-dimensional feature vectors { ϕh ( s , a ) ∈ Rd } and a set of unknown parameter matrices { µh ∈ Rd×|S| } such that ∀ ( s , a ) ∈ S ×A and 1 ≤ h ≤ H : Ph ( · | s , a ) = ( ϕh ( s , a ) ) > µh . ( 1 ) Similar assumptions are imposed on the reward function too . In words , linear MDP posits that the probability transition matrix is low-rank ( with rank at most d ) with the column space known a priori , which forms the basis for sample size saving compared to the unstructured setting . • Value-based linear realizability . Rather than relying on linear embedding of the model specification ( namely , the transition kernel and reward function ) , another class of linear representation assumes that the action-value function ( or Q-function ) can be well predicted by linear combinations of known feature vectors { ϕh ( s , a ) ∈ Rd } . A concrete framework of this kind assumes linear realizability of the optimal Q-function ( denoted by Q ? h at time step h from now on ) , that is , there exist some unknown vectors { θ ? h ∈ Rd } such that Q ? h ( s , a ) = 〈 ϕh ( s , a ) , θ ? h 〉 ( 2 ) holds for any state-action pair ( s , a ) and any step h. It is self-evident that this framework seeks to represent the optimal Q-function — which is an |S||A|H-dimensional object — via H parameter vectors each of dimension d. Throughout this work , an MDP obeying this condition is said to be an MDP with linearly realizable optimal Q-function , or more concisely , an MDP with linear Q ? . 1.2 Sample size barriers with linearly realizable Q ? . This paper focuses on MDPs with linearly realizable optimal Q-function Q ? . In stark contrast to linear MDPs that allow for sample-efficient RL ( in the sense that the sample complexity is almost independent of |S| and |A| but instead depends only polynomially on d , H and possibly a certain sub-optimality gap ) , MDPs with linear Q ? do not admit a similar level of sample efficiency in general . To facilitate discussion , we summarize key existing results for a couple of settings . Here and below , f ( d ) = Ω ( g ( d ) ) means that f ( d ) is at least on the same order as g ( d ) for d large enough . • Sample inefficiency under a generative model . Even when a generative model or a simulator is available — so that the learner can query arbitrary state-action pairs to draw samples from ( Kearns and Singh , 1999 ) — one can find a hard MDP instance in this class that requires at least min { exp ( Ω ( d ) ) , exp ( Ω ( H ) ) } samples regardless of the algorithm in use ( Weisz et al. , 2021b ) . • Sample efficiency with a sub-optimality gap under a generative model . The aforementioned sample size barrier can be alleviated if , for each state , there exists a sub-optimality gap ∆gap between the value under the optimal action and that under any sub-optimal action . As asserted by Du et al . ( 2020a , Appendix C ) , a sample size that scales polynomially in d , H and 1/∆gap is sufficient to identify the optimal policy , assuming access to the generative model . • Sample inefficiency with a large sub-optimality gap in online RL . Turning to the standard episodic online RL setting ( so that in each episode , the learner is given an initial state and executes the MDP for H steps to obtain a sample trajectory ) , sample-efficient algorithms are , much to our surprise , infeasible even in the presence of a large sub-optimality gap . As has been formalized in Wang et al . ( 2021b ) , it is possible to construct a hard MDP instance with a constant sub-optimality gap that can not be solved without at least min { exp ( Ω ( d ) ) , exp ( Ω ( H ) ) } samples . In conclusion , the linear Q ? assumption , while succinctly capturing the low-dimensional structure , still presents an undesirable hurdle for RL . The sampling mechanism commonly studied in RL literature precludes sample-efficient solutions even when a favorable sub-optimality gap exists . 1.3 Our contributions . Having observed the exponential separation between the generative model and standard online RL when it comes to the linear Q ? problem , one might naturally wonder whether there exist practically relevant sampling mechanisms — more flexible than standard online RL yet more practical than the generative model — that promise substantial sample size reduction . This motivates the investigation of the current paper , as summarized below . A new sampling protocol : sampling with state revisiting . We investigate a flexible sampling protocol that is built upon classical online/exploratory formalism but allows for state revisiting ( detailed in Algorithm 1 ) . In each episode , the learner starts by running the MDP for H steps , and is then allowed to revisit any previously visited state and re-run the MDP from there . The learner is allowed to revisit states for an arbitrary number of times , although executing this feature too often might inevitably incur an overly large sampling burden . This sampling protocol accommodates several realistic scenarios ; for instance , it captures the “ save files ” feature in video games that allows players to record player progress and resume from the save points later on . In addition , state revisiting is reminiscent of Monte Carlo Tree Search implemented in various real-world applications , which assumes that the learner can go back to father nodes ( i.e. , previous states ) ( Silver et al. , 2016 ) . This protocol is also referred to as local access to the simulator in the recent work Yin et al . ( 2021 ) . A sample-efficient algorithm . Focusing on the above sampling protocol , we propose a value-based method — called LinQ-LSVI-UCB — adapted from the LSVI-UCB algorithm ( Jin et al. , 2020 ) . The algorithm implements the optimism principle in the face of uncertainty , while harnessing the knowledge of the sub-optimality gap to determine whether to revisit states . Our algorithm achieves a sample complexity that scales polynomially in the feature dimension d , the horizon H , and the inverse sub-optimality gap 1/∆gap , but is otherwise independent of the size of the state/action space . 2 Model and assumptions . We present precise problem formulation as well as a couple of key assumptions . Here and throughout , we denote by |S| the cardinality of a set S , and adopt the notation [ n ] : = { 1 , · · · , n } . 2.1 Basics of Markov decision processes . Finite-horizon MDP . The focus of this paper is the setting of a finite-horizon MDP , as represented by the quintupleM = ( S , A , { Ph } Hh=1 , { rh } Hh=1 , H ) . Here , S : = { 1 , · · · , |S| } denotes the state space , A : = { 1 , · · · , |A| } is the action space , H is the time horizon , Ph stands for the probability transition kernel at time step h ∈ [ H ] ( namely , Ph ( · | s , a ) is the transition probability from state s upon execution of action a at step h ) , whereas rh : S × A → [ 0 , 1 ] represents the reward function at step h ( namely , we denote by rh ( s , a ) the immediate reward received at step h when the current state is s and the current action is a ) . For simplicity , it is assumed throughout that all rewards { rh ( s , a ) } are deterministic and reside within the range [ 0 , 1 ] . Note that our analysis can also be directly extended to accommodate random rewards , which we omit here for the sake of brevity . Policy , value function , and Q-function . We let π = { πh } 1≤h≤H represent a policy or action selection rule . For each time step h , πh represents a deterministic mapping from S to A , namely , action πh ( s ) is taken at step h if the current state is s. The value function associated with policy π at step h is then defined as the cumulative reward received between steps h and H under this policy : ∀ ( s , h ) ∈ S × [ H ] : V πh ( s ) : = E [ H∑ t=h rt ( st , at ) ∣∣∣ st = s ] . ( 3 ) Here , the expectation is taken over the randomness of an MDP trajectory { st } Ht=h induced by policy π ( namely , at = πt ( st ) and st+1 ∼ P ( · | st , at ) for any h ≤ t ≤ H ) . Similarly , the action-value function ( or Q-function ) associated with policy π is defined as ∀ ( s , a , h ) ∈ S ×A× [ H ] : Qπh ( s , a ) : = E [ H∑ t=h rt ( st , at ) ∣∣∣ sh = s , ah = a ] , ( 4 ) which resembles the definition ( 3 ) except that the action at step h is frozen to be a . Our normalized reward assumption ( i.e. , rh ( s , a ) ∈ [ 0 , 1 ] ) immediately leads to the trivial bounds ∀ ( s , a , h ) ∈ S ×A× [ H ] : 0 ≤ V πh ( s ) ≤ H and 0 ≤ Qπh ( s , a ) ≤ H. ( 5 ) A recurring goal in reinforcement learning is to search for a policy that maximizes the value function and the Q-function . For notational simplicity , we define the optimal value function V ? = { V ? h } 1≤h≤H and optimal Q-function Q ? = { Q ? h } 1≤h≤H respectively as follows ∀ ( s , a , h ) ∈ S ×A× [ H ] : V ? h ( s ) : = max π V πh ( s ) and Q ? h ( s , a ) : = max π Qπh ( s , a ) , with the optimal policy ( i.e. , the one that maximizes V π ) represented by π ? = { π ? h } 1≤h≤H .

Summary: Recent work in RL has been seeking ways that RL can draw the kinds of generalization that we see in supervised learning. Recently, a couple groups obtained efficient RL algorithms by assuming that the reward and transition probabilities are linear combinations of some set of basis functions. A somewhat more appealing, weaker assumption would be for merely the Q-function (state-action value function) to be linear. Unfortunately, recent works also found hard examples of environments with linear Q-functions, that require exponential sample complexity. Indeed, even with a simulator, there are exponential lower bounds -- only recently have any kind of polynomial-time guarantees, with a polynomial dependence on the smallest gap between the value of any optimal action any any sub-optimal action.  This work proposes a new model in which linear Q-functions are learnable. The model can be seen as a weakening of the simulator model: whereas the  simulator model allowed the learner to query any state-value pair to get a  sample of the reward and transition functions, this model merely allows the  learner to return to states that it previously visited on some trajectory.