Settling the Horizon-Dependence of Sample Complexity in Reinforcement Learning

Throughout this paper, for a given positive integer $H$, we use $[H]$ to denote the set $\{0,1,\ldots,H-1\}$. For a condition $\mathcal{E}$, we use $\mathbb{I}[\mathcal{E}]$ to denote the indicator function, i.e., $\mathbb{I}[\mathcal{E}]=1$ if $\mathcal{E}$ holds and $\mathbb{I}[\mathcal{E}]=0$ otherwise.

For a random variable $X$ and a real number $\varepsilon\in[0,1]$, its $\varepsilon$-quantile $\mathcal{Q}_{\varepsilon}(X)$ is defined so that

Let $C=(\mathcal{S},P,\mu)$ be a Markov chain where $\mathcal{S}$ is the state space, $P:\mathcal{S}\to\Delta(\mathcal{S})$ is the transition operator and $\mu\in\Delta(\mathcal{S})$ is the initial state distribution. A Markov chain $C$ induces a sequence of random states

where for each $s_{0}\sim\mu$ and $s_{h+1}\sim P(s_{h})$ for each $h\geq 0$.

Let $M=\left(\mathcal{S},\mathcal{A},P,R,H,\mu\right)$ be an $H$-horizon Markov Decision Process (MDP) where $\mathcal{S}$ is the finite state space, $\mathcal{A}$ is the finite action space, $P:\mathcal{S}\times\mathcal{A}\rightarrow\Delta\left(\mathcal{S}\right)$ is the transition operator which takes a state-action pair and returns a distribution over states, $R:\mathcal{S}\times\mathcal{A}\rightarrow\Delta\left(\mathbb{R}\right)$ is the reward distribution, $H\in\mathbb{Z}_{+}$ is the planning horizon (episode length), and $\mu\in\Delta\left(\mathcal{S}\right)$ is the initial state distribution. In the learning setting, $P$, $R$ and $\mu$, are unknown but fixed, and $\mathcal{S}$, $\mathcal{A}$ and $H$ are known to the learner. Throughout, we assume $H=\Omega(|\mathcal{S}|\log|\mathcal{S}|)$ since otherwise we can apply existing algorithms (e.g. [ZJD20]) to obtain a sample complexity that is independent of $H$. Throughout the paper, for a state $s\in\mathcal{S}$, we occasionally abuse notation and use $s$ to denote the deterministic distribution that always takes $s$.

We now introduce how a RL agent (or an algorithm) interacts with an unknown MDP. In the setting of Theorem 1.1, in each episode, the agent first receives an initial state $s_{0}\sim\mu$. For each time step $h\in[H]$, the agent first decides an action $a_{h}\in\mathcal{A}$, and then observes and moves to $s_{h+1}\sim P(s_{h},a_{h})$ and obtains a reward $r_{h}\sim R(s_{h},a_{h})$. The episode stops at time $H$, where the final state is $s_{H}$. Note that even if the agent decides to stop before time $H$, e.g. at time $\sqrt{H}$, this still counts as one full episode.

In the generative model setting (as in Theorem 1.3), the agent is allowed to start from any $s_{0}\in\mathcal{S}$ instead of $s_{0}\sim\mu$, and is allowed to restart at any time. The sample complexity is defined as the number of batches of queries (each batch consists of $H$ queries) the agent uses.

The final goal of RL is to output a good policy $\pi$ with respect to the unknown MDP that the agent interacts with. In this paper, we consider (non-stationary) deterministic policies $\pi$ which choose an action $a$ based on the current state $s\in\mathcal{S}$ and the time step $h\in[H]$. Formally, $\pi=\{\pi_{h}\}_{h=0}^{H-1}$ where for each $h\in[H]$, $\pi_{h}:\mathcal{S}\to\mathcal{A}$ maps a given state to an action. The policy $\pi$ induces a (random) trajectory

where $s_{0}\sim\mu$, $a_{0}=\pi_{0}(s_{0})$, $r_{0}\sim R(s_{0},a_{0})$, $s_{1}\sim P(s_{0},a_{0})$, $a_{1}=\pi_{1}(s_{1})$, $r_{1}\sim R(s_{1},a_{1})$, $\ldots$, $a_{H-1}=\pi_{H-1}(s_{H-1})$, $r_{H-1}\sim R(s_{H-1},a_{H-1})$ and $s_{H}\sim P(s_{H-1},a_{H-1})$. Throughout the paper, all policies are assumed to be deterministic unless specified otherwise.

To measure the performance of a policy $\pi$, we define the value of a policy as

Note that in general a policy does not need to be deterministic, and it can even depend on the history in which case a policy chooses an action at time $h$ based on the entire transition history before $h$. It is well-known (see e.g. [Put94]) that the optimal value of $M$ can be achieved by a non-stationary deterministic optimal policy $\pi^{*}$. Hence, we only need to consider non-stationary deterministic policies.

The goal of the agent is then to find a policy $\pi$ with $V^{\pi}_{M,H}\geq V^{\pi^{*}}_{M,H}-\epsilon$, i.e., obtain an $\epsilon$-optimal policy, while minimizing the number of episodes of environment interactions (or batches of queries if in the generative model).

For a policy $\pi$, we also define

to be the $\delta$-quantile of the visitation frequency of a state-action pair $(s,a)$.

For the sake of the analysis, we shall also consider stationary policies. A stationary deterministic policy $\pi$ chooses an action $a$ solely based on the current state $s\in\mathcal{S}$, i.e, $\pi_{0}=\pi_{1}=\ldots=\pi_{H-1}$. We use $\Pi_{\mathsf{st}}$ to denote the set of all stationary policies. Note that $\left|\Pi_{\mathsf{st}}\right|=|\mathcal{A}|^{|\mathcal{S}|}$.

We remark that when the horizon length $H$ is finite, the value of the best stationary policy and that of the best non-stationary policy can differ significantly. Consider the case when there are two states $s,s^{\prime}$ and two actions $a,a^{\prime}$. Starting from $s$, taking action $a$ will go back to $s$ and obtain a reward of $1/H$, while taking action $a^{\prime}$ will go to $s^{\prime}$ and obtain a reward of $1$. Taking any action at $s^{\prime}$ will go back to $s^{\prime}$ with reward 0. In this case, the optimal non-stationary policy has value $2-1/H$ since it can exit $s$ at time $H$, while deterministic stationary policies can only have value $\leq 1$ (even random stationary policy can only achieve value $2-\Omega(1)$).

For a stationary policy $\pi:\mathcal{S}\to\mathcal{A}$, we use $M^{\pi}=(\mathcal{S},P^{\pi},\mu)$ to denote the Markov chain induced by $M$ and $\pi$, where the transition operator $P^{\pi}$ is defined so that

Below, we formally introduce the bounded total reward assumption.

As discussed in the previous section, this assumption is more general than the standard reward uniformity assumption, where $r_{h}\in[0,1/H]$, for all $h\in[H]$. Throughout this paper, we assume that the above assumption holds for the unknown MDP that the agent interacts with.

The above assumption in fact implies a very interesting consequence on the reward values.

We also introduce another variant of MDP, discounted MDP, which is specified by $M=\left(\mathcal{S},\mathcal{A},P,R,\gamma,\mu\right)$, where $\gamma\in(0,1)$ is a discount factor and all other components have the same meaning as in an $H$-horizon MDP. Note that we consider discounted MDP only for the purpose of analysis, and the unknown MDP that the agent interacts with is always a finite-horizon MDP. The difference between a discounted MDP and an $H$-horizon MDP is that discounted MDPs have an infinite horizon length, i.e., the length of a trajectory can be infinite. Hence, to measure the value of a discounted MDP, suppose policy $\pi$ obtains a random trajectory,

we denote

as the discounted value of $\pi$. Throughout, for a (discounted or finite-horizon) MDP $M=\left(\mathcal{S},\mathcal{A},P,R,\cdot,\mu\right)$, we simply denote $V^{\pi}_{M,H}$ as the value function of $\left(\mathcal{S},\mathcal{A},P,R,H,\mu\right)$ and $V^{\pi}_{M,\gamma}$ as the value function of $\left(\mathcal{S},\mathcal{A},P,R,\gamma,\mu\right)$.

Our algorithm in the generative model is conceptually simple: for each state-action pair $(s,a)$, we draw $O(H)$ samples from $P(s,a)$ and $R(s,a)$ and then return the optimal policy with respect to the empirical model $\widehat{M}$ which is obtained by using the empirical estimators for $P$ and $R$ (denoted as $\widehat{P}$ and $\widehat{R}$). Here for simplicity, we assume $R=\widehat{R}$ which allows us to focus on the estimation error induced by the transition probabilities. Moreover, we assume that $P$ differs from $\widehat{P}$ only for a single state-action pair $(s,a)$. To further simplify the discussion, we assume that there are only two different states on the support of $P(s^{\prime}\mid s,a)$ (say $s_{1}$ and $s_{2}$).

In order to prove the correctness of the algorithm, we show that for any policy $\pi$, the value of $\pi$ in the empirical model $\widehat{M}$ is close to that in the true model $M$. However, standard analysis based on dynamic progarmming shows that the difference between the value of $\pi$ in $\widehat{M}$ and that in $M$ could be as large as $H$ times the estimation error on $P(s,a)$, which is clearly insufficient for obtaining an algorithm which uses $O(1)$ batches of queries. Our main idea here is to show that for most trajectories $T$, the probability of $T$ in the empirical model $\widehat{M}$ is a multiplicative approximation to that in the true model $M$ with constant approximation ratio.

To establish the multiplicative approximation guarantee, our observation is that one should consider $s_{1}$ and $s_{2}$, the two states on the support of $P(s,a)$, as a whole. To see this, consider the case where $P(s_{1}\mid s,a)=P(s_{2}\mid s,a)=1/2$. Again, the additive estimation errors on both $P(s_{1}\mid s,a)$ and $P(s_{2}\mid s,a)$ are roughly $O(1/\sqrt{H})$. Now, consider a trajectory that visits both $(s,a,s_{1})$ and$(s,a,s_{2})$ for $H/2$ times. Note that the multiplicative approximation ratio between $\widehat{P}(s^{\prime}\mid s,a)^{H/2}$ and $P(s^{\prime}\mid s,a)^{H/2}$ could be as large as $\exp(\sqrt{H})$, for both $s^{\prime}=s_{1}$ and $s^{\prime}=s_{2}$. However, since $\widehat{P}(s_{1}\mid s,a)+\widehat{P}(s_{2}\mid s,a)=1$ as the empirical estimator $\widehat{P}(s,a)$ is still a probability distribution, it must be the case that $\widehat{P}(s_{1}\mid s,a)/P(s_{1}\mid s,a)=1-2\delta$ and $\widehat{P}(s_{2}\mid s,a)/{P}(s_{2}\mid s,a)=1+2\delta$ where $\delta=P(s_{1}\mid s,a)-\widehat{P}(s_{1}\mid s,a)$ and thus $|\delta|\leq O(1/\sqrt{H})$. Since $(1+2\delta)^{H/2}(1-2\delta)^{H/2}=(1-4\delta^{2})^{H/2}$ is a constant, $(\widehat{P}(s_{1}\mid s,a))^{H/2}(\widehat{P}(s_{2}\mid s,a))^{H/2}$ is a constant factor approximation to the true probability $(P(s_{1}\mid s,a))^{H/2}(P(s_{2}\mid s,a))^{H/2}$ due to cancellation.

In our analysis, to formalize the above intuition, for each trajectory $T$, we take $T$ into consideration only when $|m_{T}(s,a,s^{\prime})-P(s^{\prime}\mid s,a)\cdot m_{T}(s,a)|\leq O(\sqrt{P(s^{\prime}\mid s,a)\cdot H})$ for both $s^{\prime}=s_{1}$ and $s^{\prime}=s_{2}$. Here $m_{T}(s,a)$ is the number of times that $(s,a)$ is visited on $T$ and $m_{T}(s,a,s^{\prime})$ is the number of times that $(s,a,s^{\prime})$ is visited on $T$. By Chebyshev’s inequality (see Lemma 5.7 for details), we only ignore a small subset of trajectories whose total probability can be upper bounded by a constant. For the remaining trajectories, it can be shown that $\widehat{P}(s_{1}\mid s,a)^{m_{T}(s,a,s_{1})}\cdot\widehat{P}(s_{2}\mid s,a)^{m_{T}(s,a,s_{2})}$ is a constant factor approximation to $P(s_{1}\mid s,a)^{m_{T}(s,a,s_{1})}\cdot P(s_{2}\mid s,a)^{m_{T}(s,a,s_{2})}$ so long as $|\widehat{P}(s,a,s^{\prime})-P(s,a,s^{\prime})|\leq O(\sqrt{P(s,a,s^{\prime})/H})$ for both $s^{\prime}=s_{1}$ and $s^{\prime}=s_{2}$ due to the cancellation mentioned above. The formal argument is given in Lemma 5.5. Note that using $O(H)$ samples, $|\widehat{P}(s,a,s^{\prime})-P(s,a,s^{\prime})|\leq O(\sqrt{P(s,a,s^{\prime})/H})$ holds only when $P(s,a,s^{\prime})\geq\Omega(1/H)$. On the other hand, we can also ignore trajectories that visit $(s,a,s^{\prime})$ with $P(s,a,s^{\prime})\leq O(1/H)$ since such trajectories have negligible cumulative probability by Markov’s inequality (formalized in Lemma 5.6).

The above analysis can be readily generalized to handle perturbation on the transition probabilities of multiple state-action pairs, and to handle the case when the transition operator $P(\cdot\mid s,a)$ is not supported on two states. In summary, by using $O(H)$ samples for each state-action pair $(s,a)$, the empirical model $\widehat{M}$ provides a constant factor approximation to the probabilities of all trajectories, except for a small subset of them whose cumulative probability can be upper bounded by a constant. Hence, for all policies, the empirical model provides an accurate estimate to its value and thus, the optimal policy with respect to the empirical model is near-optimal.

In the discussion above, we heavily rely on the ability of the generative model to obtain $\Omega(H)$ samples for each state-action pair. However, for the RL setting, it is not possible to reach every state-action pair freely. Although each trajectory contains $H$ state-action-state tuples (corresponding to a batch of queries in the generative model), these samples may not cover states that are crucial for learning an optimal policy. Indeed, one could use all possible deterministic non-stationary policies to collect samples, which shall then cover the whole state-action space. Unfortunately, such a naïve method introduces a dependence on the number of non-stationary policies which is exponential in $H$. The sample complexity of other existing methods in the literature also inevitably depends on $H$ as their sample complexity intrinsically depends on the number of non-stationary policies.

In this work, we adopt a completely different approach for exploration. Our new idea is to show that if there exists a non-stationary policy that visits $(s,a)$ for $f$ times in expectation, then there exists a stationary policy that visit $(s,a)$ for $f/\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$ times in expectation. This is formalized in Lemma 4.7. If the above claim is true, then intuitively, one can simply enumerate all stationary policies and sample $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$ trajectories using each of them to obtain $f$ samples of $(s,a)$. Note that there are only $|\mathcal{A}|^{|\mathcal{S}|}$ stationary policies, which is completely independent of $H$. In order to prove the above claim, we show that for any stationary policy $\pi$, its value in the infinite-horizon discounted setting is close to that in the finite-horizon undiscounted setting (up to a factor of $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$) by using a properly chosen discount factor (see Lemma 4.5). Note that this implies the correctness of the above claim since there always exists a stationary optimal policy in the infinite-horizon discounted setting.

In order to show the value of a stationary policy in the infinite-horizon discounted setting is close to that in the finite-horizon setting, we study reaching probabilities in time invariant Markov chains. In particular, we show in Lemma 4.4 that in a time invariant Markov chain, for any $H\geq|\mathcal{S}|$, the probability of reaching a specific state $s$ within $H$ steps is close the probability of reaching $s$ within $4H$ steps, up to a factor of $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$. Previous literature on time invariant Markov chains mostly focus on the asymptotic behavior, and as far as we are aware, we are the first to prove the above claim. Note that this claim directly establishes a connection between the value of a stationary policy in the infinite-horizon discounted setting and that in the finite-horizon setting. Moreover, as a direct consequence of the above claim, we can show that if $H>2|\mathcal{S}|$, the value of a stationary policy within $H$ steps is close to that of the same policy within $H/2$ steps, up to a factor of $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$. This consequence is crucial for later parts of the analysis.

The above analysis shows that if there exists a non-stationary policy that visits $(s,a)$ for $f$ times in expectation, then our algorithm, which uses all stationary policies to collect samples, visits $(s,a)$ for $f/\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$ times in expectation. However, this does not necessarily mean that one can obtain $f$ samples of $(s,a)$ by sampling $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$ trajectories using our algorithm with good probability. To see this, consider the case when our policy visits $(s,a)$ for $H$ times with probability $1/\sqrt{H}$ and does not visit $(s,a)$ with probability $1-1/\sqrt{H}$. In this case, our policy may not obtain even a single sample for $(s,a)$ unless one rollouts the policy for $O(\sqrt{H})$ times. Therefore, instead of obtaining a visitation frequency guarantee which holds in expectation, it is more desirable to obtain a visitation frequency guarantee that holds with good probability.

To resolve this issue, we establish a connection between the expectation and the $\epsilon$-quantile of the visitation frequency of a stationary policy. We note that such a connection could not hold without any restriction. To see this, consider a policy that visits $(s,a)$ for $H$ times with probability $\epsilon/2$. In this case, the expected visitation frequency is $\epsilon H/2$ while the $\epsilon$-quantile is zero. On the other hand, suppose the initial state $s_{0}=s$ almost surely, then such a connection is easy to establish by using the Martingale Stopping Theorem. In particular, in Corollary 4.9, we show that if there exists a non-stationary policy that visits $(s,a)$ for $f$ times with probability $\epsilon$, then there exists a stationary policy that visits $(s,a)$ for $\epsilon f/\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$ times with constant probability, when the initial state $s_{0}=s$ almost surely.

In general, the initial state $s_{0}$ comes from a distribution $\mu$ and could be different from $s$ with high probability. To tackle this issue, in our algorithm, we simultaneously enumerate two stationary policies $\pi_{1}$ and $\pi_{2}$. $\pi_{1}$ should be thought as the policy that visits $(s,a)$ with highest probability within $H/2$ steps starting from the initial state distribution $\mu$, and $\pi_{2}$ should be thought as the policy that maximizes the $\epsilon$-quantile of the visitation frequency of $(s,a)$ within $H/2$ steps when $s_{0}=s$. In our algorithm, we execute $\pi_{1}$ before $(s,a)$ is visited for the first time, and switch to $\pi_{2}$ once $(s,a)$ has been visited. Intuitively, we first use $\pi_{1}$ to reach $(s,a)$ for the first time and then use $\pi_{2}$ to collect as many samples as possible. As mentioned above, the value of a stationary policy within $H$ steps is close to the value of the same policy within $H/2$ steps, up to a factor of $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)$. Thus, by sampling the above policy (formed by concatenating $\pi_{1}$ and $\pi_{2}$) for $\exp\left(O(|\mathcal{S}|\log|\mathcal{S}|)\right)/\epsilon^{2}$ times, we obtain at least $f$ samples for $(s,a)$, if there exists a non-stationary policy that visits $(s,a)$ for $f$ times with probability $\epsilon$. This is formalized in Lemma 5.1.

By the above analysis, suppose $m(s,a)$ is the largest integer such that there exists a non-stationary policy that visits $(s,a)$ with probability $\epsilon$ for $m(s,a)$ times, then our dataset contains $\Omega(m(s,a))$ samples of $(s,a)$. However, $m(s,a)$ could be significantly smaller than $H$ and therefore the perturbation analysis established in the generative model no longer applies here. For example, previously we show that if $|m_{T}(s,a,s^{\prime})-P(s^{\prime}\mid s,a)\cdot m_{T}(s,a)|\leq O(\sqrt{P(s^{\prime}\mid s,a)\cdot H})$, then $\widehat{P}(s_{1}\mid s,a)^{m_{T}(s,a,s_{1})}\cdot\widehat{P}(s_{2}\mid s,a)^{m_{T}(s,a,s_{2})}$ is a constant factor approximation to $P(s_{1}\mid s,a)^{m_{T}(s,a,s_{1})}\cdot P(s_{2}\mid s,a)^{m_{T}(s,a,s_{2})}$ when $|\widehat{P}(s,a,s^{\prime})-P(s,a,s^{\prime})|\leq O(\sqrt{P(s,a,s^{\prime})/H})$ for both $s^{\prime}=s_{1}$ and $s^{\prime}=s_{2}$. However, if $m(s,a)\ll H$, it is hopeless to obtain an estimate $\widehat{P}(s,a,s^{\prime})$ with $|\widehat{P}(s,a,s^{\prime})-P(s,a,s^{\prime})|\leq O(\sqrt{P(s,a,s^{\prime})/H})$. Fortunately, our perturbation analysis still goes through so long as $m_{T}(s,a,s^{\prime})\leq P(s^{\prime}\mid s,a)\cdot m_{T}(s,a)+O(\sqrt{P(s^{\prime}\mid s,a)\cdot m(s,a)})$ and $|\widehat{P}(s,a,s^{\prime})-P(s,a,s^{\prime})|\leq O(\sqrt{P(s,a,s^{\prime})/m(s,a)})$, i.e., replacing all $H$ appearances with $m(s,a)$.

The above analysis introduces a final subtlety in our algorithm. In particular, $m(s,a)$ in the empirical model $\widehat{M}$ could be significantly larger than that in the true model. On the other hand, the number of samples of $(s,a)$ in our dataset is at most $O(m(s,a))$ where $m(s,a)$ is defined by the true model. This means the value estimated in the empirical model $\widehat{M}$ could be significantly larger than that in the true model $M$. To resolve this issue, we employ the principle of “pessimism in the face of uncertainty” and for each policy $\pi$, the estimated value of $\pi$ is set to be the lowest value among all models that lie the confidence set. Since the true model always lies in the confidence set, the estimated value is then guaranteed to be close to the true value.