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: This paper considers MDPs with linear transition models. Assuming access to a generative model, under certain anchor state assumption, the authors prove the following two results:  1. A model-based approach achieves a sample complexity of $O~(K / ((1 - \gamma)^3 \epsilon^2))$ which matches the minimax optimal lower bound up to logarithm factors;  2. A model-free approach based on Q-learning achieves a sample complexity of  $O~(K / ((1 - \gamma)^3 \epsilon^2))$.