On Reward-Free Reinforcement Learning with Linear Function Approximation
Ruosong Wang, Simon S. Du, Lin F. Yang, Ruslan Salakhutdinov
Introduction
In reinforcement learning (RL), an agent repeatedly interacts with an unknown environment to maximize the cumulative reward. To achieve this goal, RL algorithm must be equipped with exploration mechanisms to effectively solve tasks with long horizons and sparse reward signals. Empirically, there is a host of successes by combining deep RL methods with different exploration strategies. However, the theoretical understanding of exploration in RL by far is rather limited.
In this work we study the reward-free RL setting which was formalized in the recent work by Jin et al. (2020). There are two phases in the reward-free setting: the exploration phase and the planning phase. During the exploration phase, the agent collects trajectories from an unknown environment without any pre-specified reward function. Then, in the planning phase, a specific reward function is given to the agent, and the goal is to use samples collected during the exploration phase to output a near-optimal policy for the given reward function. From a practical point of view, this paradigm is particularly suitable for 1) the batch RL setting (Bertsekas and Tsitsiklis, 1996) where data collection and planning are explicitly separated and 2) the setting where there are multiple reward function of interest, e.g., constrained RL (Achiam et al., 2017; Altman, 1999; Miryoosefi et al., 2019; Tessler et al., 2018). From a theoretical point view, this setting separates the exploration problem and the planning problem which allows one to handle them in a theoretically principled way, in contrast to the standard RL setting where one needs to deal both problems simultaneously.
The sample complexity bound in (Jin et al., 2020), although being near-optimal in the tabular setting, can be unacceptably large in practice due to the polynomial dependency on the size of the state space. For environments with a large state space, function approximation schemes are needed for generalization. RL with linear function approximation is arguably the simplest yet most fundamental setting. Clearly, in order to understand more general function classes, e.g., deep neural networks, one must understand the class of linear functions first. In this paper, we study RL with linear function approximation in the reward-free setting, and our goal is to answer the following question:
Is it possible to design provably efficient RL algorithms with linear function approximation in the reward-free setting?
We obtain both polynomial upper bound and hardness result to the above question.
Our first contribution is a provably efficient algorithm for reward-free RL under the linear MDP assumption (Yang and Wang, 2019; Jin et al., 2019), which, roughly speaking, requires both the transition operators and the reward functions to be linear functions of a -dimensional feature extractor given to the agent. See Assumption 2.1 for the formal statement of the linear MDP assumption. Our algorithm, formally presented in Section 3, samples trajectories during the exploration phase, and outputs -optimal policies for an arbitrary number of reward functions satisfying Assumption 2.1 during the planning phase with high probability. Here is the feature dimension, is the planning horizon and is the required accuracy.
One may wonder whether is possible to further weaken the linear MDP assumption, since it requires the feature extractor to encode model information, and such feature extractor might be hard to construct in practice. Our second contribution is a hardness result for reward-free RL under the linear assumption, which only requires the optimal value function to be a linear function of the given feature extractor and thus weaker than the linear MDP assumption. Our hardness result, formally presented in Section 4, shows that under the linear assumption, any algorithm requires exponential number of samples during the exploration phase, so that the agent could output a near-optimal policy during the planning phase with high probability. The hardness result holds even when the MDP is deterministic.
Our results highlight the following conceptual insights.
Reward-free RL might require the feature to encode model information. Under model-based assumption (linear MDP assumption), there exists a polynomial sample complexity upper bound for reward-free RL, while under value-based assumption (linear assumption), there is an exponential sample complexity lower bound. Therefore, the linear assumption is strictly weaker than the linear MDP assumption in the reward-free setting.
Reward-free RL could be exponentially harder than standard RL. For deterministic systems, under the assumption that the optimal -function is linear, there exists a polynomial sample complexity upper bound (Wen and Van Roy, 2013) in the standard RL setting. However, our hardness result demonstrates that under the same assumption, any algorithm requires exponential number of samples in the reward-free setting.
Simulators could be exponentially more powerful. In the setting where the agent has sampling access to a generative model (a.k.a. simulator) of the MDP, the agent can query the next state sampled from the transition operator given any state-action pair as input. In the supplementary material, we show that for deterministic systems, under the linear assumption, there exists a polynomial sample complexity upper bound in the reward-free setting when the agent has sampling access to a generative model. Compared with the hardness result above, this upper bound demonstrates an exponential separation between the sample complexity of reward-free RL in the generative model and that in the standard RL model. To the best our knowledge, this is the first exponential separation between the standard RL model and the generative model for a natural question.
1 Related Work
This paper studies linear function approximation. Linear MDP is the setting where both the transition and the reward are linear functions of a given feature extractor. Recently, in the standard RL setting, many works (Yang and Wang, 2019; Jin et al., 2019; Cai et al., 2020; Zanette et al., 2019) have provided polynomial sample complexity guarantees for different algorithms in linear MDPs. Technically, our algorithm, which works in the reward-free setting, combines the algorithmic framework in (Jin et al., 2019) with a novel exploration-driven reward function (cf. Section 3). Linear is another setting where only the optimal -function is assumed to be linear, which is weaker than the assumptions in the linear MDP setting. In the standard RL setting, it is an open problem whether one can use polynomial number of samples to find a near-optimal policy in the linear setting (Du et al., 2020a). Existing upper bounds all require additional assumptions, such as (nearly) deterministic transition Wen and Van Roy (2013); Du et al. (2019b, 2020b).
Preliminaries
Throughout this paper, for a given positive integer , we use to denote the set .
When the initial distribution and the transition operators are all deterministic, we say is a deterministic system. In this case, we may regard each transition operator as a function that maps state-action pairs to a states. We note that deterministic systems are special cases of general MDPs.
A policy chooses an action based on the current state and the time step . Formally, where for each , maps a given state to an action. The policy induces a trajectory , where , , , , , , etc.
An important concept in RL is the -function. For a specific set of reward functions , given a policy , a level and a state-action pair , the -function is defined as
Similarly, the value function of a given state is defined as
For a specific set of reward functions , We use to denote an optimal policy with respect to , i.e., is a policy that maximizes
We also denote and . We say a policy is -optimal with respect to if
Throughout the paper, when is clear from the context, we may omit from , , , and .
2 Linear Function Approximation
The following linear MDP assumption, which was first introduced in (Yang and Wang, 2019; Jin et al., 2019), states that the model of the MDP can be predicted by linear functions of the given features.
An MDP is said to be a linear MDP if the followings hold:
there are unknown signed measures such that for any , ;
As in (Jin et al., 2019), we assume for all and , , , and .
The following linear assumption, which is a common assumption in the theoretical RL literature (see e.g. (Du et al., 2019b, 2020a)), states that the optimal -function can be predicted by linear functions of the given features.
We note that Assumption 2.2 is weaker than Assumption 2.1. Under Assumption 2.1, it can be shown that for any policy , is a linear function of the given feature extractor . In this paper, we show that Assumption 2.2 is strictly weaker than Assumption 2.1 in the reward-free setting, meaning that reward-free RL under Assumption 2.2 is exponentially harder than that under Assumption 2.1.
3 Reward-Free RL
In the reward-free setting, the goal is to design an algorithm that efficiently explore the state space without the guidance of reward information. Formally, there are two phases in the reward-free setting: exploration phase and planning phase.
During the exploration phase, the agent interacts with the environment for episodes. In the -th episode, the agent chooses a policy which induces a trajectory. The agent observes the states and actions as usual, but does not observe any reward values. After episodes, the agent collects a dataset of visited state-actions pairs which will be used in the planning phase.
Planning Phase.
During the planning phase, the agent is no longer allowed to interact with the MDP. Instead, the agent is given a set of reward functions where is the deterministic reward function in level , and the goal here is to output an -optimal policy with respect to using the collected dataset .
To measure the performance of an algorithm, we define the sample complexity to be the number of episodes required in the exploration phase to output an -optimal policy in the planning phase.
Reward-Free RL for Linear MDPs
In this section, we present our reward-free RL algorithm under the linear MDP assumption.
The exploration phase of the algorithm is presented in Algorithm 1, and the planning phase is presented in Algorithm 2.
During the exploration phase of the algorithm, we employ the least-square value iteration (LSVI) framework introduced in (Jin et al., 2019). In each episode, we first update the parameters that are used to calculate the -functions, and then execute the greedy policy with respect to the updated -function to collect samples. As in (Jin et al., 2019), to encourage exploration, Algorithm 1 adds an upper-confidence bound (UCB) bonus function .
The main difference between Algorithm 1 and the one in (Jin et al., 2019) is the definition of the exploration-driven reward function. Since the algorithm in (Jin et al., 2019) is designed for the standard RL setting, the agent can obtain reward values by simply interacting with the environment. On the other hand, in the exploration phase of the reward-free setting, the agent does not have any knowledge about the reward function. In our algorithm, in each episode, we design an exploration-driven reward function which is defined to be , where is the UCB bonus function defined in Line 8. Note that we divide by so that always lies in $u_{h}(\cdot,\cdot)$) is large. After sufficient number of episodes, the uncertainty of all state-action pairs should be low on average, since otherwise the agent would have visited those state-action pairs with large uncertainty as guided by the reward function.
Planning Phase.
After the exploration phase, the returned dataset contains sufficient amount of information for the planning phase. In the planning phase (Algorithm 2), for each step , we optimize a least squares predictor to predict the -function, and return the greedy policy with respect to the predicted -function. During the planning phase, we still add an UCB bonus function to guarantee optimism. However, as mentioned above and will be made clear in the analysis, since the agent has acquired sufficient information during the exploration phase, should be small on average, which implies the returned policy is near-optimal.
2 Analysis
In this section we outline the analysis of our algorithm. The formal proof is deferred to the supplementary material. We first give the formal theoretical guarantee of our algorithm.
After collecting trajectories during the exploration phase, with probability , our algorithm outputs an -optimal policy for an arbitrary number of reward functions satisfying Assumption 2.1 during the planning phase.
Now we show how to prove Theorem 3.1. Our first lemma shows that the estimated value functions are optimistic with high probability, and the summation of should be small.
With probability , for all ,
for some constant where is as defined in Algorithm 1.
Note that the definition of the exploration driven reward function used in the -th episode depends only on samples collected during the first episodes. Therefore, the first part of the proof is nearly identical to that of Theorem 3.1 in (Jin et al., 2019). To prove the second part of the lemma, we first recursively decompose (similar to the standard regret decomposition for optimistic algorithms), and then use the fact that and the elliptical potential lemma in (Abbasi-Yadkori et al., 2012) to given an upper bound on . The formal proof is provided in the supplementary material.
Our second lemma shows that with high probability, if one divides the bonus function (defined in Line 5 in Algorithm 2) by and uses it as a reward function, then the optimal policy has small cumulative reward on average.
With probability , for the function defined in Line 5 in Algorithm 2, we have
for some absolute constant .
for all , which implies the desired result.
Our third lemma states the estimated -function is always optimistic, and is upper bounded by plus the UCB bonus function . The lemma can be proved using the same concentration argument as in (Jin et al., 2019).
With probability , for an arbitrary number of reward functions satisfying Assumption 2.1 and all , we have
Now we sketch how to prove Theorem 3.1 by combining Lemma 3.2 and Lemma 3.3. Note that With probability , the events defined in Lemma 3.2 and Lemma 3.3 both hold. Conditioning on both events, we have
where the first inequality follows by Lemma 3.3, the second inequality follows by Lemma 3.3 and decomposing the -function recursively, the third inequality follows by the definition of , and the last inequality follows by Lemma 3.2.
In this section we prove lower bound for reward-free RL under the linear assumption. We show that there exists a class of MDPs which satisfies Assumption 2.2, such that any reward-free RL algorithm requires exponential number of samples during the exploration phase in order to find a near-optimal policy during the planning phase. In particular, we prove the following theorem.
Since deterministic systems are special cases of general MDPs, the hardness result in Theorem 4.1 applies to general MDPs as well. In the remaining part of this section, we describe the construction of the hard instance and outline the proof of Theorem 4.1.
In the hard instance, there are levels of states
where contains all states that can be reached in level . The action space . For each , we represent each state in by an integer in , i.e., , , , etc. We also have and . The initial states is .
Transition.
For each , for each , is fixed and thus known to the algorithm. In particular, for each , for each , we define where . We will define the transition operator for those states shortly.
Feature Extractor.
For all states , we define
Finally, for all states , we define
The Hard MDPs.
By Yao’s minimax principle (Yao, 1977), to prove a lower bound for randomized algorithms, it suffices to define a hard distribution and show that any deterministic algorithm fails for the hard distribution. We now define the hard distribution. We first define the transition operator for those states . To do this, we first pick a state-action pair from uniformly at random, and define
To define the transition function for those states , we pick a random action from uniformly at random, and define
The Reward Function.
We now define the optimal -function which automatically implies a set of reward function . During the planning phase, the agent will receive as the reward functions. By construction, there exists a unique trajectory with . For each , we define in Assumption 2.2 as . This implies that for each ,
For each , we have
which implies that for all . Now for each , for each , we define so that the Bellman equations hold. Moreover, by construction, for each , we have when , and when and thus .Note that this is slightly different from the assumption that . However, this can be readily fixed by shifting all reward values by .
Proof of Hardness.
Now we sketch the final proof of the hardness result. We define to be the event that for all where are the state-action pairs collected by the algorithm, we have . For any deterministic algorithm, we claim that if the algorithm samples at most trajectories during the exploration phase, with probability at least over the randomness of the distribution of MDPs, holds. This is because the feature extractor is fixed and thus the algorithm receives the same feedback before reaching . Since there are state-action pairs and only one of them satisfies , and the algorithm samples at most trajectories during the exploration phase, holds with probability at least .
Now during the planning phase, by construction of the optimal -function, the only -optimal policy is . However, conditioned on , any deterministic algorithm correctly output with probability at most , since conditioned on , does not contain , and the set of reward functions also does not depend on . Therefore, during the planning phase of the algorithm, a -optimal policy is found with probability at most .
Conclusion
This paper provides both positive and negative results for reward-free RL with linear function approximation. Our results imply three new exponential separations: 1) linear MDP v.s. linear , 2) standard RL v.s. reward-free RL, and 3) query with a simulator v.s. query without a simulator. An interesting future direction is to generalize our results to more general function classes using techniques, e.g., in (Wen and Van Roy, 2013; Ayoub et al., 2020; Wang et al., 2020).
Acknowledgments
RW and RS are supported in part by NSF IIS1763562, AFRL CogDeCON FA875018C0014, and DARPA SAGAMORE HR00111990016. SSD is supported by NSF grant DMS-1638352 and the Infosys Membership.
References
Appendix A Missing Proofs in Section 3
In this section, for all , we denote
Since , we have
for appropriate choices of and .
To prove Lemma 3.1, we need a concentration lemma similar to Lemma B.3 in [Jin et al., 2019].
Suppose Assumption 2.1 holds. Let be the event that for all ,
for some absolute constant . Then .
The proof is nearly identical to that of Lemma B.3 in [Jin et al., 2019]. The only deference in our case is that we have a different reward functions at different episodes. However, note that in our case
Thus our value function is of the form
In our proof, we condition on the event defined in Lemma A.1, which holds with probability at least . Since , we have
is an unknown vector. By Assumption 2.1, . Therefore,
We thus have, for all and ,
Now we prove the first part of the lemma.
Our proof is by induction on . Indeed, for , it holds that for all ,
since . Suppose for some , it holds that for all ,
Notice that for all ,
Second Part.
To prove the second part, for all , we denote
Note that for each , is a martingale difference sequence with . Define to be the even that
By Azuma–Hoeffding inequality, we have .
By Lemma D.2 in [Jin et al., 2019], we have
Conditioned on which holds with probability at least , we have
A.2 Proof of Lemma 3.2
which we condition on in the rest of the proof. Therefore, we have,
Hence we have for all ,
for some absolute constant . ∎
A.3 Proof of Lemma 3.3
Using the same argument in the proof of Lemma 3.1, with probability at least , for all and , we have
Therefore, for all and ,
Moreover, . Since , we have
Now we prove for all and , . We prove by induction on . When this is clearly true. Suppose for some , for all . We have
Since and , it suffices to prove that
A.4 Proof of Theorem 3.1
In our proof we condition on the events defined in Lemma 3.2 and Lemma 3.3 which hold with probability at least . By Lemma 3.3, for any ,
By definition of , we have
By taking for a sufficiently large constant , we have
which implies is -optimal with respect to . ∎
In this section, we present an algorithm for reward-free RL under the linear assumption (Assumption 2.2) in deterministic systems, when the agent has access to a generative model (a.k.a. simulator) of the MDP. More specifically, for each state action , for each , we assume the agent can query . We show that after querying the transition operator for polynomial number of times during the exploration phase, during the planning phase, the agent can find an optimal policy for any given reward function .
The exploration phase of our algorithm is described in Algorithm 3, while the planning phase is described in Algorithm 4.
During the exploration phase, for each level , we find such that
Notice that during the exploration phase, the algorithm query the transition operator for times in total. To prove the correctness, we prove by induction on that during the planning phase, . Note that this is clearly true when . Suppose . It is clear that , which implies by the Bellman equation. By Assumption 2.2, if ,
Appendix C Missing Proofs in Section 4
for all with .
This is a direct implication of Lemma A.1 in [Du et al., 2020a] by setting and . ∎
Note that the above lemma implies the existence of the required feature exactor, since for each , there are less than state-action pairs in . We simply define the feature of the -th state-action pair in to be in the above lemma.
In order to prove Theorem 4.1, by Yao’s minimax principle [Yao, 1977], it suffices to prove that for the hard distribution constructed in Section 4, for any deterministic algorithm that samples at most trajectories during the exploration phase, the probability (over the randomness of the hard distribution) that outputs a -optimal policy in the planning phase is at most .
We first show that for the deterministic algorithm , among all the choices for , is in the collected dataset for at most choices for during the exploration phase. Note that whenever , we must have and . Therefore, the feedback received by is always the same unless . However, since samples at most trajectories during the exploration phase, there are most choices for during the exploration phase for which is in the collected dataset .
Recall that is deterministic. For any choice of , if is not in the collected dataset , the collected dataset is always the same, no matter or . Moreover, for any fixed choice of , it can be verified that the reward function does not depend on the choice of . Note that during the planning phase, algorithm deterministically maps the collected dataset and the reward function to a policy. Furthermore, the only -optimal policy must satisfy . However, for any choice of , if is not in the collected dataset , does not depend on since both the collected dataset and the reward function do not depend on . Therefore, for those choices of , outputs a -optimal policy with probability at most . Therefore, the probability that outputs a -optimal policy is at most