Planning in Markov Decision Processes with Gap-Dependent Sample Complexity

Introduction

In reinforcement learning (RL), an agent repeatedly takes actions and observes rewards in an unknown environment described by a state. Formally, the environment is a Markov Decision Process (MDP) $\mathcal{M}=\langle\mathcal{S},\mathcal{A},p,r\rangle$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ the action space, $p=\left\{p_{h}\right\}_{h\geq 1}$ a set of transition kernels and $r=\left\{r_{h}\right\}_{h\geq 1}$ a set of reward functions. By taking action $a$ in state $s$ at step $h$, the agent reaches a state $s^{\prime}$ with probability $p_{h}(s^{\prime}|s,a)$ and receives a random reward with mean $r_{h}(s,a)$. A common goal is to learn a policy $\pi=(\pi_{h})_{h\geq 1}$ that maximizes cumulative reward by taking action $\pi_{h}(s)$ in state $s$ at step $h$. If the agent has access to a generative model, it may plan before acting by generating additional samples in order to improve its estimate of the best action to take next.

In this work, we consider Monte-Carlo planning as the task of recommending a good action to be taken by the agent in a given state $s_{1}$, by using samples gathered from a generative model. Let $Q^{\star}(s_{1},a)$ be the maximum cumulative reward, in expectation, that can be obtained from state $s_{1}$ by first taking action $a$, and let $\hat{a}_{n}$ be the recommended action after $n$ calls to the generative model. The quality of the action recommendation is measured by its simple regret, defined as $\bar{r}_{n}(\hat{a}_{n}):=V^{\star}(s_{1})-Q^{\star}(s,\hat{a}_{n}),\;\mbox{where}\;V^{\star}(s_{1}):=\max_{a}Q^{\star}(s_{1},a)$.

We propose an algorithm in the fixed confidence setting $(\varepsilon,\delta)$: after $n$ calls to the generative model, the algorithm should return an action $\hat{a}_{n}$ such that $\bar{r}_{n}(\hat{a}_{n})\leq\varepsilon$ with probability at least $1-\delta$. We prove that its sample complexity $n$ is bounded in high probability by a quantity that depends on the sub-optimality gaps of the actions that are applicable in state $s_{1}$. We also provide experiments showing its effectiveness. The only assumption that we make on the MDP is that the support of the transition probabilities $p_{h}(\cdot|s,a)$ should have cardinality bounded by $B<\infty$, for all $s$, $a$ and $h$.

Monte-Carlo Tree Search (MCTS) is a form of Monte-Carlo planning that uses a forward model to sample transitions from the current state, as opposed to a full generative model that can sample anywhere. Most MCTS algorithms sample trajectories from the current state , and are widely used in deterministic games such as Go. The AlphaZero algorithm guides planning using value and policy estimates to generate trajectories that improve these estimates. The MuZero algorithm combines MCTS with a model-based method which has proven useful for stochastic environments. Hence efficient Monte-Carlo planning may be instrumental for learning better policies. Despite their empirical success, little is known about the sample complexity of state-of-the-art MCTS algorithms.

The earliest MCTS algorithm with theoretical guarantees is Sparse Sampling , whose sample complexity is polynomial in $1/\varepsilon$ in the case $B<\infty$ (see Lemma 1). However, it is not trajectory-based and does not select actions adaptively, making it very inefficient in practice.

A first category of algorithms rely on optimistic planning , and require additional assumptions: a deterministic MDP , the open loop setting in which policies are sequences of actions instead of state-action mappings (the two are equivalent in MDPs with deterministic transitions), or an MDP with known parameters . For MDPs with stochastic and unknown transitions, polynomial sample complexities have been obtained for StOP , TrailBlazer and SmoothCruiser but the three algorithms suffer from numerical inefficiency, even for $B<\infty$. Indeed, StOP explicitly reasons about policies and storing them is very costly, while TrailBlazer and SmoothCruiser require a very large amount of recursive calls even for small MDPs. We remark that popular MCTS algorithms such as UCT are not $(\varepsilon,\delta)$-correct and do not have provably small sample complexities.

In the setting $B<\infty$, BRUE is a trajectory-based algorithm that is anytime and whose sample complexity depends on the smallest sub-optimality gap $\Delta:=\min_{a\neq a^{\star}}\left(V^{\star}(s_{1})-Q^{\star}(s_{1},a)\right)$. For planning in deterministic games, gap-dependent sample complexity bounds were previously provided in a fixed-confidence setting . Our proposal, MDP-GapE, can be viewed as a non-trivial adaptation of the UGapE-MCTS algorithm to planning in MDPs. The defining property of MDP-GapE is that it uses a best arm identification algorithm, UGapE , to select the first action in a trajectory, and performs optimistic planning thereafter, which helps refining confidence intervals on the intermediate Q-values. Best arm identification tools have been previously used for planning in MDPs and UGapE also served as a building block for StOP .

Finally, going beyond worse-case guarantees for RL is an active research direction, and in a different context gap-dependent bounds on the regret have recently been established for tabular MDPs .

Contributions

We present MDP-GapE, a new MCTS algorithm for planning in the setting $B<\infty$. MDP-GapE performs efficient Monte-Carlo planning in the following sense: First, it is a simple trajectory-based algorithm which performs well in practice and only relies on a forward model. Second, while most practical MCTS algorithms are not well understood theoretically, we prove upper bounds on the sample complexity of MDP-GapE. Our bounds depend on the sub-optimality gaps associated to the state-action pairs encountered during exploration. This is in contrast to StOP and TrailBlazer, two algorithms for the same setting, whose guarantees depend on a notion of near-optimal nodes which can be harder to interpret, and that can be inefficient in practice. In the anytime setting, BRUE also features a gap-dependent sample complexity, but only through the worst-case gap $\Delta$ defined above. As can be seen in Table 1, the upper bound for MDP-GapE given in Corollary 1 improves over that of BRUE as it features the gap of each possible first action $a_{1}$, $\Delta_{1}(s_{1},a_{1})=V^{\star}(s_{1})-Q_{1}^{\star}(s_{1},a_{1})$, and scales better with the planning horizon $H$. Furthermore, our proof technique relates the pseudo-counts of any trajectory prefix to the gaps of state-action pairs on this trajectory, which evidences the fact that MDP-GapE does not explore trajectories uniformly.

Learning Framework and Notation

where the supremum is taken over (deterministic) policies $\pi=(\pi_{1},\dots,\pi_{H})$, and the expectation is on a trajectory $s_{1},a_{1},\dots,s_{h},a_{h}$ where $s_{h}\sim p_{h-1}(\cdot|s_{h-1},a_{h-1})$ and $a_{h}=\pi_{h}(s_{h})$ for $h\in[2,H]$. With this definition, an optimal action in state $s_{1}$ is $a^{\star}\in\text{argmax}_{a\in\mathcal{A}(s_{1})}Q^{\star}(s_{1},a)$.

We assume that there is a maximal number $K$ of actions available in each state, and that, for each $(s,a)$, the support of $p_{h}(\cdot|s,a)$ is bounded by $B$: that is, $B$ is the maximum number of possible next states when applying any action. We further assume that the rewards are bounded in $$. For each pair of integers$i,h$such that$i\leq h$, we introduce the notation$[i,h]=\{i,\ldots,h\}$and$[h]=[1,h]$.

A sequential planning algorithm proceeds as follows. In each episode $t$, the agent uses a deterministic policy on the form $\pi^{t}=(\pi_{1}^{t},\ldots,\pi_{H}^{t})$ to generate a trajectory $(s_{1},a_{1}^{t},r_{1}^{t},\ldots,s_{H}^{t},a_{H}^{t},r_{H}^{t})$, where $a_{h}^{t}=\pi_{h}^{t}(s_{h}^{t})$, $r_{h}^{t}$ is a reward with expectation $r_{h}(s_{h}^{t},a_{h}^{t})$ and $s_{h+1}^{t}\sim p_{h}(\cdot|s_{h}^{t},a_{h}^{t})$. After each episode the agent decides whether it should perform a new episode to refine its guess for a near-optimal action, or whether it can stop and make a guess. We denote by $\tau$ the stopping rule of the agent, that is the number of episodes performed, and $\hat{a}_{\tau}$ the guess.

We aim to build an $(\varepsilon,\delta)$-correct algorithm, that is an algorithm that outputs a guess $\hat{a}_{\tau}$ satisfying

while using as few calls to the generative model $n=H\tau$ (i.e. as few episodes $\tau$) as possible.

Our setup permits to propose algorithms for planning in the undiscounted episodic case (in which our bounds will not blow up when $\gamma=1$) and in discounted MDPs with infinite horizon. Indeed, choosing $H$ such that $2\gamma^{H}/(1-\gamma)\leq{\varepsilon}$, an $(\varepsilon,\delta)$-correct algorithm for the discounted episodic setting recommends an action that is $2\varepsilon$-optimal for the discounted infinite horizon setting.

A (recursive) baseline

Sparse Sampling can be tuned to output a guess $\hat{a}$ that satisfies (1), as specified in the following lemma, which provides a baseline for our undiscounted episodic setting (see Appendix F). Note that Sparse Sampling is not strictly sequential as it does not repeatedly select trajectories.

If $B<\infty$, Sparse Sampling using horizon $H$ and performing $\mathcal{O}\left(({H^{5}}/{\varepsilon^{2}})\log\left({BK}/{\delta}\right)\right)$ transitions in each node is $(\varepsilon,\delta)$-correct with sample complexity $O(n_{\text{SS}})$ for $n_{\text{SS}}:={H^{5}(BK)^{H}}/{\varepsilon^{2}}$.

Structure of the optimal Q-value function

In our algorithm, we will build estimates of the intermediate Q-values, that are useful to compute the optimal Q-value function $Q^{\star}(s_{1},a)$. Defining

$Q^{\star}(s_{1},a)=Q_{1}(s_{1},a)$ and the optimal action-values $Q=(Q_{1},\ldots,Q_{H})$ can be computed recursively using the Bellman equations, where we use the convention $Q_{H+1}(\cdot,\cdot)=0$:

Let $\pi^{\star}=(\pi_{1}^{\star},\ldots,\pi_{H}^{\star})$ denote a deterministic optimal policy where, for $h\in[H]$, $\pi_{h}^{\star}(s_{h})=\arg\max_{a}Q_{h}(s_{h},a)$, with ties arbitrarily broken. Hence the optimal value in $s_{h}$ is $Q_{h}(s_{h},\pi_{h}^{\star}(s_{h}))$.

The MDP-GapE Algorithm

In this section we present MDP-GapE, a generalization of UGapE to Monte-Carlo planning. Like BAI-MCTS for games a core component is the construction of confidence intervals on $Q_{1}(s_{1},a)$. The construction below generalizes that of OP-MDP for known transition probabilities.

Let $n_{h}^{t}(s_{h},a_{h},s_{h+1}):=\sum_{s=1}^{t}\mathds{1}\left((s_{h}^{s},a_{h}^{s},s_{h+1}^{s})=(s_{h},a_{h},s_{h+1})\right)$ be the number of observations of transition $(s_{h},a_{h},s_{h+1})$, and $R_{h}^{t}(s_{h},a_{h}):=\sum_{s=1}^{t}r_{h}^{s}(s_{h},a_{h})\mathds{1}\left((s_{h}^{s},a_{h}^{s})=(s_{h},a_{h})\right)$ the sum of rewards obtained when selecting $a_{h}$ in $s_{h}$. We define the empirical transition probabilities $\hat{p}^{t}$ and expected rewards $\hat{r}^{t}$ as follows, for state-action pairs such that $n_{h}^{t}(s_{h},a_{h}):=\sum_{s}n_{h}^{t}(s_{h},a_{h},s)>0$:

As rewards are bounded in $$, we define the following Kullback-Leibler upper and lower confidence bounds on the mean rewards$r_{h}(s_{h},a_{h})$ :

In order to define confidence bounds on the values $Q_{h}$, we introduce a confidence set on the probability vector $p_{h}(\cdot|s_{h},a_{h})$. We define $\mathcal{C}_{h}^{t}(s_{h},a_{h})=\Sigma_{B}$ if $n_{h}^{t}(s_{h},a_{h})=0$ and otherwise

where $\Sigma_{B}$ is the set of probability distribution over $B$ elements, $\beta^{p}$ is an exploration function and $\operatorname{KL}(p,q)=\sum_{s\in\text{Supp}(p)}p(s)\log\tfrac{p(s)}{{q}(s)}$ is the Kullback-Leibler divergence between two categorical distributions $p$ and $q$ with supports satisfying $\text{Supp}(p)\subseteq\text{Supp}(q)$.

We now define our confidence bounds on the action values inductively. We use the convention $U_{H+1}^{t}(\cdot,\cdot)=L_{H+1}^{t}(\cdot,\cdot)=0$, and for all $h\in[H]$,

As explained in Appendix A of , optimizing over these KL confidence sets can be reduced to a linear program with convex constraints, that can be solved efficiently with Newton Iteration, which has complexity $O(B\log(d))$ where $d$ is the desired digit precision.

We provide in Section 4.1 an explicit choice for the exploration functions $\beta^{r}(n,\delta)$ and $\beta^{p}(n,\delta)$ that govern the size of the confidence intervals. Note that if the rewards or transitions are deterministic, or if we know $p$, we can adapt our confidence bounds by setting $\beta^{p}=0$ or $\beta^{r}=0$.

MDP-GapE

As any fixed-confidence algorithm, MDP-GapE depends on the tolerance parameter $\varepsilon$ and the risk parameter $\delta$. The dependency in $\varepsilon$ is explicit in the stopping rule (4), while the dependency in $\delta$ is in the tuning of the confidence bounds, that depend on $\delta$.

After $t$ trajectories observed, MDP-GapE selects the $(t+1)$-st trajectory using the policy $\pi^{t+1}=(\pi^{t+1}_{1},\dots,\pi^{t+1}_{H})$ where the first action choice is made according to UGapE:

where $b^{t}$ is the current guess for the best action, which is the action $b$ with the smallest upper confidence bound on its gap $Q_{1}^{\star}(s_{1},a^{\star})-Q_{1}(s_{1},b)$, and $c^{t}$ is some challenger:

Then for all remaining steps we follow an optimistic policy, for all $h\in[2,H]$,

and the guess output when stopping is $\hat{a}_{\tau}=b^{\tau}$. A generic implementation of MDP-GapE is given in Algorithm 1 in Appendix A, where we also discuss some implementation details. Note that, in sharp contrast with the deterministic stopping rule proposed for Sparse Sampling in Lemma 1, MDP-GapE uses an adaptive stopping rule.

Analysis of MDP-GapE

Recall that MDP-GapE uses policy $\pi^{t+1}=(\pi_{1}^{t+1},\dots,\pi_{H}^{t+1})$ to select the $(t+1)$-st trajectory, $s_{1},a_{1}^{t+1},s_{2}^{t+1},a_{2}^{t+1},\dots,s_{H}^{t+1},a_{H}^{t+1}$, satisfying $a_{h}^{t+1}=\pi_{h}^{t+1}(s_{h}^{t+1})$ and $s_{h+1}^{t+1}\sim p_{h}\left(\cdot\left|s_{h}^{t+1},a_{h}^{t+1}\right.\right)$.

To define an event $\mathcal{E}$ that holds with high probability, let $\mathcal{E}^{r}$ (resp. $\mathcal{E}^{p}$) be the event that the confidence regions for the mean rewards (resp. transition kernels) are correct:

For a state-action pair $(s_{h},a_{h})$, let $p_{h}^{\pi}(s_{h},a_{h})$ be the probability of reaching it at step $h$ under policy $\pi$, and let $p_{h}^{t}(s_{h},a_{h})=p_{h}^{\pi^{t}}(s_{h},a_{h})$. We define the pseudo-counts of the number of visits of $(s_{h},a_{h})$ as $\bar{n}_{h}^{t}(s_{h},a_{h}):=\sum_{s=1}^{t}p_{h}^{s}(s_{h},a_{h})\,.$ As $n_{h}^{t}(s_{h},a_{h})-\bar{n}_{h}^{t}(s_{h},a_{h})$ is a martingale, the counts should not be too far from the pseudo-counts. Given a rate function $\beta^{\text{cnt}}$, we define the event

Finally, we define $\mathcal{E}$ to be the intersection of these three events: $\mathcal{E}=\mathcal{E}^{r}\cap\mathcal{E}^{p}\cap\mathcal{E}^{\text{cnt}}$.

1 Correctness

One can easily prove by induction (see Appendix B) that

In Lemma 2 below, we provide a calibration of the thresholds functions $\beta^{r},\beta^{p}$ and $\beta^{\text{cnt}}$ such that this sufficient condition holds. This result, proved in Appendix C, relies on new time-uniform concentration inequalities that follow from the method of mixtures .

Moreover, the maximum of these three thresholds defined (by continuity when $B=1$) as

is such that $n\mapsto\beta(n,\delta)$ is non-decreasing and $n\mapsto\beta(n,\delta)/n$ is non-increasing.

2 Sample Complexity

In order to state our results, we define the following sub-optimality gaps. $\Delta_{h}(s_{h},a_{h})$ measures the gap in future discounted reward between the optimal action $\pi_{h}^{\star}(s_{h})$ and the action $a_{h}$, whereas $\Delta_{1}^{\star}(s_{1},a_{1})$ also takes into account the gap of the second best action and the tolerance level $\varepsilon$.

Recall that $\Delta=\min_{a\neq a^{\star}}\left[Q_{1}(s_{1},a^{\star})-Q_{1}(s_{1},a)\right]$. For all $h\in[H]$, we let

Our sample complexity bounds follow from the following crucial theorem, which we prove in Appendix D, that relates the pseudo-counts of state-action pairs at time $\tau$ to the corresponding gap.

If $\mathcal{E}$ holds, every $(s_{h},a_{h})$ is such that

The number of episodes used by MDP-GapE satisfies

The upper bound on the sample complexity $n=H\tau$ of MDP-GapE that follows from Corollary 1 improves over the $\mathcal{O}(H^{5}(BK)^{H}/\varepsilon^{2})$ sample complexity of Sparse Sampling. It is also smaller than the $\mathcal{O}(H^{4}(BK)^{H}/\Delta^{2})$ samples needed for BRUE to have a reasonable upper bound on its simple regret. The improvement is twofold: first, this new bound features the problem dependent gap $\Delta(s_{1},a_{1})\vee\Delta\vee\varepsilon$ for each action $a_{1}$ in state $s_{1}$, whereas previous bounds were only expressed with $\varepsilon$ or $\Delta$. Second, it features an improved scaling in $H^{2}$.

In particular, using that $\tau=\sum_{t_{1:H}\in\mathcal{T}}\bar{n}_{h}^{\tau}(t_{1:H})$ where $\mathcal{T}$ is the set of $(BK)^{H}$ complete trajectories leads to a sample complexity bound featuring all gaps. However, its improvement over the bound of Corollary 1 is not obvious in the general case. For $B=1$, that is for planning in a deterministic MDP with possibly random rewards, a slightly different proof technique leads to the following improved gap-dependent sample complexity bound (see the proof in Appendix E).

Numerical Experiments

We consider random discounted MDPs with infinite horizon in which the maximal number $B$ of successor states and the sparsity of rewards are controlled. The transition kernel is generated as follows: for each transition in $\mathcal{S}\times\mathcal{A}$, we uniformly pick $B$ next states in $\mathcal{S}$. The cumulative transition probabilities to these states are computed by sorting $B-1$ numbers uniformly sampled in $(0,1)$. The reward kernel is computed by selecting a proportion of the transitions to have non-zero rewards with means sampled uniformly in $(0,1)$. The values for these parameters are shown in Table 3(a).

We verify empirically that MDP-GapE is ($\varepsilon$, $\delta$)-correct while stopping with a reasonable number of oracle calls. Table 3(b) shows the choice of parameters for the algorithm. For various values of the desired accuracy $\varepsilon$ and of the corresponding planning horizon $H=\lceil\log_{\gamma}(\varepsilon(1-\gamma)/2)\rceil$ (see Section 2), we run simulations on 200 random MDPs. We report in Table 4 the distribution of the number $n=\tau H$ of oracle calls and the simple regret $\bar{r}_{n}(\hat{a}_{n})$ of MDP-GapE over these 200 runs. We first observe that MDP-GapE verifies $\bar{r}_{n}(\hat{a}_{n})<\varepsilon$ in all simulations, despite the use of smaller exploration functions compared to those prescribed in Lemma 2. We then compare its sample complexity to that of Sparse Sampling, which is deterministic and for which $n_{\text{SS}}$ given in Lemma 1 is a tight upper bond. We see that the sample complexity of MDP-GapE is an order of magnitude smaller than that of Sparse Sampling.

Scaling in ε𝜀\varepsilon

Comparison to the state of the art

In the fixed-confidence setting, most existing algorithms are considered theoretical and cannot be applied to practical cases. For instance, for our problem with $K=5$ and $\varepsilon=1$, Sparse Sampling and SmoothCruiser both require a fixed budgetIn non-regularized MDPs, SmoothCruiser has the same sample complexity as Sparse Sampling. of at least $n_{\text{SS}}=$8\text{\times}{10}^{9}$$. Likewise, Trailblazer is a recursive algorithm which did not terminate in our setting. We did not implement StOP as it requires to store a tree of policies, which is very costly even for moderate horizons. In comparison, Table 4 shows that MDP-GapE stopped after$n=$1.9\text{\times}{10}^{4}$$oracle calls in the worst case. To the best of our knowledge, MDP-GapE is the first$(\varepsilon,\delta)$-correct algorithm for general MDPs with an easy implementation and a reasonable running time in practice. The only planning algorithms that can be run in practice are in the fixed-budget setting, which we now consider.

Fixed-budget evaluation

We compare MDP-GapE to three existing baselines: first, the KL-OLOP algorithm , which uses the same upper-confidence bounds on the rewards $u_{h}^{t}$ and states values $U_{h}^{t}$ as MDP-GapE, but is restricted to open-loop policies, i.e. sequences of actions only. Second, the BRUE algorithm which explores uniformly and handles closed-loop policies. Third, the popular UCT algorithm , which is also closed-loop and performs optimistic exploration at all depths. UCT and its variants lack theoretical guarantees, but they have been shown successful empirically in many applications. For each algorithm, we tune the planning horizon $H$ similarly to KL-OLOP, by dividing the available budget $n$ into $\tau$ episodes, where $\tau$ is the largest integer such that $\tau\log\tau/(2\log 1/\gamma)\leq n$, and choose $H=\log\tau/(2\log 1/\gamma)$. The exploration functions are those of KL-OLOP and depend on $\tau$: $\beta_{r}(n_{h}^{t},\delta)=\beta_{p}(n_{h}^{t},\delta)=\log(\tau)$. Again, we perform 200 simulations and report in Figure 2 the mean simple regret, along with its $95\%$ confidence interval. We observe that MDP-GapE compares favourably with these baselines in the high-budget regime.

Conclusion

We proposed a new, efficient algorithm for Monte-Carlo planning in Markov Decision Processes, that combines tools from best arm identification and optimistic planning and exploits tight confidence regions on mean rewards and transitions probabilities. We proved that MDP-GapE attains the smallest existing gap-dependent sample complexity bound for general MDPs with stochastic rewards and transitions, when the branching factor $B$ is finite. In future work, we will investigate the worse-case complexity of MDP-GapE, that is try to derive an upper bound on its sample complexity that only features $\varepsilon$ and some appropriate notion of near-optimality dimension.

Acknowledgments

Anders Jonsson is partially supported by the Spanish grants TIN2015-67959 and PCIN-2017-082.

References

Appendix A Detailed Algorithm

In this section we provide a detailed algorithm for MDP-GapE, namely Algorithm 1.

There are different ways to store and update the confidence bounds on the $Q$-value (that is, to specify the UpdateBounds subroutine) according to how we merge information across states.

The most obvious one, suggested by previous work (and also implemented for our experiments) does not merge information at all and builds a search tree in which a node $(s_{h},a_{h})$ at depth $h$ is identified with the sequence of $h$ states and actions that leads to it. It leads to a very simple update: after each trajectory, one only needs to update the confidence bounds, $U_{h}(s_{h},a_{h})$ and $L_{h}(s_{h},a_{h})$, of the visited action-state pairs. Another option is to merge information for the same states and a fixed depth. But in this case the search tree becomes a graph and after each trajectory we need to re-compute the values $U_{h}(s_{h},a_{h})$ for all stored state action pairs $(s_{h},a_{h})$ at each depth.

Appendix B Correctness of MDP-GapE

In this section we prove the correctness of MDP-GapE under the assumption that the event $\mathcal{E}^{r}\cap\mathcal{E}^{p}$ holds. Concretely, we prove by induction that

The base case is given by $h=H+1$, in which case by our previous convention,

For the inductive case, assume that the inclusion holds at depth $h+1$. Then we have

Appendix C Concentration Events

In this section we prove that the event $\mathcal{E}$ holds with high probability. But before we need several concentration inequalities.

Let $X_{1},X_{2},\ldots,X_{n},\ldots$ be i.i.d. samples from a distribution supported over $\{1,\ldots,m\}$, of probabilities given by $p\in\Sigma_{m}$, where $\Sigma_{m}$ is the probability simplex of dimension $m-1$. We denote by $\widehat{p}_{n}$ the empirical vector of probabilities, i.e. for all $k\in\{1,\ldots,m\}$

For all $p\in\Sigma_{m}$, for all $\delta\in$,

We apply the method of mixture with a Dirichlet prior on the mean parameter of the exponential family formed by the set of categorical distribution on $\{1,\ldots,m\}$. Letting

be the log-partition function, the following quantity is a martingale:

where in the second inequality we used Lemma 3. Now we choose the uniform prior $\alpha=(1,\ldots,1)$. Hence we get

It remains to upper-bound the entropic term

Thus we can lower bound the martingale as follows

Using the fact that, for any supermartingale it holds that

which is a well-known property used in the method of mixtures (see ), we conclude that

where $\varphi_{p}(\lambda)=\log(p_{m}+\sum_{k=1}^{m-1}p_{k}e^{\lambda_{k}})$ and $p^{\lambda}=\nabla\varphi_{p_{0}}(\lambda)$.

There is a more general way than the ad hoc one below to prove the result. First note that

C.2 Deviation Inequality for Bounded Distribution

Let $X_{1},X_{2},\ldots,X_{n},\ldots$ be i.i.d. samples from a distribution $\nu$ of mean $\mu$ supported on $$. We denote by$\widehat{\mu}_{n}$ the empirical mean

It is well known, see , that we can "project" the distribution $\nu$ on a Bernoulli distribution with the same mean and then use deviation inequality for Bernoulli to concentrate the empirical mean. This method dos not lead to the sharpest confidence intervals but it provides a good trade-off between complexity computation and accuracy.

For all distribution $\nu$ of mean $\mu$ supported on the unit interval, for all $\delta\in$,

Then we can follow the proof of Proportion 1 with $m=2$ and where $M_{n}^{\lambda}$ is only a supermartingale but this does not change the result as the property (6) still holds. Thus the proposition follows by specifying Proposition 1 to the case $m=2$. ∎

C.3 Deviation Inequality for sequence of Bernoulli Random Variables

C.4 Proof of Lemma 2

We just prove that each event forming $\mathcal{E}=\mathcal{E}^{r}\cap\mathcal{E}^{p}\cap\mathcal{E}^{n}$ holds with high probability. For the first one using Proposition 2, since the reward are bounded in the unit interval we have

where we used Doob’s optional skipping in the second inequality in order to apply Proposition 2, see Section 4.1 of . Similarly for the confidence regions for the probabilities transitions, using Proposition 1 we obtain

It remains to control the counts, using Proposition 3,

where we used that by definition of the pseudo-counts

Appendix D Proof of Theorem 1

In this section we present the proof of Theorem 1, which relies on three important ingredients. The first ingredient is Lemma 5 in Appendix D.1, which provides a relationship between the state-action gaps and the diameter $D_{h}^{t}(s_{h},a_{h}):=U_{h}^{t}(s_{h},a_{h})-L_{h}^{t}(s_{h},a_{h})$ of the confidence intervals. The second ingredient is Lemma 8 in Appendix D.2, which provides an upper bound on the diameter $D_{h}^{t}(s_{h},a_{h})$. The third ingredient is Lemma 9 in Appendix D.3, which relates the actual counts of state-action pairs to the corresponding pseudo-counts. After providing these ingredients, we present the detailed proof of Theorem 1 in Appendix D.4.

Before stating Lemma 5, we prove an important property of the UGapE algorithm. We recall that $b^{t}$ and $c^{t}$ are the candidate best action and its challenger, defined as

The policy at the root is then defined as $\pi^{t+1}_{1}(s_{1})=\underset{b\in\{b^{t},c^{t}\}}{\text{argmax }}\left[U_{1}^{t}(s_{1},b)-L_{1}^{t}(s_{1},b)\right]$.

For all $t\in[\tau_{\delta}-1]$, the following inequalities hold:

$U^{t}_{1}\left(s_{1},c^{t}\right)-L^{t}_{1}\left(s_{1},b^{t}\right)\leq U^{t}_{1}\left(s_{1},\pi^{t+1}(s_{1})\right)-L^{t}_{1}\left(s_{1},\pi^{t+1}(s_{1})\right)$,

$U^{t}_{1}\left(s_{1},b^{t}\right)-L^{t}_{1}\left(s_{1},c^{t}\right)<2\left[U^{t}_{1}\left(s_{1},\pi^{t+1}(s_{1})\right)-L^{t}_{1}\left(s_{1},\pi^{t+1}(s_{1})\right)\right]$.

We show the first part by contradiction. If the inequality does not hold, we obtain

where the last inequality follows from the definition of $b^{t}$. Combining the two inequalities yields $U_{1}^{t}(s_{1},b^{t})<U_{1}^{t}(s_{1},c^{t})<\max_{a\neq c^{t}}U_{1}^{t}(s_{1},a)$, which contradicts the definition of $c^{t}$.

For the second part, if $t<\tau_{\delta}$ then the algorithm has not yet stopped, implying

As a consequence of Lemma 4, we can upper bound any confidence interval involving $b^{t}$ and $c^{t}$.

For each pair of actions $a,a^{\prime}\in\{b^{t},c^{t}\}$, it holds that

If $\mathcal{E}$ holds and $t<\tau_{\delta}$, for all $h\in[H]$ and $s_{h}\in\mathcal{S}_{h}(\pi^{t+1})$,

The proof for $h\in[2,H]$ is immediate from the correctness of the confidence bounds implied by $\mathcal{E}$, and the fact that the selection is optimistic:

For $h=1$, we prove separately that each term in the max is smaller that the right hand side of desired inequality, that is

Now, by definition of the stopping rule, if $t<\tau_{\delta}$, $U_{1}^{t}\left(s_{1},c^{t}\right)-L_{1}^{t}\left(s_{1},b^{t}\right)>\varepsilon$. Using the first property in Lemma 4 yields

Then, exploiting the fact that the action with largest UCB is either $b^{t}$ or $c^{t}$, it holds on $\mathcal{E}$ that

Using Corollary 3 to further upper bound the right hand side yields

Finally, one can also write, on the event $\mathcal{E}$,

In each of the four possible choices of $(a,a^{\prime})$, Corollary 3 implies that

Lemma 5 follows by combining (7), (8) and (9) with the definition of $\Delta_{1}^{\star}\left(s_{1},\pi^{t+1}(s_{1})\right)$. ∎

D.2 Upper bounding the diameters

In this section we state and prove Lemma 8. We use the notation $\sigma_{h}=\sum_{i=0}^{h-1}\gamma^{i}$ to upper bound the discounted reward in $h$ steps. As a first step, we prove the following auxiliary lemma.

If $\mathcal{E}$ holds, for each $h\in[H]$, each $(s_{h},a_{h})$ and each $q\in\mathcal{C}_{h}^{t}(s_{h},a_{h})$,

where we have used Pinsker’s inequality to bound the $L^{1}$-norm using the KL divergence, combined with the fact that both $q$ and $p$ are close to the empirical transition probabilities $\hat{p}^{t}$ under $\mathcal{E}$. ∎

As a consequence, we can express the upper bound $U^{t}$ in terms of the true transition probabilities $p$.

If $\mathcal{E}$ holds, for each $h\in[H]$ and each $(s_{h},a_{h})$,

We can also express the lower bound $L^{t}$ in terms of the transition probabilities $p$ and policy $\pi^{t+1}$.

If $\mathcal{E}$ holds, for each $h\in[H]$ and each $(s_{h},a_{h})$,

We exploit the fact that for each $h\in[H]$, each $(s_{h},a_{h})$ and each $q\in\mathcal{C}_{h}^{t}(s_{h},a_{h})$,

The proof is analogous to the proof of Lemma 6. We can now write

If $\mathcal{E}$ holds, for all $h\in[H]$, $s_{h}\in\mathcal{S}_{h}(\pi^{t+1})$ and $a_{h}$,

The bound on the diameter follows directly from Corollary 4 and Lemma 7:

where we used $\mathcal{E}^{r}\supseteq\mathcal{E}$ and Pinsker’s inequality to bound

To obtain the final expression in Lemma 8, we observe that it also trivially holds that

The conclusion follows by observing that one can get rid of the maximum with 1 in the denominator by using instead the convention $1/0=+\infty$. ∎

D.3 Relating counts to pseudo-counts

D.4 Detailed proof of Theorem 1

Using Lemma 5 and the fact that $p^{t+1}_{h}(s_{h},a_{h})=0$ if $a_{h}\neq\pi^{t+1}(s_{h})$ yields

By induction, one then obtains the following upper bound:

Summing for $t\in\{0,\dots,\tau-1\}$ the inequalities given by (10) yields

The rest of the proof consists in upper bounding $B_{h}^{\tau}(s_{h},a_{h})$ in terms of the pseudo counts $\overline{n}_{h}^{\tau}(s_{h},a_{h})$.

Step 4: from counts to pseudo-counts

Using Lemma 9 to relate the counts to the conditional pseudo-counts, one can write

Since $\gamma\leq 1$ and $x>1$, we can write

where $r=1/x<1$. The latter is an arithmetico-geometric sum that can be upper bounded as

Appendix E Proof of Theorem 2

The proof of Theorem 2 uses the same ingredients as the proof of Theorem 1: Lemma 5 which relates the gaps to the diameters of the confidence intervals $D_{h}^{t}(s_{h},a_{h})=U_{h}^{t}(s_{h},a_{h})-L_{h}^{t}(s_{h},a_{h})$ and a counterpart of Lemma 8 for the deterministic case, stated below.

If $\mathcal{E}$ holds, and $t_{1:H}=(s_{1},a_{1},\dots,s_{H},a_{H})$ is the $(t+1)$-st trajectory generated by MDP-GapE, for all $h\in[H]$,

It follows from Lemma 11 that for all $h\in[H]$, along the $(t+1)$-st trajectory $t_{1:H}=(s_{1},a_{1},\dots,s_{H},a_{H})$,

Using Lemma 5, if $t<\tau$, if $t_{1:H}$ is the trajectory selected at time $(t+1)$, either $n^{t}(t_{1:H})=0$ or

It follows that for any trajectory $t_{1:H}$,

The conclusion follows from Lemma 12 and from the fact that $\tau=\sum_{t_{1:H}\in\mathcal{T}}n^{\tau}(t_{1:H})$.

Appendix F Sample complexity of Sparse Sampling in the Fixed-Confidence Setting

For simplicity, and without loss of generality, assume that the reward function is known. Let $C>0$. Sparse Sampling builds, recursively, the estimates $\widehat{V}_{h}$ and $\widehat{Q}_{h}$ for $h\in[H+1]$, starting from $\widehat{V}_{H+1}(s)=0$ and $\widehat{Q}_{H+1}(s,a)=0$ for all $(s,a)$. Then, from a target state-action pair $(s,a)$, it samples $C$ transitions $Z_{i}\sim p_{h}(\cdot|s,a)$ for $i\in[C]$ and computes:

For an initial state $s$, its output is $\widehat{Q}_{1}(s,a)$ for all $a\in[K]$. For any state $s$, consider the events

defined for $h\in[H+1]$, where $\varepsilon_{h}:=(H-h+1)H\sqrt{(2/C)\log(2/\delta^{\prime})}$ for some $\delta^{\prime}>0$.

with probability at least $1-\delta$, where $\delta=2K\delta^{\prime}\left((BK)^{H}-1\right)/(BK-1)$. Finally, we let $\varepsilon:=H^{2}\sqrt{(2/C)\log(2/\delta^{\prime})}/2$ and solve for $C$, obtaining

Thus predicting $\hat{a}=\underset{a}{\text{argmax }}\widehat{Q}_{1}(s_{1},a)$ after $\mathcal{O}\left(C(BK)^{H}\right)$ sampled transitions we have

Appendix G A Technical Lemma

We state and prove below a technical result that permits to obtain an upper bound on $n$ from a condition of the form $n\Delta^{2}\leq\beta(n,\delta)$, like the one which appears in Theorem 1.

Let $n\geq 1$ and $a,b,c,d>0$. If $n\Delta^{2}\leq a+b\log(c+dn)$ then

Since $\log(x)\leq\sqrt{x}$ and $\sqrt{x+y}\leq\sqrt{x}+\sqrt{y}$ for all $x,y>0$, we have