SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The low-rank structure has been shown useful in other bandit problems. Katariya et al. [2017] considered the case when the learner chooses both "context" and arm and the reward map  $A$ , when viewed as a  $S \times K$  matrix, has rank one. This is a special case of a linear bandit setting with a fixed action set where both the features underlying the actions and the parameter vectors are "shaped" as matrices of matching dimensions, and both matrices are rank-one. Similarly to our problem, the challenge is to make the regret depend on  $S + K$  rather than on  $S \times K$  (but note that  $S$  does not have the interpretation of the number of contexts here). This paper demonstrates that this is possible when considering gap-dependent bounds.
Trinh et al. [2020] improved the result of Katariya et al. [2017] and gave an asymptotically optimal regret bound. Kveton et al. [2017] generalized the setting of Katariya et al. [2017] to the case when the reward  $A$  (viewed as matrix) is of rank  $r \leq \min(S, K)$  and is a "hott topics matrix", the learner in each round can choose  $r$  row and column indices, observes the entries of  $A$  at the resulting submatrix of  $A$  in noise and incurs the maximum reward in this submatrix. For this problem they show an instance dependent bound that depends on  $S, K, r$  only through  $(S + K)\mathrm{poly}(r)$ .
Jun et al. [2019] dropped the extra conditions on the mean reward matrix besides assuming that it has low rank. Jang et al. [2021] gave the first upper bound of the form  $\tilde{O}((S + K)r\sqrt{T})$ , though with inefficient algorithms, which was the first bound better than the naive bound  $\tilde{O}(SK\sqrt{T})$ . Later, Lu et al. [2021] dropped the condition on the action matrices and also extended the results to the generalized linear setting (they assume that both the action matrices and the parameter matrix has a constant Frobenius norm bound). The regret bound of Jun et al. [2019], Lu et al. [2021] takes the form  $\tilde{O}((S + K)^{3/2}\sqrt{rT})$ , which is still worse than the earlier mentioned naive bound. Lattimore and Hao [2021] prove that (up to logarithmic factors) when  $S = K$  and both the action and reward matrices are rank one and symmetric, the minimax regret is of order  $S\sqrt{T}$ .
Recently, Kang et al. [2022] improved the previous state-of-the-art for the generalized linear setting. The new regret bound, which they prove "under a mild" extra assumption, takes the form  $\tilde{O}((S + K)r\sqrt{T})$  and is conjectured to match the order of the minimax regret. Noticing that the minimax lower bounds developed for the finite-armed stochastic bandits (e.g., Exercise 15.2 of Lattimore and Szechesváři 2020) is applicable to the rank-one setting (as the proofs use actions and rewards that are rank one, one-hot matrices), we see that the regret is at least  $\Omega(\sqrt{SKT})$  in these problems. As  $\sqrt{S + K} \approx \sqrt{\max(S, K)} \ll \sqrt{SK} \leq \sqrt{\max(S, K)^2} = \max(S, K) \approx S + K$ , this lower bound rules out upper bounds of the form  $O(\sqrt{(S + K)\mathrm{poly}(r)T})$ , though it is compatible with the upper bound mentioned above.<sup>5</sup>
Jain and Pal [2022] study a related problem where the learner chooses one action per context in each round and incurs a reward for each of the actions. They also propose an epsilon-greedy type algorithm and show that its regret scales as  $T^{2/3} \log(S + K)$  provided some technical conditions hold, one of which is that the condition number of the reward  $A$  when viewed as a matrix is constant. The main limitation of this work is that the technical conditions are restrictive and the problem is easier as in each round an observation is available for each context, whereas in our setting the contexts arrive at random from a distribution which may be very far from the uniform distribution, which, as we shall see, will require extra care.
A.2 Bandit meta-learning
Bandit meta-learning is concerned with learning a lower dimensional subspace across bandit tasks [Kveton et al., 2020, 2021, Cella et al., 2020], while we consider only a single task.
To get this naive bound, following an argument of Jang et al. [2021], notice that the problem is an instance of  $d$ -dimensional linear bandits with  $d = SK$  with actions and parameters belonging to the unit sphere, in which case the minimax regret, up to logarithmic factors, is of order  $d\sqrt{T}$  (e.g., Theorem 24.2, Lattimore and Hao 2021). If the action set has one-hot matrices only, this bound improves to  $\sqrt{dT}$ , as discussed before.
Here,  $f(S, K) \approx g(S, K)$  means that  $f, g$  are within a constant factor of each other, while  $f(S, K) \ll g(S, K)$  means that  $f(S, K) = o(g(S, K))$  as  $S, K \to \infty$ .
A.3 Contextual bandit problems
Contextual bandit problems, introduced by Auer et al. [2002], can be seen as a special case of the prediction with expert advice problem, which is an online problem. In a prediction with expert advice problem, the learner is given access to the recommendations of  $N$  experts in the form of distributions over the  $K$  actions. The learner still needs to choose an action in each round with the goal of competing with the total reward collected by the best expert in hindsight. Auer et al. [2002] consider the adversarial case when in each round, each action is assigned a reward in an arbitrary way from a known, finite interval. The authors propose the EXP4 algorithm that is shown to achieve  $O(\sqrt{KT\log N})$  regret in  $T$  rounds of interactions. Beygelzimer et al. [2011] extended this result to control the regret with high probability, while Dudik et al. [2011] introduced a computationally efficient variant assuming a computationally efficient cost-sensitive classification oracle.
To reduce an  $r$ -lumpable contextual bandit problem to prediction with expert advice, we need to choose the experts. The obvious choice here is that an expert is a  $[S] \to [K]$  map that is the composition of an  $[S] \to [r]$  map followed by an  $[r] \to [K]$  map. Denoting by  $N$  the number of such experts, we see that  $\log(N) = S \log(r) + r \log(K)$ , and hence the regret of EXP4 is of order  $\Omega(\sqrt{SKT})$ , which is not better than not using the lumpability structure.
Another line of work focuses on oracle-based contextual bandits [Foster and Rakhlin, 2020, Simchi-Levi and Xu, 2022]. In this setting, the learner has access to function class  $\mathcal{F}$  and the optimal regret is  $O(\sqrt{KT\log(|\mathcal{F}|)})$ , where  $|\mathcal{F}|$  is the size of the function class. However, with the same argument, we can conclude  $|\mathcal{F}| = \Omega(S)$  even under the lumpable assumption. Therefore, the results in this direction also give no direct implication to the problem we consider.
A.4 Stochastic linear bandits with changing action sets
In this setting, in each round  $t = 1,2,\ldots$  the learner first receives  $K$ $d$ -dimensional vectors,  $x_{t,1},\dots,x_{t,K}$ , each corresponding to an action. Choosing action  $j$  gives a reward whose mean (given the past) is  $x_{t,j}^{\top}\theta$  for some unknown parameter vector  $\theta \in \mathbb{R}^{d}$ . For our case, one can choose  $d = SK$ :  $\theta$  can be the "flattening" of  $A$  and  $x_{t,j}$  is a unit vector so that  $x_{t,j}^{\top}\theta = A(i_t,j)$ . Applying the SupLinRel algorithm from Section 4 of Auer [2002] to this setting, we get the regret bound  $O(\sqrt{dT\log^3(KT)})$ , which shows no improvement compared to ignoring lumpability.
A.5 Matrix completion problems
The offline version of our problem is closely related to matrix completion problems, where the goal is to reconstruct a matrix with missing values under the low-rank constraint [Arora et al., 2012, Jain et al., 2013, Chen et al., 2020]. Based on these ideas, Sen et al. [2016] propose an epsilon-greedy algorithm for the contextual low-rank bandit problem and show that its regret is of order  $T^{2/3}(S \, \text{poly}(r, \log K))^{1/3}$ . This result holds under an assumption that the reward matrix has a nonnegative decomposition  $A = UV$  where entries of  $U$  and  $V$  are all nonnegative. If  $A$  is nonnegative valued itself,  $r$ -lumpability of the contexts implies that such a decomposition exist. However, the main result of this work needs additional conditions. In particular, the context distribution needs to be near-uniform, and due to their "restricted isometry property" assumptions, the context groups need to be of approximately the same size. In particular, if all context groups except one have a single member, their bound degrades to the trivial bound mentioned earlier.
A.6 Lumpable Markov decision processes
Lumpable MDPs in contemporary literature on machine learning are referred to as block MDPs [Du et al., 2019]. Learning to act in a block MDPs has been the subject a numerous recent papers [Dann et al., 2018, Du et al., 2019, Misra et al., 2020, Feng et al., 2020, Zhang et al., 2020, 2022]. Furthermore, learning to act in a block MDP is a special case of the so-called low-rank setting, which was also heavily studied [Jiang et al., 2017, Uehara et al., 2022, Modi et al., 2021, Papini et al., 2021, Zhang et al., 2021, Misra et al., 2020], both in the online and in the PAC settings. The learner in these problems is given a set of feature maps with the promise that at least one of the feature maps will allow a factored (low-rank) representation of the transition dynamics and the reward.
Despite the significant effort that went into studying this problem, in our setting none of the existing results improve upon the naive bound that one can get for non-lumpable problems. Duan et al. [2019], Ni et al. [2021] study the offline setting, where the considerations are rather different. Finally, Kwon et al. [2021] consider the problem when over multiple episodes, a learner interacts with a finite-horizon MDP, which is chosen at random from one of  $r$  such unknown MDPs. The learner is given no additional information about the hidden identity of the MDP that it faces. The challenge here is that even if the MDPs were known, the problem of acting optimally is nontrivial, while in our case this is not a problem. As such, while superficially, the problems may look similar, they are quite different.
In this section, we show some probabilistic lemmas that will be useful to prove our results. The following lemma is a high-probability version of the classical "coupon collector's problem", which says that the expected number of coupons required to draw with replacement is  $\Theta(K \log K)$  in order to get each of  $K$  coupon at least once.
Lemma 6. Given a set  $\mathcal{K}$  with  $|\mathcal{K}| \leq K$  and consider  $M$  i.i.d samples drawn from  $\mathrm{unif}(\mathcal{K})$ . Then with probability  $1 - \delta'$ , every element in  $\mathcal{K}$  appears at least once in these samples as long as  $M \geq K \log (K / \delta')$ .
Proof. Fix an element  $j \in \mathcal{K}$ . The probability that  $j$  appears in none of the  $M$  samples can be bounded by
$$
\left(1 - \frac {1}{| \mathcal {K} |}\right) ^ {M} \leq \left(1 - \frac {1}{K}\right) ^ {M} \leq \exp \left(- \frac {M}{K}\right) \leq \frac {\delta^ {\prime}}{K}.
$$
A union bound on every  $j$  in  $\kappa$  finishes the proof.
Lemma 7 (Concentration of Subgaussian random variables). Let  $X_{1} - \mu, \ldots, X_{M} - \mu$  be a sequence of independent 1-subgaussian random variables and  $\widehat{\mu} = \frac{1}{M}\sum_{m=1}^{M}X_{m}$ . Then with probability  $1 - \delta', \delta' < \frac{1}{2}$ , we have
$$
| \widehat {\mu} - \mu | \leq 2 \sqrt {\frac {\log (1 / \delta^ {\prime})}{M}}
$$
Lemma 8 (Bernstein's inequality). Let  $X_{1},\ldots ,X_{M}$  be a sequence of independent random variables with  $X_{m} - \mathbb{E}[X_{m}]\leq b$  almost surely for every  $m$  and  $v = \sum_{m = 1}^{M}\operatorname {Var}[X_m]$ . With probability  $1 - \delta^{\prime}$ , we have
$$
\sum_ {m = 1} ^ {M} X _ {m} \leq \sum_ {m = 1} ^ {M} \mathbb {E} \left[ X _ {m} \right] + \sqrt {2 v \log \left(1 / \delta^ {\prime}\right)} + \frac {2 b}{3} \log \left(1 / \delta^ {\prime}\right)
$$
The following lemma is a direct consequence of Bernstein's inequality, which states, in terms of the coupon collector's problem, that after drawing  $M$  coupons, one can expect with a high probability a particular coupon appears at least  $Mp / 2$  times if its probability of appearing is  $p$ .
Lemma 9. Let  $X_{1},\ldots ,X_{M}$  be i.i.d Bernoulli random variables so that  $\mathbb{E}[X_m] = p$  for all  $m$ . With probability  $1 - \delta^{\prime}$ , we have  $\sum_{m = 1}^{M}X_{m}\geq Mp / 2$  as long as  $Mp\geq 16\log (1 / \delta ')$ .
Proof. Let  $Y_{m} = 1 - X_{m}$  We have  $\mathbb{E}[Y_m] = 1 - p$ . By Lemma 8, with probability  $1 - \delta'$ , we have
$$
M - \sum_ {m = 1} ^ {M} X _ {m} \leq M \cdot (1 - p) + \sqrt {2 M p (1 - p) \log (1 / \delta^ {\prime})} + \log (1 / \delta^ {\prime}).
$$
Rearranging terms and using  $Mp \geq 16\log(1/\delta')$  gives
$$
\begin{array}{l} \sum_ {m = 1} ^ {M} X _ {m} \geq M p - \sqrt {2 M p (1 - p) \log (1 / \delta^ {\prime})} - \log (1 / \delta^ {\prime}) \\ \geq M p - \frac {M p}{\sqrt {8}} - \frac {M p}{1 6} \\ \geq \frac {M p}{2} \\ \end{array}
$$
We first provide an outline of the proof in Appendix C.1 and show the complete proofs of the lemma in the corresponding subsections.
To begin with, we assign adaptive target optimality (accuracy) to each block according to their block size:
$$
\varepsilon_ {b} = \max  \left\{1, \frac {1}{\sqrt {r \omega (b)}} \right\} \varepsilon , \qquad \text {f o r} b \in [ r ].
$$
The following lemma states that as long as  $\mathcal{W}$  contains an  $\widetilde{O} (\varepsilon_b)$ -optimal action for each block  $b\in [r]$ , then the output policy is  $\widetilde{O} (\varepsilon)$ -optimal with high probability.
Lemma 10. For a positive constant  $C$ , suppose for any  $b \in [r]$  with  $\omega(b) \geq \varepsilon / r$ ,  $\mathcal{W}$  contains a  $C\varepsilon_b$ -optimal action, then  $\pi^{\mathrm{out}}$  is  $2(C + 1)\varepsilon$ -optimal for the original context-lumpable bandit problem.
As a result, if we can prove that the precondition of Lemma 10 holds with high probability, then the correctness of Algorithm 1 follows immediately. To do so, we first argue that for every block  $b \in [r]$ , we have sufficiently explored the entire action set at accuracy level  $n = \lceil \log (1 / \varepsilon_b^2) \rceil$  in the data collection step, as stated in the next lemma.
Lemma 11. With probability at least  $1 - 2\delta$ , for any  $b \in [r]$  with  $\omega(b) \geq \varepsilon / r$ , we have that: for any  $j \in [K]$ , there exists  $(i, j) \in \mathcal{D}_n$  satisfying  $g(i) = b$  where  $n = \lceil \log(1 / \varepsilon_b^2) \rceil$ .
Intuitively, Lemma 11 says that we have tested all actions on each block  $b \in [r]$  at the corresponding accuracy level  $n = \lceil \log (1 / \varepsilon_b^2) \rceil$ . Based on Lemma 11, the following lemma states that in the second step we are able to filter out an  $\widetilde{O}(\varepsilon_b)$ -optimal action for block  $b$  by using the information contained in  $\mathcal{D}_n$ .
Lemma 12. With probability at least  $1 - 4\delta$ , for any  $b \in [r]$  with  $\omega(b) \geq \varepsilon / r$ , a  $10\varepsilon_b\sqrt{\log(rSK / \delta)}$ -optimal action of block  $b$  is added into  $\mathcal{W}$  in the process of shrinking  $\mathcal{D}_n$  where  $n = \lceil \log(1 / \varepsilon_b^2) \rceil$ .
Now we have proved the correctness of Algorithm 1. However, to obtain the desired sample complexity, it remains to show that the size of the optimal action candidate set  $\mathcal{W}$  is relatively small, which we prove next.
Lemma 13. With probability at least  $1 - 4\delta$ ,  $|\mathcal{W}| \leq Nr$ .
Now we are ready to prove Theorem 1 by combining all the above lemmas.
Proof of Theorem 1. The step of data collection uses
$$
8 \left(1 + \log \frac {1}{\varepsilon^ {2}}\right) \cdot \left(2 + \log \frac {1}{\varepsilon^ {2}}\right) \cdot \left(\log \frac {r S K}{\delta}\right) \cdot \frac {r (S + K)}{\varepsilon^ {2}}
$$
samples. By Lemma 13, the step of screening optimal action candidates uses
$$
2 (1 6 ^ {2}) \left(1 + \log \frac {1}{\varepsilon^ {2}}\right) \cdot \left(\log \frac {r S K}{\delta}\right) \cdot \frac {r S}{\varepsilon^ {2}}
$$
samples and solving the simplified context-lumpable bandit problem uses
$$
4 \left(1 + \log \frac {1}{\varepsilon^ {2}}\right) \cdot \left(\log \frac {S K}{\delta}\right) \cdot \frac {S r}{\varepsilon^ {2}}
$$
samples. Overall, the number of samples is bounded by
$$
5 2 4 \left(\log \frac {r S K}{\delta}\right) \cdot \left(1 + 2 \log \frac {1}{\varepsilon}\right) ^ {2} \cdot \frac {r (S + K)}{\varepsilon^ {2}}
$$
By Lemma 10 and Lemma 12,  $\pi^{\mathrm{out}}$  is  $2\left(10\sqrt{\log(rSK / \delta)} +1\right)\varepsilon$ -optimal.
C.2 Proof of Lemma 10
Proof. We control the suboptimality of  $\pi^{\mathrm{out}}$  in the following way:
$$
\begin{array}{l} \sum_ {i \in [ S ]} \nu (i) \left(\max _ {j \in [ K ]} A (i, j) - A (i, \pi^ {\mathrm {o u t}} (i))\right) \\ \leq \sum_ {b \in [ r ]: \omega (b) \geq \varepsilon / r} \sum_ {i \in \mathcal {B} (b)} \nu (i) \left(\max  _ {j \in \mathcal {W}} A (i, j) - A (i, \pi^ {\mathrm {o u t}} (i)) + C \varepsilon_ {b}\right) + \sum_ {b \in [ r ]: \omega (b) <   \varepsilon / r} \omega (b) \\ \leq \varepsilon + C \sum_ {b \in [ r ]} \omega (b) \varepsilon_ {b} + \varepsilon \leq 2 (C + 1) \varepsilon , \\ \end{array}
$$
where the first inequality uses the precondition of the lemma, the second one uses the  $\varepsilon$ -optimality of  $\pi^{\mathrm{out}}$  in the simplified context-lumpable bandit, and the final one follows from Cauchy-Schwarz inequality.
C.3 Proof of Lemma 11
Proof. Fix a block  $b \in [r]$ . For each accuracy level, recall in the step of data collection, we sample  $L = r(S + K) \widetilde{\lg} / \varepsilon^2$  contexts i.i.d. from  $\nu$ . By Lemma 9, with probability  $1 - \frac{\delta}{r}$ , at least  $L\omega(b)/2$  of them are from block  $b$  as
$$
L \omega (b) \geq \frac {L \varepsilon}{r} \geq \widetilde {\lg} \geq 1 6 \log (r / \delta).
$$
Further, recall that we define accuracy level  $n = \lceil \log (1 / \varepsilon_b^2) \rceil$ . We first show that this  $n$  is well defined, that is,  $n \geq 1$ . Indeed, we have
$$
\frac {1}{\varepsilon_ {b} ^ {2}} = \frac {1}{\varepsilon^ {2} / (r \omega (b))} \geq \frac {\varepsilon}{\varepsilon^ {2}} = \frac {1}{\varepsilon} \geq 2
$$
as long as  $\varepsilon \leq \frac{1}{2}$ . Since we add a context-action pair into  $\mathcal{D}_n$  once we have collected
$$
2 ^ {n} = 2 ^ {\lceil \log (1 / \varepsilon_ {b} ^ {2}) \rceil} <   2 ^ {\log (1 / \varepsilon_ {b} ^ {2}) + 1} = 2 / \varepsilon_ {b} ^ {2}
$$
samples for estimating its reward. Note that in the end there are at most  $|\mathcal{B}(b)|(2^n - 1)$  samples from block  $b$  unused. Thus, with probability  $1 - \frac{\delta}{r}$ , there are at least  $L\omega(b)/2 - |\mathcal{B}(b)|(2^n - 1)$  samples used. Therefore, the number of context-action pairs, where the contexts are from block  $b$ , that are added into  $\mathcal{D}_n$  is at least
$$
\begin{array}{l} \frac {L \omega (b) / 2 - | \mathcal {B} (b) | (2 ^ {n} - 1)}{2 ^ {n}} \geq \frac {L \omega (b) / 2}{2 / \varepsilon_ {b} ^ {2}} - S \quad \left(2 ^ {n} \leq 2 / \varepsilon_ {b} ^ {2} \text {a n d} | \mathcal {B} (b) | \leq S\right) \\ \geq \frac {L \varepsilon^ {2} / \left(2 r \varepsilon_ {b} ^ {2}\right)}{2 / \varepsilon_ {b} ^ {2}} - S \quad (\text {b y}) \\ = \geq 1 6 (S + K) \log (r S K / \delta) - S \quad (\text {t h e v a l u e o f} L) \\ \geq K \log (r S K / \delta). \\ \end{array}
$$
Conditioned on this event, with probability  $1 - \frac{\delta}{r}$ , for any  $j \in [K]$ , there exists  $(i,j) \in \mathcal{D}_n$  satisfying  $g(i) = b$  by Lemma 6. Therefore, the lemma holds for block  $b$  with probability at least  $1 - \frac{2\delta}{r}$ . We complete the proof by a union bound on all blocks.
C.4 Proof of Lemma 12
Proof. Denote by  $i$  the first context being removed from  $\mathcal{D}_n$ , which satisfies  $g(i) = b$ . By the rule of shrinking  $\mathcal{D}_n$ , we have
$$
\left| \widetilde {A} _ {n} \left(i, j ^ {\star}\right) - \widehat {A} _ {n} \left(i ^ {\star}, j ^ {\star}\right) \right| \leq 4 \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}}.
$$
By Lemma 7 and a union bound on all pairs in  $\mathcal{D}_n$ , with probability  $1 - 2\delta / r$ , we have
$$
\left| \widetilde {A} _ {n} \left(i ^ {\prime \prime}, j ^ {\prime \prime}\right) - A \left(i ^ {\prime \prime}, j ^ {\prime \prime}\right) \right| \leq \frac {1}{2} \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}} \quad \text {a n d} \quad | \widehat {A} \left(i ^ {\prime \prime}, j ^ {\prime \prime}\right) - A _ {n} \left(i ^ {\prime \prime}, j ^ {\prime \prime}\right) | \leq \frac {1}{2} \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}} \tag {5}
$$
for every  $(i^{\prime \prime},j^{\prime \prime})\in \mathcal{D}_n$  . Therefore, we have
$$
\left| A \left(i, j ^ {\star}\right) - A \left(i ^ {\star}, j ^ {\star}\right) \right| \leq 5 \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}}, \tag {6}
$$
By Lemma 11, we know that for any  $j \in [K]$ , there exists  $(i', j) \in \mathcal{D}_n$  satisfying  $(g(i'), j) = (b, j)$ , before we remove context  $i$ . Therefore, by the definition of  $(i^{\star}, j^{\star})$ , for any  $j \in [K]$ , there exists  $(i', j) \in \mathcal{D}_n$  satisfying  $g(i') = b$  and
$$
\widehat {A} _ {n} (i ^ {\star}, j ^ {\star}) \geq \widehat {A} _ {n} (i ^ {\prime}, j),
$$
which, again by Eq. (5) with probability  $1 - 2\delta /r$
$$
A \left(i ^ {\star}, j ^ {\star}\right) \geq \max  _ {j} \mu (b, j) - \sqrt {\frac {\log \left(r S K / \delta\right)}{2 ^ {n}}}. \tag {7}
$$
By combining Eq. (6) and Eq. (7), we conclude that  $j^{\star}$  is a  $10\sqrt{\log(rSK / \delta)} \varepsilon_{b}$ -optimal action for block  $b$  because
$$
5 \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}} \geq 1 0 \sqrt {\frac {\log (r S K / \delta)}{1 / \varepsilon_ {b} ^ {2}}} = 1 0 \sqrt {\log (r S K / \delta)} \varepsilon_ {b}
$$
This completes the proof by a union bound on all blocks.
C.5 Proof of Lemma 13
Proof. It suffices to show that for each accuracy level  $n$ , we will add at most  $r$  actions into  $\mathcal{W}$  before shrinking  $\mathcal{D}_n$  to  $\emptyset$ . Below we prove a stronger argument: each time we shrink  $\mathcal{D}_n$ , we will remove all the contexts from at least one block from  $\mathcal{D}_n$ .
Denote by  $(i,j)$  an arbitrary context-action pair from  $\mathcal{D}_n$  such that  $g(i) = g(i^{\star})$ , then by Eq. (5), we have
$$
\left| \widetilde {A} _ {n} (i, j ^ {\star}) - \widehat {A} _ {n} (i ^ {\star}, j ^ {\star}) \right| \leq \left| \widetilde {A} _ {n} (i, j ^ {\star}) - A (i, j ^ {\star}) \right| + \left| A (i ^ {\star}, j ^ {\star}) - \widehat {A} _ {n} (i ^ {\star}, j ^ {\star}) \right| \leq 4 \sqrt {\frac {\log (r S K / \delta)}{2 ^ {n}}},
$$
which implies all context-action pairs with context from block  $g(i^{\star})$  will be removed from  $\mathcal{D}_n$ .
C.6 Proof of Theorem 2
Proof. By Lemma 9 and a union bound over  $[S]$ , we have that with probability  $1 - \delta$  for all  $i \in [S]$ :
$$
\frac {1}{2} \left(\nu (i) - \frac {\log (S / \delta)}{J}\right) \leq \hat {\nu} (i) \leq 2 \left(\nu (i) + \frac {\log (S / \delta)}{J}\right),
$$
which implies that  $\nu(i) \geq \varepsilon/(4S)$  for  $\hat{\nu}(i) \geq \varepsilon/S$  and  $\nu(i) \leq 3\varepsilon/S$  for  $\hat{\nu}(i) < \varepsilon/S$ . By definition, for any  $l \in [L-1]$ ,  $\mathcal{X}_l$  consists of contexts such that  $\hat{\nu}(i) > 2^{-l-1} \geq \varepsilon/S$ . As a result, for any  $l \in [L-1]$ , by Lemma 9, we have that at least  $N\nu(\mathcal{X}_l)/2$  out of  $N$  samples corresponds to  $\mathcal{X}_l$  and thus can be used in learning  $\pi_l$  by executing Algorithm 1, where  $\nu(\mathcal{X}_l) := \sum_{i \in \mathcal{X}_l} \nu(i) \geq \varepsilon/(4S)$ .
Moreover, notice that inside each  $\mathcal{X}_l$ , the conditional context distribution is almost uniform, so we can invoke Theorem 1 to obtain that  $\pi_l$  is
$$
\frac {C \varepsilon}{\sqrt {\nu (\mathcal {X} _ {l})}} - o p t i m a l
$$