Sample-Efficient Reinforcement Learning for Linearly-Parameterized MDPs with a Generative Model

1 Introduction . Reinforcement learning ( RL ) studies the problem of learning and decision making in a Markov decision process ( MDP ) . Recent years have seen exciting progress in applications of RL in real world decision-making problems such as AlphaGo [ 46 , 47 ] and autonomous driving [ 33 ] . Specifically , the goal of RL is to search for an optimal policy that maximizes the cumulative reward , based on sequential noisy data . There are two popular approaches to RL : model-based and model-free ones . • The model-based approaches start with formulating an empirical MDP by learning the probability transition model from the collected data samples , and then estimating the optimal policy / value function based on the empirical MDP . • The model-free approaches ( e.g . Q-learning ) learn the optimal policy or the optimal ( action- ) value function from samples . As its name suggests , model-free approaches do not attempt to learn the model explicitly . ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Generally speaking , model-based approaches enjoy great flexibility since after the transition model is learned in the first place , it can then be applied to any other problems without touching the raw data samples . In comparison , model-free methods , due to its online nature , are usually memory-efficient and can interact with the environment and update the estimate on the fly . This paper is devoted to investigating the sample efficiency of both model-based RL and Q-learning ( arguably one of the most commonly adopted model-free RL algorithms ) . It is well known that MDPs suffer from the curse of dimensionality . For example , in the tabular setting where the state space S and the action space A are both finite , to obtain a near optimal policy or value function given sampling access to a generative model , the minimax optimal sample complexity scales linearly with |S| × |A| [ 5 , 2 ] . However contemporary applications of RL often encounters environments with exceedingly large state and action spaces , whilst the data collection might be expensive or even high-stake . This suggests a large gap between the theoretical findings and practical decision-making problems where |S| and |A| are large or even infinite . To close the aforementioned theory-practice gap , one natural idea is to impose certain structural assumption on the MDP . In this paper we follow the feature-based linear transition model studied in [ 63 ] , where each state-action pair ( s , a ) ∈ S×A admits aK dimensional feature vector φ ( s , a ) ∈ RK that expresses the transition dynamics P ( ·|s , a ) = Ψφ ( s , a ) for some unknown matrix Ψ ∈ R|S|×K which is common for all ( s , a ) . This model encompasses both the tabular case and the homogeneous model in which the state space can be partitioned into K equivalent classes . Assuming access to a generative model [ 31 , 32 ] , under this structural assumption , this paper aims to answer the following two questions : How many samples are needed for model-based RL and Q-learning to learn an optimal policy under the feature-based linear transition model ? In what follows , we will show that the answer to this question scales linearly with the dimension of the feature space K and is independent of |S| and |A| under the feature-based linear transition model . With the aid of this structural assumption , model-based RL and Q-learning becomes significantly more sample-efficient than that in the tabular setting . Our contributions . We focus our attention on an infinite horizon MDP with discount factor γ ∈ ( 0 , 1 ) . We use ε-optimal policy to indicate the policy whose expected discounted cumulative rewards are ε close to the optimal value of the MDP . Our contributions are two-fold : • We demonstrate that model-based RL provably learns an ε-optimal policy by performing planning based on an empirical MDP constructed from a total number of Õ ( K ( 1− γ ) 3 ε2 ) samples , for all ε ∈ ( 0 , ( 1− γ ) −1/2 ] . Here Õ ( · ) hides logarithmic factors compared to the usual O ( · ) notation . To the best of our knowledge , this is the first theoretical guarantee for model-based RL under the feature-based linear transition model . This sample complexity bound matches the minimax limit established in [ 63 ] up to logarithmic factor . • We also show that Q-learning provably finds an entrywise ε-optimal Q-function using a total number of Õ ( K ( 1− γ ) 4 ε2 ) samples , for all ε ∈ ( 0 , 1 ] . This sample complexity upper bound improves the state-of-the-art result in [ 63 ] and the dependency on the effective horizon ( 1− γ ) −4 is sharp in view of [ 34 ] . These results taken collectively show the minimax optimality of model-based RL and the suboptimality of Q-learning in sample complexity . 2 Problem formulation . This paper focuses on tabular MDPs in the discounted infinite-horizon setting [ 10 ] . Here and throughout , ∆d−1 : = { v ∈ Rd : ∑d i=1 vi = 1 , vi ≥ 0 , ∀i ∈ [ d ] } stands for the d-dimensional probability simplex and [ N ] : = { 1 , 2 , · · · , N } for any N ∈ N+ . Discounted infinite-horizon MDPs . Denote a discounted infinite-horizon MDP by a tuple M = ( S , A , P , r , γ ) , where S = { 1 , · · · , |S| } is a finite set of states , A = { 1 , · · · , |A| } is a finite set of actions , P : S ×A → ∆|S|−1 represents the probability transition kernel where P ( s′|s , a ) denotes the probability of transiting from state s to state s′ when action a is taken , r : S ×A → [ 0 , 1 ] denotes the reward function where r ( s , a ) is the instantaneous reward received when taking action a ∈ A while in state s ∈ S , and γ ∈ ( 0 , 1 ) is the discount factor . Value function and Q-function . Recall that the goal of RL is to learn a policy that maximizes the cumulative reward , which corresponds to value functions or Q-functions in the corresponding MDP . For a deterministic policy π : S → A and a starting state s ∈ S , we define the value function as V π ( s ) : = E [ ∞∑ k=0 γkr ( sk , ak ) ∣∣∣ s0 = s ] for all s ∈ S. Here , the trajectory is generated by ak = π ( sk ) and sk+1 ∼ P ( sk+1|sk , ak ) for every k ≥ 0 . This function measures the expected discounted cumulative reward received on the trajectory { ( sk , ak ) } k≥0 and the expectation is taken with respect to the randomness of the transitions sk+1 ∼ P ( ·|sk , ak ) on the trajectory . Recall that the immediate rewards lie in [ 0 , 1 ] , it is easy to derive that 0 ≤ V π ( s ) ≤ 11−γ for any policy π and state s. Accordingly , we define the Q-function for policy π as Qπ ( s , a ) : = E [ ∞∑ k=0 γkr ( sk , ak ) |s0 = s , a0 = a ] for all ( s , a ) ∈ S × A . Here , the actions are chosen by the policy π except for the initial state ( i.e . ak = π ( sk ) for all k ≥ 1 ) . Similar to the value function , we can easily check that 0 ≤ Qπ ( s , a ) ≤ 1 1−γ for any π and ( s , a ) . To maximize the value function or Q function , previous literature [ 9 , 51 ] establishes that there exists an optimal policy π ? which simultaneously maximizes V π ( s ) ( resp . Qπ ( s , a ) ) for all s ∈ S ( resp . ( s , a ) ∈ S ×A ) . We define the optimal value function V ? and optimal Q-function Q ? respectively as V ? ( s ) : = max π V π ( s ) = V π ? ( s ) , Q ? ( s , a ) : = max π Qπ ( s , a ) = Qπ ? ( s , a ) for any state-action pair ( s , a ) ∈ S ×A . Linear transition model . Given a set of K feature functions φ1 , φ2 , · · · , φK : S × A → R , we define φ to be a feature mapping from S ×A to RK such that φ ( s , a ) = [ φ1 ( s , a ) , · · · , φK ( s , a ) ] ∈ RK . Then we are ready to define the linear transition model [ 63 ] as follows . Definition 1 ( Linear transition model ) . Given a discounted infinite-horizon MDPM = ( S , A , P , r , γ ) and a feature mapping φ : S ×A → RK , M admits the linear transition model if there exists some ( unknown ) functions ψ1 , · · · , ψK : S → R , such that P ( s′|s , a ) = K∑ k=1 φk ( s , a ) ψk ( s ′ ) ( 1 ) for every ( s , a ) ∈ S ×A and s′ ∈ S. Readers familiar with linear MDP literatures might immediately recognize that the above definition is the same as the structure imposed on the probability transition kernel P in the linear MDP model [ 63 , 30 , 65 , 26 , 52 , 56 , 60 ] . However unlike linear MDP which also requires the reward function r ( s , a ) to be linear in the feature mapping φ ( s , a ) , here we do not impose any structural assumption on the reward . Example 1 ( Tabular MDP ) . Each tabular MDP can be viewed as a linear transition model with feature mapping φ ( s , a ) = e ( s , a ) ∈ R|S|×|A| ( i.e . the vector with all entries equal to 0 but the one corresponding to ( s , a ) equals to 1 ) for all ( s , a ) ∈ S ×A . To see this , we can check that Definition 1 is satisfied with K = |S| × |A| and ψ ( s , a ) ( s′ ) = P ( s′|s , a ) for each s , s′ ∈ S and a ∈ A . This example is a sanity check of Definition 1 , which also shows that our results ( Theorem 1 and 2 ) can recover previous results on tabular MDP [ 2 , 34 ] by taking K = |S| × |A| . Example 2 ( Simplex Feature Space ) . If all feature vectors { φ ( s , a ) } ( s , a ) ∈S×A fall in the probability simplex ∆K−1 , a linear transition model can be constructed by taking ψk ( · ) to be any probability measure over S for all k ∈ [ K ] . A key observation is that the model size of linear transition model with known feature mapping φ is |S|K ( the number of coefficients ψk ( s′ ) in ( 1 ) ) , which is still large when the state space S is large . In contrast , it will be established later that to learn a near-optimal policy or Q-function , we only need a much smaller number of samples , which depends linearly on K and is independent of |S| . Next , we introduce a critical assumption employed in prior literature [ 63 , 67 , 45 ] . Assumption 1 ( Anchor state-action pairs ) . Assume there exists a set of anchor state-action pairs K ⊂ S ×A with |K| = K2 such that for any ( s , a ) ∈ S ×A , its corresponding feature vector can be expressed as a convex combination of the feature vectors of anchor state-action pairs { ( s , a ) : ( s , a ) ∈ K } : φ ( s , a ) = ∑ i : ( si , ai ) ∈K λi ( s , a ) φ ( si , ai ) for K∑ i=1 λi ( s , a ) = 1 and λi ( s , a ) ≥ 0 . ( 2 ) Further , we assume that the vectors in { φ ( s , a ) : ( s , a ) ∈ K } are linearly independent . We pause to develop some intuition of this assumption using Examples 1 and 2 . In Example 1 , it is straightforward to check that tabular MDPs satisfies Assumption 1 with K = S ×A . In terms of Example 2 , without loss of generality we can assume that the subspace spanned by the features has full rank , i.e . span { φ ( s , a ) : ( s , a ) ∈ S ×A } = RK ( otherwise we can reduce the dimension of feature space ) . Then we can also check that Example 2 satisfies Assumption 1 with arbitrary K ⊆ S ×A such that the vectors in { φ ( s , a ) : ( s , a ) ∈ K } are linearly independent . In fact , this sort of “ anchor ” notion appears widely in the literature : [ 3 ] considers “ anchor word ” in topic modeling ; [ 19 ] defines “ separability ” in their study of non-negative matrix factorization ; [ 48 ] introduces “ aggregate ” in reinforcement learning ; [ 21 ] studies “ anchor state ” in soft state aggregation models . These concepts all bear some kind of resemblance to our definition of anchor state-action pairs here . Throughout this paper , we assume that the feature mapping φ is known , which is a widely adopted assumption in previous literature [ 63 , 30 , 68 , 26 , 52 , 56 , 60 ] . In practice , large scale RL usually makes use of representation learning to obtain the feature mapping φ . Furthermore , the learned representations can be selected to satisfy the anchor state-action pairs assumption by design . A useful implication of Assumption 1 is that we can represent the transition kernel as P ( ·|s , a ) = ∑ i : ( si , ai ) ∈K λi ( s , a ) P ( ·|si , ai ) , ( 3 ) This follows simply from substituting ( 2 ) into ( 1 ) ( see ( 14 ) in Appendix A for a formal proof ) .

Summary: The paper considers reinforcement learning problem in infinite-horizon, discounted reward MDPs with a generative model. The authors aim to obtain upper bounds on the number of samples required to efficiently learn the optimal policy or the optimal $Q$-function when number of states and actions are very large. In order to achieve this, the authors impose generalization by considering transition probabilities to be linear in known state-action features (a.k.a. the linear MDP setting). The authors further assume access to a number of anchor state-action pairs, whose features (essentially) form a basis of the feature space. Under these assumptions on the MDP, the authors propose both model-based and model-free algorithm for learning the optimal policy (or $Q$-function), and analyze their sample complexity guarantees. The authors prove that the sample complexities of both the algorithms scale linearly with the dimension of the feature space (which is also equal to the number of anchor state-action pairs). The sample complexity bound of the model-based algorithm is the first of its kind for linear MDPs, and matches the minimax lower bound of Yang and Wang, 2019. The sample complexity bound of the model-free algorithm, though suboptimal, improves over the previously known bound of Yang and Wang, 2019.