Title: Reward-Free Model-Based Reinforcement Learning with Linear Function Approximation
Authors: Weitong ZHANG
Dongruo Zhou
Quanquan Gu
Venue: NeurIPS 2021
Keywords: Reinforcement Learning
Reward-Free Exploration
Model Based RL

Abstract: We study the model-based reward-free reinforcement learning with linear function approximation for episodic Markov decision processes (MDPs). In this setting, the agent works in two phases. In the exploration phase, the agent interacts with the environment and collects samples without the reward. In the planning phase, the agent is given a specific reward function and uses samples collected from the exploration phase to learn a good policy. We propose a new provably efficient algorithm, called UCRL-RFE under the Linear Mixture MDP assumption, where the transition probability kernel of the MDP can be parameterized by a linear function over certain feature mappings defined on the triplet of state, action, and next state. We show that to obtain an $\epsilon$-optimal policy for arbitrary reward function, UCRL-RFE needs to sample at most $\tilde O(H^5d^2\epsilon^{-2})$ episodes during the exploration phase. Here, $H$ is the length of the episode, $d$ is the dimension of the feature mapping. We also propose a variant of UCRL-RFE using Bernstein-type bonus and show that it needs to sample at most $\tilde O(H^4d(H + d)\epsilon^{-2})$ to achieve an $\epsilon$-optimal policy. By constructing a special class of linear Mixture MDPs, we also prove that for any reward-free algorithm, it needs to sample at least $\tilde \Omega(H^2d\epsilon^{-2})$ episodes to obtain an $\epsilon$-optimal policy. Our upper bound matches the lower bound in terms of the dependence on $\epsilon$ and the dependence on $d$ if $H \ge d$. 

TL;DR: A reward free approach for model based reinforcement learning with function approximation, lower bound provided

Review discussion:
The paper is incremental in terms of the originality and significance: (1) the importance of the study on the proposed model should have been more discussed, since a very similar model (linear MDP) has already been studied quite a bit as a benchmark for theoretical studies on function approximation, and (2) techniques (optimistic value-iteration, Bernstein-type bonus) are not particularly new. Pseudo value function seems not much more than some optimistic value function for linear mixture MDP, which is new only because the considered model is different from linear MDP which have slightly different representation for transition and reward models. Main results are also somewhat predictable given many previous works on linear function approximation [3, 7, 10, 12, 17, 26, 28].
Overall, I recommend rejection for this paper.
Comments
1. The claim for matching upper and lower bounds sounds misleading since it only holds for H \ge d (and in this case, dependence on H is sub-optimal anyway) and computationally inefficient algorithm.
2. Line 43-44 (They also proved ...): In which setting [17] proved exponential lower bound?
3. Line 46-50 (However, both works ...): I do not understand why the model considered in previous work is a very larger class than linear mixture MDP.
4. Eq (4.4): Why do you need regularizer on $\theta$? On which line in Algorithm 2 uses $\lambda$?
5. Theorem 4.2: What is B? Any value range for B?
See also the limitations.
No
Thank you for your detailed feedback, here's our response to your comments.
Q1: Pseudo value function seems not much more than some optimistic value function for linear mixture MDP
A1: We respectfully disagree. Pseudo value function is quite different from the optimistic value function.
Instead, the pseudo value function is crafting the maximum uncertainty level for each state-action pair $(s, a)$. It is not maintaining the optimistic estimation of the value function given the reward function.
Q2: Linear mixture MDP, which is new only because the considered model is different from linear MDP which has a slightly different representation for transition and reward models.
A2: We respectfully disagree that Linear Mixture MDP is slightly different from Linear MDP with different representations.
Indeed this is a fundamental difference in reinforcement learning with linear function approximation. In detail, Linear MDP is a model-free assumption by assuming the linear representation is built on the state-action $(s, a)$.
However, Linear Mixture MDP is a model-based assumption assuming the linear representation is built on the state-action-next state pair.
Thus Linear Mixture MDP is quite different from Linear MDP and cannot be claimed as ‘slightly different’ from Linear MDP
Q3: The claim for matching upper and lower bounds sounds misleading since it only holds for H \ge d (and in this case, dependence on H is sub-optimal anyway) and computationally inefficient algorithm.
A3: We acknowledge that there’s still a gap between the lower bound and the upper bound and we hope this work can inspire further work.
However, our result is still the first lower bound based on the same settings.
Furthermore, the proposed algorithm is computationally efficient when the state space is finite.
When the state space is infinite, we can use epsilon-net to cover the continuous state space, and reduce the state space from infinite to finite. Note that this is not the drawback of our algorithm, but the drawback of linear mixture MDPs (the feature mapping $\phi_f(s, a)$ is defined as the integration of the original feature mapping $\phi$ with respect to a certain function $f$).
Q4: Line 43-44 (They also proved ...): In which setting [17] proved exponential lower bound?
Q5: Line 46-50 (However, both works ...): I do not understand why the model considered in previous work is a very larger class than linear mixture MDP.
A4, 5: For the lower bound part, the previous work [17] is only assuming the optimal Q function $Q*(s, a)$ is a linear function of the given representation $\phi(s, a)$. It would be obvious that the previous work does not provide any assumption on the MDP transition kernel.
Therefore, the previous work [17] is considering a very larger class than linear mixture MDP.
Q6: Eq (4.4): Why do you need a regularizer on θ? On which line in Algorithm 2 uses λ?
A6: As in the previous bandit literature, we use the regularizer to make the covariance matrix non-singular. We use regularizer $\lambda$ on Line 2 in Algorithm 2 to initialize the `covariance matrix’ $U$
Q7: Theorem 4.2: What is B? Any value range for B?
A7: According to Definition 3.1, B is the norm bound for $\theta^*$, we don’t need any value range for $B$ since the regularizer $\lambda$ is set to $B^{-2}$ thus $B$ only appears in the logarithmic terms in our result.
Thanks to the authors for thoughtful response. As the authors explained and promised to emphasize technical challenge (particularly, comparison to [17]) in the revision, I increased my score.
However, I am still not convinced on significance of the result, because their matching upper bound holds only for H > d and computationally inefficient algorithm in continuous state spaces. Please see the limitations.
Thank you for your positive feedback! We acknowledge that there’s still some gap between the upper and lower bounds. Since our work provides the first polynomial lower bound in the linear function approximation setting, we hope our work can inspire some follow-up works to close this gap.
Regarding computational efficiency, we want to emphasize that this inefficiency is a common issue when applying model-based algorithms to continuous state space.
In fact, in real-world RL applications, the continuous state space often needs to be discretized, in order to apply the RL algorithms.
Thus, it would not be a big concern that our algorithms cannot handle continuous state space efficiently.
This papers develops a lower bound specifically for linear MDPs, which matches the upper bound in terms of $\epsilon$.
The major concern is that the MDP studied in this paper is time homogeneous, which may be restrictive in many practical cases. Is this homogeneity assumption fundamental to the proposed algorithm?
In other words, is it possible to extend the analysis to the case where transition probabilities are different from step to step?
There are a few typos, such as Page 1 "the algorithm is given a specific reward function and use the collected data" where "use" should be replaced by "uses".
No
Thank you for your constructive feedback. We will revise the typos you have pointed out.
Q1: Is it possible to extend the analysis to the case where transition probabilities are different from step to step?
A1: We believe it would be very easy to extend the analysis to the time-inhomogeneous settings. In detail, we can replace the estimation of $\theta$ with $\theta_1, \theta_2, \cdots, \theta_H$.
Then we need to do the value-targeted regression for $\theta_h$ using only the data collected from the $h$-th step. In this case, Lemma A.3 will be modified to have an $\tilde O(H^3\sqrt{K})$ bound and Lemma A.5 should be modified to have an $O(H^{2.5}\sqrt{K}) bound. We will also definitely add this in the final version.
Thank the authors for spending time responding to my reviews. The response partly addresses my major concern. I've updated the score accordingly.
Thank you for your positive feedback and for increasing the score!
The paper is well written and the proofs seem correct (see specific comments, I did not check the proofs relative to Algorithm 3). I think that in absolute terms it is interesting to study the optimal problem independent rate for reward-free exploration problem with linear kernel assumption.
But on the other hand, the implementation seems limited to the tabular case and no simple examples are provided.
Thus the proposed solution seems rather theoretical.
Main comments:
-Can you provide intuition on the different sources of "optimism" in your algorithm. Because they are: the exploration driven reward function, additional bonuses in PLAN, and pseudo value function which should all help to explore.
In particular why we cannot use the value V_h^k for the target and in the exploration-driven rewards instead of the max (this would save computations)?
-Can you elaborate on the implementation of UCRL-RFE when S is large or infinite?
In particular, what is the time complexity of the algorithm in this case?  E.g. do we actually need to do a complete pass over all state-action pairs in Planning Module, Algorithm 1?
-What do you think is the optimal dependence on the horizon H and how to obtain it? More generally would it be possible to match the lower bound?
-What are the main technical differences/challenges with respect to [17] in linear-kernel MDP?
Specific comments:
-L49: Can you detail why the bound is too conservative?
-53: maybe here you should explain a bit what is the difference between linear MDP and linear kernel MPD.
-L93, Table 1: For the results, you should be precise with the setting otherwise the table is wrong. The papers [9,12,14] consider episodic MDP with bounded rewards in [0,1] and step-dependent transitions. In this setting, the optimal rate is of order O(H^3S^2A/\epsilon^2) and O(H^2S^2A/\epsilon^2) if the transitions are step-independent. [27] consider an additional assumption where the sum of the rewards is equal to 1 (and step-independent transitions) where they obtain a rate of O(S^2A/\epsilon^2) which they prove to be optimal given this extra assumption. Note that in particular, their algorithm has no guarantees in the setting considered by  [9,12,14].
-L20: It is not clear if the expectation or the variance are well defined for any function V no?
-L126: Do you have any example of MDP that satisfies this assumption except the trivial example of tabular MDP?
In particular, an MDP where it is easy to compute the summation in L129.
-L138 Maybe you should describe how the agent interacts with the mdp during the exploration phase.
-L186: Intuitively why do you need to take the maximum over f.  Because the "exploration driven
reward function"  is not enough to explore?
-L206: can you elaborate on this point because the main focus of the paper is to deal with large state space?
-L264: Since \tiled{\theta} and \hat{\theta} estimate both the same parameter \theta^* why you need two separate estimators?
-L315:  Thus do you think that the lower bound is sub-optimal in this case?
-L445: For all s?
-L477: Conditioned on the events of lemma A.2 and A.3  do we still have that conditionally the expectation of the random variable in the sum is zero in order to apply Hoeffding inequality?
-L538: there are extra }
-L578: the first equality is inequality.
-L584: can you detail this part, what is E_{s_h}?
-L688: K should be in the sum and you should keep the sum in the middle inequality.
-L740: you also need to choose \alpha such that the probabilities are in (0,1).
-L745: bound on the cardinal is missing in the statement since taking for \cM a singleton is sufficient. You mean for all M\leq ... there exits ...
-L774: x and x' independent.
-L788: Ok you can in fact invoke the chain rule here and you do not have to consider a deterministic algorithm. P_\theta(Y_k) and E_{Y_k} are not defined.
-L811: could you define clearly \hat{\theta}
No
Thank you for your detailed and constructive comments. We will address all writing suggestions you pointed out and here’s our response to your comments.
Q1: Can you provide intuition on the different sources of "optimism" in your algorithm, why we cannot use the value V_h^k for the target and in the exploration-driven rewards instead of the max?
A1: Firstly speaking, the exploration driven reward function and pseudo value function can not be simply interpreted as ‘optimism’.
In fact, they are providing the maximum uncertainty level of the state-action pair $(s, a)$ given any possible value function. The additional bonus in PLAN is the ‘optimism’, which maintains an optimistic estimation of the value function given reward function $r$. We cannot use the value $V_h^k$ for the target and in the exploration-driven reward, because the value function $V_h^k$ is built on a given reward function $r_h$.
However, due to the nature of reward-free exploration, we cannot get a reward function during the exploration phase.
Q2: Can you elaborate on the implementation of UCRL-RFE when S is large or infinite? Do we actually need to do a complete pass over all state-action pairs in Planning Module, Algorithm 1?
A2: We do not need to do a complete pass over all state-action pairs in the Planning Module (Algorithm 1).
In particular, we can calculate the estimated Q function only to the state we encountered.
Thus the planning module can be implemented efficiently.
Q3: What do you think is the optimal dependence on the horizon H and how to obtain it?
A3: We think the optimal dependence on horizon $H$ should be quadratic, according to our lower bound and the result in tabular setting [27, 9]. We hope our work can inspire follow-up work to obtain this bound.
Q4: What are the main technical differences/challenges with respect to [17] in linear-kernel MDP?
A4: The main technical challenges of UCRL-RFE is the pseudo value function $u_h^k$, which is designed to craft the maximum uncertainty of the state-action pair $(s, a)$ caused by the MDP transition kernel. This pseudo value function for the exploration phase is never used in previous literature [17].
Secondly, comparing our work with the traditional analysis routine of linear-kernel MDP with Bernstein-type bonus [28], the main difference of our algorithm is we not only need to maintain $\hat \theta, \tilde \theta$ and $\check \theta$ using the target value function $V_{h+1}^k$, but also need to do the regression for $\theta$ using the pseudo value function $u_h^k$. This additional regression will indeed introduce more challenges compared to the previous literature.
Recently there has been great in understanding the sample complexity of reward-free exploration in reinforcement learning. In this paper, the authors study reward-free exploration in linear mixture models. The main methodology is to combine the idea of defining exploration-driven reward functions as the uncertainty bonus (such idea was proposed in [17]) and mechanisms for quantifying uncertainty in linear mixture models.
Moreover, the authors show how to improve the sample complexity by using a Bernstein-type uncertainty bonus. The authors also prove a lower bound by constructing a hard instance with three states and exponential number of actions.
My main concern is the lack the novelty. The algorithms presented in this paper follow the same methodology developed in [17], with modifications on the exploration-driven reward function so that it works for linear mixture models (see Remark 1). Is that true that the main algorithm in this paper is a direct combination of the methodology in [17] and uncertainty quantification techniques for linear mixture models? If not, the authors should consider adding a section highlighting the technical challenges they face when adapting the methodology in [17] to linear mixture models. Such a discussion will significantly improve the paper.
Minor comment
Table 1: the algorithm in [25] does not work in linear MDPs.
Rather, it requires linear completeness and certain explorability assumption.
No
We sincerely thank the reviewer for the comments and we will address the assumption needed by [25] in Table 1 in the revision.
Q1: “the main algorithm in this paper is a direct combination of the methodology in [17] and uncertainty quantification techniques for linear mixture models?”
A1: We respectfully disagree.
In fact, to guarantee the convergence of our Algorithms, we need to design the pseudo value function $u_h^k$ which has never appeared in previous works [17]. This unique exploration-driven value function characterizes the **maximum** uncertainty level of the state-action pair $(s, a)$ caused by the MDP transition kernel.
Moreover, in Algorithm 3, we construct the estimators $\hat \theta, \check \theta, \tilde \theta$ by solving the regression problem over the value function $V_{h+1}^k$, but the estimator $\theta$ over the pseudo value function $u_h^k$. This is different from the traditional uncertainty quantification technique used for linear mixture models [28, 29], which only depend on the value function. This new estimation $\theta$ indeed brings in additional challenges in the analysis. We will definitely add another subsection to highlight these difficulties in the camera-ready version.

Meta-review:
role: PC
data: comment: The paper provides the first analysis of reward free in linear mixture MDPs. I think this is a quite natural complement to the literature of RFE. While the paper is well-executed, I find it hard to identify the technical challenges and why certain technical tools are necessary. This concern was also raised by several reviewers. The authors should explain the challenges specific to this setting and why the paper is not a "simple" combination of techniques available in the literature. Due to the theoretical nature of the paper, I think these are quite important points. In addition, during the discussion, no reviewer was willing to champion the paper, which currently stands in a very borderline position. I decided to vote for rejection due to the lack of clarity on the algorithmic choices and novelty to a field that has seen numerous advances in the last couple of years. 

Decision: Accept (Poster)