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: The paper propose and analyze a model-based algorithm for the reward-free exploration setting. The analysis shows that the proposed algorithm is (nearly) minimax optimal. The proposed algorithm is agnostic of the planning algorithm that is used to solve an MDP constructed with the estimated model. This property is of independent interest because it allows the computation of the value of the policies, and the estimated model can be used for any down-stream tasks.