Tractable Optimality in Episodic Latent MABs
Authors: Jeongyeol Kwon
Yonathan Efroni
Constantine Caramanis
Shie Mannor
Venue: NeurIPS 2022 Accept (2022)
Keywords: multi-armed bandits
partially observable MDPs
experimental design
latent variable models
method-of-moments
maximum likelihood estimation

Abstract: We consider a multi-armed bandit problem with $M$ latent contexts, where an agent interacts with the environment for an episode of $H$ time steps. Depending on the length of the episode, the learner may not be able to estimate accurately the latent context. The resulting partial observation of the environment makes the learning task significantly more challenging. 
Without any additional structural assumptions, existing techniques to tackle partially observed settings imply the decision maker can learn a near-optimal policy with $O(A)^H$ episodes, but do not promise more. 
In this work, we show that learning with {\em polynomial} samples in $A$ is possible. We achieve this by using techniques from experiment design. Then, through a method-of-moments approach, we design a procedure that provably learns a near-optimal policy with $O(\poly(A) + \poly(M,H)^{\min(M,H)})$ interactions. In practice, we show that we can formulate the moment-matching via maximum likelihood estimation. In our experiments, this significantly outperforms the worst-case guarantees, as well as existing practical methods.


Tractable Optimality in Episodic Latent MABs
Jeongyeol Kwon University of -aaonnnn-Maaooon jeongyeol www.wwwsw...e..ooooo.w......w..e edu
Yonathan Efroni
Meta, New York jonathan efroni @gmail. commail..omm. com
Constantine Caramanis The University of Texas at Austin constantine@utexass edu
Shie Mannor
Technion, NVIDIA
shie@ee technion. ac i1, smannor@nvidis com
Abstract
We consider multi-armed bandit problem with A actions and M latent contexts, where an agent interacts with the environment for an episode of time steps. Depending on the length of the episode, the learner may not be able to estimate accurately the latent context. The resulting partial observation of the environment makes the learning task significantly more challenging. Without any additional structural assumptions, existing techniques to tackle partially observed settings imply the decision maker can learn near-optimal policy with O(A) episodes, but do not promise more. this work, we show that learning with polynomial samples in A is possible. We achieve this by using techniques from experiment design. Then, through a method-of -ommmmnss-sss approach, we design a procedure that provably learns a near-optimal policy with (ppoly(A) + poly(M, H) min(M,H)) interactions. In practice, we show that we can formulate the moment-matchnng via maximum likelihood estimation. our experiments, this significantly outperforms the worst-case guarantees, as well as existing practical methods.
Introduction
In Multi-Armed Bandits (MABs), an agent learns act optimally by interacting with an unknown environment. In many applications, interaction sessions are short, relative to the complexity of the overall task. As motivating example, we consider a large content website that interacts with users arriving to the site, by making sequential recommendations. In each such episode, the system has only a few chances to recommend items before the user leaves the website an interaction time typically much smaller than the number of actions or the types of users. In such settings, agents often have access short horizon episodes, and have to learn how to process observations from different episodes to learn the best adaptation strategy.
Motivated by such examples, consider the problem of learning near-optimal policy in Latent Multi-Armed Bandits (LMABs). In each episode with time-horizon H, the agent interacts with one of M possible MAB environments (e.g., type of user) randomly chosen nature. Without knowing the identity of the environment (we call this the latent context), an agent aims to maximize the expected cumulative reward per episode (see Definition 2.1 for formal description). The LMAB framework is different from the setting considered in [31] where multiple episodes proceed in parallel without limiting the horizon of an episode H. For long horizons H. we show that possible find near-optimal policy and to determine near-optimal actions for each episode, as if we knew the hidden context. H & where A the total number of actions, however, this no longer possible. Instead, we aim to learn the best history-dependent policy for H time steps.
36th Conference on Neural Information Processing Systems (NeurIPS 2022).

---

1.1 Our Results and Contributions
In this work, we study the problem of learning a near optimal policy of an LMABs for both long and short horizon H. In the long horizon case, we show the problem's 's latent parameters are learnable. However, for short horizon, the latent parameters may not be identifiable and different perspective is required.
A naive approach to learn near optimal policy of an LMAB starts by casting the problem Markov Decision Process (MDP). Assume that the reward values are discrete and defined over a finite alphabet of cardinality Z. By defining the state space as the set of all sequences of history observations an LMAB can be formulated as an MDP with O((AZ) " states. Appealing to standard reinforcement learning algorithms (e.g.. [18, 6)) we can learn an e-optimal policy with O(AZ) /e² samples.
natural way to improve upon the naive approach through using the unique structure of the LMAB setting. In this work we focus on the following question: can we learn near-optimal policy with fewer than (AA)) /22 samples as the naive approach?
Our work answers the above question affirmatively. Specifically, our main contributions are as follows. We show that the dependence of our algorithm is poly (A) + poly H, M) mi),HH,MM and thus tractable when either H or M is small, even for very large A. That is, are particularly focused on the setting with a few contexts or relatively short episodes (in comparison to A), Le., M = O(1) or H = O(1), where a natural objective is to learn near-optimal history-dependent policy for H time steps (see also Figure 3 in Appendix 1.
1.2
Related Work
Due to space constraints, we discuss only the most closely related work here, and defer a lengthier discussion on related work to Appendix B.
Learning priors in multi-armed bandit problems Several recent works have considered the
Bayesian learning framework with short time-horizon [29, 35, 21]. The focus in this line of work is on the design of algorithms that learn the prior, while acting with a fixed Bayesian algorithm (e.g.. Thompson sampling). While Bayesian learning with short time-horizon may be viewed as a special case of LMABs, the baseline policy we compare ourselves to is the optimal H-step policy, which is a harder baseline than considering fixed Bayesian algorithm.
Some prior work considers the framework of Latent MDPs (LMDPs), which is the MDP generalization of LMAB [16, 37, 24, 23]. In particular, [24] has shown the cooomation-theorete limit of learning in LMDPs with no assumptions on MDPs, i.e., an exponential number of sample episodes {2AA²)) is necessary learn a near-optimal policy in LMDPs. In contrast, we show that in LMABs, the required number of episodes can be polynomial in A. This does not contradict the result in [24], since their lower bound construction comes from the challenge in exploraxpoootiion with latent contexts. In contrast, there is no state-exploration issue for bandits, which enables the polynomial sample complexity in A. Furthermore, our upper bound does not require any assumptions or additional information such as good initialization or separations as in [24]. To the best of our knowledge, no existing results are known for learning a near optimal policy of LMAB instances for M 2 3 without further assumptions.
Latent MDPs
Learning in POMDPs with full-rank observations One popular learning approach in partially observable systems is the lensor-decomposition method, which extracts the realization of model parameters from third-order tensors [1. 7, 15]. However, the recovery of model parameters require specific geometric assumptions on the full-rankness of a certain test-observation matrix. Furthermore, most prior work requires a uniform reachability assumption, i.e., all latent spaces should be reached with non-negligible probabilities by any exploration policy for the parameter recovery. Recent results in [19, 30] have shown that the uniform reachability assumption can be dropped with the optimism principle. However, they still require the full-rankness of test-observation matrix to keep the volume of a confidence set explicitly bounded. Since L.MAB instances do not necessarily satisfy the full-rankness assumption, their results do not imply an upper bound for learning LMABs.

---

2 Preliminaries
We define the problem of episodic latent multi-armed bandits with time-horizon H 2 2 as follows:
Definition 2.1 (Latent Multi-Armed Bandit (LMAB)) Let LMAB be a tuple where A is sel of actions, {wm J are the mixing weights (
(A, {wm jol m=11 (mm Jm=1) such that latent context m randomly chosen with probability Wm. and pmm is the model parameter that describes a reward distribution, i.e. IP Pam.0 (r a) = P m, a.). according a E A conditioning on a latent context m (each pm parameterizes the probability model, and not necessarily mean reward vector).
action
We do not assume priori knowledge of mixing weights. The bulk of this paper considers discrete reward realizations, when support of the reward distribution is finite and bounded. Appendix D, we also show that our results can be adapted to Gaussian reward distributions.
Assumption 2.2 (Discrete Rewards) The reward distribution has finite and bounded support. The reward attains a value in the set 2. We assume that for all 2 E Z we have |z| S 1. We denote the cardinality of Z Z and assume that Z O().
As an example, Bernoulli distribution satisfies Assumption 2.2 with 2 (0, and 2 = 2. We denote
the probability of observing a reward value z by playing an action a as jmm (a, 2) = P(r = m, a) in context m. We often use jm hewaar-probabiliyy vector in gaA2 indexed by tuple (a, =) E Ax
At the beginning of every episode, latent context m [M] sampled from a mixing distribution {wm JMM 1 and fixed for H time steps, however we cannot observe m directly. We consider a policy class II which contains all History-dependen policies
(A Z)" + A. Our goal to find a near optimal policy E II that maximizes the expected cumulative reward V for each episode
V max (TT) = E" E
where the expectation is taken over latent contexts and rewards generated by an LMAB instance, and actions following policy म.
Tt
Definition 2.3 (Approximate Planning Oracle) A planning orucle receives an LMAB instance and returns e-approximate policy n such that V - ((=) S.
Concretely, the point-based value-iteration (PBVI) algorithm [33] is an e-approximate planning algorithm which runs in time ((HMAZ((FF/)00))
Additional notation. For any quantity q with respect to the true LMAB, B, we use to refer to its corresponding empirical estimate. use Wmin := mi Wm for the minimum mixing weight. For any 1 -order tensor 7, E Txxx...x. (n repeated times), we denote ||Ti|l for the element-wise largest absolute value. For any vector in E IR in dimension n E N4. we use n to denote degree tensorization of v. We often use I| (a, pi fim(a, (a,-))l1 to mean the statistical distance between reward distributions: Exez z |mm (a, 2) - pmm (a, z)|. Lastly, we denote by and r. as the action and reward realizations observed at time step E [H]
Sample-Complexity of Learning LMABS
In this section, develop our main algorithm that learns near-optimal policy of an LMAB instance with O(poly(A) + poly (MM mfnm((M..))) samples. Towards achieving this goal, we first elaborate on a low rank representation of the LMAB problem. Then, we show how experimental design techniques can assist in utilizing this structure to improve over the naive upper bound to the problem.
3.1 Dimensionality Reduction via Experimental Design
When = 1, estimating the latent context is not possible, and executing a standard MAB strategy (e.g., UCB) is optimal. However, to obtain a near-optimal history-dependent policy with longer time-horizons H > 1, tracking the mean reward of an action is not enough, since rewards from previous time steps are correlated with current and future rewards.

---

We consider a model- Since there are A actions, this implies that we have O(4') quantities to estimate, which would incur OAA') sample complexity if we separately measure every correlation between / pairs of actions. On the other hand, when M < A, the distributions pm (recall these are reward probabilities for each action, in each context) occupy only M -d-dimensional subspace of AZ-dimensional space: U := span f1. P..... Mas) }. Thus if we could estimate U, and project all observations to U, we could remove the unfavorable dependence on A.
model-based based learning approach from the 1th -odder moments of reward observations.
However, even when the subspace U is known a priori, dimensionality reduction with bandit feedback non-trivial. We need to compute directly the projection onto U estimates of the reward probability
vectors. To see the challenge, let (3, ) C AZ be an orthonormal basís of U. To compute the
projected estimate 33. we need a sampling policy such that rTa) Σx (a, =)I. The variance of this sampling policy can be as much as ((1,,, which in general scales with || (()). Therefore, reliable estimation for any statistics in the reduced dimension would still have a dependence on A. In particular, since our approach is based on higher-order method-of -offoomfttstss, an estimation of 1th -order statistics for 1 > 2 would require (244)) samples.
The general idea of dimensionality reduction is not fatally flawed. Instead, we need to avoid the pitfalls of the approach outlined above. First, the calculation above shows that we need to control the - II of the vectors 8y. ideally, || Jm such that maxae | (a, :) jim (a, I11- I to compute an c-optimal policy.
O(1). Second, we only need good estimate for each
The key is to show that we can choose 3z 2s's to be subset of the standard basis in RAZ (hence they will have ||ل,ك|ı = 1), in such a way that guarantees the approximation quality for jmm from estimating 3y ss. In terms the original problem, the existence of such subset the standard basis is equivalent to the existence of small set of informative action-value pairs, that are sufficiently correlated with all other action-value pairs. This is called a core set in the experimental design literature, and its existence in our context, is an important consequence of the Hiefer-Wolfowir theorem. Specifically, we can select a core set of coordinates with the following crucial lemma:
Lemma 3.1 Let U be a given k-dimensional linear subspace in TRd where d >> k and let u, ū E U. There exists an algorithm that runs in time Ódd??) and returns the following.
1. A core set of at most n 4k log log k + 16 coordinates {iz}i=1 S [d] such that || allo S 22k maxje(n] (u(t,) - u(tz))-
2. A linear transformation that maps (ū(i),.... ū(ix)] to its corresponding ū E U.
Note that in our setting, d = AZ and A = M. We prove Lemma 3.1 in Section C.I, essentially as corollary of the Seffer-Wolfowitz theorem (and its geometric interpretation) for (near)-optimal experimental design [20,3] 39].
In the context of bandits and RI,, the Kiefer-Wolfowitz theorem has been used to study how the misspecification in linear representation changes the problem landscape (e.g., [27, 32]). However, experimental design has not been previously used for dimensionality reduction for problems with bandit feedback (as far as we know). We believe it is a powerful tool. We review the basics of experimental design in Appendix B.2.
A direct consequence of Lemma 3.1 is an algorithm to find a set of coordinates of size Ô(M) that are sufficient to reconstruct the latent reward model {um MM where pmm E I AZ We refer to this set of coordinates as the set of core action-value pairs.
Corollary 3.2 Suppose the subspace U span(pi,.... pas) is given. Then for any is E U, there exists an algorithm that runs in time A(ZA2)) and returns the following.
1. A set of core action-value pairs {(az, 2,)17 n = 4M log log M + 16, such that for all m E [M]
S
A
x
of size at most Z
max || aEA
(a, 111 S 2Z /2M
max (mm Dm (i)1, jEIn)
where Vm, Pm E R such that V. ( = pmm (aj, 2z) and fm = pm (ag. 23).

---

2. A linear transformation that maps [D(ii)...... ((t)) to its corresponding ji E U.
Assuming access to the subspace U (we remove this assumption in Section 3.4), Corollary 3.2 implies a strategy for estimating the latent model parameters [jmml=e obtain core eetion-value pairs {(a;, Pj) 23)13=11 11=1+ estimate {v/m /mmi+ and construct latent model parameters via a linear transformation of {um MM m=1- In the following sections, we show that by matching higher-order moments we can IM
estimate [Ym) m= to sufficiently good accuracy with only polynomial dependence in A.
3.2 H 2 2M- - 1: Identifiable Regime in Wasserstein Metric
In the regime where that H 2 2M - we can measure moments up to order (2M 1)th, Then, we can leverage recent advances in learning parameters of mixture distributions from higher-order moments [40, 12]. That is, we can recover {Pm //m/m=1 by estimating 2M 1 higher-ordes moments M and find {im } m=1 that matches them. This further emphasizes the importance of the dimensionality reduction we take to obtain the set of core action-values pairs; otherwise, moment-based approach M would have an exponential dependence in A.
to recover { (Am) m=1
We now elaborate on the estimation procedure of {vm J from higher-order moments. We define the l-order tensor as Ti == Σm== Wmn Wmm @1 For an LMAB instance, we can access the tensor Ti using observational data by simply estimating correlations between core action-value pairs within each episode. Specifically, for any I = (i1, 12 ..... it) E [n]] the tensor [i(1,, i2.... it) is also given by
Ti i2, (2,,....,)) in) = T=* [
{r. = zi, }]
where म a policy that performs the sequence of actions (at, (12: ..., ai,) for ====,,. I,....I.
Suppose find estimators {(wmm )m))m m= such that moments match Ti up error d for all / e [2M - 1]. The results in [40, 12] imply that the closeness in moments of up to (2M 1)th degree implies the closeness in model parameters:
①m 2 1100 < & all ( = 1,2, ......MMMMM.......... 2M- - 1 for some sufficiently Lemma 3.3 Suppose || - EM
small 6 > 0. Then,
Σ T(m, IIum 1me I 0 ((MM---11(2M-11)
inf (m,'²)E[MMp
(2)
where I is a joint distribution over (m, m') [M]] satisfying
m(m,m²)
RR MRMM
=I I'm, m²) = Wm.: EM T'm, mm) Wm-
The form of guarantee given for {wm, Pm JM is in the Wasserstein distance between two latent model parameters. A useful property of the Wasserstein metric is that the distance measure is invariant to permutation of individual components, and flexible with arbitrarily small mixing probabilities or arbitrarily close components [40] (see also Appendix B.3 for the review on Wasserstein distance).
Once we obtain estimates of (um JM 1= we can estimate (jmm } m=1 as implied by Corollary 3.2 (see Appendix C.8 for further details). We then show that for any history-dependent policy the expected cumulative rewards of ι are approximately the same for any close LMAB instances.
LM Proposition 3.4 Let B- (A, (wm) 1 {pm 1 and B (A, { ( J (ma JM be any two LMABs. Then, for any distory-dependee policy x (A 2)" A, we have
(V(T)- V < H inf
Σ (EL]))
(r(m,m')
max I| (a, ii (a-))))))))) (4) aed
where the infimum over T is taken over joint distributions over mm,,m)) satisfying (3).
From Corollary 3.2 we have max Illm (a, -) iimm (a, 1] S 22.22M Im D lloo for any m, m' E [M]. Plugging this into Proposition 3.4, we have
(VT) (T)) - poly H, 2,M, m). 8 122
Thus, using Lemma 3.3 with 5 < (poly(H, Z, M, m)) m2)c))2-1, we can conclude that any coppimal policy for B is O (e)-optimal for the underlying LMAB B.
5

---

3.3 H < 2M- - 1: Unidentifiable Regime with Short Time-Horizon
more interesting regime the one which the time-horizon is smaller than the required degree of moments 2.M- 1. For such setting we cannot measure moments of degree higher than H. Therefore, H 2.M- 1, we cannot rely on the identifiablility of the underlying LMAB model.
Instead, we make the following claim: to compute an optimal policy only for time horizon H < 2M-1, we only need match the 17th order moments. This is formalized in the next lemma.
Lemma 3,5 Suppose that || Em== pendent policy π E 11, we have
23 H Wm/fin
-mM= Wm fm
2 H
1o S 8, then for any history de-
-(T)-((T) S HZH
Therefore, it sufficient find estimates of mixing weights and latent models that match the measured Hh moment, and then compute an optimal policy for the estimated model. Lemma 3.5 is natural for discrete reward distributions with bounded support. Interestingly, we extend this result to Gaussian rewards, which are continuous and unbounded, in Appendix D.
Hence, natural idea is to estimate parameters {u } MM {pm JM that satisfy the condition in Lemma 3.5 without incurring (A* sample complexity. This is possible if we have good estimates of moment- matching parameters for set of core action-values.
Proposition 3.6 For any given 1, if ΣM
Wm MiT
2 - -mM
Dm Drr
'
11 S 6, then
EM
Wmmmm
Σ
m=1
tm
@
||gg S (2M) -8.
By Proposition 3.6, conclude that it is sufficient estimate Tи >m= Wmmm and find
X H
{(wm, pm) M that matches Tи element-wise up to accuracy 5 = (Z/H)
3.4 Main Result
The sections above have outlined the key ideas we need assuming knowledge of U. In this section, we show how we can estimate U, and we describe the complete procedure that learns a near-optimal policy of an LMAB instance. We give the details in Algorithm 1. Our algorithm is divided to three steps as detailed below.
Step 1: Estimating U. Algorithm 1 first estimates the second-order moments M₂
Em=n Wm.lmHi. Let Ü be the top-M cigenvectors of an empirical estimate of M2, where M, is constructed by collecting samples reward correlations by taking random actions at first two time steps. Note that Ü may not be a good proxy for some (mm with small mixing probability wm Rs and thus all elements U are not necessarily close to U. We defer the details of subspace recovery procedure the proof of the following lemma in Appendix C.7.
Lemma 3.7 Ler Ú be a subspace spanned by top-M eigenvectors of M2 After we estimate M2 using No = O(A*Z2 (og(ZA/m) / episodes, with probability at least 1 - 1) for all m E [M], there exists Am || 1o I S "mn such that pm Am Ü.
Our choice of dub differs two regimes as the following:
8zub
if: H 2M -
22MH2'
(Winin + c/(MHH(ZZ2MM)" c/(HVM)
min
8aub
else: H < 2M - 1, (5)
22 H
where Wiinn = minme(M] Wm-- The main difference two regimes that for 2 2M 1, when the parameters are identifiable, latent contexts with small mixing probabilities can be ignored

---

Algorithm
1: Input: Accuracy level danh, Pear > 0, model parameters M, A, 2, H.
2: // Step I: Estimate subspace U.
3: Construct M2, an estimate of M2 = E MM
WmHmpPm using No = O(A4222/88 episodes.
4: Calculate U, the span of top-M eigenvectors of M2.
5: Step Estimate {(wm. Vm.) m== )m,))
6: Get core action-value pairs {(az, b by calling Corollary 3.2.
7: Construct (Ti) min(H.2MM11) using N = mn / episodes.
8: Find valid empirical parameters {wm,, Pm } satisfying (7).
9: Step 3: Use the core-action value pairs construct an empirical LMAB.
Construct empirical model B (A,, [ {Pm)m=))
11: Output: optimal policy of B by calling a planning oracle (Definition 2.3) for B.
once Wm = o(e/M) to guarantee the closeness in distributions of observations. However, in the parameter unidentifiable regime, total variation distance is bounded only through errors in the moment space. Therefore, the estimated subspace needs to be accurate even for contexts with small mixing probabilities to keep the higher-order moments well approximated. Note that for instances with well- bal-nbaldcedd mixing probabilities, i.e., if Wmin = QLL/M), the order of duub remains the same as in the H 2 2M 1 case.
Step 2: Moment Matching, Given U, the subspace spanned by top-M eigenvectors of My, Algo- rithm 1, follows the procedure described in Sections 3.1.3.3. constructs the set of (approximate) core action-values pairs {aaj 2,) 2 (see Appendix C.8.1 1 for the detailed algorithm). Then, we construct higher-order moments of the action-value pairs. For every multi-index (i1, d2, it) [n]', using N1 Olog(Inn) /n)/ episodes, we execute at as for / 1...., and estimate higher-order moments (as also described in equation (1))
N 1 frt =
fiii, un)
N
Using standard concentration inequalities and applying the union bounds over all elements in tensors, we get || oo ōtar with probability least 1-r, r/. Then we find empirical parameters {1mm Pm) MM m=1 that satisfy
EE
th 1
ri
< 6irr:
AI [min HM - 1)].
2M We set derr O(C((ZHHHM///))) 1)) when H 2 2M - 1. In the parameter unidentifiable regime H 2M - 1, we set òisr O (Cc/H)/(Z222)))))
Step 3: Constructing Empirical LMAB. Finally, Algorithm uses the estimates {(wn, pm ) JM to construct an empirical model B = (A,, {m 1m ( after proper clipping and normal- ization. For this step to succeed, we require Wm and im to be valid parameters for the reconstruction of a valid empirical model. We state details on the recovery procedure in Appendix C.8.1.
m==
Once a valid empirical model B obtained, the remaining step is to call the planning oracle that gives an e-approximate optimal policy for B. We conclude this section with an end-to-end guarantee.
Theorem 3.8 For any LMAB instance with M latent contexts, with probability at least - 7, Algo- rithm 1 returns an e-optimal policy given total number of episades of at most
if 2 2M 2M poly(H, M, Z, A, 1/c) log(AZ/ + poly H, M) 2M-1 log(M/n)/MMM2
poly wmin: Z, A, 1/c) log(AZ/m) H² log( log(M/n) (2M7²) /2, otherwise.
Note that the dependency on A is polynomial. This polynomial term is needed to control the subspace estimation error. The exponential term is derived from the closeness in moments as discussed carlier.
7

---

Remark 3.9 (Note on Wmin) Ideally. mixing weights are well balanced such that Wmin ((1/M). However, general we may have arbitrarily small Wmin- Note that when H 2 2M 1, the upper bound does not depend on Wmin- However when H < 2M - with arbitrarily small Wminn effectively we have Wmin X M)) as can be seen in equation (5). This is mainly because in the parameter unidentifiable regime, we give guarantees from moment-closente very accurate to approximate higher-order moments well. Nevertheless, we are not aware whether this is light, and believe this to be an interesting question for future research.
and thus the subspace has to be
Continuous Rewards Gaussian Case. So far we have focused on rewards with finite support Z = O(1). However, some steps in the algorithm cannot be straightforwardly extended to continuous reward distributions. In Appendix D, we show a similar upper bound that most
poly(H, M, A, 1/e) log(A/n) + poly (H, M, log(l/ (7r))) min(H.M) /
assuming the rewards are Gaussian with unknown mean (see Theorem D.2 for the exact upper bound).
Maximum Likelihood Implementation
In the previous section, we derived procedure that learns near-optimal policy of an LMAB instance. However, our procedure relies on matching higher-order moments, which, general, not suitable for practical implementation. In this section, describe a maximum likelihood (MLE) method, motivated by our previous results. Importantly, we can use the Expectation Expectation-Maximizze Maximizatiinnn (EM) [11] heuristic to find an approximate solution the MLE optimization problem.
We can start from the set of core action-value pairs {(a,. 3,)};=1 given by Corollary 3.2. Now for every time step f. we choose is randomly from uniform distribution over [na], play action as,, and observe b, := 1l {r+- ). After repeating this for N episodes, we formulate the log-likelihood function with parameterization 0 = {(wm, )m))
m= as follows:
N
M
IN (0) =
Σ log
Σ m==1
Wm UI (brlm (1t) + (1- - (im())))
(8)
N
To prevent the confusion with searching parameters, we use q" to denote any quantity q constructed with ground truth parameters, e.g., 0" == {(win "r,, " JM Let ON be the maximum likelihood estimator, i.e., ON arg maxsge (0) in some valid parameter set O.
We can recover the sample-complexity guaranteed by moment-matching methods studied in the previous section. Specifically, we show that the maximum likelihood estimator with sufficiently many samples have nearly matching moments:
Lemma 4.1 Consider the maximum likelihood estimator O2 {(wm: Pm)) with N episodes +1,4M
for large enough N. N > C. m Jog(N/n)/6 for some sufficiently large constant C' > 0, then with probability at least 1 - 17.
|T n px S 8,
Vl E [min (H, 2M - 1]]]-
Therefore, by setting the same accuracy parameter å used Algorithm 1, we can obtain the required IN {(mm, fm)) >m)) M for matching moments. We only replace the moment-matchin step (Step
in Algorithm 1) with solving the MLE optimization problem (8).
Implementation Details While the lik-likelihood formulation (8) is still non-convex and thus intractable, we can rely on powerful heuristics: we initialize the parameters with clustering methods [2] spectral methods [1, 7J. and run the EM algorithm to improve the accuracy [11] (regardless of conditions or guarantees). Therefore, while the complexcy upper bound for both methods remains the same, benefit from the likelihooo formulation despite the non-convexity.
Specifically, we follow the same steps in Algorithm 1. We first find the core action-event pairs from second-order moments as described (Step 1). Then we form the maximum log-likelihood objective (8) and find an approximate MLE solution with the EM algorithm (Step 2). After obtaining 0yy, we recover an empirical model which we use as input to an approximate planning oracle (Step 3). We
8

---

M = 4, A 20
M 7, A 50
0.90
0.90
Per-Step Expected Rewards
Per-Step Expected Rewards
0.85
0.85
LUCRL+EM ED+MLE Tensor
LUCRL+EM
ED+MLE
Tensor
Naive UCB
Genie
Genie
0.70
0.70
0.65
0.65
0.60
20
50
15
(H)
(H)
Figure time-step rewards increasing lengths of episodes with history-dependent policies returned after
the exploration phase
compare our MLE-based implementation with experimental design (ED MLE) to other baselines applicable to learning in LMABs (e.g., L-UCRL + EM [24], tensor-decomposition methods [1, 7, 41], naive-UCB [6]). our experimental results, we can observe that in practice, our method (ED+MLE) outperforms other baselines and the worst case guarantee given in Theorem 3.8.
Provable Benefits of MLE In Appendix E.2, we also show that MLE solutions can automatically adapt to mild separation conditions. Specifically, we show that if H = ((M2/7² for some separation parameter 7 > 0 (see Assumption E.I), then learning ncar-optimal policy with the MLE solution 0x-
achieve the polynomial sample complexity for
In Figure 1, compare the averaged per time-step rewards obtained with returned policies from different algorithms after some exploration episodes. Note that with increasing H, it becomes easier to identify contexts, i.e., have bigger separations, and thus MLE based approaches benefit from larger separations. Due to space constraints, we defer the remaining details to Appendix A.
5 Conclusion and Future Work
In this work, we designed an algorithm that learns a near-optimal policy for an LMAB instance that requires only poly(A) + poly (M, -i.. number of samples. We achieved this with the experimental design and method-of-momeett and the maximum likelihood estimation which is more suitable practice. We discuss few limitations this work and future directions below.
Lower Bounds The question of a lower bound on the sample complexity of the LMAB problem remains unresolved. For LMDP, an MDP extension of LMAB, a lower bound of ( ) is known due [24]. While we conjecture that some exponential dependence in M unavoidable, characterizing the minimax dependence of the exponent is left. as an interesting future research question.
We believe that the moment-mateen matching based approach we took in this work offers a promising way for designing RL algorithms in the presence of latent contexts. Specifically, we believe these techniques can be used to design algorithm that finds near-optimal policy of LMDPs [24] or reward-mixing MDPs [23] with M O(1) and without further separation assumptions.
Latent MDPs
Linear Bandits/ / Continuous Rewards Our work has focused on the tabular setting, where all arms are independent. will be an interesting future work consider the same objective in linear bandit settings. Also, while we only considered Gaussian rewards, we believe approach can also be extended to broader class of parametric distributions. Investigating more general classes rewards is an important future research direction from practical perspective.

---

Acknowledgement
This research was partially funded by the NSF IFML Institute (NSF 2019844), the NSF AI-EDGE Institute (NSF 2112471), and NSF 1704778, and was also supported by the Israel Science Foundation (grant No. 2199/20).
References
[1] A. Anandkumar, Ge, D. Hsu, S. M. Kakade, and M. Telgarsky. Tensor decompositions for learning latent variable models. The Journal of Machine Learning Research, 15(1): 15(1):2773-2832, 2014.
[2] D. Arthur and S. Vassilvitskii. k-means++ the advantages of careful seeding. In Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithmss, pages 1027-1035, 2007.
[3] J.-Y. Audibert, Bubeck, et al. Minimax policies for adversarial and stochastic bandits. In COLT, volume 7, pages 1-122, 2009.
[4] Auer, N. Cesa-Bianchis and P. Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-256, 2002.
[5] P. Auer, N. Cesa-Bianchi, Y. Freund, and R. E. Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-77, 2002.
[6] M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. arXiv preprint arXiv: 701.00333..50 704405449, 2017.
[7] K. Azizzadencsheli, A. Lazaric, and A. Anandkumar. Reinforcement learning pomdps using spectral methods. In Conference Learning Theory, pages 193-256, 2016.
[8] A. Bhaskara, M. Charikar, A. Moitra, and Vijayaraghavan. Smoothed analysis of tensor de- compositions. In Proceedings of the forty-sixth annual ACM symposium Theory of computing, pages 594-6004400 ACM, 2014.
[9] E. Brunskill and L. Li. Sample complexity of multi-task reinforcement learning. In Uncertainty in Artificial Intelligence, page 122. Citeseer, 2013.
[10] R. Chawla, A. Sankararaman, A. Ganesh, and S. Shakkottai. The gossiping insert-eliminate algorithm for multi-agent bandits. In International Conference on Artificial Intelligence and Statistics, pages 3471-3481. PMLR, 2020.
[11] A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the Royal Statistical Society: Series B (Methodological), 39(1):1-22, 1977.
[12] N. Doss, Y. Wu, P. Yang, and H. H. Zhou. Optimal estimation of high-dimensional gaussian mixtures. arXiv preprint arXiv:2002.058/8, 2020.
[13] A. Garivier and E. Moulines. On upper-confidence bound policies for switching bandit problems. In International Conference on Algorithmic Learning Theory, pages 174-188. Springer, 2011. 1.
[14] C. Gentile, S. Li, and G. Zappella. Online clustering of bandits. International Conference on Machine Learning, pages 757-765, 2014.
[15] A. Gopalan, O.-A. Maillard, and M. Zaki. Low-rank bandits with latent mixtures. arXiv preprint arXiv: 01000000150000 508, 2016.
[16] A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv preprint arXiv:/502.000 /502.02259 2015.
[17] J. Hu, X. Chen, C. Jin, L. Li, and L. Wang. Near-optimal representation learning for linear bandits and linear RL. In International Conference on Machine Learning, pages 4349-4358 PMLR, 2021.
10

---

[18] T. Jaksch, R. Ortner, and P. Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11:1563-1600, 2010.
[19] C. Jin, S. Kakade, A. Krishnamurthy, and Q. Liu. Sample-efficient reinforcement learning of undercomplete pomdps. Advances in Neural Information Processing Systems, 33:1 13:18530-1853, 2020.
[20] J. Kiefer and J. Wolfowitz. The equivalence of two extremum problems. Canadian Journal of Mathematics, 12:363-366, 1960.
[21] B. Kveton, M. Konobeev, M. Zaheer, C.-w. Hsu, M. Mladenov, C. Boutilier, and C. Szepesvari. thompson sampling. In International Conference on Machine Learning, pages 5884-5893 PMLR, 2021.
[22] Kwon and C. Caramanis. The EM algorithm gives sample-optimality for learning mixtures of well-separated gaussians. In Conference Learning Theory, pages 2425-2487, 2020.
[23] Kwon, Y. Efroni, C. Caramanis, and Mannor. Reinforcement learning in reward-mixin mdps. Advances in Neural Information Processing Systems, 34, 2021.
[24] Kwon, Y. Efroni, C. Caramanis, and S. Mannor. RL for latent mdps: Regret guarantees and lower bound. Advances in Neural Information P Processing Systems, 34, 2021.
[25] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Coordinated attacks against contextual bandits: Fundamental limits and defense mechanisms. In Proceedings of the 39th International Conference Machine Learning, pages 11772 PMLR, 2022.
[26] T. Lattimore and C. Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
[27] T. Lattimore, C. Szepesvari, and G. Weisz. Learning with good feature representations in bandits and in with a generative model. International Conference on Machine Learning, pages 3662-5670. PMLR, 2020.
[28] M. L. Littman, Cassandra, and L. P. Kaelbling. Learning policies for partially observable environments: Scaling up. In Machine Learning Proceedings 1995, pages 362-370. Elsevier, 1995.
[29] C.-Y. Liu and I. Li. On the prior sensitivity of thompson sampling. In International Conference on Algorithmic Learning Theory, pages 321-336. Springer, 2016.
[30] Q Liu, A. Chung, C. Szepesvári, and C. Jin. When is partially observable reinforcement learning not scary? arXiv preprint wwxiv:2204.08661 2022.
[31] O.-A. Maillard and S. Mannor. Latent bandits. In International Conference on Machine Learning, pages 136-144, 2014.
[32] A. Modi, N. Jiang, A. Tewari, and S. Singh. Sample complexity of reinforcement learning using linearly combined model ensembles. In International Conference on Artificial Intelligence and Statisties, pages 2010-2020. PMLR, 2020.
[33] J. Pineau, G. Gordon, and S. Thrun. Anytime point-based approximations for large pomdps. Journal of Artificial Intelligence Research, 27:335-380, 2006.
[34] W. Schudy and M. Sviridenko. Concentration and moment inequalities for polynomials of independent random variables. In Proceedings of the twenty-third annual ACM-SIAM symposium on Discrete Algorithms, pages 437-446. SIAM, 2012.
[35] M. Simchowitz, C. Tosh, A. Krishnamurthy, D. J. Hsu, T. Lykouris, M. Dudik, and R. E. Schapire. Bayesian decision makinng under misspecified priors with applications to meta-lcarning. Advances in Neural Information Processing Systems, 34, 2021.
[36] A. Slivkins and E. Upfal. Adapting to a changing environment: the brownian restless bandits. In COLT, pages 343-354, 2008.

---

[37] L. N. Steimle, D. L. Kaufman, and B. T. Denton. Multi-model markov decision processes. Optimization Online URL http://www. optimization-onlite org/DB F1LE/2018/1/64553 pdf. 2018.
[38] W. R. Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285 294, 1933.
[39] M. J. Todd. Minimum-volume ellipsoids: Theory and algorithms. SIAM, 2016.
[40] Y. Wu, P. Yang, et al. Optimal estimation of gaussian mixtures via denoised method of moments.
Annals of Statistics, 18(4))1981-2007, 2020.
[41] X. Zhou, Y. Xiong, N. Chen, and X. Gao. Regime switching bandits. In Thirty- Fifrtt Conference on Neural Information Processing Systems, 2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This work is primarily focused on the theoretical aspect and its purpose is too general for us speculate about broader societal impacts.
(d) Did you describe the limitations your work? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]: in Supplementary Materials
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] in Supplementary Materials
(c) Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [No] We did average errors over multiple independent
experiments, but error bars are not reported due to computational burden.
(d) Did you include the amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
The remaining issues are not applicable us.
12


Reviews:
[{"review_id": "q4_0Bty_tyb", "reviewer": "Reviewer_E6HY", "summary": "This paper studies latent MABs with discrete reward. First, the paper suggested that in order to identify a near-optimal policy, it suffices to identify the high-order moments of the action-reward distribution, which essentially recover the distribution over all possible action-reward trajectory. Second, since this tensor is low rank, the paper proposed to decouple the estimation into three steps: first estimate the subspace, then use experimental design to find a core-set, and finally only need to estimate a sub high-order tensor corresponding to the core set. By doing so, they avoid the $A^{\\mathcal{O}(\\min\\{M,H\\})}$ factor in sample complexity.", "questions": "Please see my questions above.", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The result is almost assumption-free. Even if this paper only focuses on the finite MAB setting, it is nice to see they don't require any observable or revealing condition. \n\n2. The algorithm design is cute. First estimate the subspace, and then instead of directly estimating the $\\ell_2$ projection onto the subspace, they use the subspace to compute a $\\ell_1$ core set of much smaller size than $AZ$.\n\n\nWeakness:\n\n1. It is hard for me to tell how significant it is to improve from $(AZ)^M$ to $Z^M$. If $M$ is large, the result obtained is still exponential in $M$. If $M$ is small, then $A^M$ is still fine?\n\n2. I found the technical writing in Section 3 is a little bit over-detailed. For example, in Section 3.1, the paper spent much words motivating the core-set trick before giving an overview of the overall algorithm flow. Maybe give an algorithmic overview first and then explain each module in detail? Anyway this is just my own feeling and not an important issue.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "oYapmr3hve", "reviewer": "Reviewer_rtQX", "summary": "The paper studies episodic latent MABs, where the learner interacts with a sequence of bandit problems in episodes of H steps and each bandit problem is associated with a latent context m \\in [M] which characterizes the arms' reward distributions. The authors present an algorithm based on moment matching that achieves a sample complexity polynomial in A and exponential in M or H, hence improving over a trivial reduction to a partially-observable MDP which would yield A^H complexity (at least in the regime where M or H is small).", "questions": "Overall I think the paper provides an interesting and significant contribution (though, as mentioned above, I am not very familiar with the related literature). I am thus more inclined to vote for acceptance. Here are a few additional comments/questions besides the strengths/weaknesses reported above.\n\n1. While the episodic latent MAB setting introduces many challenges, I think it should be better motivated. For instance, in the recommender system example given in the first paragraph of the introduction, it is typically reasonable to assume that one knows some information about the user interacting with the system, in which case it is possible to use standard contextual-bandit strategies. Even when no information about the latent contexts is available, if H is large, one can simply run a no-regret bandit algorithm from scratch at each episode. While this would completely ignore the structure of the problem (and probably yield poor performance), it would result in a much simpler and tractable strategy. Could the authors characterize other relevant problems that fit into this formalism and where one really cares about competing with the optimal history-dependent policy (eg where running a contextual bandit algorithm is not possible)?\n\n2. One of the main assumptions here is that the learner knows the number of contexts M. What if M is unknown? Would there be any hope to estimate it?\n\n3. Is w_min, the minimum probability of sampling a context, assumed to be strictly positive? What if it is zero?\n\n4. Another line of related works that should be probably mentioned is the one of \"structured bandits\" (see eg [1,2,3] below and references therein). As in the setting studied here, it is assumed that the arms' reward distributions dependent on some latent parameter. Differently from this paper, the learner knows the mappings but does not know the latent parameter\n\n5. A minor comment on wording: a sample complexity of O(A^H) is still polynomial in A. So when stating the sample complexity as depending on poly(A), it would be good to clarify that the exponent of A is neither H nor M.\n\n6. One downside of the approach is that it must explicitly distinguish between the large and small H regimes. Would it be possible to do a more unified treatment? \n\n\n[1] Lattimore, Tor, and Rémi Munos. \"Bounded regret for finite-armed structured bandits.\" Advances in Neural Information Processing Systems 27 (2014).\n[2] Tirinzoni, Andrea, Alessandro Lazaric, and Marcello Restelli. \"A novel confidence-based algorithm for structured bandits.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2020.\n[3] Degenne, Rémy, Han Shao, and Wouter Koolen. \"Structure adaptive algorithms for stochastic bandits.\" International Conference on Machine Learning. PMLR, 2020.", "limitations": "The authors discussed the limitations of their work. I do not see any societal impact.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Novelty: the method proposed here, its analysis, and the results are, to my knowledge, novel (though I am not super familiar with the related literature on latent bandits/MDPs)\n- The main result seems interesting and significant as it advances our understanding of the episodic latent MAB setting (in particular by confirming that the latter is provably easier than latent MDPs)\n- While notation is quite complicated and there are a lot of technical details, I think the paper is very well written. In particular, the authors provide both intuitive and formal explanations of their results and clearly discuss the results while properly comparing them with related literature\n\nWeaknesses:\n- While the problem studied here is definitely very complicated and requires quite technical methods to be solved, I feel that the setting itself should be better motivated (see below)\n- The results are only for rewards with finite support or Gaussian distributions and it is not clear whether they extend beyond that (eg to general sub-Gaussian rewards). This might limit the overall scope\n- The \"theoretical\" algorithm based on tensor decomposition seems to be quite computationally inefficient. While the authors propose a simpler maximum-likelihood implementation, that still requires solving a non-convex optimization problem for which one has to rely on heuristics (hence loosing theoretical guarantees).\n- While the main sample complexity result improves over the naive A^H baseline and advances our understanding of the problem at hand, it is still exponential in either M or H. Since no lower bound is provided, one still wonders whether such an exponential dependence is unavoidable\n- While some numerical simulations have been reported in appendix, I feel that they should be given more importance and be moved to the main text", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "jWYh5fHLYU3", "reviewer": "Reviewer_RY7y", "summary": "The present paper considers the challenging policy learning problem of latent multi-armed bandits. The setting as follows. There are M latent contexts, and corresponding to each latent context there is a distribution over multi-armed bandit problems. The objective of the learner is to identify a policy for sequentially choosing actions in the context of a random instance of such a multi-armed bandit problem. Performance is quantified in terms of the cumulative regret over the horizon H, when averaged over a distribution on the set of M contexts. In the setting, both the individual bandit problems for each context, and the distribution over contexts are unknown a priori. The key question is how many episodes the learner must explore before an epsilon-optimal policy is identified. In particular, is this number of the same order of the cardinality of the set of all action sequences (i.e. A^H, where A is the number of actions and H is the horizon).\nThe present paper demonstrates that it is possible to learn an epsilon-optimal policy in poly(A)+ploy(H,M)^min(H,M) episodes. \nA key insight is that when the number of latent contexts M is significantly less than the number of actions A then one can reduce the scope for unfavorable dependencies upon the number of actions by projecting onto the M dimensional space of reward vectors, for each contexts. This motivates an approach where the goal is efficiently project sequences of actions and rewards onto this space. This requires an interesting adaptation of the Keifer-Wolfowitz theorem. The key difficulty then lies in identifying the space itself.\nThe analysis considers two separate regimes. The first is where the horizon is at least twice the number of latent contexts. If this is the case then the dimension reduction approach is used to infer the context. In the second regime, the horizon is to small to identify the context. If this is the case then a more direct approach is taken which directly matches the Hth order moments.\nThe main result is Theorem 3.8 which gives a bound on the required sample complexity which is polynomial in A, the number of actions. \nThe paper also considers a more traditional maximum likelihood based variant of the methodology based on the EM algorithm, which performs well in practice.\n", "questions": "Would anything be lost in just assuming that all of the arm-dependent reward distributions are Bernoulli? It seems that this would simplify the presentation.", "limitations": "The primary limitations are discussed above in the “strengths and weaknesses” section.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper addresses a very challenging and interesting problem, which seems well-motivated by practical applications. In particular, the setting goes substantially beyond the closely related area of learning priors in Bayesian multi-armed bandits. Moreover, the algorithmic approach based on dimensionality reduction and Keifer-Wolfowitz theory seems novel in this setting.\nI think a major limitation with the paper is the quality of the writing. I found the paper quite difficult to read in many places. This is partly due to my lack of expertise in this area. That being said, I think there were also issues with the clarity of presentation. As an example, the abstract refers to “A”. This denotes the number of actions, but we don’t actually learn this until line 32 of the text. There are also various strange uses of brackets. As another example, consider Assumption 2.2. In particular, its not at all clear what Z=O(1) is with respect to at this stage. \nI would also argue that Figure 1 gives relatively little value over and above Section 1.2. It would be preferable to use the space to include a figure on empirical results for Section 4.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YQyQy4qUJs", "reviewer": "Reviewer_XPh1", "summary": "This paper studies the problem called Episodic Latent Multi-Armed Bandits, which includes $M$ MABs with $A$ actions. In each episode, an MAB is sampled with some probability, then the agent interacts with the MAB for $H$ steps, without knowing the true latent context. Previous works derived algorithms to learn a near-optimal policy with $O(A)^H$ episodes. This paper exploits the low-rank structure of the latent reward distributions and uses a moment-matching method to derive an algorithm with a $\\mathrm{poly}(A)+\\mathrm{poly}(H, M)^{\\min \\\\{H, M \\\\}}$ sample complexity. The analysis depends on a finite, discrete support set of reward distributions, and can be extended to a Gaussian. This paper further provides a Maximum Likelihood implementation to approximate high-order moment matching and it shows relatively good results.", "questions": "None.", "limitations": "None.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This paper gives nearly unconditional (without strong separability assumptions) results for LMAB problems, which are all polynomial sample complexities in $A$. This shows a separating gap between LMAB, Latent MDP, and reward-mixing MDPs.\n2. To my knowledge, the use of the Kiefer-Wolfowitz theorem (Lemma 3.1 and Corollary 3.2) for dimension reduction with bandit feedback is novel.\n3. This paper provides simulation to corroborate theories.\n4. This paper is clearly written, with clear notations and logic.\n5. The moment-matching technique may be further extended to LMDPs. \n\n\nWeaknesses:\n1. The assumption that either $H = O(1)$ or $M = O(1)$ is not negligible. I hope the authors can either provide a $\\mathrm{poly}(H, M)^{\\min \\\\{H, M \\\\}}$ lower-bound or further improve on this exponential term.\n2. Some polynomial terms in Theorem 3.8 have large constant exponents. I don't know if this is tight or not. But according to the reward-mixing MDP paper [22], I guess they have room to be improved.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept