Year: 2021

Abstract:
Although model-based reinforcement learning (RL) approaches are considered more sample efficient, existing algorithms are usually relying on sophisticated planning algorithm to couple tightly with the model-learning procedure. Hence the learned models may lack the ability of being re-used with more specialized planners. In this paper we address this issue and provide approaches to learn an RL model efficiently without the guidance of a reward signal. In particular, we take a plug-in solver approach, where we focus on learning a model in the exploration phase and demand that any planning algorithm on the learned model can give a near-optimal policy. specifically, we focus on the linear mixture MDP setting, where the probability transition matrix is a (unknown) convex combination of a set of existing models. We show that, by establishing a novel exploration algorithm, the plug-in approach learns a model by taking Õ(dH/ ) episodes with the environment and any -optimal planner on the model gives an O( )-optimal policy on the original model. This sample complexity matches our lower bound for non-plug-in approaches and is statistically optimal. We achieve this result by leveraging a careful maximum total-variance bound using Bernstein inequality and properties specified to linear mixture MDP.

heading: 1 INTRODUCTION
text: 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) cannot 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.

heading: 2 RELATED WORK
text: 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 cannot 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 cannot be applied to the setting with linear function approximation.

heading: 3 PRELIMINARIES

heading: 3.1 EPISODIC MDPS
text: 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).

heading: 3.2 LINEAR MIXTURE MDPS
text: We study a special class of MDPs called 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). This model is defined as follows in the previous literature. Definition 1. (Ayoub et al., 2020) Let φ(s, a, s′) : S×A×S → Rd be a feature mapping satisfying that for any bounded function V : S → [0, 1] and any tuple (s, a) ∈ S ×A, we have ‖φV (s, a)‖2 ≤ 1, where φV (s, a) = ∑ s′∈S φ (s, a, s
′)V (s′) . An MDP is called a linear mixture MDP if there exists a parameter vector θh ∈ Rd with ‖θh‖2 ≤ B for a constant B and a feature vector φ(·, ·, ·),
such that Ph(s′|s, a) = θ>h φ(s, a, s′) for any state-action-next-state triplet (s, a, s′) ∈ S × A × S and step h ∈ [H].

heading: 3.3 REWARD-FREE REINFORCEMENT LEARNING
text: We study the problem of reward-free exploration with plug-in solver. Our setting is different from the reward-free exploration setting studied in the previous literature. Formally, there are two phases in this setting: exploration phase and planning phase.
During the exploration phase, the agent interacts with the environment forK episodes. In episode k, the agent chooses a policy πk which induces a trajectory. The agent observes the states and actions sk,1, ak,1, · · · , sk,H , ak,H as usual, but does not observe any rewards. After K episodes, the agent calculates the estimated model {P̃h = θ̃>h φ}h∈[H], which will be used in the planning phase to calculate the optimal policy.
During the planning phase, the agent is no longer allowed to interact with the MDP. Also, it cannot directly observe the history samples obtained in the exploration phase. Instead, the agent is given a set of reward function {Rh}h∈[H], where Rh : S ×A → [0, 1] is the deterministic reward in step h. For notation convenience, we occasionally use R as a shorthand of {Rh}h∈[H] during the analysis. We define V̂ π,P̃h (s,R) as the value function of transition P̃ and reward R, i.e.
V̂ π,P̃h (s,R) = E
[ H∑
h′=h
Rh (sh′ , πh′ (sh′)) | sh = s, P̃ ] .
In the planning phase, the agent calculates the optimal policy π̂R with respect to the reward function R in the estimated model P̃ using any opt-optimal model-based solver. That is, the returned policy π̂R satisfies:
V̂ ∗,P̃1 (s1, R)− V̂ π̂R,P̃ 1 (s1, R) ≤ opt.
The agent’s goal is to output an accurate model estimation P̃ after the exploration phase, so that the policy π̂R calculated in planning phase can be + opt-optimal w.r.t. any reward function R. Compared with the reward-free setting studied in the previous literature (Jin et al., 2020a; Wang et al., 2020; Ménard et al., 2020; Kaufmann et al., 2021), the main difference is that we require that the algorithm maintains a model estimation instead of all the history samples after the exploration phase, and can use any model-based solver to calculate the near-optimal policy in the planning phase.

heading: 4 REWARD-FREE RL WITH PLUG-IN SOLVER

heading: 4.1 ALGORITHM
text: The exploration phase of the algorithm is presented in Algorithm 1. Recall that for a given value function Ṽk,h+1, we have PhṼk,h+1(sk,h, ak,h) = θ>h φṼk,h+1(sk,h, ak,h) for any k, h. Therefore, Ṽk,h+1(sk,h+1) and φṼk,h+1(sk,h, ak,h) can be regarded as the stochastic reward and the linear feature of a linear bandits problem with linear parameter θh. We employ the standard least-square regression to learn the underlying parameter θh. In each episode, we first update the estimation θ̂k based on history samples till episode k − 1. We define the auxiliary rewards Rk,h to guide exploration. We calculate the optimistic Q-function using the parameter estimation θ̂k, and then execute the greedy policy with respect to the updated Q-function to collect new samples.
The main problem is how to define the exploration-driven reward Rk,h(s, a), which measures the uncertainty for state-action pair (s, a) at the current step. In the setting of linear mixture MDPs, the linear feature φV (s, a) is a function of both state-action pair (s, a) and the next-step value function V . This is not a big deal in the standard exploration setting (Ayoub et al., 2020; Zhou et al., 2020a). However, in the reward-free setting, since the reward function is not given beforehand, we need to upper bound the estimation error of value functions for any possible rewards. To tackle
Algorithm 1 Reward-free Exploration: Exploration Phase Input: Failure probability δ > 0 and target accuracy > 0 λ← B−2, β ← H √ d log(4H3Kλ−1δ−1) + √ λB, V ← {V : S → [0, H]}
for episode k = 1, 2, · · · ,K do Qk,H+1(·, ·) = 0, Vk,H+1(·) = 0 5: for step h = H,H − 1, · · · , 1 do Λk,h ← ∑k−1 t=1 φt,h(st,h, at,h)φt,h(st,h, at,h) > + λI
θ̂k,h ← (Λk,h)−1 ∑k−1 t=1 φt,h(st,h, at,h)Ṽt,h+1,s,a(st,h+1)
Ṽk,h+1,s,a ← arg maxV ∈V ‖ ∑ s′ φ(s, a, s ′)V (s′)‖ (Λk,h)−1
,∀s, a φk,h(s, a)← ∑ s′ φ(s, a, s ′)Ṽk,h+1,s,a(s ′),∀s, a
10: uk,h(s, a)← β √ φk,h(s, a)> (Λk,h) −1 φk,h(s, a),∀s, a
Define the exploration-driven reward function Rk,h(s, a) = uk,h(s, a),∀s, a Qk,h(s, a)← min { θ̂>k,h ( ∑ s′ φ(s, a, s ′)Vk,h+1(s ′)) +Rk,h(s, a) + uk,h(s, a), H } Vk,h(s)← maxa∈AQk,h(s, a), πk,h(s) = arg maxa∈AQk,h(s, a)
end for 15: for step h = 1, 2, · · · , H do
Take action ak,h = πk,h(sk,h) and observe sk,h+1 ∼ Ph(sk,h, ak,h) end for
end for Find P̃h such that the transition P̃h(·|·, ·) = θ̃>h φ(·, ·, ·) is well-defined and ∥∥∥θ̃h − θ̂K,h∥∥∥ ΛK,h ≤
β for h ∈ [H] 20: Output: {P̃h}Hh=1
this problem, we use maxV ∈V β‖φV (s, a)‖Λ−1k,h = maxV ∈V β √ φV (s, a)> (Λk,h) −1 φV (s, a) as a measure of the maximum uncertainty for the state-action pair (s, a), where V = {V : S → [0, H]} is the set of all possible value functions, and Λk,h is the summation of all the history samples {st,h, at,h, st,h+1}k−1t=1 with feature φt,h(st,h, at,h) = arg maxφV β‖φV (s, a)‖Λ−1t,h . In each episode, we define Rk,h(s, a) = uk,h(s, a) = maxV ∈V β‖φV (s, a)‖Λ−1k,h , where Rk,h(s, a) is the exploration-driven reward used to guide exploration, and uk,h(s, a) is the additional bonus term which helps to guarantee that Qk,h(s, a) is an optimistic estimation. After sufficient number of episodes, the uncertainty along all possible trajectories should be low in average, since otherwise the agent would have visited those state-action pairs with large uncertainty as guided by the reward function. Finally, the algorithm returns the model estimation {P̃h}h that is well defined (i.e.∑ s′ P̃h(s ′|s, a) = 1, P̃h(s′|s, a) ≥ 0,∀s, a, s′), and satisfy the constraints ∥∥∥θ̃h − θ̂K,h∥∥∥
ΛK,h ≤ β.
Algorithm 1 involves two optimization problems in line 8 and line 19. These problems can be formulated as the standard convex optimization problem with a slight modification. We discuss the implementation details in Appendix A.2.

heading: 4.2 ANALYSIS
text: In this subsection, we present the sample complexity bound of our Algorithm 1. Theorem 1. With probability at least 1 − δ, after collecting K = Õ ( d2H4
2
) trajectories , Algo-
rithm 1 returns a transition model P̃ , then for any given reward in the planning phase, a policy returned by any opt-optimal plug-in solver on (S,A, P̃ , R,H, ν) isO( + opt)-optimal for the true MDP, (S,A, P,R,H, ν).
We also propose a lower bound for reward-free exploration in Appendix C. Our lower bound indicates that Ω(d2H3/ 2) episodes are necessary to find an -optimal policy with constant probability. Compared with this bound, our result matches the lower bound w.r.t. the dimension d and the pre-
cision except for logarithmic factors. We achieve this by connecting the sample complexity lower bound with the regret lower bound of certain constructed learning algorithms in the standard online exploration setting.
There is also a recent paper of Wang et al. (2020) studying reward-free exploration in the setting of linear MDPs. Their sample complexity is Õ(d3H6/ 2), thus our bound is better than theirs by a factor of H2. This is achieved by setting the exploration-driven reward Rk,h(s, a) = uk,h(s, a) and a more careful analysis of the value estimation error, while they set Rk,h(s, a) = uk,h(s, a)/H in their algorithm. Though the setting is different, we find that our parameter choice and the more refined analysis is applicable to their setting, which can help to further improve their bound by a factor of H2. Please see Appendix D for the detailed discussion.
Now we briefly explain the main idea in the proof. The detailed explanation is deferred to Appendix A. Firstly, we introduce the value function Ṽ ∗h (s,R), which is recursively defined from step H + 1 to step 1:
Ṽ ∗H+1(s,R) = 0,∀s ∈ S (1)
Ṽ ∗h (s,R) = max a∈A
{ min { Rh(s, a) + PhṼ ∗ h+1(s, a,R), H }} ,∀s ∈ S, h ∈ [H] (2)
Compared with the definition of V ∗h (s,R), the main difference is that we take minimization over the value and H at each step. We state the following lemma, which gives an upper bound on the sub-optimality gap of π̂ in the planning phase. Lemma 1. With probability at least 1 − δ, the sub-optimality gap of the policy π̂R for any reward function R in the planning phase satisfies V ∗1 (s1, R)− V π̂R 1 (s1, R) ≤ 4Ṽ ∗1 (s1, RK) + opt, where RK is the exploration-driven reward used for episode K in the exploration phase.
This lemma connects the sub-optimality gap with the value function of the auxiliary reward in episode K. So the remaining problem is how to upper bound Ṽ ∗1 (s1, RK). Since the explorationdriven reward Rk is non-increasing w.r.t. k, it is not hard to prove that KṼ ∗1 (s1, RK) ≤∑K k=1 Ṽ ∗ 1 (s1, Rk). We use the following two lemmas to upper bound ∑K k=1 Ṽ ∗ 1 (s1, Rk). Lemma 2. With probability at least 1− δ, Ṽ ∗h (s,Rk) ≤ Vk,h(s) holds for any (s, a) ∈ S ×A, h ∈ [H] and k ∈ [K]. Lemma 3. With probability at least 1− δ,
K∑ k=1 Vk,1(s1) ≤ 6H2d √ K log(4H3KB2/δ) log(1 +KH2B2/d).
With the help of the optimistic bonus term uk,h, we can prove that the estimation value Vk,h is always optimistic w.r.t Ṽ ∗h , which is illustrated in Lemma 2. Therefore, we have ∑K k=1 Ṽ
∗ 1 (s1, Rk) ≤∑K
k=1 Vk,1(s1). By adapting the regret analysis in the standard RL setting to the reward-free setting, we can upper bound the summation of Vk,1(s1) in Lemma 3. Combining the above two lemmas, we derive the upper bound of Ṽ ∗1 (s1, RK), then we bound the sub-optimality gap of π̂ by Lemma 1.

heading: 5 IMPROVING THE DEPENDENCE ON H
text: In Section 4, we propose an efficient algorithm with O(d2H4/ 2) sample complexity. There is still a gap of H between this result and the lower bound. In this section, we close this gap with a careful maximum total-variance bound using Bernstein inequality.
In the previous results studying regert minimization in online RL setting (Azar et al., 2017; Zhou et al., 2020a; Wu et al., 2021b), one commonly-used approach to obtain the minimax rate is to upper bound the regret using the total variance of the value function by the Bernstein’s concentration inequalities, and finally bound the summation of the one-step transition variance by the law of total variance (Lattimore & Hutter, 2012; Azar et al., 2013; 2017). However, the situation becomes much more complicated in the reward-free setting with linear function approximation. Recall that Ṽk,h+1,s,a(s′) defined in Line 9 of Algorithm 1 is the next-step value function that maximizes the uncertainty ‖φV (s, a)‖Λ−1k,h for state-action pair (s, a). One naive approach is to still use
maxV ‖φV (s, a)‖Λ−1k,h as the uncertainty measure for (s, a) and upper bound the error rate by the summation of the one-step transition variance of Ṽk,h+1,sk,h,ak,h(s
′). However, we can not upper bound the variance summation of Ṽk,h+1,sk,h,ak,h(s
′) in H steps by O(H2) similarly by the law of total variance, since {Ṽk,h,sk,h,ak,h}h∈[H] is not the value functions induced from the same policy and transition dynamics. To tackle the above problem, we need to define the exploration-driven reward and the confidence bonus for each state-action pair in a more refined way. Our basic idea is to carefully measure the expected total uncertainty along the whole trajectories w.r.t. each MDPs in the confidence set. We will explain the detail in the following subsections.

heading: 5.1 ALGORITHM
text: Algorithm 2 Reward-free Exploration: Exploration Phase Input: Failure probability δ > 0 and target accuracy > 0 λ← B−2, β̂ ← 16 √ d log(1 +KH2/(dλ)) log(32K2H/δ) + √ λB
β̌ ← 16d √ log(1 +KH2/(dλ)) log(32K2H/δ) + √ λB
β̃ ← 16H2 √ d log(1 +KH4/(dλ)) log(32K2H/δ) + √ λB
5: Set Λi,k,h ← λI , θ̂i,k,h ← 0 for k = 1, h ∈ [H], i = 1, 2, 3, 4, 5 Set U1,h to be the set containing all the θ̃h that makes P̃h well-defined, h ∈ [H]. for episode k = 1, 2, · · · ,K do
Calculate πk, θ̃k, Rk = arg maxπ,θ̃h∈Uk,h,R V π,P̃ k,1 (s1, R), where V is defined in Eqn 7.
for step h = 1, 2, · · · , H do 10: Take action according to the policy πk,h and observe sk,h+1 ∼ Ph(·|sk,h, ak,h)
end for for step h = 1, 2, · · · , H do
Update {Λi,k+1,h}5i=1 using Eqn 67, 68, 71, 72 and 77 Update the model estimation {θ̂i,k+1,h}5i=1 using Eqn 69, 70, 73, 74 and 78
15: Add the constraints (Eqn 3) to the confidence set Uk,h, and obtain Uk+1,h end for
end for Output: { P̃K,h(·|·, ·) = θ̃>K,hφ(·, ·, ·) }H h=1 .
Our algorithm is described in Algorithm 2. At a high level, Algorithm 2 maintains a high-confidence set Uk,h for the real parameter θh in each episode k, and calculates an optimistic value function V π,P̃k,1 (s1, R) for any reward function R and transition P̃h = θ̃ > h φ with θ̃h ∈ Uk,h. Roughly speaking, the value function V π,P̃k,1 (s1, R) measures the expected uncertainty along the whole trajectory induced by policy π in the MDP with transition P̃ and reward R. To collect more “informative” samples and minimize the worst-case uncertainty over all possible transition dynamics and reward function, we calculate πk = arg maxπ maxθ̃h∈Uk,h,R V π,P̃ k,1 (s1, R), and execute the policy πk to collect more data in episode k. To ensure that the model estimation θ̃k+1,h is close to the true model θh w.r.t. features φV (s, a) of different value functions V , we use the samples collected so far to calculate five model estimation {θ̂i,k,h}5i=1 and the corresponding constraints at the end of the episode k. Each constraint is an ellipsoid in the parameter space centered at the parameter estimation θ̂i,k,h with covariance matrix Λi,k,h and radius βi, i.e.
‖θ̃h − θ̂i,k,h‖Λi,k,h ≤ βi. (3)
We update Uk+1,h by adding these constraints to the the confidence set Uk,h.
The remaining problems are how to define the value function V π,P̃k,1 (s1, R) that represents the expected uncertainty for policy π, transition P̃ and reward R, and how to update the model estimation θi,k,h and the confidence set Uk,h.
Uncertainty Measure Instead of using maxV ‖φV (s, a)‖Λ−1k,h to measure the maximum uncertainty for the state-action pair (s, a) in Algorithm 1, we separately define the uncertainty along the trajectories induced by different policy π, reward function R and transition dynamics P̃h = θ̃>h φ. Specifically, Recall that V̂ π,P̃h (s,R) is the value function of policy π in the MDP model with transition P̃ and reward R. We define the following exploration-driven reward for transition P̃ , reward R and policy π:
uπ,P̃1,k,h(s, a,R) = β̂ ∥∥∥∥∥∑ s′ φ(s, a, s′)V̂ π,P̃h+1 (s ′, R) ∥∥∥∥∥ (Λ1,k,h) −1 . (4)
Suppose Ṽ π,P̃k,H+1(s,R) = 0,∀s ∈ S , we define the corresponding value function recursively from step H + 1 to step 1. That is,
Ṽ π,P̃k,h (s,R) = min { uπ,P̃1,k,h(s, πh(s), R) + PhṼ π,P̃ k,h+1(s, πh(s), R), H } ,∀s ∈ S. (5)
Ṽ π,P̃k,h (s,R) can be regarded as the expected uncertainty along the trajectories induced by policy π for the value function V̂ π,P̃h (s,R). In each episode k, we wish to collect samples by executing the policy π that maximizes the uncertainty measure maxθ̃∈Uk,h,R Ṽ π,P̃ k,1 (s,R). However, the definition of Ṽ π,P̃k,1 (s,R) explicitly depends on the real transition dynamics P , which is unknown to the agent. To solve this problem, we construct an optimistic estimation of Ṽ π,P̃k,h (s,R). We define V π,P̃ k,h (s,R) from step H + 1 to step 1 recursively. Suppose V π,P̃k,H+1(s,R) = 0. We calculate
uπ,P̃2,k,h(s, a,R) = β̂ ∥∥∥∥∥∑ s′ φ(s, a, s′)V π,P̃k,h+1(s ′, R) ∥∥∥∥∥ (Λ2,k,h) −1 , (6)
V π,P̃k,h (s,R) = min { uπ,P̃1,k,h(s, πh(s), R) + u π,P̃ 2,k,h(s, πh(s), R) + P̃hV π,P̃ k,h+1(s, πh(s), R), H } ,
(7)
where uπ,P̃2,k,h is the confidence bonus which ensures that the optimism V π,P̃ k,h (s,R) ≥ Ṽ π,P̃ k,h (s,R) holds with high probability. After calculating V π,P̃k,h (s,R), we take maximization over all θ̃h ∈ Uk,h and reward function R, and calculate πk = arg maxπ maxθ̃h∈Uk,h,R V π,P̃ k,1 (s1, R). We execute πk to collect more samples in episode k.
Weighted Ridge Regression and Confidence Set In Algorithm 2, we maintain five model estimation {θ̂i,k,h}5i=1 and add five constraints to the confidence set Uk,h in each episode. These constraints can be roughly classified into three categories. The first and the third constraints are applied to ensure that P̃k,hV̂ π,P̃ h+1 (s, a,R) is an accurate estimation of Pk,hV̂ π,P̃ h+1 (s, a,R) for any reward function R and P̃k,h = θ̃>k,hφ satisfying θ̃k,h ∈ Uk,h, while the second and the forth constraints are used to guarantee that P̃k,hV π,P̃ h+1 (s, a,R) is an accurate estimation of PhV π,P̃ h+1 (s, a,R) for any reward function R and P̃k,h = θ̃>k,hφ satisfying θ̃k,h ∈ Uk,h. The last constraint is applied due to technical issue and will be explained in Appendix B.2. In this subsection, we introduce the basic idea behind the construction of the first and the third constraints. The second and the forth constraints follow the same idea but consider the different value function V π,P̃h+1 (s, a,R). The formal definition of the parameters {θ̂i,k,h}5i=1 and the constraints are deferred to Appendix B.2.
For notation convenience, we use Vk,h(s) as a shorthand of V̂ πk,P̃k k,h (s,Rk) in this part. The construction of our confidence sets is inspired by a recent algorithm called UCRL-VTR+ proposed by Zhou et al. (2020a). Recall that we use (sk,h, ak,h) to denote the state-action pair that the agent
encounters at step h in episode k. We use the following ridge regression estimator to calculate the corresponding θ̂1,k,h in episode k:
θ̂1,k,h = arg min θ∈Rd λ‖θ‖22 + k−1∑ t=1 [ θ>φVt,h+1/σ̄1,t,h(st,h, at,h)− Vt,h+1(st,h+1, Rt)/σ̄1,t,h ]2 , (8)
where σ̄21,k,h = max { H2/d, [V̄k,hVk,h+1](sk,h, ak,h) + E1,k,h } is an optimistic estimation of the one-step transition variance:
VhVk,h+1(sk,h, ak,h) = Es′∼Ph(·|sk,h,ak,h) [ (Vk,h+1(s ′)− PhVk,h(sk,h, ak,h)) 2 ] .
In the definition of σ̄21,k,h, V̄k,hVk,h+1(sk,h, ak,h) is an empirical estimation for the variance VhVk,h+1(sk,h, ak,h), and E1,k,h is a bonus term defined in Eqn. 61 which ensures that σ̄k,h is an optimistic estimation of VhVk,h+1(sk,h, ak,h). One technical issue here is how to estimate the variance VhVk,h+1(sk,h, ak,h). Recall that by definition,
VhVk,h+1(sk,h, ak,h) =PhV 2k,h+1(sk,h, ak,h)− [PhVk,h+1(sk,h, ak,h)]2 (9) =θ>h φV 2k,h+1(sk,h, ak,h)− [θ > h φVk,h+1(sk,h, ak,h)] 2. (10)
We use θ̃>k,hφV 2k,h+1(sk,h, ak,h) − [ θ̃>k,hφVk,h+1(sk,h, ak,h) ]2 as our variance estimator V̄k,hVk,h+1(sk,h, ak,h), where θ̃k,h ∈ Uk,h is the parameter which maximizes the value V π,P̃k,1 (s1, R) in episode k. To ensure that V̄k,hVk,h+1(sk,h, ak,h) is an accurate estimator, we maintain another parameter estimation θ̂3,k,h using history samples w.r.t. the feature φV 2k,h+1(sk,h, ak,h).
θ̂3,k,h = arg min θ∈Rd λ‖θ‖22 + k−1∑ t=1 [ θ>φV 2t,h+1(st,h, at,h)− V 2 t,h+1(st,h+1) ]2 . (11)
After calculating θ̂1,k,h and θ̂3,k,h, we add the first and the third constraints (Eqn. 3) to the confidence set Uk,h, where Λ1,k,h and Λ3,k,h is the corresponding covariance matrix of all history samples, i.e. Λ1,k,h = ∑k−1 t=1 σ̄ −2 1,t,hφVt,h+1(st,h, at,h)φ > Vt,h+1
(st,h, at,h) and Λ3,k,h =∑k−1 t=1 φV 2t,h+1(st,h, at,h)φ > V 2t,h+1 (st,h, at,h).

heading: 5.2 REGRET
text: We present the regret upper bound of Algorithm 2 in Theorem 2. In the regime where d ≥ H and ≤ H/ √ d, we can obtain Õ(d2H3/ 2) sample complexity upper bound, which matches the sample complexity lower bound except logarithmic factors.
Theorem 2. With probability at least 1 − δ, after collecting K = Õ ( d2H3+dH4 2 + d2.5H2+d2H3 ) trajectories, Algorithm 2 returns a transition model P̃K , then for any given reward in the planning phase, a policy returned by any opt-optimal plug-in solver on (S,A, P̃K , R,H, ν) is O( + opt)optimal for the true MDP, (S,A, P,R,H, ν).

heading: 6 CONCLUSION
text: This paper studies the sample complexity of plug-in solver approach for reward-free reinforcement learning. We propose a statistically efficient algorithm with sample complexity Õ ( d2H4/ 2 ) . We further refine the complexity by providing an another algorithm with sample complexity Õ ( d2H3/ 2 ) in certain parameter regimes. To the best of our knowledge, this is the first minimax sample complexity bound for reward-free exploration with linear function approximation. As a side note, our approaches provide an efficient learning method for the RL model representation, which preserves values and policies for any other down-stream tasks (specified by different rewards). Our sample complexity bound matches the lower bound only when d ≥ H and ≤ H/ √ d. It is unclear whether minimax rate could be obtained in a more broader parameter regimes. We plan to address this issue in the future work.

heading: A OMITTED DETAILS IN SECTION 4

heading: A.1 NOTATIONS
text: In this subsection, we summarize the notations used in Section A.
Symbol Explanation E1 The high-probability event for Theorem 1 sk,h, ak,h The state and action that the agent encounters in episode k and step h Ṽk,h+1,s,a The value function with maximum uncertainty: arg maxV ∈V ‖ ∑ s′ φ(s, a, s ′)V (s′)‖ (Λk,h)−1
φk,h(s, a) ∑ s′ φ(s, a, s ′)Ṽk,h+1,s,a(s ′)
θ̂k,h The estimation of θh in episode k: (Λk,h) −1∑k−1 t=1 φt,h(st,h, at,h)Ṽt,h+1,s,a(st,h+1)
Λk,h The covariance matrix in (k, h): ∑k−1 t=1 φt,h(st,h, at,h)φt,h(st,h, at,h) > + λI
uk,h(s, a) The uncertainty measure: β √ φk,h(s, a)> (Λk,h) −1 φk,h(s, a) Rk,h(s, a) The exploration-driven reward which equals uk,h(s, a) Qk,h(s, a) The Q function defined in line 12 of Algorithm 1 Vk,h(s) The value function defined in line 13 of Algorithm 1 Ṽ ∗h (s,R) The value function defined in Eqn 1 Q̃∗h(s,R) The Q value similarly defined as Ṽ ∗ h (s,R) Ṽ πh (s,R) The value function of policy π similarly defined as Ṽ ∗ h (s,R) θ̃h The parameter estimation returned at the end of the exploration phase P̂k,h(s
′|s, a) θ̂>k,hφ(s, a, s′) P̃h(s ′|s, a) θ̃>h φ(s, a, s′)
A.2 IMPLEMENTATION DETAILS OF ALGORITHM 1
The optimization problem in line 8 of Algorithm 1 can be formulated in the following way:
max V ∥∥∥∥∥∑ s′ φ(s, a, s′)V (s′) ∥∥∥∥∥ (Λk,h)−1 s.t. 0 ≤ V (s) ≤ H,∀s ∈ S (12)
In general, solving this optimization problem is hard. For the case of finite state space (S ≤ ∞), a recent work of Zhang et al. (2021a) relaxed the problem to the following linear programming problem:
max V ∥∥∥Σ−1/21,k Φ(s, a)f∥∥∥ 1 s.t.‖V ‖∞ ≤ H, (13)
where Φ(s, a) = ( φ (s, a, S1) , · · · ,φ ( s, a, S|S| )) and f = ( f (S1) , · · · , f ( S|S| ))> . As discussed in Zhang et al. (2021a), the sample complexity will be worse by a factor of d if we solve the linear programming problem as an approximation. For the case where the state apace is infinite, we can use state aggregation methods (Ren & Krogh, 2002; Singh et al., 1995) to reduce the infinite state space to finite state space and apply the approximation approaches to solve it.
The optimization problem in line 8 of Algorithm 1 is to find parameter θ̃h satisfying several constraints. For the case where the state space is finite, we can solve this problem in the following way:
min θ̃h ∥∥∥θ̃h − θ̂K,h∥∥∥2 ΛK,h
s.t. ∑ s′ θ̃>φ(s, a, s′) = 1
θ̃>φ(s, a, s′) ≥ 0
The above problem can be regarded as a quadratic programming problem and can be solved efficiently by the standard optimization methods. By Lemma 4, we know that the true parameter
θh satisfies ‖θ̂K,h − θh‖ΛK,h ≤ β with high probability. Therefore, the solution θ̃h satisfies the constraint ∥∥∥θ̃h − θ̂K,h∥∥∥ ΛK,h ≤ β with high probability.
For the case where the state space is infinite, we can also solve the above problem using state aggregation methods. In particular, if the linear mixture MDP model can be regarded as a linear combination of several base MDP models (i.e. φ(s, a, s′) = (P1(s′|s, a), P2(s′|s, a), · · · , Pd(s′|s, a))> where Pi(s′|s, a) is the transition probability of certain MDP model), then we can formulate the optimization problem in the following way:
min θ̃h ∥∥∥θ̃h − θ̂K,h∥∥∥2 ΛK,h
s.t. θ̃>1 = 1
θ̃ 0, which can also be solved efficiently when the state space is infinite.

heading: A.3 HIGH-PROBABILITY EVENTS
text: We firstly state the following high-probability events. Lemma 4. With probability at least 1−δ/2, the following inequality holds for any k ∈ [K], h ∈ [H]:
‖θ̂k,h − θh‖Λk,h ≤ β. (14)
Proof. For some step h ∈ [H], the agent selects an action with feature xk,h = φk,h(sk,h, ak,h). The noise ηk,h = Ṽk,h,sk,h,ak,h(s
′) − θ>h φk,h(sk,h, ak,h) satisfies H-sub-Gaussian. By Lemma 20, the following inequality holds with probability 1− δ2H :
‖θ̂k,h − θh‖Λk,h ≤ H √ d log( 4H3K
λδ ) + λ1/2B. (15)
By taking union bound over all h ∈ [H], we can prove the lemma.
Lemma 5. With probability at least 1− δ/2, we have K∑ k=1 H∑ h=1 (PhVk,h+1(sk,h, ak,h)− Vk,h+1(sk,h+1)) ≤ √ 2H3K log(4/δ). (16)
Proof. This lemma follows directly by Azuma’s inequality.
During the following analysis, we denote the high-probability events defined in Lemma 4 and Lemma 5 as E1.

heading: A.4 PROOF OF LEMMA 1
text: Lemma 6. (Restatement of Lemma 1) Under event E1, the sub-optimality gap of the policy π̂R for any reward function R in the planning phase can be bounded by
V ∗1 (s1, R)− V π̂R 1 (s1, R) ≤ 4Ṽ ∗1 (s1, RK) + opt. (17)

heading: Proof.
text: V ∗1 (s1, R)− V π̂R 1 (s1, R) (18) = ( V ∗1 (s1, R)− V̂ π∗R,P̃ 1 (s1, R) ) + ( V̂ π̂R,P̃1 (s1, R)− V π̂R 1 (s1, R) ) + ( V̂ π∗R,P̃ 1 (s1, R)− V̂ π̂R,P̃ 1 (s1, R) ) (19)
≤ ( V ∗1 (s1, R)− V̂ π∗R,P̃ 1 (s1, R) ) + ( V̂ π̂R,P̃1 (s1, R)− V π̂R 1 (s1, R) ) + opt. (20)
The inequality is because that the policy π̂R is the opt-optimal policy in the estimated MDP M̂ , i.e. V̂ π∗R,P̃ 1 (s1, R) ≤ V̂ π̂R,P̃ 1 (s1, R) + opt .
For the notation convenience, for a certain sequence {Rh}Hh=1, we define the function Wh({Rh}) recursively from stepH+1 to step 1. Firstly, we defineWH+1({Rh}) = 0. Wh({Rh}) is calculated recursively from Wh+1({Rh}):
Wh({Rh}) = min {H,Rh +Wh+1({Rh})} . (21)
Similarly with the definition of Ṽ ∗h (s,R), we introduce the value function Ṽ π h (s,R), which is recursively defined from step H + 1 to step 1:
Ṽ πH+1(s,R) = 0,∀s ∈ S (22) Ṽ πh (s,R) = { min { Rh(s, π(s)) + PhṼ π h+1(s, π(s), R), H }} ,∀s ∈ S, h ∈ [H]. (23)
We use traj ∼ (π, P ) to indicate that the trajectory {sh, ah}Hh=1 is sampled from transition P with policy π. For any policy π, we have
∣∣∣Es1∼µ (V̂ π,P̃1 (s1, R)− V π1 (s1))∣∣∣ (24) = ∣∣∣Etraj∼(π,P )W1 ({(P̃h − Ph)V̂ π,P̃Kh+1 (sh, ah, R)})∣∣∣ (25)
= ∣∣∣∣∣Etraj∼(π,P )W1 ({ (θ̃h − θh) ∑ s′ φ(sh, ah, s ′)V̂ π,P̃h+1 (s ′, R) })∣∣∣∣∣ (26) ≤Etraj∼(π,P )W1 ∥∥∥θ̃h − θh∥∥∥ΛK,h ∥∥∥∥∥∑ s′ φ(sh, ah, s ′)V̂ π,P̃h+1 (s ′, R) ∥∥∥∥∥ (ΛK,h) −1   (27)
≤Etraj∼(π,P )W1 2β ∥∥∥∥∥∑ s′ φ(sh, ah, s ′)V̂ π,P̃h+1 (s ′, R) ∥∥∥∥∥ (ΛK,h) −1   (28)
≤2Etraj∼(π,P )W1 ({uK,h(sh, ah)}) (29) =2Ṽ π1 (s1, RK) (30) ≤2Ṽ ∗1 (s1, RK). (31)
The second inequality is due to lemma 4 and the definition of θ̃. Plugging this inequality back to Inq. 18, we can prove the lemma.

heading: A.5 PROOF OF LEMMA 2
text: Lemma 7. (Restatement of Lemma 2) Under event E1, Ṽ ∗h (s,Rk) ≤ Vk,h(s) holds for any (s, a) ∈ S ×A, h ∈ [H] and k ∈ [K]. .
Proof. We prove the lemma by induction. Suppose Ṽ ∗h+1(s,Rk) ≤ Vk,h+1(s), Q̃∗h(s, a,Rk)−Qk,h(s, a) (32)
≤− uk,h(s, a) + ( θk,h − θ̂k,h )> · ∑ s′ φ(s, a, s′)Vk,h(s ′) + Ph(Ṽ ∗ h+1 − Vk,h+1)(s, a) (33)
≤− uk,h(s, a) + ( θk,h − θ̂k,h )> · ∑ s′ φ(s, a, s′)Vk,h(s ′) (34)
≤− uk,h(s, a) + ∥∥∥θk,h − θ̂k,h∥∥∥
Λk,h ∥∥∥∥∥∑ s′ φ(s, a, s′)Vk,h(s ′) ∥∥∥∥∥ (Λk,h)−1
(35)
≤− uk,h(s, a) + ∥∥∥θk,h − θ̂k,h∥∥∥
Λk,h max V ∈V ∥∥∥∥∥∑ s′ φ(s, a, s′)V (s′) ∥∥∥∥∥ (Λk,h)−1
(36)
≤− uk,h(s, a) + βmax V ∈V ∥∥∥∥∥∑ s′ φ(s, a, s′)V (s′) ∥∥∥∥∥ (Λk,h)−1
(37)
=0. (38)
The first inequality is due to induction condition Ṽ ∗h+1(s,Rk) ≤ Vk,h+1(s). The last inequality is due to Lemma 4. Since Q̃∗h(s, a) ≤ Qk,h(s, a,Rk) for any a ∈ A, we have Ṽ ∗h (s,Rk) ≤ Vk,h(s).

heading: A.6 PROOF OF LEMMA 3
text: Lemma 8. (Restatement of Lemma 3) Under event E1, ∑K k=1 Vk,1(sk,1) ≤
6H2d √ K log(4H3KB2/δ) log(1 +KH2B2/d).

heading: Proof.
text: Qk,h(sk,h, ak,h) (39)
= min { H,Rk,h(sk,a, ak,h) + uk,h(sk,a, ak,h) + P̂k,hVk,h+1(sk,h, ak,h) } (40)
≤min { H, 2β ‖φk,h(sk,h, ak,h)‖Λ−1k,h + P̂k,hVk,h+1(sk,h, ak,h) } (41)
≤min { H, 2β ‖φk,h(sk,h, ak,h)‖Λ−1k,h } + min { H, ( P̂k,h − Ph ) Vk,h+1(sk,h, ak,h) } (42)
+ (PhVk,h+1(sk,h, ak,h)− Vk,h+1(sk,h+1)) + Vk,h+1(sk,h+1) (43) ≤3βmin { 1, ‖φk,h(sk,h, ak,h)‖Λ−1k,h }
+ (PhVk,h+1(sk,h, ak,h)− Vk,h+1(sk,h+1)) + Vk,h+1(sk,h+1). (44)
The last inequality is from( P̂k,h − Ph ) Vk,h+1(sk,h, ak,h) =(θ̂k,h − θh)φVk,h+1(sk,h, ak,h) (45)
≤ ∥∥∥θ̂k,h − θh∥∥∥
Λk,h ∥∥φVk,h+1(sk,h, ak,h)∥∥Λ−1k,h (46) ≤β ∥∥φVk,h+1(sk,h, ak,h)∥∥Λ−1k,h (47)
≤β ‖φk,h(sk,h, ak,h)‖Λ−1k,h . (48)
From Inq 39, we have K∑ k=1 Vk,1(s1) = K∑ k=1 H∑ h=1 3βmin { 1, ‖φk,h(sk,h, ak,h)‖Λ−1k,h }
(49)
+ K∑ k=1 H∑ h=1 (PhVk,h+1(sk,h, ak,h)− Vk,h+1(sk,h+1)) . (50)
For the first term, by Lemma 21, we have for any h ∈ [H],
K∑ k=1 min { 1, ‖φk,h(sk,h, ak,h)‖Λ−1k,h } ≤ √√√√K K∑ k=1 min { 1, ‖φk,h(sk,h, ak,h)‖2Λ−1k,h }
(51)
≤ √ 2dK log(1 +KH2/(dλ)). (52)
Since λ = B−2, we have K∑ k=1 H∑ h=1 3βmin { 1, ‖φk,h(sk,h, ak,h)‖Λ−1k,h } ≤ 3H2d √ 2K log(4H3KB2/δ) log(1 +KH2B2/d).
(53)
For the second term, by Lemma 5,
K∑ k=1 H∑ h=1 (PhVk,h+1(sk,h, ak,h)− Vk,h+1(sk,h+1)) ≤ √ 2H3K log(4/δ). (54)
Therefore, we have
K∑ k=1 Vk,1(sk,1) ≤ 6H2d √ K log(4H3KB2/δ) log(1 +KH2B2/d). (55)
Lemma 9. Under event E1, Ṽ ∗1 (s1, RK) ≤ 6H2d √ log(4H3KB2/δ) log(1+KH2B2/d) K .
Proof. Since Ṽ ∗1 (sk,1, Rk) ≤ Vk,1(sk,1) by lemma 7, we have
K∑ k=1 Ṽ ∗1 (s1, Rk) ≤ K∑ k=1 Vk,1(s1) ≤ 6H2d √ K log(4H3KB2/δ) log(1 +KH2B2/d). (56)
By the definition of Rk, we know that Rk(s, a) is non-increasing w.r.t k. Therefore,
K∑ k=1 Ṽ ∗1 (s1, RK) ≤ K∑ k=1 Ṽ ∗1 (s1, Rk) ≤ 6H2d √ K log(4H3KB2/δ) log(1 +KH2B2/d). (57)
The lemma is proved by dividing both sides by K.

heading: A.7 PROOF OF THEOREM 1
text: Combining the results in Lemma 9 and Lemma 6, we know that for any reward function R,
V ∗1 (s1, R)− V π̂R 1 (s1, R) ≤ 24H2d
√ log(4H3KB2/δ) log(1 +KH2B2/d)
K + opt. (58)
Choosing K = C1H 4d2 log(4H3KB2/δ) log(1+KH2B2/d)
2 for some constant C1 suffices to guarantee that V ∗1 (s1, R)− V π̂R 1 (s1, R) ≤ + opt.

heading: B OMITTED DETAILS IN SECTION 5

heading: B.1 NOTATIONS
text: In this subsection, we summarize the notations used in Section B.
Symbol Explanation E2 The high-probability event for Theorem 2 sk,h, ak,h The state and action that the agent encounters in episode k and step h φV (s, a,R) ∑ s′ φ(s, a, s
′)V π,P (s′, R) for certain value function V V̂ π,P̃h (s,R) The value function of policy π in the MDP model with transition P̃ and reward R Ṽ π,P̃k,h (s,R) The expected uncertainty along the trajectories induced by policy π for V̂ π,P̃ h (s,R) (Defined in Eqn 5) V π,P̃k,h (s,R) The optimistic estimation of Ṽ π,P̃ k,h (s,R) (Defined in Eqn 7) uπ,P̃1,k,h(s, a,R) The exploration-driven reward for transition P̃ , reward R and policy π (Defined in Eqn 4) uP̃ ,π2,k,h(s, a,R) The confidence bonus ensuring the optimism V π,P̃ k,h (s,R) ≥ Ṽ π,P̃ k,h (s,R) (Defined in Eqn 6) VhV (s, a) The one-step transition variance w.r.t. certain value function V V̄1,k,h (s, a) The empirical variance estimation w.r.t. the value V̂ πk,P̃kh (s,Rk) (Defined in Eqn 59) V̄2,k,h (s, a) The empirical variance estimation w.r.t. the value V πk,P̃kh (s,Rk) (Defined in Eqn 60) E1,k,h, E2,k,h The confidence bonus for the variance estimation V̄1,k,h (s, a) and V̄2,k,h (s, a), respectively σ̄21,k,h, σ̄ 2 2,k,h The optimistic variance estimation for V̄1,k,h (s, a) and V̄2,k,h (s, a), respectively θ̂i,k,h The parameter estimation w.r.t. certain value function (Defined in Section B.2) Λi,k,h The empirical covariance matrix w.r.t. certain value function (Defined in Section B.2) Uk,h The confidence set containing θh with high probability πk, θ̃k, Rk arg maxπ,θ̃h∈Uk,h,R V π,P̃ k,1 (s1, R) Ỹk,h(s) The “value function” for the MDP with transition P̃k and reward V̄1,k,h τ log(32K2H/δ) log2(1 +KH4B2)

heading: B.2 OMITTED DETAILS OF ALGORITHM 2
text: In algorithm 2, we maintain five different parameter estimation θ̂1,k,h, θ̂2,k,h, θ̂3,k,h, θ̂4,k,h and θ̂5,k,h, and the corresponding covariance matrix Λ1,k,h,Λ2,k,h,Λ3,k,h,Λ4,k,h and Λ5,k,h. θ̂1,k,h, θ̂2,k,h are the parameter estimation using history samples w.r.t the features of variancenormalized value function V̂ πk,P̃kk,h+1 /σ̄1,k,h and V πk,P̃k k,h+1 /σ̄2,k,h, respectively. θ̂3,k,h, θ̂4,k,h are the
parameter estimation using samples w.r.t the features of value function ( V̂ πk,P̃kk,h+1 )2 and ( V πk,P̃kk,h+1 )2 .
θ̂5,k,h is somewhat technical and will be explained later. for any i ∈ [5] and h ∈ [H], θ̂i,1,h are initialized as 0, and Λi,1,h are initialized as λI in Algorithm 2. U1,h is the set containing all the θ̃h that makes P̃h well-defined, i.e. ∑ s′ θ̃ > h φ(s ′|s, a) = 1 and θ̃>h φ(s′|s, a) ≥ 0,∀s, a.
After observing {sk,h, ak,h}Hh=1 in episode k, we calculate V̄1,k,h (s, a) and V̄2,k,h (s, a) as the corresponding variance in the empirical MDP with transition dynamics P̃k.
V̄1,k,h(s, a) =θ̃>k,h ∑ s′ φ(s, a, s′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2 − [ θ̃>k,h ∑ s′ φ(s, a, s′)V̂ πk,P̃kk,h+1 (s ′, Rk) ]2 ,
(59)
V̄2,k,h(s, a) =θ̃>k,h ∑ s′ φ(s, a, s′) ( V πk,P̃kk,h+1 (s ′, Rk) )2 − [ θ̃>k,h ∑ s′ φ(s, a, s′)V πk,P̃kk,h+1 (s ′, Rk) ]2 .
(60)
To guarantee the variance estimation is optimistic, we calculate the confidence bonus E1,k,h and E2,k,h for variance estimation V̄1,k,h (sk,h, ak,h) and V̄2,k,h (sk,h, ak,h):
E1,k,h = min H2, 4Hβ̌k ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ (Λ1,k,h)−1  (61) + min
H2, 2β̃k ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2∥∥∥∥∥
(Λ3,k,h)−1  , (62) E2,k,h = min
H2, 4Hβ̌k ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ (Λ2,k,h)−1  (63) + min
H2, 2β̃k ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′) ( V πk,P̃kk,h+1 (s ′, Rk) )2∥∥∥∥∥
(Λ4,k,h)−1
 . (64)
We add the bonuses to V̄1,k,h(sk,h, ak,h) and V̄2,k,h(sk,h, ak,h) and maintain the optimistic variance estimation σ̄21,k,h and σ̄ 2 2,k,h:
σ̄21,k,h = max { H2/d, V̄1,k,h(sk,h, ak,h) + E1,k,h } , (65)
σ̄22,k,h = max { H2/d, V̄2,k,h(sk,h, ak,h) + E2,k,h } . (66)
We use σ̄i,k,h to normalize the value obtained in each step. By setting λ = B−2 and solving the ridge regression defined in Eqn. 8, we know that θ̂i,k+1,h and Λi,k+1,h are updated in the following way:
Λ1,k+1,h = Λ1,k,h + (σ̄1,k,h) −2 (∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) )(∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) )> ,
(67)
Λ2,k+1,h = Λ1,k,h + (σ̄2,k,h) −2 (∑ s′ φ(sk,h, ak,h, s ′)V πk,P̃kk,h+1 (s ′, Rk) )(∑ s′ φ(sk,h, ak,h, s ′)V πk,P̃kk,h+1 (s ′, Rk) )> ,
(68)
θ̂1,k+1,h = (Λ1,k+1,h) −1 k∑ t=1 (σ̄1,t,h) −2 (∑ s′ φ(st,h, at,h, s ′)V̂ πt,P̃tt,h+1 (s ′, Rt) ) V̂ πt,P̃tt,h+1 (st,h+1, Rt),
(69)
θ̂2,k+1,h = (Λ2,k+1,h) −1 k∑ t=1 (σ̄2,t,h) −2 (∑ s′ φ(st,h, at,h, s ′)V πt,P̃tt,h+1 (s ′, Rt) ) V πt,P̃tt,h+1 (st,h+1, Rt).
(70)
Similarly, we update the estimation w.r.t ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2 and ( V πk,P̃kk,h+1 (s ′, Rk) )2 :
Λ3,k+1,h = Λ3,k,h + (∑ s′ φ(sk,h, ak,h, s ′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2)(∑ s′ φ(sk,h, ak,h, s ′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2)> ,
(71)
Λ4,k+1,h = Λ4,k,h + (∑ s′ φ(sk,h, ak,h, s ′) ( V πk,P̃kk,h+1 (s ′, Rk) )2)(∑ s′ φ(sk,h, ak,h, s ′) ( V πk,P̃kk,h+1 (s ′, Rk) )2)> ,
(72)
θ̂3,k+1,h = (Λ3,k+1,h) −1 k∑ t=1 (∑ s′ φ(st,h, at,h, s ′) ( V̂ πt,P̃tt,h+1 (s ′, Rt) )2)( V̂ πt,P̃tt,h+1 (st,h+1, Rt) )2 ,
(73)
θ̂4,k+1,h = (Λ4,k+1,h) −1 k∑ t=1 (∑ s′ φ(st,h, at,h, s ′) ( V πt,P̃tt,h+1 (s ′, Rt) )2)( V πt,P̃tt,h+1 (st,h+1, Rt) )2 .
(74)
We define Ỹk,h to be the “value function” with transition P̃k and reward V̄1,k,h.
Ỹk,H+1(s) = 0, (75)
Ỹk,h(s) = V̄1,k,h(s, πk(s)) + P̃k,hỸk,h+1(s, πk(s)). (76)
θ̂5,k,h is the parameter estimation using samples w.r.t. Ỹk,h(s):
Λ5,k+1,h = Λ5,k,h + (∑ s′ φ(sk,h, ak,h, s ′)Ỹk,h+1(s ′) )(∑ s′ φ(sk,h, ak,h, s ′)Ỹk,h+1(s ′) )> ,
(77)
θ̂5,k+1,h = (Λ5,k+1,h) −1 k∑ t=1 (∑ s′ φ(st,h, at,h, s ′)Ỹt,h+1(s ′) ) Ỹt,h+1(sk,h+1). (78)
We update the high-confidence set Uk,h by adding following five constraints:
‖θ̃h − θ̂1,k,h‖Λ1,k,h ≤ β̂, (79)
‖θ̃h − θ̂2,k,h‖Λ2,k,h ≤ β̂, (80)
‖θ̃h − θ̂3,k,h‖Λ3,k,h ≤ β̃, (81)
‖θ̃h − θ̂4,k,h‖Λ4,k,h ≤ β̃, (82)
‖θ̃h − θ̂5,k,h‖Λ5,k,h ≤ β̃. (83)
Note that we add new constraints to this confidence set Uk,h in each episode, instead of updating Uk,h as the intersection of the new constraints. This operation is designed to ensure that the cardinality of the confidence set Uk,h is non-increasing w.r.t. the episode k, so that V π,P̃k,1 (s1, R) is always non-increasing w.r.t. k.
During the proof, we use φV π,P (s, a,R) as a shorthand of ∑ s′ φ(s, a, s ′)V π,P (s′, R).

heading: B.3 HIGH-PROBABILITY EVENTS
text: For notation convenience, we use V1,k,h(s, a) and V2,k,h(s, a) to denote the one-step transition variance with regard to V̂ πk,Pkk,h+1 and V πk,Pk k,h+1 , i.e.
V1,k,h(s, a) = Es′∼Ph(·|s,a) [( V̂ πk,Pkk,h+1 (s ′, Rk)− PhV̂ πk,Pkk,h (s, a,Rk) )2] , (84)
V2,k,h(s, a) = Es′∼Ph(·|s,a) [( V πk,Pkk,h+1 (s ′, Rk)− PhV πk,Pkk,h (s, a,Rk) )2] . (85)
(86)
Lemma 10. With probability at least 1− δ, the following event holds for any h ∈ [H], k ∈ [K]:∥∥∥θ̃k,h − θh∥∥∥ Λ1,k,h
≤ 2β̂, (87)∥∥∥θ̃k,h − θh∥∥∥ Λ2,k,h
≤ 2β̂, (88)∣∣V̄1,k,h(sk,h, ak,h)− V1,k,h(sk,h, ak,h)∣∣ ≤ E1,k,h, (89)∣∣V̄2,k,h(sk,h, ak,h)− V2,k,h(sk,h, ak,h)∣∣ ≤ E2,k,h, (90)∥∥∥θ̃k,h − θh∥∥∥ Λ5,k,h ≤ 2β̃. (91)
Proof. We firstly prove the first and the third inequality, the second and the fourth inequality can be proved following the same idea. By the definition of V̄1,k,h(sk,h, ak,h) and V1,k,h(sk,h, ak,h):∣∣V̄1,k,h(sk,h, ak,h)− V1,k,h(sk,h, ak,h)∣∣ (92)
≤min { H2, 2H ∥∥∥θ̃k,h − θh∥∥∥ Λ1,k,h ∥∥∥∥φV̂ πk,Pkk,h+1 (sk,h, ak,h, Rk) ∥∥∥∥
(Λ1,k,h)−1
} (93)
+ min { H2, ∥∥∥θ̃k,h − θh∥∥∥ Λ3,k,h ∥∥∥∥φ(V̂ πk,Pkk,h+1 )2 (sk,h, ak,h, Rk) ∥∥∥∥
(Λ3,k,h)−1
} . (94)
Let xk = (σ̄1,k,h)−1φV̂ πk,Pkk,h+1 (sk,h, ak,h), and the noise
ηk = (σ̄1,k,h) −1V̂ πk,Pkk,h+1 (sk,h+1, Rk)− (σ̄1,k,h) −1 〈 φ V̂ πk,Pk k,h+1 (sk,h, ak,h, Rk) , θh 〉 .
Since σ̄1,k,h ≥ √ H2/d, we have ‖xk‖2 ≤ √ d, E [ η2k | Gk ] ≤ d. By Lemma 22, we have with prob at least 1− δ/(8H), for all k ≤ K,∥∥∥θh − θ̂1,k,h∥∥∥ Λ1,k,h ≤ 16d √ log(1 +KH2/(dλ)) log(32K2H/δ) + √ λB = β̌. (95)
Similarly, we can prove that∥∥∥θh − θ̂3,k,h∥∥∥ Λ3,k,h ≤ 16H2 √ d log(1 +KH4/(dλ)) log(32K2H/δ) + √ λB = β̃. (96)
Since ∥∥∥θ̃k,h − θ̂1,k,h∥∥∥
Λ1,k,h ≤ β̂ ≤ β̌ by Inq 79, we have ∥∥∥θh − θ̃k,h∥∥∥ Λ1,k,h ≤ 2β̌. Similarly,
since ∥∥∥θ̃k,h − θ̂3,k,h∥∥∥
Λ3,k,h ≤ β̃ by Inq 81, we have ∥∥∥θh − θ̃k,h∥∥∥ Λ3,k,h
≤ 2β̃. Plugging the above inequalities back to Inq. 92, we have∣∣V̄1,k,h(sk,h, ak,h)− V1,k,h(sk,h, ak,h)∣∣ ≤ E1,k,h. (97) The above inequality indicates that σ̄1,k,h ≥ V1,k,h, which is an optimistic estimation. As a result, the variance of the single-step noise ηk = (σ̄1,k,h)−1V̂ πk,Pk k,h+1 (sk,h+1, Rk) −
(σ̄1,k,h) −1 〈 φ V̂ πk,Pk k,h+1 (sk,h, ak,h, Rk) , θh 〉 satisfies E [ η2k | Gk ] ≤ 1. By Lemma 22, we can prove
a tighter confidence guarantee for θ̂1,k,h:∥∥∥θh − θ̂1,k,h∥∥∥ Λ1,k,h ≤ 16 √ d log(1 +KH2/(dλ)) log(32K2H/δ) + √ λB = β̂. (98)
Combining with ∥∥∥θ̃1,k,h − θ̂1,k,h∥∥∥
Λ1,k,h ≤ β̂ by Inq 79, we have ∥∥∥θh − θ̃1,k,h∥∥∥ Λ1,k,h ≤ 2β̂.
Now we prove the last inequality in this lemma. Recall that we define
Ỹk,H+1(s) = 0, (99)
Ỹk,h(s) = V̄1,k,h(s, πk(s)) + P̃k,hỸk,h+1(s, πk(s)). (100)
Note that Ỹk,h(s) ≤ H2 by the law of total variance (Lattimore & Hutter, 2012; Azar et al., 2013). Therefore, the variance on the single-step noise in at most H2. By Lemma 22, we have∥∥∥θh − θ̂5,k,h∥∥∥
Λ5,k,h ≤ β̃. (101)
Combining with ∥∥∥θ̃h − θ̂5,k,h∥∥∥
Λ5,k,h ≤ β̃ by Inq 83, we have ∥∥∥θh − θ̃5,k,h∥∥∥ Λ5,k,h ≤ 2β̃.
Lemma 11. With probability at least 1− δ/2, we have
K∑ k=1 H∑ h=1 ( PhV πk,P̃k k,h+1 (sk,h, ak,h, Rk)− V πk,P̃k k,h+1 (sk,h+1, ak,h+1, Rk) ) ≤ √ 2H3K log(8/δ), (102)
K∑ k=1 H∑ h=1 ( PhỸk,h+1(sk,h, ak,h, Rk)− Ỹk,h+1(sk,h+1, ak,h+1, Rk) ) ≤ √ 2H5K log(8/δ). (103)
Proof. This lemma follows directly by Azuma’s inequality for martingale difference sequence and union bound.
During the following analysis, we denote the high-probability events defined in Lemma 10 and Lemma 11 as E2.

heading: B.4 PROOF OF THEOREM 2
text: Lemma 12. Under event E2, we have
K∑ k=1 V πk,P̃kk,1 (s1, Rk) (104)
≤O (√ (dH4 + d2H3)K log(KH/δ) log2(KH4B2)) + (d2.5H2 + d2H3) log(KH/δ) log2(KH4B2) ) .
(105)
This lemma can be proved with the technique for regret analysis. We will explain it in Appendix B.5.
Lemma 13. (Optimism) Under event E2, for any h ∈ [H], k ∈ [K], s ∈ S , π ∈ Π, R and P̃ satisfying P̃h = θ̃>h φ, we have
V π,P̃k,h (s,R) ≥ Ṽ π,P̃ k,h (s,R). (106)
Proof. This lemma is proved by induction. Suppose V π,P̃k,h+1(s,R) ≥ Ṽ π,P̃ k,h+1(s,R), we have
V π,P̃k,h (s,R)− Ṽ π,P̃ k,h (s,R) (107)
≥uπ,P̃2,k,h(s, π(s), R) + (P̃k,h − Ph)V π,P̃ k,h+1(s, π(s), R) + Ph ( V π,P̃k,h+1 − Ṽ π,P̃ k,h+1 ) (s, π(s), R)
(108)
≥uπ,P̃2,k,h(s, π(s), R) + (P̃k,h − Ph)V π,P̃ k,h+1(s, π(s), R) (109)
≥uπ,P̃2,k,h(s, π(s), R)− ∥∥∥θ̃k,h − θh∥∥∥
Λ2,k,h ∥∥∥∥∥∑ s′ φ(s, π(s), s′)V π,P̃k,h+1(s ′, R) ∥∥∥∥∥ (Λ2,k,h)−1
(110)
≥0. (111)
This indicates that V π,P̃k,h (s,R) ≥ Ṽ π,P̃ k,h (s,R) holds for step h. Lemma 14. under event E2, we have V πK ,P̃KK,1 (s1, RK) ≤ ∑K k=1 V πk,P̃k k,1 (s1,Rk) K .
Proof. Firstly, we prove that V π,P̃k,1 (s,R) is non-increasing w.r.t k for any fixed π, P̃ , s and R. This can be proved by induction. Suppose for any k1 ≤ k2, V π,P̃k1,h+1(s,R) ≥ V π,P̃ k2,h+1
(s,R) for any s. Recall that
V π,P̃k,h (s,R) = min { uπ,P̃1,k,h(s, πh(s), R) + u π,P̃ 2,k,h(s, πh(s), R) + P̃hV π,P̃ k,h+1(s, πh(s), R), H } .
(112)
Since Λ1,k1,h 4 Λ1,k2,h, Λ2,k1,h 4 Λ2,k2,h and V π,P̃ k1,h+1 (s,R) ≥ V π,P̃k2,h+1(s,R) for any s, we can prove that
uπ,P̃1,k1,h(s, πh(s), R) ≥ u π,P̃ 1,k2,h (s, πh(s), R), (113)
uπ,P̃2,k1,h(s, πh(s), R) ≥ u π,P̃ 2,k2,h (s, πh(s), R), (114)
P̃hV π,P̃ k1,h+1 (s, πh(s), R) ≥ P̃hV π,P̃k2,h+1(s, πh(s), R). (115)
Therefore, V π,P̃k1,h(s,R) ≥ V π,P̃ k2,h (s,R) holds for step h and any s.
Since the cardinality of the transition set Uk,h is non-increasing (We add more constraints in each episode), we know that θ̃k2,h ∈ Uk1,h. By the optimality of πk1 , θ̃k1 and Rk1 in episode k1, we have V πk1 ,P̃k1 k1,1 (s1, Rk1) ≥ V πk2 ,P̃k2 k1,1 (s1, Rk2).
Combining the above two inequalities, we have
V πk1 ,P̃k1 k1,1 (s1, Rk,1) ≥ V πk2 ,P̃k2 k1,1 (s1, Rk2) ≥ V πk2 ,P̃k2 k2,1 (s1, Rk2). (116)
This indicates that the value V πk,P̃kk,1 (s1, Rk) is non-increasing w.r.t. k. Therefore,
KV πK ,P̃KK,1 (s1, RK) ≤ K∑ k=1 V πk,P̃kk,1 (s1, Rk). (117)
Lemma 15. Under event E2, the sub-optimality gap of the policy π̂R returned in the planning phase for any reward function R can be bounded by
V ∗1 (s1, R)− V π̂R 1 (s1, R) ≤ 4V πK ,P̃K K,1 (s1, RK) + opt. (118)
Proof. Recall that we use V̂ πh to denote the value function of policy π on the estimated model θ̃K .
V ∗1 (s1, R)− V π̂R 1 (s1, R) (119) = ( V ∗1 (s1, R)− V̂ π∗R,P̃K 1 (s1, R) ) + ( V̂ π̂R,P̃K1 (s1, R)− V π̂R 1 (s1) ) + ( V̂ π∗R,P̃K 1 (s1, R)− V̂ π̂R,P̃K 1 (s1, R) ) (120)
≤V ∗1 (s1, R)− V̂ π∗R,P̃K 1 (s1, R) + V̂ π̂R,P̃K 1 (s1, R)− V π̂R 1 (s1, R) + opt. (121)
The inequality is because that the policy π̂R is the opt-optimal policy in the estimated MDP M̂ .
For notation convenience, we use traj ∼ (π, P ) to denote that the trajectory ({sh, ah}Hh=1) is sampled with transition P and policy π. For a certain sequence {Rh}Hh=1, we define the function Wh({Rh}) recursively from step H + 1 to step 1. Firstly, we define WH+1({Rh}) = 0. Wh(R) is calculated from Wh+1(R):
Wh({Rh}) = min {H,Rh +Wh+1({Rh})} . (122)
With the above notation, we can prove that for any policy π ∈ Π,
∣∣∣V̂ π,P̃K1 (s1, R)− V π1 (s1)∣∣∣ (123) ≤ ∣∣∣Etraj∼(π,P )W1 ({(P̃K,h − Ph)V̂ π,P̃Kh+1 (sh, ah, R)})∣∣∣ (124)
= ∣∣∣∣∣Etraj∼(π,P )W1 ({ (θ̃K,h − θh) ∑ s′ φ(sh, ah, s ′)V̂ π,P̃Kh+1 (s ′, R) })∣∣∣∣∣ (125) ≤Etraj∼(π,P )W1 ∥∥∥θ̃K,h − θh∥∥∥Λ1,K,h ∥∥∥∥∥∑ s′ φ(sh, ah, s ′)V̂ π,P̃Kh+1 (s ′, R) ∥∥∥∥∥ (Λ1,K,h) −1   (126)
≤Etraj∼(π,P )W1 2β̂ ∥∥∥∥∥∑ s′ φ(sh, ah, s ′)V̂ π,P̃Kh+1 (s ′, R) ∥∥∥∥∥ (Λ1,K,h) −1   (127)
=2Etraj∼(π,P )W1 ({ uπ,P̃K1,K,h(sh, ah, R) }) (128)
=2Ṽ π,P̃KK,1 (s1, R) (129)
≤2V π,P̃KK,1 (s1, R) (130)
≤2V πK ,P̃KK,1 (s1, RK). (131)
The second inequality is due to lemma 10 . The third inequality is due to Lemma 13. The last inequality is due to the optimality of πK and RK . Plugging this inequality back to Inq. 119, we can prove the lemma.
Proof. (Proof of Theorem 2) Combining Lemma 12, Lemma 14 and Lemma 15, we know that K = C2(d 2H3+dH4) log(dH/(δ )) log2(KHB/(δ )) 2 + C2(d 2.5H2+d2H3) log(dH/(δ )) log2(KHB/(δ ))
for some constant C2 suffices to guarantee that V ∗1 (s1, R)− V π̂1 (s1, R) ≤ + opt.

heading: B.5 PROOF OF LEMMA 12
text: Proof. From the standard value decomposition technique (Azar et al., 2017; Jin et al., 2018), we can decompose V πk,P̃kk,1 (s1, Rk) into following terms:
V πk,P̃kk,1 (s1, Rk) (132)
≤ H∑ h=1 min { H,uπk,P̃k1,k,h (sk,h, ak,h, Rk) } + H∑ h=1 min { H,uπk,P̃k2,k,h (sk,h, ak,h, Rk) } (133)
+ H∑ h=1 min { H, (P̃k,h − Ph)V πk,P̃kk,h+1 (sk,h, ak,h, Rk) } (134)
+ H∑ h=1 ( PhV πk,P̃k k,h+1 (sk,h, ak,h, Rk)− V πk,P̃k k,h+1 (sk,h+1, ak,h+1, Rk) ) (135)
≤ H∑ h=1 min { H,uπk,P̃k1,k,h (sk,h, ak,h, Rk) } (136)
+ 2 H∑ h=1 min { H,uπk,P̃k2,k,h (sk,h, ak,h, Rk) } (137)
+ H∑ h=1 ( PhV πk,P̃k k,1 (sk,h, ak,h, Rk)− V πk,P̃k k,1 (sk,h+1, ak,h+1, Rk) ) , (138)
where the second inequality is derived from
(P̃k,h − Ph)V πk,P̃kk,h+1 (sk,h, ak,h, Rk) =(θ̃k,h − θh) ∑ s′ φ(sk,h, ak,h, s ′)V πk,P̃kk,h+1 (s ′, Rk) (139)
≤ ∥∥∥θ̃k,h − θh∥∥∥
Λ2,k,h ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ Λ−12,k,h
(140)
≤uπk,P̃k2,k,h (sk,h, ak,h, Rk). (141)
Eqn 138 is a martingale difference sequence. By Lemma 11, the summation over all k ∈ [K] is at most √ 2H3K log(4/δ). We mainly focus on Eqn 136 and Eqn 137.
Upper bound of Eqn 136 Firstly, we bound the summation of min { H,uπk,P̃k1,k,h (sk,h, ak,h, Rk) } .
K∑ k=1 H∑ h=1 min { H,uπk,P̃k1,k,h (sk,h, ak,h, Rk) } (142)
= K∑ k=1 H∑ h=1 min H, β̂ ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ Λ−11,k,h  (143) ≤
K∑ k=1 H∑ h=1 β̂σ̄1,k,h min 1, ∥∥∥∥∥∥ ∑ s′ φ(sk,h, ak,h, s ′) V̂ πk,P̃kk,h+1 (s ′, Rk) σ̄1,k,h ∥∥∥∥∥∥ Λ−11,k,h  (144)
≤β̂ √√√√√√∑ k,h σ̄21,k,h ∑ k,h min 1, ∥∥∥∥∥∥ ∑ s′ φ(sk,h, ak,h, s′) V̂ πk,P̃kk,h+1 (s ′, Rk) σ̄1,k,h ∥∥∥∥∥∥ 2
Λ−11,k,h  (145) ≤β̂ √∑
k,h
σ̄21,k,h · 2dH log(1 +KH/λ). (146)
The first inequality is due to Cauchy-Schwarz inequality. The second inequality is due to Lemma 21.
By the definition of σ̄21,k,h, we have
σ̄21,k,h ≤ H2/d+ E1,k,h + V̄1,k,h(sk,h, ak,h). (147) Now we bound ∑ k,hE1,k,h and ∑ k,h V̄1,k,h(sk,h, ak,h) respectively. For ∑ k,hE1,k,h, we have
∑ k,h E1,k,h ≤ H∑ h=1 K∑ k=1 min H2, 4Hβ̌ ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ (Λ1,k,h)−1  (148) +
H∑ h=1 K∑ k=1 min H2, 2β̃ ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2∥∥∥∥∥
(Λ3,k,h)−1  . (149)
For the first part,
H∑ h=1 K∑ k=1 min H2, 4Hβ̌ ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk) ∥∥∥∥∥ (Λ1,k,h)−1  (150) ≤4H
H∑ h=1 K∑ k=1 β̌σ̄1,k,h min 1, ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′)V̂ πk,P̃kk,h+1 (s ′, Rk)/σ̄1,k,h ∥∥∥∥∥ (Λ1,k,h)−1  (151)
≤4Hβ̌ √√√√√√∑ k,h σ̄21,k,h ∑ k,h min 1, ∥∥∥∥∥∥ ∑ s′ φ(sk,h, ak,h, s′) V̂ πk,P̃kk,h+1 (s ′, Rk) σ̄1,k,h ∥∥∥∥∥∥ 2
Λ−11,k,h  (152) ≤4Hβ̌
√∑ k,h σ̄21,k,h · 2dH log(1 +KH/λ), (153)
where the last inequality is due to Lemma 21. Similarly, for the second part, by Lemma 21, we have
H∑ h=1 K∑ k=1 min H2, 2β̃ ∥∥∥∥∥∑ s′ φ(sk,h, ak,h, s ′) ( V̂ πk,P̃kk,h+1 (s ′, Rk) )2∥∥∥∥∥
(Λ3,k,h)−1  (154) ≤2β̃ √ 2dH2K log(1 +KH/λ). (155)
Note that V̄1,k,h(sk,h, ak,h) is the empirical variance of V̂ πk,P̃kk,h+1 (s′, Rk) with transition P̃k(s
′|sk,h, ak,h). We bound the summation of V̄1,k,h(sk,h, ak,h) by the law of total variance (Lattimore & Hutter, 2012; Azar et al., 2013).
Recall that we define
Ỹk,H+1(s) = 0, (156)
Ỹk,h(s) = V̄1,k,h(s, πk(s)) + P̃k,hỸk,h+1(s, πk(s)). (157)
By the law of total variance, we have Ỹk,1(s) ≤ H2 holds for any s, a and k. We now bound the difference between ∑K k=1 Ỹk,1(sk,1) and ∑K k=1 ∑H h=1 V̄1,k,h(sk,h, ak,h).
Ỹk,1(sk,1)− H∑ h=1 V̄1,k,h(sk,h, ak,h) (158)
=P̃k,1Ỹk,2(sk,1, ak,1)− H∑ h=2 V̄1,k,h(sk,h, ak,h) (159)
≤min { H2, ( P̃k,1 − P1 ) Ỹk,2(sk,1, ak,1) } + ( P1Ỹk,2(sk,1, ak,1)− Ỹk,2(sk,2) ) (160)
+ ( Ỹk,2(sk,2)−
H∑ h=2 V̄1,k,h(sk,h, ak,h)
) . (161)
Therefore, we have K∑ k=1 Ỹk,1(sk,1)− K∑ k=1 H∑ h=1 V̄1,k,h(sk,h, ak,h) (162)
≤ K∑ k=1 H∑ h=1 min { H2, ( P̃k,h − Ph ) Ỹk,h+1(sk,h, ak,h) } (163)
+ H∑ h=1 K∑ k=1 ( PhỸk,h+1(sk,h, ak,h)− Ỹk,h+1(sk,h+1) ) . (164)
For Eqn 164, this term can be regarded as a martingale difference sequence, thus can be bounded by H2 √ KH by Lemma 11. For Eqn 163, we can bound this term in the following way:
K∑ k=1 H∑ h=1 min { H2, ( P̃k,h − Ph ) Ỹk,h+1(sk,h, ak,h) } (165)
≤ K∑ k=1 H∑ h=1 min { H2, ∥∥∥θ̃k,h − θh∥∥∥ Λ5,k,h ‖φỸ (sk,h, ak,h, R)‖Λ−15,k,h } (166)
≤β̃ K∑ k=1 H∑ h=1 min { 1, ‖φỸ (sk,h, ak,h, R)‖Λ−15,k,h }
(167)
≤β̃ H∑ h=1 √√√√K K∑ k=1 min { 1, ‖φỸ (sk,h, ak,h, R)‖ 2 Λ−15,k,h } (168)
≤β̃ √ 2dH2K log(1 +KH/λ). (169)
The second inequality is due to Lemma 10. The third inequality is due to Cauchy-Schwarz inequality. The last inequality is due to Lemma 21.
From the above analysis, we have∑ k,h V̄1,k,h(sk,h, ak,h) ≤ H2K + β̃ √ 2dH2K log(1 +KH/λ). (170)
Therefore, the summation of σ̄21,k,h can be bounded as∑ k,h σ̄21,k,h (171)
≤H3K/d+ ∑ k,h E1,k,h + ∑ k,h V̄1,k,h (172)
≤H3K/d+ 4Hβ̌ √∑
k,h
σ̄21,k,h · 2dH log(1 +KH/λ) + 3β̃ √ dH2K log(1 +KH/λ) +H2K.
(173) We define τ = log(32K2H/δ) log2(1 +KH4B2). Solving for ∑ k,h σ̄
2 1,k,h, we have∑
k,h
σ̄21,k,h ≤c1 ( H3K/d+H2K +H3d3τ + √ H2d3Kτ ) (174)
≤2c1 ( H3K/d+H2K +H3d3τ ) , (175)
where c1 denote a certain constant. The last inequality is due to 2 √ ab ≤ a+b for a, b ≥ 0. Plugging the above inequality back to Inq 142, we have for a constant c2, K∑ k=1 H∑ h=1 uπk,P̃k1,k,h (sk,h, ak,h, Rk) ≤ c2 √ (dH4 + d2H3)Kτ + c2d 2.5H2τ. (176)
Upper bound of Eqn 137 Now we focus on the summation of uπk,P̃k2,k,h (sk,h, ak,h, Rk). The proof follows almost the same arguments as the summation of uπk,P̃k1,k,h (sk,h, ak,h, Rk), though the upper
bound of ∑K k=1 ∑H h=1 V̄2,k,h(sk,h, ak,h) is derived in a different way. Following the same proof idea, we can show that K∑ k=1 H∑ h=1 uπk,P̃k2,k,h (sk,h, ak,h, Rk) ≤ β̂ √∑ k,h σ̄22,k,h2dH log(1 +KH/λ), (177)
σ̄22,k,h ≤ H2/d+ E2,k,h + V̄2,k,h(sk,h, ak,h), (178) K∑ k=1 H∑ h=1 E2,k,h ≤ 4Hβ̌ √∑ k,h σ̄22,k,h · 2dH log(1 +KH/λ) + 2β̃ √ 2dH2K log(1 +KH/λ).
(179)
Combining Inq 178 and Inq 179 and solving for ∑ k,h σ̄
2 2,k,h, we have∑
k,h
σ̄22,k,h (180)
≤2H3K/d+ 16dH3β̌2 log(1 +KH/λ) + 4β̃ √ dH2K log(1 +KH/λ) + 2 ∑ k,h V̄2,k,h(sk,h, ak,h)
(181) ≤6H3K/d+ 16dH3β̌2 log(1 +KH/λ) + 4β̃2d2/H log(1 +KH/λ) + 2 ∑ k,h V̄2,k,h(sk,h, ak,h) (182) Plugging this inequality back to Inq 177, we have
K∑ k=1 H∑ h=1 uπk,P̃k2,k,h (sk,h, ak,h, Rk) (183)
≤β̂ √ 6H4K log(1 +KH/λ) + 4β̌β̂dH2 log(1 +KH/λ) + 2β̃β̂d3/2 log(1 +KH/λ) (184)
+ √ 32d2H log(32K2H/δ) log2(1 +KH4B2) ∑ k,h V̄2,k,h(sk,h, ak,h). (185)
We now bound the last term in the above equation. Define τ = log(32K2H/δ) log2(1 +KH4B2), we have: √
32d2Hτ ∑ k,h V̄2,k,h(sk,h, ak,h) (186)
≤ √√√√32d2Hτ K∑ k=1 H∑ h=1 Es′∼P̃k,h(·|sk,h,ak,h) [ (V πk,P̃kk,h+1 ) 2(s′, Rk) ]
(187)
≤ √√√√32d2H2τ K∑ k=1 H∑ h=1 Es′∼P̃k,h(·|sk,h,ak,h) [ V πk,P̃kk,h+1 (s ′, Rk) ]
(188)
≤4 √ 2c̃3d 2H3τ + 1
c̃3H K∑ k=1 H∑ h=1 Es′∼P̃k,h(·|sk,h,ak,h) [ V πk,P̃kk,h+1 (s ′, Rk) ]
(189)
=4 √ 2c̃3d 2H3τ + 1
c̃3H K∑ k=1 H∑ h=1 V πk,P̃kk,h+1 (sk,h+1, Rk) (190)
+ 1
c̃3H K∑ k=1 H∑ h=1 ( PhV πk,P̃k k,h+1 (sk,h, ak,h, Rk)− V πk,P̃k k,h+1 (sk,h+1, Rk) ) (191)
+ 1
c̃3H K∑ k=1 H∑ h=1 ( P̃k,h − Ph ) V πk,P̃kk,h+1 (sk,h, ak,h, Rk), (192)
where c̃3 ≥ 1 is a constant to be defined later. The last inequality is due to 2 √ ab ≤ a + b for any a, b ≥ 0. Note that by Lemma 10, Eqn 192 is upper bounded by 1 c3H ∑ k,h u πk,P̃k 2,k,h (sk,h, ak,h, Rk). Eqn 191 is a martingale difference sequence, and can be bounded
by √ HK log(1/δ) with high probability. Plugging the above inequality back to Inq 183 and solving
for ∑K k=1 ∑H h=1 u πk,P̃k 2,k,h (sk,h, ak,h, Rk), we have
K∑ k=1 H∑ h=1 uπk,P̃k2,k,h (sk,h, ak,h, Rk) (193)
≤c3 ( √ dH4Kτ + (d2.5H2 + d2H3)τ + 1
c̃3H K∑ k=1 H∑ h=1 V πk,P̃kk,h+1 (sk,h+1, Rk)
) , (194)
where c3 ≥ 1 is a constant here. We set c̃3 = c3, then we have
K∑ k=1 H∑ h=1 uπk,P̃k2,k,h (sk,h, ak,h, Rk) (195)
≤c3 (√ dH4Kτ + (d2.5H2 + d2H3)τ ) + 1
H K∑ k=1 H∑ h=1 V πk,P̃kk,h+1 (sk,h+1, Rk), (196)
Bounding the summation of V πk,P̃kk,1 (s1, Rk) In the above analysis, we derive the upper bound of Eqn 136 and Eqn 137 (Inq 176 and Inq 195). Now we can finally bound ∑K k=1 V πk,P̃k k,1 (sk,1, Rk). For a constant c4, K∑ k=1 V πk,P̃kk,1 (sk,1, Rk) (197)
≤c4 √ (dH4 + d2H3)Kτ) + c4(d 2.5H2 + d2H3)τ + 1
H K∑ k=1 H∑ h=1 V πk,P̃kk,h+1 (sk,h+1, Rk). (198)
Following the same analysis, the above inequality actually also holds for any step h ∈ [H]: K∑ k=1 V πk,P̃kk,h (sk,h, Rk) (199)
≤c4 √ (dH4 + d2H3)Kτ) + c4(d 2.5H2 + d2H3)τ + 1
H K∑ k=1 H∑ h1=h V πk,P̃kk,h1+1(sk,h1+1, Rk). (200)
Define G = c4 (√ (dH4 + d2H3)Kτ) + c4(d 2.5H2 + d2H3)τ ) and ah =∑K
k=1 V πk,P̃k k,h (sk,h, Rk). The above inequality can be simplified into the following form:
ah ≤ G+ 1
H H∑ h1=h ah1 , (201)
and we have aH+1 = 0. From the elementary calculation, we can prove that a1 ≤ (1+ 1H ) HG ≤ eG. Therefore, we have K∑ k=1 V πk,P̃kk,1 (sk,1, Rk) (202)
≤c5 (√ (dH4 + d2H3)K log(KH/δ) log2(KH4B2)) + (d2.5H2 + d2H3) log(KH/δ) log2(KH4B2) ) .
(203)
for a constant c5 = ec4.

heading: C THE LOWER BOUND FOR REWARD-FREE EXPLORATION
text: In this section, we prove that even in the setting of non-plug-in reward-free exploration, the sample complexity in the exploration phase to obtain an -optimal policy is at least Ω(d 2H3
2 ). We say an algorithm can ( , δ)-learns the linear mixture MDPM if this algorithm returns an -optimal policy forM with probability at least 1 − δ in the planning phase after receiving samples for K episodes in the exploration phase. Our theorem is stated as follows.
Theorem 3. Suppose ≤ min ( C1 √ H,C2 √ dH4B ) , B > 1, H > 4, d > 3 for certain positive
constants C1 and C2. Then for any algorithm ALG there exists an linear mixture MDP instance M = (S,A, P,R,H, ν) such that if ALG ( , δ)-learns the problemM, ALG needs to collect at least K = Cd2H3/ 2 episodes during the exploration phase, where C is an absolute constant, and 0 < δ < 1 is a positive constant that has no dependence on ,H, d,K.
Compared with the lower bound proposed by Zhang et al. (2021a), we further improve their result by a factor of d. This lower bound also indicates that our sample complexity bound in Theorem 2 is statistically optimal. We prove Theorem 3 in Appendix C.1.

heading: C.1 PROOF OF THEOREM 3
text: Our basic idea to prove Theorem 3 is to connect the sample complexity lower bound with the regret lower bound of another constructed learning algorithm in the standard online exploration setting.
To start with, we notice that the reward-free exploration is strictly harder than the standard nonreward-free online exploration setting (e.g. Jin et al. (2018); Zhou et al. (2021; 2020a)), where the reward function is deterministic and known during the exploration. Therefore, if we can prove the sample complexity lower bound in the standard online exploration setting, this bound can also be applied to the reward-free setting.
For readers who are not familiar with the formulation of online exploration setting studied in this section, we firstly introduce the preliminaries. Compared with the reward-free exploration, the only
difference is that the rewardR(s, a) is fixed and known to the agent. In each episode, 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 the deterministic 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. The agent’s goal is to find a -optimal policy π after K episodes. We say a policy π is -optimal if
E [ H∑ h=1 Rh (sh, ah) | π ] ≥ E [ H∑ h=1 Rh (sh, ah) | π∗ ] − ,
where π∗ is the optimal policy for the MDP (S,A, P,R,H, ν). Theorem 4. Suppose ≤ min ( C1 √ H,C2 √ dH4B ) , B > 1, H > 4, d > 3. Then for any
algorithm ALG1 solving the non-reward-free online exploration problem, there exists an linear mixture MDPM = (S,A, P,R,H, ν) such that ALG1 needs to collect at least K = Cd2H3/ 2 episodes to output an -optimal policy for the linear mixture MDPM with probability at least 1−δ, where C is an absolute constant, and 0 < δ < 1 is a positive constant that has no dependence on ,H, d,K.
From the above discussion, Theorem 3 can be directly proved by reduction if Theorem 4 is true.
Proof. (Proof of Theorem 3) The theorem can be proved by contradiction. Suppose there is an algorithmALG which can ( , δ)-learns any linear mixture MDP instanceMwith onlyK ′ ≤ Cd2H3/ 2 episodes. Then we can use this algorithm to solve the online exploration problem by simply ignoring the information about the reward function and directly calling the exploration algorithm ofALG during the exploration phase. Then in the planning phase, we use the planning algorithm of ALG to output a policy based on the reward function as well as samples collected in the exploration phase. Therefore, this indicates that ALG can output an -optimal policy for the non-reward-free online exploration problem with probability at least 1 − δ after only K ′ ≤ Cd2H3/ 2 episodes. This contradicts the sample complexity lower bound in Theorem 4.
Now we discuss on how to prove Theorem 4.
Proof. (Proof of Theorem 4) Set K1 = cK for a positive constant c ≥ 2. We construct another algorithm ALG2 for any possible ALG1 in the following way: In ALG2, the agent firstly runs ALG1 for K1/c = K episodes. After K episodes, suppose ALG1 outputs a policy π̂ according to certain policy distribution νπ . ALG2 executes the policy π̂ in the following c−1c K1 episodes. The interaction ends after K1 episodes. ALG2 can be regarded as an algorithm which firstly runs the online exploration algorithm ALG1 for K episodes, and then evaluates the performance of the policy π̂ in the following episodes. We study the total regret ofALG2 from episodeK+1 to episode K1.
Recently, Zhou et al. (2020a) proposed the following regret lower bound for linear mixture MDPs. In order to avoid confusion, we use the notation K2 instead of K to denote the number of total episodes of the regret minimization problem.
Lemma 16. (Theorem 5.6 in Zhou et al. (2020a)) Let B > 1 and suppose K2 ≥ max { (d− 1)2H/2, (d− 1)/(32H(B − 1)) } , d ≥ 4, H ≥ 3. Then for any algorithm there exists an episodic, B-bounded linear mixture MDP parameterized by Θ = (θ1, · · · , θH) such that the expected regret is at least Ω(dH √ HK2).
To prove the above regret lower bound, Zhou et al. (2020a) construct a class of hard instances which is an extension of the hard instance class for linear bandits (cf. Theorem 24.1 in Lattimore & Szepesvári (2020)). They show that for any algorithm, there exists a hard instance in the instance class such that the regret is at least Ω(dH √ HK2). By Theorem 16 and setting K2 = K1, we know that for any possible algorithm ALG2, there exists a hard instance M such that ∑K1 k=1 E (V ∗(s1)− V πk(s1)) ≥ c2dH √ HK1 for a positive constant c2, where πk denotes the policy used in episode k for algorithm ALG2. The expectation is over all randomness of the algorithm and the environment.
Note that in the hard instance constructed in Zhou et al. (2020a), the per-step regret is at most E (V ∗(s1)− V πk(s1)) ≤ dH4√2 √ H K1 for any episode k ≤ K1. By choosing c = max { 1 2 √ 2c2 , 2 }
, we know that for the instanceM,
K1∑ k=K1/c+1 E (V ∗(s1)− V πk(s1)) ≥ ( c2 − 1 4 √ 2c ) dH √ HK1 ≥ c2 2 dH √ HK1. (204)
By the definition of ALG2, we have πk = π̂ if k > K1/c. Therefore, we have
(c− 1)KEπ̂∼µπ,s1∼µ ( V ∗(s1)− V π̂(s1) ) ≥ c2 2 dH √ HK1 = c2c 2 dH √ HK. (205)
Dividing both sides by (c− 1)K, we have
Eπ̂∼µπ,s1∼µ ( V ∗(s1)− V π̂(s1) ) ≥ c2c 2(c− 1) dH √ H/K. (206)
The above inequality indicates that, for any algorithm ALG1, there exists an instanceM such that the expected sub-optimality gap of the policy returned by ALG1 is at least c2c2(c−1)dH √ H/K after collecting samples for K episodes.
SupposeALG1 returns an -optimal policy π̂ with probability at least 1− δ. Recall that the per-step sub-optimality gap is at most E (V ∗(s1)− V πk(s1)) ≤ dH4√2 √ H cK in the constructed hard instance. We have
(1− δ) · + δ · dH 4 √ 2
√ H
cK ≥ Eπ̂∼µπ,s1∼µ
( V ∗(s1)− V π̂(s1) ) ≥ c2c 2(c− 1) dH √ H/K. (207)
We set δ to be a constant satisfying 0 < δ < min { 1, 2 √
2c1.5c2 c−1
} . Solving the above inequality, we
have K ≥ Cd 2H3 2 for a positive constant C.
D IMPROVED BOUND FOR REWARD-FREE EXPLORATION IN LINEAR MDPS
In this section, we explain how our choice of the exploration-driven reward can be used to improve the sample complexity bound for reward-free exploration in linear MDPs (Wang et al., 2020). We study the same reward-free setting as that in Wang et al. (2020), which is briefly explained in Appendix D.1. We describe our algorithm and bound in Appendix D.2, and prove our theorem in Appendix D.3.

heading: D.1 PRELIMINARIES
text: For the completeness of explanation, we briefly restate the reward-free setting studied in Wang et al. (2020). Compared with the setting in this work, the main differences are twofold: Firstly, they study the linear MDPs setting instead of linear mixture MDPs in this work. Secondly, they study the standard reward-free exploration setting without the constraints of the plug-in solver.
The linear MDP assumption, which was first introduced in Yang & Wang (2019); Jin et al. (2020b), states that the model of the MDP can be represented by linear functions of given features.
Assumption 1. (Linear MDP) an MDPM = (S,A, P,R,H, ν) is said to be a linear MDP if the following hold:
• There are d unknown signed measures µh = (µ (1) h , µ (2) h , · · · , µ (d) h ) such that for any
(s, a, s′) ∈ S ×A× S , Ph (s′ | s, a) = 〈µh (s′) , φ(s, a)〉.
• There exists H unknown vectors η1, η2, · · · , ηH ∈ Rd such that for any (s, a) ∈ S × A, Rh(s, a) = 〈φ(s, a), ηh〉.
We assume for all (s, a) ∈ S ×A and h ∈ [H], ‖φ(s, a)‖ ≤ 1, ‖µh(s)‖2 ≤ √ d and ‖η‖2 ≤ √ d.
For the reward-free exploration studied in Wang et al. (2020), the agent can collects a dataset of visited state-action pairsD = {( skh, a k h )} (k,h)∈[K]×[H] which will be used in the planning phase. Then during the planning phase, the agent can follow a certain designed learning algorithm to calculate an -optimal policy w.r.t. the reward function R using the dataset D.

heading: D.2 ALGORITHM AND THEOREM
text: Our algorithm can be divided into two parts. The exploration phase of the algorithm is presented in Algorithm 3, and the planning phase is presented in Algorithm 4.
Algorithm 3 Reward-free Exploration for Linear MDPs: Exploration Phase Input: Failure probability δ > 0 and target accuracy > 0 β ← cβ · dH √ log (dHδ−1ε−1) for some cβ > 0
K ← cK · d3H4 log ( dHδ−1ε−1 ) /ε2 for some cK > 0
for episode k = 1, 2, · · · ,K do 5: Qk,H+1(·, ·)← 0, Vk,H+1(·)← 0
for step h = H,H − 1, · · · , 1 do Λk,h ← ∑k−1 t=1 φ(st,h, at,h)φ(st,h, at,h) > + I
uk,h(·, ·)← β √ φ(·, ·)> (Λk,h)−1 φ(·, ·)
Define the exploration-driven reward function Rk,h(·, ·) = uk,h(·, ·) 10: ŵk,h ← (Λk,h)−1 ∑k−1 t=1 φ(st,h, at,h)Vt,h+1(st,h+1)
Qk,h(·, ·)← min { ŵ>k,hφ(·, ·) +Rk,h(·, ·) + uk,h(·, ·), H } Vk,h(s)← maxa∈AQk,h(s, a), πk,h(s) = arg maxa∈AQk,h(s, a)
end for for step h = 1, 2, · · · , H do
15: Take action ak,h = πk,h(sk,h) and observe sk,h+1 ∼ Ph(sk,h, ak,h) end for
end for Output: D ← {(sk,h, ak,h)}(k,h)∈[K]×[H].
Algorithm 4 Reward-free Exploration for Linear MDPs: Planning Phase Input: Dataset D ← {(sk,h, ak,h)}(k,h)∈[K]×[H], reward functions R = {Rh}h∈[H] Qk,H+1(·, ·) = 0, Vk,H+1(·) = 0 for step h = H,H − 1, · · · , 1 do
Λh ← ∑K t=1 φ(st,h, at,h)φ(st,h, at,h) > + I
5: uh(·, ·)← min { β √ φ(·, ·)> (Λk,h)−1 φ(·, ·), H } ŵh ← (Λh)−1 ∑K t=1 φ(st,h, at,h)Vt,h+1(st,h+1)
Qh(·, ·)← min { ŵ>h φ(·, ·) +Rh(·, ·) + uh(·, ·), H } Vh(s)← maxa∈AQh(s, a), πh(s) = arg maxa∈AQh(s, a)
end for 10: Output: π = {πh}h∈[H]
Compared with the algorithm in Wang et al. (2020), the main difference is that we set Rk,h = uk,h in the exploration phase, instead of Rk,h = uk,h/H . The following theorem states the complexity bound of our algorithms.
Theorem 5. With probability at least 1− δ, after collecting O(d3H4 log(dHδ−1 −1)/ 2) trajectories during the exploration phase, our algorithm outputs an -optimal policy for any given reward during the planning phase.

heading: D.3 PROOF OF THEOREM 5
text: The proof of Theorem 5 follows the proof framework in Wang et al. (2020) with a slight modification. Therefore, we only sketch the proof and mainly focus on explaining the differences. Firstly, we introduce the value function Ṽ ∗h (s,R), which is recursively defined from step H + 1 to step 1:
Ṽ ∗H+1(s,R) = 0,∀s ∈ S (208)
Ṽ ∗h (s,R) = max a∈A
{ min { Rh(s, a) + PhṼ ∗ h+1(s, a,R), H }} ,∀s ∈ S, h ∈ [H] (209)
Compared with the definition of V ∗h (s,R), the main difference is that we take minimization over the value and H at each step. We can similarly define Q̃∗h(s, a,R), Ṽ π h (s,R) and Q̃ π h(s, a,R). Lemma 17. With probability 1− δ/2, for all k ∈ [K],
Ṽ ∗1 (s1,k, Rk) ≤ Vk,1(s1,k) (210) and
K∑ k=1 Vk,1(sk,1) ≤ c √ d3H4K log(dKH/δ) (211)
for some constant c > 0 where Vk,1 is as defined in Algorithm 3.
This lemma corresponds to Lemma 3.1 in Wang et al. (2020). The main difference is that we replace V ∗1 with Ṽ ∗ 1 . Note that Ṽ ∗ h (s,R) ≤ H by the definition of Ṽ ∗h (s,R). Lemma 17 can be similarly proved following the proof of Lemma 3.1 in Wang et al. (2020) and replacing V ∗h by Ṽ ∗ h during the proof. Lemma 18. With probability 1− δ/4, for the function uh(·, ·) defined in Line 5 in Algorithm 4, we have
Es∼µ [ Ṽ ∗1 (s, uh) ] ≤ c′ √ d3H4 · log(dKH/δ)/K (212)
Compared with Lemma 3.2 in Wang et al. (2020), we replace the term V ∗1 (s, uh/H) with Ṽ ∗1 (s, uh). Note that with our choice of exploration-driven reward in the exploration phase (i.e. Rk,h = uk,h), we replace the term uh/H with uh in the expectation. This lemma can be proved by following the proof of Lemma 3.2 in Wang et al. (2020). Lemma 19. With probability 1 − δ/2, for any reward function satisfying Assumption 1 and all h ∈ [H], we have
Q∗h(·, ·, r) ≤ Qh(·, ·) ≤ Rh(·, ·) + ∑ s′ Ph (s ′ | ·, ·)Vh+1 (s′) + 2uh(·, ·). (213)
Since our algorithm in the planning phase is exactly the same with that of Wang et al. (2020), this lemma shares the same idea of Lemma 3.3 in Wang et al. (2020).
Now we can prove Theorem 5 with the help of Lemma 18 and Lemma 19.
Proof. (Proof of Theorem 5) We condition on the events defined in Lemma 18 and Lemma 19, which hold with probability at least 1− δ. By Lemma 19, we have for any s ∈ S,
V1(s) = max s Q1(s, a) ≥ max a Q∗1(s, a,R) = V ∗ 1 (s,R), (214)
which implies
Es1∼µ [V ∗1 (s1, R)− V π1 (s1, r)] ≤ Es1∼µ [V1 (s1)− V π1 (s1, R)] . (215) Note that 0 ≤ Vh(s) ≤ H and 0 ≤ V πh (s,R) ≤ H since 0 ≤ R(s, a) ≤ 1. Therefore, we always have Vh(s)− V πh (s,R) ≤ H . For any sh, we have
Vh(sh)− V πh (sh, R) (216) ≤min { H,PhVh+1(sh, πh(sh)) + 2uh(sh, πh(sh))− PhV πh+1(sh+1, R) } (217)
By recursively decomposing Vh(sh)− V πh (sh, R) from step H to 1, we have Es1∼µ [V1 (s1)− V π1 (s1, r)] ≤ Es∼µ [ Ṽ π1 (s, u) ] . (218)
By definition of Ṽ ∗1 (s, u), we have Es∼µ [ Ṽ π1 (s, u) ] ≤ 2Es∼µ [ Ṽ ∗1 (s, u) ] . By Lemma 18,
Es∼µ [ Ṽ ∗1 (s, u) ] ≤ c′ √ d3H4 · log(dKH/δ)/K. (219)
By taking K = cKd3H4 log(dHδ−1 −1)/ 2 for a sufficiently large constant cK > 0, we have Es1∼µ [V ∗1 (s1, R)− V π1 (s1, r)] ≤ c′ √ d3H4 · log(dKH/δ)/K ≤ . (220)

heading: E AUXILIARY LEMMAS
text: Lemma 20. (Self-Normalized Bound for Vector-Valued Martingales, Theorem 1 and 2 in AbbasiYadkori et al. (2011)) Let {Ft}∞t=0 be a filtration. Let {ηt}∞t=1 be a real-valued stochatic process such that ηt is Ft-measurable and ηt is conditionally R-sub-Gaussian for some R ≥ 0, i.e.
∀λ ∈ R,E[eληt |Ft−1] ≤ exp ( λ2R2
2
) . (221)
Let {Xt}∞t=1 be an Rd-valued stochastic process such that Xt is Ft−1-measurable. Assume that V is a d× d positive definite matrix. For any t ≥ 0, define
V̄t = V + t∑ s=1 XsX > s , St = t∑ s=1 ηsXs. (222)
Then, for any δ > 0, with probability at least 1− δ, for all t ≥ 0,
‖St‖2V̄ −1t ≤ 2R 2 log
( det ( V̄t )1/2 det(V )−1/2
δ
) . (223)
Further, let V = Iλ, λ > 0. Define Yt = 〈Xt, θ∗〉 + ηt and assume that ‖θ∗‖2 ≤ S, ‖Xt‖2 ≤ L,∀t ≥ 1. Then for any δ > 0,with probability at least 1− δ, for all t ≥ 0, θ∗ satisfies∥∥∥θ̂t − θ∗∥∥∥
V̄t ≤ R
√ d log ( 1 + tL2/λ
δ
) + λ1/2S, (224)
where θ̂t is the l2-regularized least-squares estimation of θ∗ with regularization parameter λ > 0 based on history samples till step t. Lemma 21. (Lemma 11 in Abbasi-Yadkori et al. (2011)) Let {Xt}∞t=1 be a sequence in Rd, V a d× d positive definite matrix and define V̄t = V + ∑t s=1XsX > s . Then, we have that
log
( det ( V̄n )
det(V )
) ≤
n∑ t=1 ‖Xt‖2V̄ −1t−1 . (225)
Further , if ‖Xt‖2 ≤ L for all t, then n∑ t=1 min { 1, ‖Xt‖2V̄ −1t−1 } ≤ 2 ( log det ( V̄n ) − log detV ) ≤ 2 ( d log (( trace(V ) + nL2 ) /d ) − log detV ) ,
(226)
and finally, if λmin(V ) ≥ max(1, L2), then n∑ t=1 ‖Xt‖2V̄ −1t−1 ≤ 2 log det ( V̄n ) det(V ) . (227)
Lemma 22. (Bernstein inequality for vector-valued martingales, Theorem 4.1 in Zhou et al. (2020a)) Let {Gt}∞t=1 be a filtration, {xt, ηt}t≥1 a stochastic process so that xt ∈ Rd is Gtmeasurable and ηt ∈ R is Gt+1-measurable. Fix R,L, σ, λ > 0, µ∗ ∈ Rd. For t ≥ 1, let yt = 〈µ∗, xt〉+ ηt and suppose that ηt, xt also satisfy
|ηt| ≤ R,E [ηt | Gt] = 0,E [ η2t | Gt ] ≤ σ2, ‖xt‖2 ≤ L. (228)
Then, for any 0 < δ < 1, with probability at least 1− δ, we have
∀t > 0, ∥∥∥∥∥ t∑ i=1 xiηi ∥∥∥∥∥ z−1t ≤ βt, ‖µt − µ∗‖zt ≤ βt + √ λ ‖µ∗‖2 , (229)
where for t ≥ 1, µt = Z−1t bt, Zt = λI + ∑t i=1 xix > i , bt = ∑t i=1 yixi and
βt = 8σ √ d log (1 + tL2/(dλ)) log (4t2/δ) + 4R log ( 4t2/δ ) . (230)

Review summaries:
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.

This paper studies reward-free exploration with linear function approximation (i.e., under the setting of linear mixture MDPs). Specifically, this work introduces two algorithms: the first one has a simple form with the sample complexity $\tilde{\mathcal{O}}(d^2H^4 / \varepsilon^2)$, and the second one is much complicated (i.e., with 5 estimators) but has an improved sample complexity $\tilde{\mathcal{O}}(d^2H^3 / \varepsilon^2)$. To improve the dependence on $H$, the main technical challenge is that the traditional Bernstein-type analysis could not be directly applied because the variance summation of $\{\tilde{V}\}_{h \in [H]}$ is not induced from the same policy and transition function.  This paper addresses this technical challenge and makes a solid contribution to reward-free exploration.

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.