Near-Optimal Reward-Free Exploration for Linear Mixture MDPs with Plug-in Solver

1 INTRODUCTION . In reinforcement learning , an agent repeatedly interacts with the unknown environment in order to maximize the cumulative reward . To achieve this goal , an RL algorithm must be equipped with effective exploration mechanisms to learn the unknown environment and find a near-optimal policy . Efficient exploration is critical to the success of reinforcement learning algorithms , which has been widely investigated from both the empirical and the theoretical perspectives ( e.g . Stadie et al . ( 2015 ) ; Pathak et al . ( 2017 ) ; Azar et al . ( 2017 ) ; Jin et al . ( 2018 ) ) . Model-based RL is one of the important approaches to solve for the RL environment . In model-based RL , the agent learns the model of the environment and then performs planning in the estimated model . It has been widely applied in many RL scenarios , including both online setting ( Kaiser et al. , 2019 ; Luo et al. , 2019 ; Azar et al. , 2017 ) and offline setting ( Yu et al. , 2020 ; Kidambi et al. , 2020 ) . It is also believed that modelbased RL is significantly more sample-efficient than model-free RL , which has been justified by many recent empirical results ( e.g . Kaiser et al . ( 2019 ) ; Wang et al . ( 2019 ) ) . Though the theoretical model-based learning in small scale problems has been studied extensively ( Azar et al. , 2017 ; Zhou et al. , 2020a ; Jin et al. , 2020a ) , it is still far from complete , especially with the presence of a function approximator . As an important implication of model-based approaches , the power of plug-in approach have been studied in several works ( Cui & Yang , 2020 ; Agarwal et al. , 2020 ) . The idea of plug-in approach is rather simple : We construct an empirical Markov Decision Process ( MDP ) using maximum likelihood estimate , then return the ( approximate ) optimal policy with efficient planning algorithm in this empirical model . The significance of plug-in approaches is two-folded . For one thing , it preserves an empirical model that keeps the value of the policies , which is of independent interests . For another , the empirical model can be used for any down-stream tasks , which makes the application much more flexible . It is shown that the plug-in approach achieves the minimax sample complexity to compute the -optimal policies with a generative model in the tabular ( Agarwal et al. , 2020 ) and linear settings ( Cui & Yang , 2020 ) . In this paper , we aim to understand the power of plug-in approach in the reward-free exploration with linear function approximation . We study the linear mixture MDPs , where the transition probability kernel is a linear mixture of a number of basis kernels ( Ayoub et al. , 2020 ; Zhou et al. , 2020b ; a ) . We first build an empirical model with an estimation of the transition dynamics in the exploration phase , and then find a near-optimal policy by planning with the empirical model via arbitrary plug-in solver in the planning phase . Our setting is different from the reward-free exploration with linear function approximation without plug-in model ( Wang et al. , 2020 ; Zanette et al. , 2020b ) , in which the agent can directly observe all history samples and design specialized model-free algorithm in the planning phase . Compared to our results , the standard plug-in approach ( Agarwal et al. , 2020 ; Cui & Yang , 2020 ) can not tackle the reward-free setting and may not predict an accurate value function under an arbitrarily given reward function after the exploration phase . Our results shows that the plug-in approach can achieve near-optimal sample complexity in the reward-free setting . In particular , we proposed a statistically efficient algorithm for reward-free exploration . Our algorithm samples Õ ( d2H4/ 2 ) trajectories during the exploration phase , which suffices to obtain O ( ) -optimal policies for an arbitrary reward function with an -optimal pluging solver in the planning phase . Here d is the feature dimension , and H is the planning horizon . Furthermore , with a more refined trajectory-wise uncertainty estimation , we further improve the sample complexity bound to Õ ( d2H3/ 2 ) in the regime where d > H and ≤ H/ √ d. This matches our lower bound Ω ( d2H3/ 2 ) for reward-free exploration in linear mixture MDPs , which indicates that our upper bound is near-optimal except for logarithmic factors . To the best of our knowledge , this is the first work that obtains minimax sample complexity bounds for the plug-in approach in reward-free exploration with linear function approximation . 2 RELATED WORK . RL with Linear Function Approximation Reinforcement learning with linear function approximation has been widely studied in the recent few years ( e.g . Jiang et al . ( 2017 ) ; Yang & Wang ( 2019 ; 2020 ) ; Jin et al . ( 2020b ) ; Modi et al . ( 2020 ) ; Du et al . ( 2019 ) ; Zanette et al . ( 2020a ) ; Cai et al . ( 2020 ) ; Ayoub et al . ( 2020 ) ; Weisz et al . ( 2021 ) ; Zhou et al . ( 2020b ; a ) ) . The linear mixture MDPs model studied in our work assumes the transition probability function is parameterized as a linear function of a given feature mapping over state-action-next-state triple ( Ayoub et al. , 2020 ; Zhou et al. , 2020b ; a ) . Based on the Bernstein inequality for vector-valued martingales , Zhou et al . ( 2020a ) proposed an efficient algorithm that obtains minimax regret in the regime where d > H . Besides linear mixture MDPs , linear MDPs is another category of RL with linear function approximation , which assumes both the transition probability function and reward function are parameterized as a linear function of a given feature mapping over state-action pairs . The algorithms with best regret bound are proposed by Jin et al . ( 2020b ) and Yang & Wang ( 2020 ) , which study model-free algorithm and model-based algorithm respectively . The minimax regret bound for linear MDPs is still unclear . Reward-Free Reinforcement Learning In contrast to the standard RL setting , reward-free reinforcement learning separates the exploration problem and the planning problem which allows one to handle them in a theoretically principled way . For tabular setting , reward-free reinforcement learning has been well-exploited in many previous results ( Jin et al. , 2020a ; Kaufmann et al. , 2021 ; Ménard et al. , 2020 ; Zhang et al. , 2020 ; 2021b ; Wu et al. , 2021a ; Bai & Jin , 2020 ; Liu et al. , 2021 ) , where the minimax rate is obtained by Ménard et al . ( 2020 ) . For reward-free exploration with linear function approximation , Wang et al . ( 2020 ) proposed the first efficient algorithm that obtains O ( d3H6/ 2 ) sample complexity for linear MDPs . However , their algorithm is model-free in nature and can not guarantee good performance with any plug-in solver . Further , Qiu et al . ( 2021 ) proposed the first provably efficient reward-free algorithm with kernel and neural function approximation . We also noticed that there is a concurrent work which also studied reward-free exploration for linear mixture MDPs ( Zhang et al. , 2021a ) . Compared with their results , we focus on the setting of the reward-free exploration with plug-in solver , which is a more general setting that covers the standard reward-free exploration studied in the previous results . Furthermore , our sample complexity bounds are tighter than theirs by a factor of H2 1 . The above two differences introduce new challenges in 1When transformed to the time-homogeneous MDPs setting studied in Zhang et al . ( 2021a ) , our algorithms can achieve sample complexity bounds Õ ( d2H3/ 2 ) and Õ ( ( d2H2 + dH3 ) / 2 ) , respectively . both the algorithmic design and the complexity analysis in this work , which makes our algorithms much more complicated than theirs . Besides , our lower bound is tighter than theirs in the dependence on d. Plug-in Approach The plug-in approach has been studied in tabular/linear case in restrictive settings . E.g. , Agarwal et al . ( 2020 ) and Cui & Yang ( 2020 ) studied the standard plug-in approach with a generative model , where the algorithm is allowed to query the outcome of any state action pair from an oracle . They showed the plug-in approach also achieved the minimax optimal sample complexity to find an -optimal policy in both tabular MDPs and linear MDPs . The reward-free algorithms proposed by Jin et al . ( 2020a ) are model-based in nature , thus can be regarded as a solution in the plug-in solver setting . However , their algorithms are restricted to the tabular case and can not be applied to the setting with linear function approximation . 3 PRELIMINARIES . 3.1 EPISODIC MDPS . We consider the setting of episodic Markov decision processes ( MDPs ) , which can be denoted by a six-tuple ( S , A , P , R , H , ν ) , where S is the set of states , A is the action set , P is the transition probability matrix so that Ph ( ·|s , a ) gives the distribution over states if action a is taken on state s at step h , Rh ( s , a ) is the deterministic reward function of taking action a on state s with support [ 0 , 1 ] in step h , H is the number of steps in each episode , and ν is the distribution of the initial state . In episode k , the agent starts from an initial state sk,1 sampled from the distribution ν . At each step h ∈ [ H ] , the agent observes the current state sk , h ∈ S , takes action ak , h ∈ A , receives reward Rh ( sk , h , ak , h ) , and transits to state sk , h+1 with probability Ph ( sk , h+1|sk , h , ak , h ) . The episode ends when sH+1 is reached . A deterministic policy π is a collection of H policy functions { πh : S → A } h∈ [ H ] . We use Π to denote the set of all deterministic policies . For a specific reward function R , we use V πh : S × R → R to denote the value function at step h under policy π w.r.t . reward R , which gives the expected sum of the remaining rewards received under policy π starting from sh = s , i.e . V πh ( s , R ) = E [ ∑H h′=hR ( sh′ , πh′ ( sh′ ) ) | sh = s , P ] . Accordingly , we define Qπh ( s , a , R ) as the expected Q-value function at step h : Q π h ( s , a , R ) = E [ R ( sh , ah ) + ∑H h′=h+1R ( sh′ , πh′ ( sh′ ) ) | sh = s , ah = a , P ] . We use π∗R to denote the optimal policy w.r.t . reward R , and we use V ∗ h ( · , R ) and Q∗h ( · , · , R ) to denote the optimal value and Q-function under optimal policy π∗R at step h. We say a policy π is -optimal w.r.t . R if E [ H∑ h=1 Rh ( sh , ah ) | π ] ≥ E [ H∑ h=1 Rh ( sh , ah ) | π∗R ] − . For the convenience of explanation , we assume the agent always starts from the same state s1 in each episode . It is straightforward to extend to the case with stochastic initialization , by adding a initial state s0 with no rewards and only one action a0 , and the transition probability of ( s0 , a0 ) is the initial distribution µ . We use PhV ( s , a , R ) as a shorthand of ∑ s′ Ph ( s ′|s , a ) V ( s′ , R ) .

Summary: This paper addresses the reward-free exploration problem with function approximation under linear mixture MDP assumption. It proposes a pair of model-based exploration algorithms, one with a lightweight methodology and sample complexity rate $\tilde{O}(H^4 d^2 \epsilon^{-2})$ and the other as a refined and more convoluted version with rate $\tilde{O}(H^3 d^2 \epsilon^{-2})$, which is matching the reported lower bound up to logarithmic factors. Crucially, the proposed approach can work with any planning solver to provide an $(\epsilon + \epsilon_{opt})$-optimal policy for any reward function.