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 reinforcement learning in infinite horizon discounted Markov decision processes with access to a generative model with linearly parameterized transitions.  At a high level, the authors consider an infinite horizon discounted Markov decision process with discount factor $\gamma$, finite state and action spaces, and known reward functions.  The high level goal is to query the transition kernel at various states in order to learn a policy $\pi$ mapping states $s$ to actions $a$ which maximizes the cumulative reward experienced across the trajectory.  It is widely known that learning in general depends on $|S| \times |A|$ which can be prohibitively large whenever $S$ or $A$ is large.  As such, the authors impose additional assumptions in the paper different from prior work with respect to the transition kernel (the only uncertainty present in the model) in order to avoid this worst-case scaling.  In particular, they assume that the algorithm designer has access to a feature function $\phi : S \times A \rightarrow R^K$ such that there are unknown functions $\psi_1, \ldots, \psi_k$ where $P(s' | s,a) = \sum_k \phi_k(s,a) \psi_k(s').$  This definition is the same structure imposed on the probability transition kernel in the linear MDP model (but doesn't also require that the reward function is a linear mapping of the features $\phi$).  This model encompasses tabular MDPs with no additional assumptions where $K$ can be taken to be $S \times A$.  In addition to the linear assumption, the authors assume access to a set of anchor state action pairs of size $K$.  With this assumption, they assume that the feature vectors can be expressed as a convex combination of the feature vectors of the anchor state action pairs (which additionally imposes that the transition kernel can be written as $P(s' | s,a) = \sum_i \lambda_i(s,a) P(s' | s_i, a_i))$ where $i$ indexes over the anchor state action pairs. The authors assume that the $\lambda$ function and the anchor pairs are known in advance to the algorithm designer.  This assumption differentiates the model from the typical infinite horizon discounted MDP to (hopefully) obtain improved regret guarantees.  In particular, the theoretical contributions for the authors follows via two separate strands.  First, the authors consider a model based reinforcement learning algorithm using the generative model.  In particular, their algorithm queries all of the anchor state-action pairs a given number of times to obtain independent samples from the transition $P( | s_i, a_i)$.  Using this, they can then construct an empirical transition kernel for all state action pairs using the linear combination arising from the anchor state action pair assumption.  With this estimate of the transition kernel, the authors can construct an empirical MDP where the true transition is replaced by its estimate (which can be done as the reward function is known), and use planning to find an optimal policy in this perturbed MDP.  Via this method the authors show that the sample complexity scales via $K / (1 - \gamma)^3 \epsilon^2$, which is minimax optimal up to logarithmic factors.  Second, the authors consider a model free reinforcement learning algorithm using the generative model.  This algorithm proceeds across time steps, where in each time step they sample a transition for each anchor state action pair, and update an estimate for the $Q$ function using the 'anchored' Bellman equations.  Via this method the authors show that the sample complexity scales via $K / (1 - \gamma)^4 \epsilon^2$, which is minimax optimal or 'sharp' under naive $Q$ learning based algorithms.  Lastly, the authors complement the theoretical results by a detailed discussion comparing the theoretical guarantees on the performance of their approaches to that in [61].