Context-lumpable stochastic bandits
Authors: Chung-Wei Lee
Qinghua Liu
Yasin Abbasi-Yadkori
Chi Jin
Tor Lattimore
Csaba Szepesvari
Venue: NeurIPS 2023 poster (2023)
Keywords: Contextual bandits
low-rank bandits
latent bandits
clustering bandits
stochastic bandit problems
context-lumpable bandits

Abstract: We consider a contextual bandit problem with $S $ contexts and $K $ actions. In each round $t=1,2,\dots$ 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\le \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 $\epsilon$-optimal policy after using at most $\widetilde O(r (S +K )/\epsilon^2)$ samples with high probability and provide a matching $\widetilde\Omega(r (S +K )/\epsilon^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{\text{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.

Context-lumpable stochastic bandits
Chung-War Wei Lee" University Southern California leechung@usc.ed
Qinghua Liu
Princeton University
ginghua @priicetoo..nneeeee
Yasin Abbasi-Yadaa
Chi Jin
Tor Lattimore
Google Deep yadkori8google..com com
Princeton University
Google DeepMind
chi wal.jeprnncetoneed
lattimore@@ooglee...
Csaba Szepesvári
Google DeepMind and University Alberta s szepesva@ualberta.co
Abstract
We consider contextual bandit problem with contexts and K actions. each round /. = 1, 2, the learner observes random context and chooses 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 lumped into r S 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 s-optimal policy after using at most (r((S + K))=²) samples with high probability and provide a matching (rr(S K)/E²) lower bound. In the regret minimization setting, give algorithm whose cumulative regret up to time T is bounded by (( KYT). To the best of our knowledge, we are the first to show the near-optimal sample complexity in the PAC setting and Ö(v (p)p))(())((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.
Introduction
Consider a recommendation platform that interacts with finite set of users in 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 reward, otherwise reward is obtained. Assume that the users can be partitioned into small number 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 fewew groups, even when the identity of the group user belongs to unknown and only learnable because
they share preferences with other users in the group.
slightly more general version of this problem can be formalized as follows: Viewing users as "contexts" and recommendations "actions" (or arms), assume that there are S contexts and K actions. round 1,2,... the learner first receives a context i,, sampled from an unknown
most works were done when interning at DeepMind.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

distribution on the set [S] := {l....,,t S} of possible contexts. The learner then chooses an action
de E [K] := {I..... .......K }} and observes reward y/: A(i + The where given the past, T)e has subgaussian tail (precise definitions are postponed to Section 2 and A [S] x [K] - IR. is an unknown function of mean rewards (RR denotes the set reals). We consider two settings when the
goal of the learner is either to identify near-optimal policy x [S] + [K]. to keep its regret small.
Policy ı is called s-optimal if
[i, T (b)) 2 max ([A((,, (11))
x'
while the regret of the learner for horizon of T is
T Σ
T max () - Σ A(in. JE[K]
Regr- = E
(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., Lattimooe and Szepesvári (2020) that there algorithms such that e-optimal policy will be discovered after
Ō SK 2
(3)
interactions, and there are also algorithms for which the regret satisfies
Regr Ö(VSKT)
Here, the notation Ō(-) 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-lumpabes
if there is a partitioning of [S] into r I min [S, K} groups such that
A(i.-) = A(t', holds whenever i, E [S] belong the same group. It is not hard to see that any algorithm needs at least Q(r(S + K))f²) interactions to discover an e-optimal policy (Theorem 3). Indeed, if we view A as an S K matrix, the lumpability condition states that A UV where U is an S X r binary matrix where each row has a single nonzero element, and V is an r x 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 S. in Eq. (3) and Eq. (4) can be replaced with (S + K ())ol((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 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 more versatile as the users are modeled as mixtures of r preference types instead of belonging to one type only.
1.1 Related works
Our problem can be seen an instance of contextual bandit problems introduced by Auer et al. [2002). For a good summary the history of the contextual bandit problem, the reader 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. [20TT]. However, lumpablility 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- -ummpebtttuubbeeeeeeee stochastic bandit problem. The earliest work is probably that of Maillard and Mannor [2014) who named this problem latent bundits: 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 Iearner. Further, the results derived are instance dependent.

---

Altogether, while this paper makes interesting contributions, it does not help in answering our questions [
The earlier work of Gentile et al. [014 considered another generalization of our problem where in round f 1, ... the learner gets action set ( C IR and the mean payoff of choosing action j S K2 given the past and the current context &, is 1 (o(i)) with 0 ......,6 0, I unknown. When specializing this to our setting, their result depends on the minimum separation ming, ( ||A --A()--)-))) [l. which, according to the authors may be 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, 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- -on-lumpable problems. [L 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 the previous paper. [enttle 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 be partitioned such that states within group of partition behave identically, both in terms of the transitions and the rewards received [e.g., Ren and We put survey of this line of work in Appendix A.6
Krogh 2002
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 ((ppoly())(S + K)T), which we believe is fundamental question in this area.
2 Notation and problem definition
For an positive integer n, we use [n]] to denote the set {l...., n). Further, for a finite set U, A(U) denotes the set of probability distributions over U and unif(U) denotes the uniform distribution over U. The set of real numbers is denoted by R.
As introduced earlier, we consider context-lumpabbe 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 E [K] is used while the context is i E [S] depends only on the identity of the group that i belongs to and the action j. That is, for some g : [S] + [r] map
A) = A((
for any j E [K],i, E [S] such that g(i) = g(())
Neither A, nor is known to the learner who interacts with the bandit instance in the following way: In rounds 2,..... the learner first observes a context in E [S] randomly chosen from a fixed, context distribution v E A((S)), independently of the past. The distribution v/ is also unknown to the learner. Next, the learner chooses an action jr E [K] to receive a reward of yl = A(t + T/ where t/2 is 1-subgaussian given the past: For any X E R,
E(exp(>m) i ft.J.................... ... Lt. Jol I exp(>²/2) almost surely (a.s.).
As described earlier, we are interested in two problem settings. the PAC setting, the learner is given a target suboptimality level æ > 0 and a target confidence ó The learner is then tasked to discover a policy क. [S] + [K] that is e-optimal (cf. Eq. (L). with probability 8. 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 á : [S] + [K]. If T is the (random) time when the learner stops, then the efficiency of the learner is characterized by E In the online setting, the goal of the learner 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
SSoeewat confusingly, Hong et al. 20200 define latent bandit problems differently from the definition given by Maillard and Mannor [2044]: They assume that in addition to the context, 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 answering our question.
' Lumpability was first introduced first by Kemeny and Snell 1976 in the context of Markov chains.
3

---

action for In the remainder of the paper, we will also need some further notation. The map induces partitioning p [1((),. B(r) of [S] in the natural way:
(b))= = (i: =
We call each subset B P** a block and call ((b) block for every b [r]. We also define block reward j) = A(i, j) of block b for every context E B(() and arm j E [K]. Finally, we define
the block distribution w E (r)) so that w(b) = Eieso) (i)).
3 Near-optimal PAC Learning in Context-lumpall 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 v(i) [1/(85), 8/5] for all i E [S] throughout this subsection. The pseudocode of our algorithm is provided in Algorithm I which consists of three main modules. Below we elaborate on each of them separately.
The high-level goal of this step is to collect sufficient amount of data for every block/action pair so that in later steps we have sufficient information to select opar-optimal arm for each block. One naive approach is to try every action on every context for sufficient amount of time (e.g., O(1/=))) 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) j) E [r] x [K], only need to sufficiently try action j on a single context i from block (i.c., E B(6)). However, the block structure is unknown priori. We circumvent this problem by assigning random action ()(i) each context for them to explore (Line 4). And after action vb()) has been tried on context sufficiently many times, we update (i) by reassigning context i with new random action (Line 10
Data collection (Line 3 and Algorithm 2
However, there still one key problem unresolved yet: how many samples to use for estimating each A(i, (i)) before reassigning context i with another random action, given fixed budget of total sample complexity. On one hand, if we only try each (i, (i)) pair a very few numbers of times before switching. then can potentially explore 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, (i)) with huge number of samples resolves the poor accuracy issue but could potentially cause the problem of under-exploraatio, especially for those blocks consisting of very few contexts. In the extreme case, if a block only contains 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 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 [log(1/e²)], and consider accuracy levels [N]. For accuracy level n, we collect a dataset D, by adding a context-action pair after the pair is played for 2n timesteps (Line 9.
Screening optimal action candidates (Line 1 Equipped with the dataset collected from Step 1, we want to identify a subset W of [K] that (i) for every block b E [r]. W contains near-optimal action of block b, and (ii) the size of W is as small as possible. We construct such W in an iterative way. For each accuracy level n [N]. we first compute the context-action pair (i" ((*)j") with the
highest reward estimate in Dn (Line 7) Intuitively, as (i*,j*) has been tried for 2" timesteps in
constructing Dn,, is natural to guess that jr could potentially be

---

certain blocks so we add it into optimal action candidate set W (Line 8). To identify the contexts from those blocks, we sample new data to estimate the reward of j* on each context E [S]. This can be done by calling Algorithm and setting the exploration set K. to be j (L.nee If context can achieve reward close enough to ( jJ") ) This is because we have found j* as an optimal action candidate for i at this accuracy level and don't t need to consider other j. We repeat the above process, and add new arm to W in each iteration, until D,, becomes empty.
action j" then peel off every (t,j) Dn (Line 11].
Solving the simplified context-lumpable bandit (Line 12 After obtaining the optimal action candidate set W, we can discard all actions not W and solve simplified -umpxx-luuppble bandit problem by replacing the original action set [K] with W. Note that ((S(W)/e)) samples suffice for learning an s-optimal policy for this simplified problem. For example,
directly try each action j E W on each context E [S] for (1/e²2 times, and define xou (i) arg maxjew A(i,j) where
A(i,)) is the empirical estimate of A(i,j) based on @(1/=²) samples.
Algorithm Algorithm for Almost-uniforn Context Distribution
Let i denote the current time and initialize N <- [log(1/e")], W <- 0, Ig 16 16log(rSK/0)
2 Define [[(i) - [ for E [S], L + r(S - K)lg/E²
Step 1. Data collection
3 for accuracy level = ......... do
L
Execute Algorithm with input L, n, K, and receive Dn and An
Step 2. Screening optimal action candidates for accuracy level n = ......NN do while D,, * 0 do
Compute (i* < arg (.)DD (i.jj)
Update optimal action candidates W < WU[j"}
Update K((i) [i+} for [S] and T/' <- 8lg2"
Execute Algorithm 2 with input L/, n, K., and reassign output to D, and Ān
10
Shrink Dn < {(i,j) E Dr An(i.,)) - An
11
Step 3. Solving the simplified context-lumpable bandit
Use 48/W]
gx log SK
8 samples to find out s.t. A(i, out (i)) maxjew Ali.j) - for all e [S]
Output out
Algorithm 2 Data Collection
Input L, n, I.
2 Let t denote the current time and initialize D,, + for context i [S] do
L
Assign ehe to be an arm drawn from unif(XC(i))
for timesteps do
Receive context it. play arm je = (tt)) and receive reward Preset (1) (i) for every i E (S)
if (E 10, 12))2" 0 then
D, + D, Ui(e. (i.)) and Å, ((+)) + Erce = Reassign (1r) ~ uni(())(i))):
Enge 111.=h.,,v(4)
10
Output Dn and A.
Now we present the theoretical guarantee for Algorithm I bound can be found in Appendix
The proof and exact constants in the
Theorem 1. Let å E (0, 1) and assume v(i E [L/(8S), (/S) for all i e [S], Algorithm outputs an
-optimal policy within Ö + K) log L/6) samples with probability least - 8.
5

---

3.2 Extension to general context distribution
Algorithm 3 Algorithm for General Context Distribution
Let 1S log( (S/o), L = [log( (, N 524 (log ) (1+ - 2 log 1 =lst z
c²
Estimate the context distribution by sampling J contexts, and denote the estimate by 0
3 Split the context set into L disjoint subsets (2i )e14) where
{i E [S] ((i)= E (2 (2-1-1 12
[ L --1]
X
- {i [5] 1(1) E (1,
1
4 for 1 [L - 1] do
L Execute Algorithm Ito learn policy गı for subset Ai from N time steps
Output out such that nut equals to over X₁ for 1 E [0 L 1]. and arbitrary over Xr.
In this subsection, we show how to generalize Algorithm 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
(SSE)) samples to obtain an empirical estimate of the context distribution, denoted by D. Then we divide the context set into many disjoint subsets he such that inside each subset X₁, the conditional empirical context distribution almost uniform. As result, we can invoke Algorithm to find a near-optimal policy Ti for each subset X₁ (Line 5) requires Ö(r(S + K))=²) time steps for every but only use samples where contexts are from X₁ and ignore the rest. Finally, we glue all Ny together to get policy gour that is near-optimal for the original problem. Formally, have the following theoretical guarantee for Algorithm 3 The proof and exact constants are in Appendix([]
Theorem 2. Let 8 1), Algorithum 3 outputs an Oss)-mitimal policy within Ö(r(S K) log(1 log(1/))/e²) samples with probability at least 1-5.
Note that we can always learn e-optimal policy for any context-lumpable bandit within Ö(SK/s²) samples by invoking any existing near-optimal algorithm for finite contextual bandits. As a result, by combining Theorem 2 with the OSSK/ O(SK/²2) upper bound, we obtain that ((min(r(S) + K), SK)////) samples suffice for learning any context- lompext-lumpable bandit, which according to the following theorem is minimax optimal up to a logarithmic factor.
Theorem 3. Learning an s-optimal policy for a context-lumpatte bandit with probability no smaller than 1/2 requires at least Qimin K),SK)/E") samples in the worst case.
4 Regret Minimization in Context-lumpahle Bandits
In this section, we extend the idea from the PAC setting to the online setting. To better introduce the key ideas, we first consider a special case when both context and block distributions are uniform (Section 4.11. Then we consider the most general case in Section 4.2
4.1 Special Case: Uniform Context and Block Distribution
In this section, we assume that distributions and are uniform, and thus, g evenly splits the contexts into r blocks so that there are S/r contexts in each block and every context appears with the same probability at each timestep. We will relax these assumptions and consider the general case in the next subsection.
For this case, we introduce Algorithm which uses phased elimination in combination with a clustering procedure. The algorithm runs in phases h = 2,2.... that are specified by error tolerance parameter ε = 2h/2 Like phased elimination algorithms for multi-armed bandits, we need to ensure at phase h actions the algorithms play are all O(Ex)-optiall Thus, at the end of each phase h, eliminate all actions that are not each block is different. Therefore, we also perform clustering on contexts and perform elimination for each subset of the partition. Specifically, we maintain a partition of contexts Pa at each phase h and initialize p1 {[]} For each cluster C E Ph. we maintain set of good arms (C),, which we will prove is ((E)) ooptimll for contexts in the cluster with high probability.
However, the set of ) actions of

---

We use Algorithm 2 to collect data similar to Algorithm 1 (Line 6p. At phase h, we try to estimate mean reward up error Ō(Ea) with probability 1 6%. The difference is that assume w is uniform, so we don't need different accuracy levels n, which will be required for the algorithm that handles the general case. Also, for every context E C, only explore arms good for now, that is, in GOOD,(C) instead of exploring all the arms [ This reflects that in the online setting, we need to minimize regret and cannot afford to explore bad arms too many times.
Based on the data we collect during the exploration stage, we then check if there is large gap across contexts in the same subset for any arms (Line 7. A large gap suggests that (i) the subset contains at least two blocks (ii) the mean reward of the arm is significantly different these blocks and we can use this arm to partition the contexts by running Algorithm (clustering stage), We repeatedly do clustering (Line 9) and split heterogeneous subsets (Line 10 until we cannot find a large gap within the same subset. If no large gap is detected, then each arm has similar mean rewards (up to error ((E)) )) across all blocks in the same subset. Then eliminate arms that significantly worse than the empirical best arm (elimination stage) from (OOD,(C) for every subset C. and start a new episode (Line I4p.
Algorithm 5 plays j for each context i C for Õ(1/²²) rounds and calculates empirical means of arm j for each context in C. It then sorts the contexts by the empirical means and performs clustering (Line 4 Specifically, the algorithm enumerates contexts in descending order of empirical means and splits contexts until a large gap between consecutive values is detected (Line 7 As we call Algorithm[]jnl if large gap is detected, we prove that in the appendix correctly separates the subset into at least two parts without putting any contexts in the same block into different parts by setting 5' = Eh/T.
Algorithm Al Algorithm for Uniform Block Distribution
Initialize Pı +- {[S]), (COD,(C) < [K] for C Pi 2 Let t denote the current time
for phase h- 1,2,.... do
Let En + 2-/2 8n + / Ig, < 641oo(rSKK6)), + V EE
Step 1. Data collection
Define Li. + r(S-K)lss nh + log(
GOOD C Ph, Ph. Caii
Execute Algorithm with input Ln, Rh, Kr. and receive DD, and Ån Step 2. Test homogeneity and perform clustering on heterogeneous subsets while 3C E Ph, P E [K] such that (i,j) - An (i.j) 2 En do 8' - EA K(i) +- {5} if C else GOOD, for C Ph. ii Eh 4r
Define 5'
Execute Algorithm with input s', 8', KC, C, and j. and get P. partition of C Initialize GOOD, (C') GOOD, (C). VC' P and update Ps. + (PUPA)
10
Step 3. Eliminate suboptimal actions in each subset
11
PA+1 + P/ for C E Ph+1 do
12
13
Calculate T4, (C,j) <- (ax))E (n,,, Vj E (O)D,
14
Update GOOD, (C) + {j:3 j E (O)D, maxy Th (C,5)) - (C, & 25,}
Similar the analysis of other phased elimination algorithms, have show that in phase specified by error level cτ. with high probability, (i) the optimal arm is not eliminated and (ii) all ((bb)tsuba arms are eliminated, that is, all arms in GOOD, (C) for all C are Ö(Ex)) We show the final regret here and defer the details to Appendix [
Theorem 4. Under the assumption that context distribution v and block distribution w are uniform, regret of Algorithm 6 is bounded as Reg, = (( KYT).
7

---

Algorithm 5 Split Contexts Into Multiple Blocks
Input: error e', confidence 8', exploration sets K. subset C, separating arm j
Initialize Ig < 64 (S/8"), I/ +- Sig/e?, n' + [1/621
3 Execute Algorithm 2 with input I/, , K, and reassign output to 'D and À
Sort contexts in C and label them as i,, djj so that A(ir.j) 2 A(i2.j) 2...2 A(i
5 Initialize P1 + {i1}. b 4-1
for k 2,3, ...
: |C] do
if A(ix-1.5) ( A(ix, 2 Vig then L Update b+ + b 1 and initialize Po+ +- {
Pa + Pb Luul0 }
10 Output (Pi....,PP) Pb)
Compared to our PAC result, we get an extra dependency on r. This is because we pay Ö(S/=²) S(e)) samples to do clustering instead of ((1/²)) in order to ensure a "perfect" partition, that is, we never separate contexts in the same block with high probability. This is crucial for our algorithm as we may call Algorithm | different phases. We left getting better than (( (S + K)T" regret upper bounds or better than @(vr(S+ + KYT) regret lower bounds (even for this uniform special case) as an important future direction.
4.2 Non-uniform Context and Block Distribution
Similar to the PAC learning setting, we can use Algorithm to reduce general context distributions to (nearly) uniform ones. We provide more details in Appendix As a result, without loss of generality, we may assume y is damost-unifonne and we focus on how to handle non-uniform block distribution w. In the extreme, there may only be one context for some blocks, which becomes challenging to estimate their mean rewards.
For this case, we introduce Algorithm 6 Intuitively, based on Algorithm different accuracy levels as Algorithm|] I However, as our goal becomes minimizing regret, it is difficult to control regret for smaller and larger blocks. Specifically, for smaller n, we explore more actions (with fewer samples) for each context in a single phase. Since fewer samples are used, we may unavoidably play suboptimal actions and suffer large regret. This becomes worse for large block as more contexts in the block suffer large regret. We note that this is not a problem in the PAC setting because only need to control sample complexity.
4 we can further employ
We fix the issue by setting different lengths L for different accuracy levels. Recall in Algorithm I we use the same length = Ö (((r(SSK//// e²) for every n. Intuitively, since we allow less accurate estima- tions for smaller n, we may use fewer data. It turns out indeed we can set L = Ör(( + K ())(m=bi/2) for level n at phase h, and with more refined analysis, provides similar guarantees as before. More importantly, since smaller n. uses (exponentially) fewer samples, the overall regret is well controlled.
In addition to setting the lengths sophisticatedly, also need to carefully maintain sets of good actions GOOD,,,. not only at each phase h but also at each accuracy level n. The partition of contexts Pr. however, only evolves with phases and is shared between different levels at the same phase. Similar to Algorithm 4 we also do clustering (Step 2) and elimination (Step 3) but do the procedures in parallel for every accuracy level n. In the clustering stage, use different thresholds to define "large gap". This reflects that n represents different accuracy levels and thus has different widths of confidence intervals. For the elimination stage, we enforce the inclusion relation GOOD, (n, C) C GOOD, (n', C) for n' < (Line (7), which will be useful in the regret analysis. We now present the main theorem and put the complete proof in Appendix G
Theorem 5. Regret of Algorithm bounded as Reg = 0 (S KKT)
We remark that Algorithm 6 is anytime, that is, it doesn't require the time horizon T' as input. However, does require the number of blocks as prior knowledge. Removing the knowledge of is an interesting future direction. One promising idea is to apply a doubling trick on r. Specifically,
we have initial guess Ō(1) and run Algorithm 6 when [Pa] > double and restart
8

---

Algorithm 6 Algorithm for Non-uniform Block Distribution
Initialize P1 < {[S]), 1 +- [K], E Pi
Let t denote the current time
for phase
1, 2........ Let En + 2 h/2 No +- h, 6n + (2/(195K), IgA < 128 128
Step 1. Data collection
for accuracy level = I,..., ......,NN do
Define Ln,n r(8 + K) 2 Kn, (3) GOOD Ph. C. ii vi [S] Execute Algorithm 2 with input Ln n, Kr,n and receive a,, and Ap,
Step 2. Test homogeneity and perform clustering heterogeneous, subsets
while 3C E Ph,
[K], n E [Na]] such that An, (i.j) (i.)) 8' _, (() + {j} if C clse GOOD,, (C) for Ph. C a Define E' + 4..
do
10
Execute Algorithm with input s', 8', K, and j. and get P, partition of Initialize GOOD,,,. + G.OD (C). AC' E P and update Ph + (PUPWMCC)
11
Step 3. Eliminate suboptimal actions in each subset
12
Ph+1 + Ph. GOODOOOO,,, (C) t [K]. Ph+1 for accuracy level 2, .... N₂, do for C Ph-1 do
13
14
15
Calculate n (C, < (A)))E. ( Vj E GOOD)..(C)
16
Let Gh. (C) < j (O)), maxy Uh j') Pa, (C, S 2/ Update GOOD + GOOD Gh,
1gh
17
the algorithm. The analysis, though, may be much more complicated. Finally, note that when knowing r, calculate in advance if VSKT is less than vr KýT and switch standard algorithm achieving Ö(VSKT) in that case. This modification guarantees the regret is never worse than the standard bound.
5 From Context-lumpable Bandits to Contextual Low-rank Bandits
Finally, we consider the more general contextual low-rank bandit problem. Specifically, we allow to have rank r, that A UV for some r matrix U and r K matrix V. Lumpability is special case in the sense that U is binary where each row has single nonzero element.
To solve the problem, We show reduction from contextual low-rank bandits bandits. Consider an a-covering of rows of U, and notice that the covering number, R., is 1/a", The context-lumpable bandits can be seen a-approximate context-lumpable bandits with R., blocks, where the reward of contexts on the same block differs at most a. Ignoring this misspecification, we may run Algorithm T and Algorithm 6 for the PAC and online settings, respectively. Moreover, turns out that our algorithms are robust to this misspecification when a is sufficiently small (O(s) in the PAC setting, for example). Therefore, the sample complexity and the regret bounds of these algorithms will be in terms of R,, despite having an exponential dependency on r. The resulting bounds can still be smaller than the naive SK /22 and VSKT bounds in some scenarios, for example, when S and K are super large. We put the details in Appendix
-untext-lumpable
Conclusions and Future Directions
We consider contextual bandit problem with S contexts and K actions. Under the assumption that the context-action reward matrix has S min{S, K) unique rows, we show an algorithm that outputs e-optimal policy and has the optimal sample complexity of Ö(r(S + K)/=²) with high

---

probability. In the regret minimization setting, we show an algorithm whose cumulative regret up to time T' is bounded as (S KYT).
An immediate next question is whether regret bound of order O(vr(5+ K)T) is achiev- able in the regret minimization setting. A second open question is concerned with obtaining Ō( D(Vpolyy (p + K)T) regret bound in contextual low-rank bandits. Our regret analysis heavily relies on the assumption that contexts arrive in an L.I.D fashion. Extending our results to the setting with adversarial contexts remains another important future direction.
References
Yasin Abbasi-Yadkord Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. Advances in neural information processing systems, 24, 2011.
Sanjeev Arora, Rong Ge, and Ankur Moitra. Learning topic models-going beyond svd. In 2012 IEEE 53rd annual symposium on foundations of computer science, pages 1-10. IEEE, 2012.
Auer. Using confidence bounds for exploitation-exploraiion trade-offs. Journal of Machine Learning Research, 3:397-422, 2002.
Peter Auer, Nicolo Cesa- Cesa-Bianchi, Yoav Freund, and Robert E Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-71 2002.
Alina Beygelzimer, John Langford, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandit algorithms with supervised learning guarantees. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 19-26. JMLR Workshop and Conference Proceedings, 2011.
Leonardo Cella, Alessandro Lazaric, and Massimiliano Pontil. Meta-learning with stochastic linear bandits. Arxiv, 2020.
Yuxin Chen, Yuejie Chi, Jianqing Fan, Cong Ma, and Yuling Yan. Noisy matrix completion: Understanding statistical guarantees for convex relaxation via nonconvex optimization. SIAM journal on optimization, 30(4):3098-31211 2020.
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff functions. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 208 214. JMLR Workshop and Conference Proceedings, 2011.
Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert Schapire. On oracle-effiieent PAC RL with rich observations. Advances in neural information processing systems, 31, 2018.
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient RL with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665-1674. PMLR, 2019.
Yaqi Duan, Tracy Ke, and Mengdi Wang. State aggregation learning from Markov transition data. Advances in Neural Information Processing Systems, 32, 2019.
Miroslav Dudik, Daniel Hsu, Satyen Kale, Nikos Karampatziakis, John Langford, Lev Reyzin, and Tong Zhang. Efficient optimal learning for contextual bandits. In Proceedings of the Twenty-Sevents Conference on Uncertainty in Artificial Intelligence, pages 169-178, 2011.
Fei Feng, Ruosong Wang, Wotao Yin, Simon S Du, and Lin Yang. Provably efficient exploration for reinforcement learning using unsupervised learning. Advances in Neural Information Processing Systems, 33:22492-225044 2020.
Dylan Foster and Alexander Rakhlin. Beyond ucb: Optimal and efficient contextual bandits with regression oracles. In International Conference on Machine Learning, pages 3199-3210. PMLR, 2020.
10

---

Claudio Gentile, Shuai Li, and Giovanni Zappella. Online clustering of bandits. In International Conference on Machine Learning, pages 757-765. PMLR, 2014.
Claudio Gentile, Shuai Li, Purushottam Kar, Alexandros Karatzoglou, Giovanni Zappella, and Evans Etrue. On context-dependen clustering of bandits. In Proceedings of the 34th International Conference on Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, pages 1253-1262, 2017.
Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, and Craig Boutilier, Latent bandits revisited. In NeurIPS, 2020.
Prateek Jain and Soumyabrata Pal. Online low rank matrix completion, 2022.
Prateek Jain, Praneeth Netrapalli, and Sujay Sanghavi. Low-rank matrix completion using alternating minimization. In Proceedings of the forty-fifth anmual ACM symposium on Theorr of computing. pages 665 674, 2013.
Kyoungseok Jang, Kwang- Swang-Sung Jun, Se- Young Yun, and Wanmo Kang. Improved regret bounds of bilinear bandits using action space analysis. In International Conference Machine Learning. pages 4744 4754. PMLR, 2021.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contex- tual decision processes with low Bellman rank are PAC-learnable. In International Conference on Machine Learning, pages 1704-1713 PMLR, 2017.
Kwang- Sung Jun, Rebecca Willett, Stephen Wright, and Robert Nowak. Bilinear bandits with low-rank structure. International Conference on Machine Learning, pages 3163-3172. PMLR, 2019.
Yue Kang, Cho-Jui Hsich, and Thomas Chun Man Lee, Efficient frameworks for generalized low-rank matrix bandit problems. In Advances in Neural Information Processing Systems, 2022.
Sumeet Katariya, Branislav Kveton, Csaba Szepesvari, Claire Vernade, and Zheng Wen. Stochastic rank-1 bandits. In Artificial Intelligence and Statistics, pages 392-401. PMLR, 2017.
John G Kemeny and James Laurie Snell. Finite Markov chains. Springer, 1976.
Branislav Kveton, Csaba Szepesvári, Anup Rao, Zheng Wen, Yasin Yadkori, and S Muthukr- ishnan. Stochastic low-rank bandits. arXiv preprint /7Xiv: 2017.
Branislav Kveton, Martin Mladenov, Chih-War Hsu, Manzil Zaheer, Csaba Szepesvári, and Craig Boutilier. Differentiable meta-learning contextual bandits. arXhv:2006.05094v1 2020.
Branislav Kveton, Mikhail Konobeev, Manzil Zaheer, Chih wei Hsu, Martin Mladenov, Craig Boutilier, and Csaba Szepesvari. Meta- Meta-Thompson sampling. Arxiv, 2021.
Jeongyeol Kwon, Yonathan Efroni, Constantine Caramanis, and Shie Mannor. RL for latent MDPs: Regret guarantees and lower bound. Advances in Neural Information Processing Systems, 34: 24523-24534, 2021.
Tor Lattimore and Botao Hao. Bandit phase retrieval. Advances in Neural Information Processing Systems, 34:18011-18811, 2021.
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
Shuai Li, Alexandros Karatzoglou, and Claudio Gentile. Collaborative filtering bandits. In Proceed- ings of the 39th International ACM SIGIR conference on Research and Development in Information Retrieval, pages 539-548, 2016.
Yangyi Lu, Amirhossein Meisami, and Ambuj Tewari. Low-rank generalized linear bandit problems. In International Conference on Artificial Intelligence and Statistics, pages 460-468. PMLR, 2021.
Odalric- Odalric-Ambrym Maillard and Shie Mannor. Latent bandits. In International Conference on Machine Learning, pages 136-144. PMLR, 2014.

---

Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstrac- tion and provably efficient rich- observation reinforcement learning. In International conference on machine learning, pages 6961-6971. PMLR, 2020.
Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank MDPs. arXiv preprint arXiv:2102.0700 2021.
Chengzhuo Ni, Anru R Zhang, Yaqi Duan, and Mengdi Wang. Learning good state and action representations via tensor decomposition. In 2021 IEEE International Symposium on Information Theory (ISIT), pages 1682-1687. IEEE, 2021.
Matteo Papini, Andrea Tirinzoni, Aldo Pacchiano, Marcello Restelli, Alessandro Lazaric, and Matteo Pirotta. Reinforcement learning in linear MDPs: Constant regret and representation selection. Advances in Neural Information Processing Systems, 15:16371-16385 2021.
Zhiyuan Ren and B.H. Krogh. State aggregation Markov decision processes. In Proceedings of the s1 IEEE Conference Decision and Control, volume 4, pages 3819-3824, 2002.
Rajat Sen, Karthikeyan Shanmugam, Murat Kocaoglu, Alexandros G. Dimakis, and Sanjay Shakkottai. Contextual bandits with latent confounders: An NMF approach, 2016.
David Simchi-Levi and Yunzong Xu. Bypassing the monster: A faster and simpler optimal algorithm for contextual bandits under realizability. Mathematics of Operations Research, 4703):1904-1931,
2022.
Aleksandrs Slivkins. Introduction to multi aedti-armed bandits. Foundations and Trends@ in Machine Learning, 12(1-2):1-288 2019.
A. Tewari and S. A. Murphy. From ads to interventions: Contextual bandits in mobile health. In Mobile Health, pages 495-517. Springer, 2017.
Cindy Trinh, Emilic Kaufmann, Claire Vernade, and Richard Combes. Solving Bernoulli rank-one bandits with unimodal Thompson sampling. In Algorithmic Learning Theory, pages 862-889.
PMLR, 2020.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline RI. in low-rank MDPs. In International Conference on Learning Representations, 2022.
Amy Zhang, Shagum Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstrac- tions for hidden-pard parammeer block MDPs. In International Conference on Learning Representations, 2020.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient representation learning in low-rank Markov decision processes. arXiv preprint arXiv:2106.11935 2021.
Xuezhou Zhang, Yuda Song. Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. Efficient reinforcement learning in block MDPs: A model-free representation learning approach. In International Conference on Machine Learning, pages 26517-26547, PMLR, 2022.
12

---

A More Related Works
A.1 From bilinear to low-rank bandits
The low-rank structure has been shown useful in other bandit problems. Katariya et al. [017]] considered the case when the learner chooses both "context" and arm and the reward map A, when viewed as S x K matrix, has rank one. This is a special case of linear bandit setting with 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 x 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. veeon et al. [017] generalized the setting Katariya et al. 2017] to the case when the reward A (viewed as matrix) is of rank r S 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 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 + polyy(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 Ö((S + (T), though with inefficient algorithms, which was the first bound better than the naive bound Ö(SK Later, Lu et al.] [202T] 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
constant Frobenius norm bound). The regret bound of Jun et al.) (2019). Lu et al. [202]] takes the
form ( + (x)33//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 SVT.
Recently, Kang et al. [022] improved the previous state-of-the-arr for the generalized linear setting. The new regret bound, which they prove "under a mild" extra assumption, takes the form Ö(SS K) /T) and is conjectured to match the order of the minimax regret. Noticing that the minimax lower bounds developed for the finite-armed stochastic handits (e.g., Exercise 15.2 of Lattimore and Szepesvári 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 (KVSKT) in these problems. As VS+ K ~ /max(S, K) < VSK= < / //max(S,FX) K)² = max(S, (S K) = S + K, this lower bound rules out upper bounds of the form O(v + ()),ll(r)TT, though it is compatible with the upper bound mentioned a
Jain and Pal 2022) study related problem where the learner chooses 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 7o/ppoyyogg(S + K') provided some technical conditions hold, one of which is that the condition number of the reward A when viewed as 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 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 leta-learnig is concerned with learning a lower dimensional subspace across bandit tasks Kveton et al.) 2020 2021 Cella 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 lincar 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 (e.g.. Theorem 24.2. Lattimore and Hao 2021). If the action set has one-hot matrices only, this bound improves to v VdT. as discussed before.
Heree f(S, K) z g(S, K) means that f. 9 are within a constant factor of each other, while ((S << 9(S K) means that f(S, K) o(g(S, K)) S, K 00.
13

---

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 online problem. In 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 choose an action in each round with the goal of competing with the total reward collected by the best expert in hindsight. |uuer et al. (2002) consider the adversarial case when in each round, each action assigned reward in an arbitrary way from a known, finite interval. The authors propose the EXP4 algorithm that is shown to achieve O(VKT log N) regret in 7 rounds of interactions. [Beygelzimer et al. [201]] 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 expert is a [S] - [K] map that is the
composition of an [S] - [r] map followed by an [r] + [K] map. Denoting by N the number of
such experts, see that log(N) S llog(r) + log( K), and hence the regret of EXP4 is of order
(VVSKT), 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 2022]. In this setting, the learner has access to function class F and the optimal regret Levi and Xu
is O(V KT log( [F) ) where [F] is the size of the function class. However, with the same argument, we can conclude [F| = ((S) even under the lumpable assumption. Therefore, the results in this direction also give 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,... the learner first receives K d-dimensional vectors, Ft,1,---,,f,K6 each corresponding to action. Choosing action gives a reward whose mean (given the past) is o for some unknown parameter vector A E Rd. For case, one can choose SK: 0 can be the "flattening" of A and zt.,j is a unit vector so that = 0 = A(i,, Applying the SupLinnel algorithm from Section 4 of Auer [2002] to this setting, we get the regret bound O( d'Tlog' (KT)), which shows improvement compared 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 AAoor et al. 2012 Jain al., 2013, Chen et al., 2020). Based on these ideas. [Sen et al.) [0166 propose an epsilon-greedy algorithm for the contextual low-rank bandit problem and show that its regret is of order (S 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 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 block MDPs Duu et al. 2019 Learning to act in a block MDPs has been the subject numerous recent papers
Dann et 2018 Du et al. [011 Misra et al., 2020 Feng et al.) 2020 Zhang al., 2020 2022] Furthermore, learning to act in 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.
14

---

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 al.) [201 Ni et al. [02T] study the offline setting, where the considerations are rather different. Finally, Kwon al. [202]] consider the problem when over multiple episodes, a learner interacts with finite-horizon MDP, which is chosen at random from one of such unknown MDPs. The learner is given no additional information about the hidden identity of the MDP that 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 problem. As such, while superficially, the problems may look similar, they are quite different.
B Auxiliary Lemmas
In this section, we show some probabilistic lemmas that will be useful 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 O(K log K) in order to get each of K coupon at least once.
Lemma 6. Given set K. with [K] S K and consider M i.i.d samples drawn from unif Then with probability 1 - 8', every element K. appears at least once in these samples as long as M 2 K (log(K/8))
Proof. Fix an element j E K. The probability that j appears in none of the M samples can be bounded by
(1-1e)
1
1 M
1
M
M
8'
S
exp
S
K
K
K
union bound on every j in IC finishes the proof.
Lemma 7 (Concentration of Subgaussian random variables). Let X, ......, Xas μ be sequence
of independent 1 I-subgaussian random variables and 2m=1 Xm. Then with probability 1 8', 8' 4. we have
M
log(1/01 (1/8))
j ml S 2V
M
Lemma (Bernstein's inequality). Let XI...., XM be a sequence of independent random variables with Xm - [E) S almost surely for every m and >M Var[Xm)) With probability 8',
we have
M
2b
E[Xm] (oggo01/8)) + 3 log(1/8')
The following lemma is direct consequence of Bernstein's inequality, which states, in terms of the coupon collector's 's problem, that after drawing M coupons, one can expect with a high probability particular coupon appears at least Mp/2 times if its probability of appearing is p.
Lemma 9. Let X1. XM be i.i.d Bernoulli random variables that [[Xm] p for all With probability 8', we have EM Xm 2 Mp/2 as long as Mp 16log(1/8))
Proof. Let Ym = 1 Xm We have E|Ym] = -1- p. By Lemma with probability 1 - d', we have
M
M - Σ X < M (1 --p)) + 22 p) (og((188) log (1/8")
15

---

Rearranging terms and using Mp 2 16 16log(1/8') gives
M
Xm 2 Mp v 2M - log(1/8") - log(1/8")
m=1
Mp
Mp
2 Mp
v8
16
Mp
2
C Proofs for Section
We first provide an outline of the proof in Appendix C.I and show the complete proofs of the lemma in the corresponding subsections.
C.1 Proof of Theorem j
To begin with, we assign adaptive target optimality (accuracy) to each block according to their block size:
Eb max { 1,
1
for b E [r].
(b)
The following lemma states that long W contains an Ô (et)-optimal action for each block b E [r]. then the output policy is Ô (e)-optimal with high probability.
Lemma 10. For positive constant C, suppose for any b [r] with w(b) 2 e/r, W contains Ce -ettoptimal action, then π out is 2 + 1)e-optimal for the original context-lumpable bandit problem.
As result, if we can prove that the precondition of Lemma T0 holds with high probability, then the correctness of Algorithm follows immediately. To do so, we first argue that for every block b E [r]. we have sufficiently explored the entire action set at accuracy level data collection step, as stated in the next lemma.
[log(1/e()] in the
Lemma 11. With probability least 1 26, for any b [r] with (b) > e/r, we have that: for any j E [K], there exists (i,j) E D, satisfying g(i) where n [loo(1/E]]]
Intuitively, Lemma 11 says that we have tested all actions on each block b E [r] at the corresponding accuracy level n Tog e/eg)] Based on Lemma 11 the following lemma states that in the second
step we are able to filter out an Ô (Eo) ooptimal action for block b by using the information contained D.,.
Lemma 12. With probability at least 1 46, for any b E [r] with w(b) 2 e/r, 10e. optimal action of block b is added into W in the process of shrinking D, where n = [log(1 [
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 W relatively small, which we prove next.
Lemma 13. With probability at least 1 - 45, [W] < Nr.
Now are ready to prove Theorem by combining all the above lemmas.
Proof of Theorem
7 The step of data collection uses
2 + log 2
(log rSK
r(S + K)
8 I log _z
e2
samples. By Lemma 13 the step of screening optimal action candidates
rSK
r
2(162)
[ log şz
log
2
16

---

samples and solving the simplified context-lumpable bandit problem uses
( (1+10g log E²
SK
Sr
log 8
g²
samples. Overall, the number samples is bounded by
rSk)
(1 1 2 log
1 !)
2
r(s + K)
524 log
8
e²
By Lemma and Lemma
nout is 2 (10/log(rSK/0) e-opimmal.
C.2 Proof of Lemma 10
Proof. We control the suboptimality of out in the following way:
Σ VE[s]
v()
max A(i,j) A(i, o0 (i))
jE(K]
<
Σ
Σ
v(i)
max A(i,)) A(i, out (i)) JEW
(())) + Ces
+
Σ
w(b)
hE(p]: ((b))e/riiB(())
belr): w(t) <e/r
<e +
Σ bE(r)
w(b)s
2(C 1)s,
where the first incquality uses the precondition of the lemma, the second one uses the e-optimality of ouu in the simplified context-lunne lumpatellumpabe bandit, and the final one follows from Kauchy-Schwar inequality.
C.3 Proof of Lemma
Proof. Fix a block b E [r]. For each accuracy level, recall in the step of data collection, we sample I- r(S K))g//² contexts i.i.d. from By Lemma with probability 1 -. at least (b) of them are from block b as
Le 2 1g 2 16log(r/6).
(b)
r
Further, recall that we define accuracy level n = [log(1/<))]- We first show that this n well defined, that is, > 1. Indeed, we have
E 2
c <2/(rw(b))
as long as < 3. Since we add a context-action pair into D, once we have collected
2 [og(1/<()) 2103(1/63)+1
2/e
samples for estimating its reward. Note that the end there are most (2" 1) samples from
block i unused. Thus, with probability s. there are at least (b) /2 [2( 1) samples used. Therefore, the number of context-action pairs, where the contexts are from block b, that are added into D, is at least
(b) [B()) 1)
2 (b)
(2 2/c3 and |B(b)] < S)
2
2/E3
> LEa (2re
(by definition of Eb)
2/e8
=> 16(S + K) log(rSK/()) S
(the value of L)
> K log(rSK/6).
Conditioned this event, with probability 1 4. for any [K]. there exists (i,j) E D, satisfying g(i) = by Lemma Therefore, the lemma holds for block with probability at least We complete the proof by union bound all blocks.
0
17

---

C.4 Proof of Lemma 12
Proof. Denote by the first context being removed from Dn, which satisfies g(1) By the rule of shrinking Dn, we have
log(rSK/0)
(1,5)))) - An("", s VV
2
By Lemma a union bound on all pairs in Dn, with probability - 28/r, we have
(") ,")" 2
1 log(rSK/o) 2
and AA(i", An(i"" 2
1
log(rSK/o)
2
(5)
for every (i", E Dn Therefore, we have
log(rSK/8) 2n
[A(i,") - A(T, bV
(6)
By Lemma know that for any j E [K]. there exists (i', E D, satisfying (g(i),.) (b, before we remove context i. Therefore, by the definition of (i* j"', for any j E [K]. there exists ((,,) Dn satisfying g(i') = b and
An 2 (i,
which, again by Eq. with probability 25/r
log(rSK/8)
A(i ( 2 max ((b.))
2n
By combining Eq. 6 and Eq.
conclude that is 10 10 /ogg::/K/00ss-optiiaa action for
block 6 because
log(rSK/8)
log(r 10 1og(rSSK/8))
10p
2
1/E3
This completes the proof by a union bound on all blocks.
C,5 Proof of Lemma 1
Proof. It suffices to show that for each accuracy level we will add at most r actions into W before shrinking D, Below we prove a stronger argument: each time we shrink Dn,, we will remove all the contexts from at least one block from D,,.
Denote by (i,5) an arbitrary context-action pair from D, such that g(i) g(it), then by Eq. 5 have
log( SK//)
(i.5") - An 5+) < Ai A(i,)))) A(T,,5 Att)" S 44
2n
which implies all context-action pairs with context from block g(²*) will be removed from Dn..
C.6
Proof of Theorem
Proof. By Lemma
9 and union bound over [S], we have that with probability 1- 8 for all [S]:
2 (()) log(S/0)
1
< ((1) 2 u(i) + log(S/0)
which implies that v11) 2 e/(4S) for (11) 2 E/S and vii) < 3e/S for i(1) < <E/S. By definition, for any E [T. 1), 2 consists of contexts such that (()) > 2-1-1 > E/S. As result, for any I e [T. 1], by Lemmag we have that at least (2(2)) /22 out of N samples corresponds X; and thus can be used in learning by executing Algorithm | where v(Xi) Σiex, v(i) e/(4S).
18

---

Moreover, notice that inside each Aú the conditional context distribution is almost uniform, so we can invoke Theorem too obtain that T
Ce
optimal
(())
over context X₁, where C' 222(21112/) (og(r/K/0 1). This means that,
v(1)
Ce
Σ
max A(i,)) A(()
vv(Xi
iEAL
As result, the total suboptimality of x ano can be upper bounded as following
Σ
(i)
max A(i,j) JELK]
A(i, (O)))
1E[S]
Σ 1E i
Σ
v(i)
max A(i,)) A(c.m(0))) + Σ v(i)
1E
SC Σ LE[L-I]
u(Xi) x +c= CVLE,
Vv(Xi
where the final equality uses Cauchy-Schwarc inequality and L S log(S/s) + 1.
C.7
Proof of Theorem
Proof. Since S S. have
r2 K, M(r)), K & S 2 K,
Q(SK),
((rK), r< K & K.
Case (1)
We construct the following hard instance:
The reward after pulling action E [K] in block b [r] sampled from distribution Bernoulli(1/2 + el = j)).
The context distribution is uniform over (S). For each i [S]. we sample g(i) uniformly at random from [K].
The above instance is the standard one commonly used in proving lower bounds for contextual bandit with S contexts and K arms [e.g.. Lattimore and Szepesvári 2020). And learning e-optimal policy for the above hard instance requires at least ((5K/ e samples.
Case (2) We only need to slightly modify the above hard instance:
The reward after pulling action j E [K] in block b [r]] is sampled from distribution Bernoulli e el =0=))) j)).
The context distribution is uniform over [S]. For each i [S], we sample g(i) uniformly at random from [r].
The above modified problem is equivalent to the original one with S contexts but arms. As result, lower bound of form 2(Sr/e²) holds.
19

---

Case (3) Similarly, we slightly modify the first hard instance:
For each block 6 [r], sample jī uniformly at random from [K]. The reward after pulling action j E [K] in block E [r] is sampled from distribution Bernoulli1122 + e = j
The context distribution is uniform over [r] and g(t) = min
The above modified problem is equivalent to the original one with K arms but contexts, As result, lower bound of form (2(rr==2) holds.
D Proof of Theorem
In this section, first show key lemmas prove Theorem 4 The proofs of the lemmas are deferred
to Appendix E
In the following discussion, we consider a single phase (i.e. fix an error ex). Similar to the analysis of other phased elimination algorithms, we have to show that in a phase specified by error level ene with high probability, (i) the optimal arm is not eliminated and (ii) all ((x))ssubpptimll arms are eliminated, that is, all arms in GOOD,(C) (C) for all C are ((E))-optimal
The following lemma is the counterpart of Lemma 11 which ensures that every arm is played in every block.
Lemma 14. With probability at least - 20 for any cluster C. Ph and any block b so that B(b) C C, we have that: for any j E GOOD,(() there exists (i, E DD satisfying g(i) = b.
The above lemma ensures that every block has a least one context assigned to include (i.j) in Dr- Thus, every arm is explored for every block.
Next, we define an event & under which the estimates AG are good:
En = An(i.)) - A(i, < Eh 4 v(i.)) E DA
The level of precision Ep is more accurate than the elimination step and will be helpful in the analysis. The following lemma states that E,, is high probability event.
Lemma 15. Event En holds with probability -8
Next, let ja = arg m K j) be the optimal arm in block b. The next lemma says that the optimal arm ja is not eliminated during the execution of the algorithm.
Lemma 16. Assume action jo E GOOD, (C) for every block E [r], and its corresponding partition Pr. 3 C 2 B(b). Then for every block b E [r]. jn is not eliminated from GOOD (C) with probability at least 30
The high-level idea of the proof is that the error of the estimated mean is smaller than Th. so Az jb) will not be much worse than other arms, given that its true mean is largest. The next lemma shows that arms in GOOD are all O(sx)-optima
Formally, we say an arm j in a block b is s-
optimal if maxy p(b, j') p(b,j) S e. Similarly, we say an arm in block b is e-suboptimal if maxz' ((b,j') (00)-((b.)) 2.
Lemma 17. For any block b E [r] and its corresponding cluster PA 3 C B(b). all subboptimal arms in block b are eliminated in GOOD with probability at least - 36,. Consequently, only 6E 6ot-optimal arms in block b are played in phase + 1 for every context in block b.
Finally, need to argue that Algorithm 5 is not called too many times. To show this, first provide a general guarantee of Algorithm
Lemma 18. Suppose we have context in E B(bw) S C in block by and context is E B(b)) C C in block by that
3r V Vig.e..
A(br. A(ir.)) = p(bn. p(b.. >
2
20

---

Then Algorithm 5 separates B((b)) and B(bb) perfectly with probabilty at least 1 1 28'. In other words, there exist indices C2 and ct2 Cu + C₁ such that (b,,)) S Pa.. and [(bb) - Pap
Consequently, with high probability, Algorithm makes progress in terms of clustering contexts every time calling AlgorithmB5 and the number of calls is bounded by r, as shown in the following lemma.
Lemma 19. When Algorithu is called by Algorithm 4 it separates at least two blocks and never separates contexts in the same block with probability at least 1 58g. Consequently, Algorithm called most r times.
We are now ready to bound regret.
Proof of Theorem Since there are at most (logr) phases, it suffices to bound regret at a single phase h. Conditioned on all the high probability good events in the previous lemmas, bound the total number of timesteps spent in phase In the data collection stage, we use at most
first
LA r(S+ samples; as for the clustering stage, by Lemma 19
we know the total length of
o7
executing Algorithm is at most
(1r)))
rL' - 1695///166 =
(=((5+K)) <7
Thus, the length of phase & is the minimum o' and regret is at most (recall that En. Vlogn Eh Ö(En))
Therefore, by Lemma 17 the
6 min { [T, (5 + K)
6Tēn, Ō }} min
( K)
Eh
(TTE)),
73 + K)
min
Eh
3 (V
On the other hand, the bad events happen with probability (( = regret contributes at most Õ(1)
In this case the
E
Missing Proofs in Appendix D
E.1 Proof of Lemma 14
Proof. Note that under the uniform block assumption, we have w(b) 2 Et. Thus, the proof
follows directly from the proof of Lemma 11 when E = Eh, = oi and n = nl..- The only difference
is when applying Lemma Lemma I1 uses I. = [K] but we need K GOOD, (C) here.
E.2 Proof of Lemma 15
Proof. Fix an (1,j) pair. Applying Lemma
we have with probability
SK
log(S (og(SK/6))
Eh
(n(i.) A(i,) < 2
2(mk)
4
We complete the proof by applying union bound on all (i.j) pairs in Di.
E.3 Proof of Lemma 16
Proof. We prove by contradiction. Assume that jo is removed from GOOD, Let
j = arg max Tr, (C, jCG
21

---

be the action achieving the highest empirical mean. By Lemma 14 there exists a context i' E ((6) so that (i', (i', Dr. Moreover, the assumption that jo eliminated implies that the condition at Line of Algorithm 4 does not hold for the current partition Ph. which further implies that there exists context E so that
Aa(i,') An 14 (C, (n)) En.
On the other hand, by the assumption that ja eliminated, again by Lemma 14 we have that there exists context i" E B(b). (i", (.6) E Dr such that
Pr J') An( (6) 2 (n(C,)) -,,(C, 25h
Combining two inequalities, we have
An (i" (b) + < F,,C,,") < (n',,'''
This means that pbb, jb) + F < j') under En, which contradicts the optimality of j.. Therefore, we conclude that j6 not eliminated.
E.4 Proof of Lemma
17
Proof. By Lemma 16 jo is not eliminated for any block b E [r] with probability at least 36
Thus, for any block be [r] and arm j GOOD,(() so that ((b.j.) (u6, З we have
(6) - (n(C, 2 (bb - (bb, 2 Eh 25h
max ma, (,, (C, 2
3'
Therefore, j is eliminated Line [4
E.5 Proof of Lemma 18
Proof. With probability 1 - 2g. context receives at least 2/s"² 2 2" samples by Lemma Thus,
A(i.j) is well defined for every i C with probability 6n. Moreover, by Lemma and union bound, we have for every context i, with probability at least 1 - 8n.
((L A(i.5) v V /g
Clearly, we have A(i.)) 2 A(i)) under Eq. (8) Consequently, to simplify the notation, we do the following modification on labels of contexts and blocks. First, we restrict the game to C, where there are S' contexts and r' blocks; also, we relabel contexts so that i, = 1, is = S", and
A(i..) A(1) > A(2) 2... Ā(S' - 1) 2 A(S") = A(úi).
Finally, given context i E [S'], we define
((b)=
max A(i'.) j) and m(b)
min
A(H',1)
v (S ( -
i' "E ((()")
for its block b- g(i) and we relabel blocks so that b... = 1 S g(i) S r' = by and
p(bb)) = (T1) 2 (bb. - ... 2 (b(b + 1) 2 (((r) = (b(b)).
It is not hard to that this modification is without loss of generality. We show next that there exists a context i, in < < it. such that Line 7 of Algorithm 5 holds, that is,
A(i 1,J) A(i,j)
Vig e'_
We prove this by contradiction. Assume Eq.
g doesn't 't hold for any k.
22

---

Then we have
Aliu: A(m,,)) (A(ba:)) [((b)) +A((i))) ta(bi) + T((bw) TT(b))
Lg
E
V
Igg
+
T((i-1)) - -T((ib)
4
4
bmn b
bj
6
((ih-1) v((o)) + Vig.d
2
b=b
3
(r 1) Vg
2
2
3
TV
ig.
2
which contradicts the condition that A(tu, - A(ii.) 2 3r Vl..e. Therefore, we conclude that Eq. (9) holds for some k:
E.6 Proof of Lemma
Proof. The proof follows directly from Lemma I8 and the fact that Algorithm 4 use e' and thus
A(i,j) A(i,j) 2 hh lgh
Ek
3
2 r gg e',
which satisfies the condition of Lemma 18
F Non-uniform Context Distribution
We show a reduction to problems with approximately uniform context distributions. The cost of this reduction an extra Ö(VST) additive regret and an extra Ollg(ST)) multiplicative factor in regret. The idea is to learn the context distribution in the first O( O(VST) timesteps. With high probability, we can estimate v(i) for any context with a constant multiplicative error as long as (1) = 2211 (VS1/STT).. bucket have the same probability up to constant factor. For contexts i with v(i) o(1/VST) we can not estimate the probability properly but we can simply ignore such contexts and suffer regret most O S L/V O(VST). We then run the algorithm that handles uniform context distribution for each bucket separately. Since there are O(log(ST) buckets, the overall regret is (logg(T)) times maximum regret over all subsets (or Vlog(ST) with refined analysis using Mauchy-Schward inequality).
J.hen Then we split the contexts into several buckets so that contexts within the same
G Proof of Theorem
Define
hb max
Eb.
for b E [r].
(b)
Also, for every phase and every level n, we define
lgA
Ek,n
2n
Next, we show the counterpart of Lemma 14 in the non-uniform case.
Lemma 20, With probability at least 1 -25,, for any cluster C E Ph, any block b with B(b) C C, we have that: for any level n S [log(1/=7,)) action j E GOOD,.n(C), (C), there exists (i.5) E Dr satisfying g(i)= b.
23

---

Proof. Fix a block b E [r]. For each accuracy level, recall in the step of data collection, we sample L = ==r(S- + K) klgg,2" 2 contexts i.d. from By Lemma with probability r of them are from block as
on.
least Lw (wbb//2 (b)/2
8T. 2 2g
(w(b)
16 16log(r/6z),
where the first incquality comes from the near-uniform context distribution assumption. Since we add context-action pair into D, once we have collected
2 2 [og(1/ 2log(1/28
2/
samples for estimating its reward. Note that in the end, there are at most [2 - samples from block b unused. Thus, with probability - 6g the number of the context-action pairs, where
S the contexts are from block that are added into D, is at least
(b) 22- [B(6)((2"
(b)/22 S
(2 2/e and [B(b]] S S)
2/Eh,b
LEn/
(by definition of En,b)
2/Ehnb
16(S + K) log(rSK/0n)
(the value of L)
2 K log(rSK/6,)
Conditioned on this event, with probability 1 بي for any j GOOD, (C), there exists (i,j) E Dz satisfying (i) by Lemma [ Therefore, the lemma holds for block b with probability least 1 26p, We complete the proof by union bound on all blocks.
Now define a good event Eg, as
En = = (()nnii) A(..)) s
1 4 Sh,m, V(i,)) E Dr, [Nal
Lemma 21. Event & holds with probability -p-
Proof. Fix (i,j) pair and level n. Applying Lemma
we have with probability SKN,
log( (N(SSNN//6))
1 Eh.n.
An, (i.j) A(i,)) S 22
2
We complete the proof by applying union bound on all (i,j) pairs in D,, and all levels n E [Na]-
In the following, we present and then prove the counterpart of Lemma 16 for Algorithm 6
Lemma 22. Assume action ju GOOD,..(C) for b [r] with (bb) 2 25n, and its corresponding
cluster Pr. 3 2 B(b). Then ja is not eliminated from GOOD (C) with probability at least 1 36r.frr any level n < [og
Proof. We prove by induction on n. The base case n 1, which satisfies the condition of Lemma a0 as
2b))2 E
log
> log
>
log (4) 2 n.
For the inductive step we prove by contradiction. Assume that jo is climinated in GOOD but not climinated in GOOD, (C) for 2 2. This means that the following inequality hold:
max j' VEGOOD,... (C)
T (C, Th, > 25
Let = arg maw F (C,j). (C, By Lemma 20 there exists context /' E B(b) so that (n, 7') jr < log(1/s{,,) Moreover, the assumption that jo is climinated implies that the condition Line does not hold for the current partition Pr. which further implies that
Th (C. Ann (i'.j)) En.
24

---

On the other hand, by the assumption that ju is eliminated, again by Lemma 20 we have that there exists context i" B()), (n, :")) 9 Jn such that
Th. (C, An, (i" .60) > 2r
Combining two incqualities, we have (n", (b6) Eh,n < (i, j'). This means that (16, j6) En < j") by Lemma 21 which contradicts the optimality of jb. Therefore, we conclude that jb 2
not eliminated.
Now we present key lemma similar to Lemma 17
Lemma 23. For any block b E [r] with (bb) 2 25h and its corresponding cluster Ph. 3 C B(b), all a suboptimal arms in block b are eliminated in GOOD (C) for any level S. [log(1/e with probability at least 1 35,. Consequently, only 6Ēh, Ghhx:-pttmal arms in block b are played in phase h 1 for every context block b.
Proof. By Lemma 22 Jo is not eliminated with high probability, so for every pair block e [r]
with (b) 2 2ea. and arm j i (C) that j6) (ubb,)) > 3 we have
max max,n,n(C.j)) Ph,n (C,5) 2 Pn,,, Th,n (C, (b) - m(b,)) 2
hh.
2 2.5 Eh,n 2.n
Therefore, j is eliminated in GOOD,,, (C), and thus eliminated in GOOD (C) for né' >
G.1 Proof of Theorem S
Proof. Fix a phase h and a level n. It suffices to bound regret within a single pair (h, n). For b E [r].
let ni = [log [ogg[]=(... 1. By Lemma 23 if n > ni. all 6,,,n,- -suboptimal arms are eliminated from GOOD,,.(C) for any cluster C Pr. Therefore, regret of playing action from GOOD,,. (C) is
65 In this case, regret is bounded by
Iga
Σ r(S K)2(n+h)/2 w(b) 65h,a BEIr)
Σ bE(rl
r(S K)2( w(b) BV 2
< Gy lh
r(S + k))(n+h) (b) En.b
bEIr]
$ 6y lg
Σ
r(S + K)2(m+h)/2 w(b) En
rb)
bE[r]
6 rg ( ()2(m+h)/² = (r(S K)21/22
If S nl. similar argument shows that regret of playing action from GOOD (C) is 6 Therefore, regret is bounded by
Σ r(S + k2) 6 6 6 Σ w(b) r ((5))r(((+k))) K)2(n+h)/2 6E
1gh
2n
belrl
S 6v Ts K)2" = -ō (r(S+ KY2*/2)
We conclude the overall regret is Ō r KYT) ( for the data collection stage by noting that T > r(S + K)2(1-1) phase ht is executed completely. The same argument holds for analyzing the clustering stage when replacing with r",, so we conclude the regret is bounded by
δ (V +KY++
25

---

H Omitted Details in Section
In this section we discuss how to solve the more general low-rank bandit problem. Note that A has rank if and only if there exist vectors Wiy.... Wg E I'' and Vi...., UK E IR" so that A(i,j) - w Vy. We first consider r = and assume that every w; is non-negative. In this case, we have
arg max A(i, j) arg max Well = arg max 1/j- jE[K]
(10)
JE[K]
jE[K]
In other words, there exists an arm that is optimal for all contexts. The problem becomes simple as we can run the EXP4 algorithm [Auer et al. 2002] using constant experts that recommend the same arm for all contexts. Thus, we will have only K experts and have the following proposition:
Proposition 24. Consider the rank-1 bandit problem with r = and ui 2 0 0oror every i [S]. The regret of the EXP4 algorithu run with K constant experts is bounded by OgyKTlog K).
However, the idea seems hard to generalize when wy E R, Eq. I0 does not hold anymore and need exponentially many experts for EXP4. Next we introduce a new idea based on a reduction to context-lumpable bandits.
In the following we define constant max; | better illustrate the idea we first assume r = 1 and consider the PAC setting. We create an a-covering of [-B, B) and "cluster" each i into one of the segment. Specifically, we have 2B intervals
B, B + 1 [-B+-1 a a
2
1 α
and each w; is assigned to the interval that contains it. Given a, let R. denote the number of intervals that have at least one context. Conceptually in the same block in context-lumpable bandits, and we have R. blocks analogously. For contexts i' in the same interval, they indeed similar in the sense that have A(i;)) A(i, O(a) for every arm j. Intuitively, if a is much smaller than E, then Algorithm can proceed normally as Ra context-lumpabe bandit problem. Consequently, we have the following theorem.
view contexts in the same interval if they
Theorem 25. Forr- = 1, by choosing a = (s), Algorithm uses Ō (R., (S + K)/s²) samples and outputs an Ö (optmmal policy.
Clearly we have R2 < min { S). so the above theorem leads to a Ō + K) sample complexity the worst case. The idea can be generalized regret minimization and r 1.
Specifically, we construct an a-grid of [-B, BJ" so that Ra ((+) and run Algorithm 6 for regret minimization. Consequently, have the following theorem:
Theorem 26. Letp
3542 and choose (S K. Then regret of Algorithm 6 bounded as Reg- = Ô ((S + P)P
Proof. Similar to previous analysis, the regret single phase is
R3 (S
} min (Ō δ R2 (
( K)
}}
δ (EN) min Ō
(11)
2
Recall that R., 0(1/a')) Therefore, when = ((E)), we have the above regret is bounded by
δ (S + K) a
of = (S+ a3r+1 + aT)) =
(S+K) K)
when choosing optimally (S+ K. Otherwise, = o(En) and Eq. 11 can be bounded by
R3 δ
R2(
R3 (R)( =
Eh
We finish the proof by noting that there are most log I phases.
26

---

The bound becomes non-trin non-trivial when S and K are large. For example, when S = K = VT, the bound T (T) for any r while the Ô (VSKT) bound given by EXP4 vacuous. The factor comes from the r term in the regret bound Theorem 5 improve the factor in context-lumpabl bandits and extend it to low-rank bandits using the reduction introduced here.
It is a promising direction to first
27


Reviews:
[{"review_id": "ApqLTi1OVG", "reviewer": "Reviewer_ihhL", "summary": "This paper studies the context-lumpable stochastic bandits, which can be understood as stochastic contextual bandits with low rank structure. A specific case of this problem where the contexts are linear was studied in previous literature, but there wasn't much study on this general problem. For this problem, this paper provides novel pure exploration lower and upper bounds. These two bounds match each other so the results in this paper are minimax optimal.  For the upper bound, this paper provides an algorithm for it based on phased elimination. ", "questions": "1. Does the setting in (https://arxiv.org/pdf/1712.04644.pdf) apply here?\n2. In most phased elimination algorithm, the sets $\\mathcal D_n$ need be computed from $\\mathcal D_{n-1}$, could you explain why it is not the case in Algorithm 1?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The bounds in this paper are minimax optimal. \n2. The problem proposed in this paper is a natural extension of linear bandits with low-rank structure. ", "weaknesses": "1. At Line 9, $\\Omega$ is better than $\\widetilde \\Omega$ because the latter could be understood as $\\Omega(\\cdot) / \\mathrm{polylog}(\\cdot)$. \n2. In Algorithm 1 Line 1, $\\mathrm{lg}$ is oftenly understood as an operator $\\log_{10}$, so I suggest use another symbol to avoid confusion.\n2. This paper can only solve the contextual bandits for stochastic contexts, while generally the contexts are considered to be adversarial.\n3. The algorithm in this paper is for pure exploration and does not lead to an optimal algorithm for regret minimization.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback. Below we reply to the reviewer’s comments and questions.\n\n- **$\\widetilde{\\Omega}$ and $\\lg$:** We thank the reviewer’s suggestion on the notations. We will change them accordingly. \n- **Stochastic contexts and suboptimal regret:** We agree that extending to adversarial contexts and finding optimal regret are important future directions. We believe our work will facilitate the development of these directions. \n- **The Stochastic Low-Rank Bandits paper:** We discuss Kveton et al. [2017] in Appendix A (see Supplementary Material). Their setting is noticeably different from ours, where the learner can pick multiple rows and columns in the reward matrix each time. We don’t see an immediate application of their algorithm to our setting or our algorithm to theirs. \n- **The question regarding $\\mathcal{D}_n$:** In fact, we don’t consider Algorithm 1 as a “phased elimination algorithm.” Unlike Algorithm 6, which loops over phases, Algorithm 1 can be seen as only running one phase given an error tolerance level of epsilon."}, {"review_id": "6zcoCe8nDE", "reviewer": "Reviewer_nWDZ", "summary": "This paper considers the problem of context-lumpable bandits, which consists of contextual bandits with $K$ actions and $S$ contexts that can be partitioned into $r \\\\le \\\\min\\\\{S, K\\\\}$ blocks. The authors consider the case when mean rewards of the actions are fixed over different contexts that belong to the same block in the partition. In the setting of PAC learning, the authors design an algorithm that uses only $\\\\tilde{O}(r(S+K)/\\\\varepsilon^2)$ samples to find an $\\\\varepsilon$-optimal policy. The authors additionally prove a nearly matching lower bound (up to logarithmic factors) on the sample complexity of such an algorithm, thus proving that it achieves near optimal. performance. The authors further focus on regret minimization for the problem of context-lumpable bandits and design a strategy with $\\\\tilde{O}(\\\\sqrt{r^3(S+K)T})$ regret after $T$ rounds. These results are shown to be more general than prior work, as they allow for arbitrary context distributions and generalize to contextual low-rank bandits.", "questions": "See weaknesses. In particular, do the authors believe that it would require a significant amount of work in order to devise and analyze an algorithm with a satisfactory gap-dependent regret bound?", "limitations": "The authors did address the limitations of their work, such as the open problem of improving the current results in the regret minimization setting.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The results provide a significant contribution and improve upon prior work.\n- There is a careful adoption of multiple techniques and interesting ideas.\n- A thorough overview on similar problems in the literature is provided to better contextualize the work, thus justifying the originality of the work.\n- The techniques employed in both the design of the algorithms and their analysis are nontrivial and can be of independent interest.", "weaknesses": "No major weaknesses were identified.\nRelative to the setting of regret minimization, the proposed algorithm essentially estimates the mean loss for each context in order to somehow recover the block partition by identifying large enough gaps between such estimates. Then, one may wonder whether an gap-dependent regret bound could be achievable, but this point does not appear to be directly addressed in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and insightful comments. Getting a gap-dependent regret bound is undoubtedly an interesting question. As the reviewer points out, our algorithm detects large gaps and eliminates suboptimal arms accordingly. Therefore, assuming the rank is constant and there is a “global” gap $\\Delta$, it is promising that our algorithm can get $\\widetilde{O}((S+K)/\\Delta)$ gap-dependent regret bounds compared to a native $\\widetilde{O}(SK/\\Delta)$ bound. However, things may become complicated when there are various gaps in the environment. There are essentially two dimensions: different gaps across arms and different gaps across blocks. It is an interesting direction to investigate if our algorithm or other algorithms can achieve a satisfactory gap-dependent regret bound or the problem of what is a satisfactory gap-dependent regret bound itself. \n"}, {"review_id": "JSIpnMDMot", "reviewer": "Reviewer_bwz6", "summary": "This work proposes a PAC learning algorithm and a regret minimization algorithm in the case when the contexts can be lumped into smaller partitions of contexts. The PAC algorithm achieves minimax sample complexity, while the regret minimization algorithm achieves a regret upper bound that could be less than the number of context groups before being lumped.", "questions": "Q1. The algorithm seems to have a lot of computation. Could the authors elaborate on the computational efficiency of the proposed algorithm?", "limitations": "Yes. This work is mostly theoretical and potential negative societal impact is not applicable.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This work provides a principled algorithm for contextual bandits for the case when contexts can be lumped into a smaller group. This case happens in real applications and is important to research. Compared to the previous works, this work relaxes the requirement of the knowledge of grouping of the contexts. This makes the algorithm more utilizable in the real world.", "weaknesses": "1. Could the authors give an example of known $r$?  Assuming $r$ is known to the algorithm seems not realistic since $r$ is usually unknown in practice. Or is there a way to relax the assumption on the knowledge of $r$?\n2. The practicality and efficiency of the algorithm are not validated through experiments. An experiment that shows whether the algorithm finds the correct lumped groups for near-optimal arms would be a significant result.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing their valuable feedback and raising important questions. Below we answer the reviewers’ questions. \n\n- **Knowledge of $r$:** We agree that the requirement of knowing r is not very satisfactory. We have some ideas on removing this requirement (see the paragraph under Theorem 5). However, given the current complication of the algorithm, we decide to put this part into future work. That being said, we don’t think this is a big issue when applying our algorithm in practice. One can use a fairly reasonable \n“upper bound” of it or treat it as a tunable hyperparameter. \n- **Computational complexity:** The algorithm has amortized constant run time per iteration. Note that some computation, such as elimination, only happens at the end of each phase, and there are only $\\log(T)$ phases over $T$ rounds."}, {"review_id": "dTVpl450bg", "reviewer": "Reviewer_15yr", "summary": "The paper considers contextual bandit problem with S contexts and K actions, under\n 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. The paper considers the task of learning $\\varepsilon$-optimal policy and prove sample complexity upper bound and a matching lower bound. The authors also present regret upper bound that depends on $r^{1.5}$.\n\nAlthough the results are original, the importance of the setting is unclear to me give the rich literature on CMAB, and the tight regret bounds known under minimal assumption of function approximation.\n(See e.g.,[1] Bypassing the Monster: A Faster and Simpler Optimal Algorithm for Contextual Bandits under Realizability, Simchi-Levi and Xu, MoR 2021).\nThe assumption that the contexts can be lumped into $r \\leq \\min\\{S,K\\}$ groups essentiality means that in practice there are only $r$ different contexts (although the cluster of each context is unknown), and the sample complexity and regret bounds depends in $r$ and $S$ (instead of $\\log(S)$ which is the SOTA under function approximation).", "questions": "1)\tEquation (3): the sample complexity mentioned there seems trivial – basically run MAB algorithm for each context separately. The comparison is unclear to me.\n2)\tPlease consider [1] (SOTA for CMAB) and compare your results to those in [1]. Explain the difference between the settings and results.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1)\tTight sample complexity bound for the problem.\n2)\tNice applications.", "weaknesses": "1)\tGiven previous literature, the significant of the results is unclear to me.\n2)\tInsufficient previous literature review. Mentioned results in equations (3) and (4) are not SOTA. Need to add a general review of CMAB literature and specifically [1].\n3)\tThe setting of almost-uniform distribution seems trivial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments. Below we address the concerns raised by the reviewer. \n\n- **Comparison with [1], and (3) and (4) seem trivial:** The paper [1] the reviewer mentioned falls into the category of “oracle-based” contextual bandits. In oracle-based contextual bandits, the near-optimal regret is $\\widetilde{O}(\\sqrt{KT\\log|\\mathcal{F}|})$, where $|\\mathcal{F}|$ is the size of the function class. (We are not aware of any algorithms achieving $\\log(S)$ dependence in the literature) Under the lumpability assumption, one may further bound $|\\mathcal{F}|$ by $O(r^S)$. However, the resulting regret bound that can be achieved from this line of work is exactly what we write in (4) (or (3) in the PAC setting). In summary, although we can apply algorithms in oracle-based contextual bandit algorithms, including one in [1], to our setting, the regret bound we can get is the same as (4), which can also be achieved by running an MAB algorithm for each context separately. In our results, we show that we can take advantage of the lumpability of contexts into a few groups, even when the identity of the group a context belongs to is unknown, and achieve much smaller regret (and sample complexity) bounds. That being said, we will ensure we adequately compare our work with this line of work in our next version. \n\n- **The setting of almost-uniform distribution seems trivial:** We strongly disagree with this comment. In fact, the problem can be reduced to this setting by our reduction, so the almost-uniform distribution setting essentially captures the problem’s main difficulty. "}]

Author rebuttal:

Meta-review: The paper studies the context-lumpable stochastic bandit problem, where contexts could be lumped into $r$ groups, and all contexts in each group share the same arm means. The authors prove tight upper and lower bound for the PAC setting and near-optimal regret bounds.

The reviewers agree that the model is very applicable and that the paper improves the existing results for this model. In addition, the algorithms suggested in the paper are nontrivial and the approach could be of independent interest.
Decision: Accept (poster)