SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chung-Wei Lee*
leechung@usc.edu
Qinghua Liu
qinghual@princeton.edu
Yasin Abbasi-Yadkori
yadkori@google.com
Chi Jin
chij@princeton.edu
Tor Lattimore
lattimore@google.com
Csaba Szepesvári
Google DeepMind and University of Alberta
szepesva@ualberta.ca
We consider a contextual bandit problem with  $S$  contexts and  $K$  actions. In each round  $t = 1,2,\ldots$  the learner observes a random context and chooses an action based on its past experience. The learner then observes a random reward whose mean is a function of the context and the action for the round. Under the assumption that the contexts can be lumped into  $r\leq \min \{S,K\}$  groups such that the mean reward for the various actions is the same for any two contexts that are in the same group, we give an algorithm that outputs an  $\varepsilon$ -optimal policy after using at most  $\widetilde{O} (r(S + K) / \varepsilon^2)$  samples with high probability and provide a matching  $\Omega (r(S + K) / \varepsilon^2)$  lower bound. In the regret minimization setting, we give an algorithm whose cumulative regret up to time  $T$  is bounded by  $\widetilde{O} (\sqrt{r^3(S + K)T})$ . To the best of our knowledge, we are the first to show the near-optimal sample complexity in the PAC setting and  $\widetilde{O} (\sqrt{\mathrm{poly}(r)(S + K)T})$  minimax regret in the online setting for this problem. We also show our algorithms can be applied to more general low-rank bandits and get improved regret bounds in some scenarios.
Consider a recommendation platform that interacts with a finite set of users in an online fashion. Users arrive at the platform and receive a recommendation. If they engage with the recommendation (e.g., they "click") then the platform receives a reward, otherwise no reward is obtained. Assume that the users can be partitioned into a small number of groups such that users in the same group have the same preferences.
We ask whether there exist algorithms that can take advantage of the lumpability of users into a few groups, even when the identity of the group a user belongs to is unknown and only learnable because they share preferences with other users in the group.
A slightly more general version of this problem can be formalized as follows: Viewing users as "contexts" and recommendations as "actions" (or arms), assume that there are  $S$  contexts and  $K$  actions. In round  $t = 1,2,\ldots$  the learner first receives a context  $i_{t}$ , sampled from an unknown
*most works were done when interning at DeepMind.
distribution on the set  $[S] := \{1, \dots, S\}$  of possible contexts. The learner then chooses an action  $j_t \in [K] := \{1, \dots, K\}$  and observes a reward  $y_t = A(i_t, j_t) + \eta_t$ , where given the past,  $\eta_t$  has a subgaussian tail (precise definitions are postponed to Section 2) and  $A: [S] \times [K] \to \mathbb{R}$  is an unknown function of mean rewards ( $\mathbb{R}$  denotes the set of reals). We consider two settings when the goal of the learner is either to identify a near-optimal policy  $\pi: [S] \to [K]$ , or to keep its regret small. Policy  $\pi$  is called  $\varepsilon$ -optimal if
$$
\mathbb {E} \left[ A \left(i _ {1}, \pi \left(i _ {1}\right)\right) \right] \geq \max  _ {\pi^ {\prime}} \mathbb {E} \left[ A \left(i _ {1}, \pi^ {\prime} \left(i _ {1}\right)\right) \right] - \varepsilon , \tag {1}
$$
while the regret of the learner for a horizon of  $T$  is
$$
\operatorname {R e g} _ {T} = \mathbb {E} \left[ \sum_ {t = 1} ^ {T} \max  _ {j \in [ K ]} A \left(i _ {t}, j\right) - \sum_ {t = 1} ^ {T} A \left(i _ {t}, j _ {t}\right) \right]. \tag {2}
$$
The expectations are taken with respect to the randomness of both the learner and environment, including contexts and rewards. It is well known (e.g., Lattimore and Szepesvári [2020]) that there are algorithms such that an  $\varepsilon$ -optimal policy will be discovered after
$$
\widetilde {O} \left(\frac {S K}{\varepsilon^ {2}}\right) \tag {3}
$$
interactions, and there are also algorithms for which the regret satisfies
$$
\operatorname {R e g} _ {T} = \widetilde {O} (\sqrt {S K T}). \tag {4}
$$
Here, the notation  $\widetilde{O}(\cdot)$  hides polylogarithmic factors of the variables involved. We say that the stochastic, finite, contextual bandit problem specified by  $A$  is  $r$ -lumpable (or, in short, the bandit problem is context-lumpable) if there is a partitioning of  $[S]$  into  $r \leq \min\{S, K\}$  groups such that  $A(i, \cdot) = A(i', \cdot)$  holds whenever  $i, i' \in [S]$  belong to the same group. It is not hard to see that any algorithm needs at least  $\Omega(r(S + K) / \varepsilon^2)$  interactions to discover an  $\varepsilon$ -optimal policy (Theorem 3). Indeed, if we view  $A$  as an  $S \times K$  matrix, the lumpability condition states that  $A = UV$  where  $U$  is an  $S \times r$  binary matrix where each row has a single nonzero element, and  $V$  is an  $r \times K$  matrix, which gives the unique mean rewards given the group indices. Hence, crude parameter counting suggests that there are  $r(S + K)$  parameters to be learned.
The question is whether  $SK$  in Eq. (3) and Eq. (4) can be replaced with  $(S + K)\mathrm{poly}(r)$  without knowing the grouping of the contexts.
More generally, we can ask the same questions for contextual bandits with the low-rank structure, where the matrix  $A$  has rank  $r$ . Equivalently, the low-rank condition means that we have the same decomposition  $A = UV$  as above but no more constraints on  $U$ . In the example of recommendation platforms, this assumption is more versatile as the users are modeled as mixtures of  $r$  preference types instead of belonging to one type only.
Our problem can be seen as an instance of contextual bandit problems introduced by Auer et al. [2002]. For a good summary of the history of the contextual bandit problem, the reader is advised to consult the article by Tewari and Murphy [2017]. Further review of prior works can also be found in the books of Slivkins [2019], Lattimore and Szepesvári [2020]. Another way to approach context-lumpable stochastic bandits is to model them as stochastic linear bandits with changing action sets Auer [2002], Chu et al. [2011], Abbasi-Yadkori et al. [2011]. However, lumpability does not give improvements on the regret bound when running the standard algorithms in these settings (EXP4 and SupLinRel, respectively). We provide more details in the appendix.
We are not the first to study the  $r$ -context-lumpable stochastic bandit problem. The earliest work is probably that of Maillard and Mannor [2014] who named this problem latent bandits: they consider the group identity of the context as latent, or missing information. While the paper introduces this problem, the main theoretical results are derived for the case when the reward distributions given the group information are known to the learner. Further, the results derived are instance dependent.
Altogether, while this paper makes interesting contributions, it does not help in answering our questions.2
The earlier work of Gentile et al. [2014] considered another generalization of our problem where in round  $t = 1,2,\ldots$  the learner gets an action set  $\{x_{t,1},\dots,x_{t,K_t}\} \subset \mathbb{R}^d$  and the mean payoff of choosing action  $1\leq j\leq K_{t}$  given the past and the current context  $i_{t}$  is  $x_{t,j}^{\top}\theta_{g(i_{t})}$  with  $\theta_{1},\ldots ,\theta_{r}\in \mathbb{R}^{d}$  unknown. When specializing this to our setting, their result depends on the minimum separation  $\min_{i,j:g(i)\neq g(j)}\| A(i,\cdot) - A(j,\cdot)\|$ , which, according to the authors may be an artifact of their analysis. An extra assumption, which the authors suggest could be removed, is that the distribution of the contexts is uniform (as we shall see, removing this assumption is considerable work). Further, as the authors note, the worst-case for their algorithm is when the contexts are equipartitioned, in which case their bound reduces to the naive bound that one gets for non-lumpable problems. Li et al. [2016] considers the variation of the problem when the set of arms (and their feature vectors) are fixed and is known before learning begins and separation is defined with respect to this fixed set of arms (and hence could be larger than before). From our perspective, their theoretical results have the same weaknesses as the previous paper. Gentile et al. [2017] further generalizes these previous works to the case when the set of contexts is not fixed a priori. However, the previously mentioned weaknesses of the results remain.
Our problem is also a special case of learning in lumpable Markov Decision Processes. In such Markov Decision Processes (MDPs) the states can be partitioned such that states within a group of a partition behave identically, both in terms of the transitions and the rewards received [e.g., Ren and Krogh, 2002].<sup>3</sup> We put a survey of this line of work in Appendix A.6.
In summary, although the problem is extensively studied in the literature, we are not aware of any previous results that have implications for the existence of the minimax regret bounds in terms of  $\widetilde{O} (\sqrt{\mathrm{poly}(r)(S + K)T})$ , which we believe is a fundamental question in this area.
2 Notation and problem definition
For an positive integer  $n$ , we use  $[n]$  to denote the set  $\{1, \ldots, n\}$ . Further, for a finite set  $U$ ,  $\Delta(U)$  denotes the set of probability distributions over  $U$  and  $\mathrm{unif}(U)$  denotes the uniform distribution over  $U$ . The set of real numbers is denoted by  $\mathbb{R}$ .
As introduced earlier, we consider context-lumpable bandits with  $S$  contexts and  $K$  actions such that the  $S$  contexts can be lumped into  $r$  groups so that the mean payoff  $A(i,j)$  given that action  $j \in [K]$  is used while the context is  $i \in [S]$  depends only on the identity of the group that  $i$  belongs to and the action  $j$ . That is, for some  $g: [S] \to [r]$  map
$$
A (i, j) = A \left(i ^ {\prime}, j\right) \quad \text {f o r a n y} j \in [ K ], i, i ^ {\prime} \in [ S ] \text {s u c h t h a t} g (i) = g \left(i ^ {\prime}\right).
$$
Neither  $A$ , nor  $g$  is known to the learner who interacts with the bandit instance in the following way: In rounds  $t = 1,2,\ldots$  the learner first observes a context  $i_t\in [S]$  randomly chosen from a fixed, context distribution  $\nu \in \Delta ([S])$ , independently of the past. The distribution  $\nu$  is also unknown to the learner. Next, the learner chooses an action  $j_{t}\in [K]$  to receive a reward of  $y_{t} = A(i_{t},j_{t}) + \eta_{t}$  where  $\eta_t$  is 1-subgaussian given the past: For any  $\lambda \in \mathbb{R}$ ,
$$
\mathbb {E} \left[ \exp \left(\lambda \eta_ {t}\right) \mid i _ {1}, j _ {1}, \dots , i _ {t}, j _ {t} \right] \leq \exp \left(\lambda^ {2} / 2\right) \text {a l m o s t s u r e l y (a . s .)}.
$$
As described earlier, we are interested in two problem settings. In the PAC setting, the learner is given a target suboptimality level  $\varepsilon > 0$  and a target confidence  $\delta$ . The learner is then tasked to discover a policy  $\pi : [S] \to [K]$  that is  $\varepsilon$ -optimal (cf. Eq. (1)), with probability  $\delta$ . In this setting, in addition to deciding what actions to choose, the learner also needs to decide when to stop collecting data and when it stops it has to return a map  $\hat{\pi} : [S] \to [K]$ . If  $T$  is the (random) time when the learner stops, then the efficiency of the learner is characterized by  $\mathbb{E}[T]$ . In the online setting, the goal of the learner is to keep its regret (cf. Eq. (2)) low. In Section 3, we will consider the PAC setting, while the online setting will be considered in Section 4.
<sup>2</sup>Somewhat confusingly, Hong et al. [2020] define latent bandit problems differently from the definition given by Maillard and Mannor [2014]: They assume that in addition to the context, a latent (unobserved) variable, influences the rewards, however, they also assume that the distributions of the reward given an action, context and latent state are all known. Their results therefore can not help us in answering our question.
<sup>3</sup>Lumpability was first introduced first by Kemeny and Snell [1976] in the context of Markov chains.
In the remainder of the paper, we will also need some further notation. The map  $g$  induces a partitioning  $\mathcal{P}^{\star} = \{\mathcal{B}(1),\ldots ,\mathcal{B}(r)\}$  of  $[S]$  in the natural way:
$$
\mathcal {B} (b) = \{i: g (i) = b \}.
$$
We call each subset  $\mathcal{B} \in \mathcal{P}^{\star}$  a block and call  $\mathcal{B}(b)$  block  $b$  for every  $b \in [r]$ . We also define block reward  $\mu(b, j) = A(i, j)$  of block  $b$  for every context  $i \in \mathcal{B}(b)$  and arm  $j \in [K]$ . Finally, we define the block distribution  $\omega \in \Delta(r)$  so that  $\omega(b) = \sum_{i \in \mathcal{B}(b)} \nu(i)$ .
3 Near-optimal PAC Learning in Context-lumpable Bandits
In this section, we present an algorithm for PAC learning in context-lumpable bandits, and prove that its sample complexity guarantee matches the minimax rate up to a logarithmic factor.
3.1 Special case: almost-uniform context distribution
To better illustrate the key ideas behind the algorithm design, we first consider a special case of almost-uniform context distribution. Formally, we assume  $\nu(i) \in [1/(8S), 8/S]$  for all  $i \in [S]$  throughout this subsection. The pseudocode of our algorithm is provided in Algorithm 1, which consists of three main modules. Below we elaborate on each of them separately.
Data collection (Line 3-4 and Algorithm 2). The high-level goal of this step is to collect a sufficient amount of data for every block/action pair so that in later steps we have sufficient information to select a near-optimal arm for each block. One naive approach is to try every action on every context for a sufficient amount of time (e.g.,  $\widetilde{O}(1 / \varepsilon^2)$ ). However, this will cause an undesired factor of  $SK$  in the sample complexity. To address this issue, note that contexts from the same block share the same reward function, which means that for every block/action pair  $(b,j) \in [r] \times [K]$ , we only need to sufficiently try action  $j$  on a single context  $i$  from block  $b$  (i.e.,  $i \in \mathcal{B}(b)$ ). However, the block structure is unknown a priori. We circumvent this problem by assigning a random action  $\psi(i)$  to each context  $i$  for them to explore (Line 4). And after action  $\psi(i)$  has been tried on context  $i$  sufficiently many times, we update  $\psi(i)$  by reassigning context  $i$  with a new random action (Line 8-10).
However, there is still one key problem unresolved yet: how many samples to use for estimating each  $A(i,\psi (i))$  before reassigning context  $i$  with another random action, given a fixed budget of total sample complexity. On one hand, if we only try each  $(i,\psi (i))$  pair a very few numbers of times before switching, then we can potentially explore a lot of different actions for each block but the accuracy of the reward estimate could be too low to identify a near-optimal policy. On the other hand, estimating each  $A(i,\psi (i))$  with a huge number of samples resolves the poor accuracy issue but could potentially cause the problem of under-exploration, especially for those blocks consisting of very few contexts. In the extreme case, if a block only contains a single context, then our total budget may only afford to test a single action on that context (block).
To address this issue, we choose different levels of accuracy for different blocks adaptively depending on their block size. Specifically, contexts from larger blocks can afford to try each random action before switching for a larger number of times to achieve higher estimate accuracy because larger blocks consist of more contexts, which means that the average number of random actions in each context inside a larger block needs to try is smaller. And for smaller blocks, the case is the opposite. Finally, we remark that the above scheme of using adaptive estimate accuracy can be implemented without any prior knowledge of the block structure. We only need to iterate over different accuracy levels using an exponential schedule (Line 3), and each block will be sufficiently explored at its appropriate accuracy level. Specifically, let  $N = \lceil \log (1 / \varepsilon^2) \rceil$ , and consider accuracy levels  $n \in [N]$ . For accuracy level  $n$ , we collect a dataset  $\mathcal{D}_n$  by adding a context-action pair after the pair is played for  $2^n$  timesteps (Line 9).
Screening optimal action candidates (Line 5-11). Equipped with the dataset collected from Step 1, we want to identify a subset  $\mathcal{W}$  of  $[K]$  so that (i) for every block  $b\in [r]$ ,  $\mathcal{W}$  contains a near-optimal action of block  $b$ , and (ii) the size of  $\mathcal{W}$  is as small as possible. We construct such  $\mathcal{W}$  in an iterative way. For each accuracy level  $n\in [N]$ , we first compute the context-action pair  $(i^{\star},j^{\star})$  with the highest reward estimate in  $\mathcal{D}_n$  (Line 7). Intuitively, as  $(i^{\star},j^{\star})$  has been tried for  $2^{n}$  timesteps in constructing  $\mathcal{D}_n$ , it is natural to guess that  $j^{\star}$  could potentially be a  $\widetilde{O}(2^{-n/2})$ -optimal action for
certain blocks so we add it into optimal action candidate set  $\mathcal{W}$  (Line 8). To identify the contexts from those blocks, we sample new data to estimate the reward of  $j^{\star}$  on each context  $i\in [S]$ . This can be done by calling Algorithm 2 and setting the exploration set  $\mathcal{K}$  to be  $j^{\star}$  (Line 10). If a context  $i$  can achieve reward close enough to  $\widehat{A}_n(i^\star ,j^\star)$  at action  $j^{\star}$ , then we peel off every  $(i,j)\in \mathcal{D}_n$  (Line 11). This is because we have found  $j^{\star}$  as an optimal action candidate for  $i$  at this accuracy level and don't need to consider other  $j$ . We repeat the above process, and add a new arm to  $\mathcal{W}$  in each iteration, until  $\mathcal{D}_n$  becomes empty.
Solving the simplified context-lumpable bandit (Line 12). After obtaining the optimal action candidate set  $\mathcal{W}$ , we can discard all actions not in  $\mathcal{W}$  and solve a simplified context-lumpable bandit problem by replacing the original action set  $[K]$  with  $\mathcal{W}$ . Note that  $\widetilde{O}(S|\mathcal{W}| / \varepsilon^2)$  samples suffice for learning an  $\varepsilon$ -optimal policy for this simplified problem. For example, we can directly try each action  $j \in \mathcal{W}$  on each context  $i \in [S]$  for  $\widetilde{\Theta}(1 / \varepsilon^2)$  times, and define  $\pi^{\mathrm{out}}(i) \in \arg \max_{j \in \mathcal{W}} \bar{A}(i,j)$  where  $\bar{A}(i,j)$  is the empirical estimate of  $A(i,j)$  based on  $\widetilde{\Theta}(1 / \varepsilon^2)$  samples.
Algorithm 1 Algorithm for Almost-uniform Context Distribution
1 Let  $t$  denote the current time and initialize  $N\gets [\log (1 / \varepsilon^2)],\mathcal{W}\gets \emptyset ,\widetilde{\lg}\gets 16\log (rSK / \delta)$    
2 Define  $\kappa (i)\gets [K]$  for  $i\in [S],L\gets r(S + K)\widetilde{\lg} /\varepsilon^{2}$  Step 1. Data collection   
3 for accuracy level  $n = 1,\ldots ,N$  do   
4 Execute Algorithm 2 with input  $L,n,\mathcal{K},$  and receive  $\mathcal{D}_n$  and  $\widehat{A}_n$  Step 2. Screening optimal action candidates   
5 for accuracy level  $n = 1,\dots ,N$  do   
while  $\mathcal{D}_n\neq \emptyset$  do   
7 Compute  $(i^{\star},j^{\star})\gets \arg \max_{(i,j)\in \mathcal{D}_n}\widehat{A}_n(i,j)$    
8 Update optimal action candidates  $\mathcal{W}\gets \mathcal{W}\bigcup \{j^{\star}\}$    
9 Update  $\kappa (i)\gets \{j^{\star}\}$  for  $i\in [S]$  and  $L^{\prime}\gets 8\widetilde{\lg} 2^{n}S$    
10 Execute Algorithm 2 with input  $L^{\prime},n,\mathcal{K},$  and reassign output to  $\widetilde{D}_n$  and  $\widetilde{A}_n$    
11 Shrink  $\mathcal{D}_n\gets \{(i,j)\in \mathcal{D}_n:|\widetilde{A}_n(i,j^\star) - \widehat{A}_n(i^\star ,j^\star)|\geq \sqrt{\frac{\lg}{2^n}}\}$
Step 3. Solving the simplified context-lumpable bandit
12 Use  $\frac{4S|\mathcal{W}|}{\varepsilon^2}\log \frac{SK}{\delta}$  samples to find  $\pi^{\mathrm{out}}$  s.t.  $A(i,\pi^{\mathrm{out}}(i))\geq \max_{j\in \mathcal{W}}A(i,j) - \varepsilon$  for all  $i\in [S]$  13 Output  $\pi^{\mathrm{out}}$
Algorithm 2 Data Collection
11 Output  $\mathcal{D}_n$  and  $\hat{A}_n$
1 Input  $L, n, \mathcal{K}$   
2 Let  $t$  denote the current time and initialize  $\mathcal{D}_n \gets \emptyset$   
3 for context  $i \in [S]$  do  
4 Assign  $\psi_t(i)$  to be an arm drawn from  $\text{unif}(\mathcal{K}(i))$   
5 for L timesteps do  
6 Receive context  $i_t$ , play arm  $j_t = \psi(i_t)$ , and receive reward  $y_t$   
7 Preset  $\psi_{t+1}(i) = \psi_t(i)$  for every  $i \in [S]$   
8 if  $(\sum_{\tau=1}^{t} \mathbb{1}(i_\tau = i_t)) \% 2^n = 0$  then  
9  $\mathcal{D}_n \gets \mathcal{D}_n \bigcup \{(i_t, \psi(i_t))\}$  and  $\widehat{A}_n(i_t, \psi(i_t)) \gets \frac{\sum_{\tau \leq t} y_\tau \mathbb{1}[i_\tau = i_t, j_\tau = \psi(i_t)]}{\sum_{\tau \leq t} \mathbb{1}[i_\tau = i_t, j_\tau = \psi(i_t)]}$   
10 Reassign  $\psi_{t+1}(i_t) \sim \text{unif}(\mathcal{K}(i_t))$ ;
Now we present the theoretical guarantee for Algorithm 1. The proof and exact constants in the bound can be found in Appendix C.
Theorem 1. Let  $\delta \in (0,1)$  and assume  $\nu(i) \in [1/(8S), 8/S]$  for all  $i \in [S]$ . Algorithm 1 outputs an  $\widetilde{O}(\varepsilon)$ -optimal policy within  $\widetilde{O}(r(S + K) \log(1/\delta)/\varepsilon^2)$  samples with probability at least  $1 - \delta$ .
3.2 Extension to general context distribution
Algorithm 3 Algorithm for General Context Distribution
1 Let  $J = \frac{4S}{\varepsilon} \log(S / \delta)$ ,  $L = \lceil \log(S / \varepsilon) \rceil$ ,  $N = 524 \left( \log \frac{rSK}{\delta} \right) \cdot \left( 1 + 2 \log \frac{1}{\varepsilon} \right)^2 \cdot \frac{r(S + K)}{\varepsilon^2}$ .
2 Estimate the context distribution by sampling  $J$  contexts, and denote the estimate by  $\hat{\nu}$
3 Split the context set into  $L$  disjoint subsets  $\{\mathcal{X}_l\}_{l\in [L]}$  where
$$
\mathcal {X} _ {l} \leftarrow \left\{ \begin{array}{l l} \{i \in [ S ]: \hat {\nu} (i) \in (2 ^ {- l - 1}, 2 ^ {- l} ] \}, & l \in [ 0: L - 1 ] \\ \{i \in [ S ]: \hat {\nu} (i) \in [ 0, 2 ^ {- l} ] \}, & l = L \end{array} \right.
$$
4 for  $l \in [L - 1]$  do
5 Execute Algorithm 1 to learn policy  $\pi_l$  for subset  $\mathcal{X}_l$  from  $N$  time steps
6 Output  $\pi^{\mathrm{out}}$  such that  $\pi^{\mathrm{out}}$  equals to  $\pi_l$  over  $\mathcal{X}_l$  for  $l\in [0:L - 1]$ , and arbitrary over  $\mathcal{X}_L$
In this subsection, we show how to generalize Algorithm 1 to handle general context distributions. We present the pseudo-code in Algorithm 3. The algorithm consists of two key steps. In the first step, we use  $J = \widetilde{O}(S / \varepsilon)$  samples to obtain an empirical estimate of the context distribution, denoted by  $\hat{\nu}$ . Then we divide the context set into many disjoint subsets  $\{\mathcal{X}_l\}_{l \in [0:L]}$  such that inside each subset  $\mathcal{X}_l$ , the conditional empirical context distribution is almost uniform. As a result, we can invoke Algorithm 1 to find a near-optimal policy  $\pi_l$  for each subset  $\mathcal{X}_l$  (Line 5). It requires  $\widetilde{O}(r(S + K) / \varepsilon^2)$  time steps for every  $\ell$  but we only use samples where contexts are from  $\mathcal{X}_l$  and ignore the rest. Finally, we glue all  $\pi_l$  together to get a policy  $\pi^{\mathrm{out}}$  that is near-optimal for the original problem. Formally, we have the following theoretical guarantee for Algorithm 3. The proof and exact constants are in Appendix C.
Theorem 2. Let  $\delta \in (0,1)$ . Algorithm 3 outputs an  $\widetilde{O}(\varepsilon)$ -optimal policy within  $\widetilde{O}(r(S + K) \log(1/\delta)/\varepsilon^2)$  samples with probability at least  $1 - \delta$ .