Tractable Optimality in Episodic Latent MABs
Anonymous Author(s)
A  liation Address email
Abstract
We consider a multi-armed bandit problem with M latent contexts, where an agentinteracts with the environment for an episode of H time steps.
Depending on thelength of the episode, the learner may not be able to estimate accurately the latentcontext.
The resulting partial observation of the environment makes the learningtask signiﬁcantly more challenging.
Without any additional structural assumptions,existing techniques to tackle partially observed settings imply the decision makercan learn a near-optimal policy with O ( A ) H episodes, but do not promise more.
In this work, we show that learning with polynomial samples in A is possible.
Weachieve this by using techniques from experiment design.
Then, through a methodof-moments approach, we design a procedure that provably learns a near-optimalpolicy with O (poly( A )+poly( M, H ) min( M,H ) ) interactions.
In practice, we showthat we can formulate the moment-matching via maximum likelihood estimation.
In our experiments, this signiﬁcantly outperforms the worst-case guarantees, aswell as existing practical methods.
Introduction
In Multi-Armed Bandits (MABs), an agent learns to act optimally by interacting with an unknownenvironment.
In many applications, interaction sessions are short, relative to the complexity of theoverall task.
As a motivating example we consider a large content website that interacts with usersarriving to the site by making sequential recommendations.
In each such episode , the system hasonly a few chances to recommend items before the user leaves the website – an interaction timetypically much smaller than the number of actions or the types of users.
In such settings, agents oftenhave access to short horizon episodes and have to learn how to process observations from di  erentepisodes to learn the best adaptation strategy.
Motivated by such examples we consider the problem of learning a 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 by nature.
Without knowing theidentity of the environment (we call this the latent context ), an agent aims to maximize the expectedcumulative reward per episode (see Deﬁnition 2.1 for a formal description).
The LMAB frameworkis di  erent from the setting considered in [29] where multiple episodes proceed in parallel withoutlimiting the horizon of an episode H .
For long horizons H , we show that it is possible to ﬁnd anear-optimal policy and to determine near-optimal actions for each episode, as if we knew the hiddencontext.
If H ⌧ A where A is the total number of actions, however, this is no longer possible.
Instead,we aim to learn the best history-dependent policy for H time steps.
1.1 Our Results and Contributions
In this work, we consider LMABs for both long and short horizon H .
In the long horizon case, weshow we can learn the problem latent parameters; yet in the short horizon case, a di  erent perspectiveis required, as learning parameters is impossible.
Partially observed problems are not a new problem, and have attracted considerable attention fordecades.
Thus one natural thought is to resort to more general modeling: LMABs can be viewed asa special case of partially observable Markov decision processes (POMDPs). From this viewpoint,we can reformulate the problem as learning in a very large MDP with O
( A ) H states (the number ofall possible histories).
Appealing to standard algorithms for MDPs ( e.g., [16, 6]) we can learn an✏ -optimal policy with O ( A ) H / ✏ 2 samples. For general POMDPs, this prohibitive dependence on Aand H is (provably) unavoidable [19].
A natural way to sidestep these bottlenecks is through searching for additional structure.
Accordingly,much past work has relied on assumptions such as separability, or reachability (see Section 1.2 formore details and references).
Methods developed in prior work critically rely on such assumptions. In contrast, our goal is to leverage directly the fact that LMABs have signiﬁcantly more structure thana general POMDPs, and to use this, rather than rely on reachability, separability, etc. Indeed, withoutsuch assumptions, no current work (see Section 1.2 for more) has succeeded in treating LMABs withfewer than O ( A ) H / ✏ 2 samples. Our goal in this paper is to answer the following question: for M =
O (1) ⌧
A contexts, can we learna near-optimal policy with fewer than O ( A ) H / ✏ 2 samples, in the absence of further assumptions ?
Our work answers the above question a  rmatively.
Speciﬁcally, our main contributions are as follows.
We show that the dependence of our algorithm is poly( A )+poly( H, M ) min
( H,M ) , and thus tractablewhen either H or M is small, even for very large A .
That is, we are particularly focused on the settingwith a few contexts or relatively short episodes (in comparison to A ), i.e., M = O (1) or H
= O (1) ,where a natural objective is to learn a near-optimal history-dependent policy for H time steps (seealso Figure 1 and Appendix B.1).
1.2 Related Work
Due to space constraints, we discuss only the most closely related work here, and defer a lengthierdiscussion on related work to Appendix B.1.
Learning priors in multi-armed bandit problems Several recent works have considered the Bayesian learning framework with short time-horizon [27, 33, 20].
The focus in this line of work ison the design of algorithms that learn the prior, while acting with a ﬁxed Bayesian algorithm (e.g.,
Thompson sampling).
While Bayesian learning with short time-horizon may be viewed as a specialcase of LMABs, the baseline policy we compare ourselves to is the optimal H -step policy, which is aharder baseline than considering a ﬁxed Bayesian algorithm.
Latent MDPs Some prior work considers the framework of Latent MDPs (LMDPs), which is the
MDP generalization of LMAB [15, 35, 23, 22].
In particular, [23] has shown the information-theoreticlimit of learning in LMDPs with no assumptions on MDPs, i.e., an exponential number of sampleepisodes ⌦ ( A M ) is necessary to learn a near-optimal policy in LMDPs.
In contrast, we show thatin LMABs, the required number of episodes can be polynomial in A .
This does not contradict theresult in [23], since their lower bound construction comes from the challenge in state-exploration with latent contexts.
In contrast, there is no state-exploration issue for bandits, which enables thepolynomial sample complexity in A .
Furthermore, our upper bound does not require any assumptionsor additional information such as good initialization or separations as in [23].
To the best of ourknowledge, no existing results are known for learning a near optimal policy of LMAB instances for
M  3 without further assumptions.
Multitask RL with explicit clustering We can cluster observations from each episode if we are givensu  ciently long time horizons H =
˜ ⌦
( A/  2 ) in each episode [9, 13, 15].
Here,  > 0 is the amountof ‘separation’ between contexts such that for all m 6 = m 0 , max a 2 A k µ m ( a, · )  µ m 0 ( a, · ) k 2  where µ m is a mean-reward vector of actions in the m th context.
We focus on signiﬁcantly moregeneral cases where there is no obvious way of clustering observations, e.g., when H ⌧ A or µ mcan be arbitrarily close to some other µ m 0 for m 6 = m 0 .
If we are given a similar separation condition,we also show that a polynomial sample complexity is achievable as long as H = ˜ O ( M 2 /  2 ) (seealso section E.2). Note that this could still be in H ⌧ A regime with large number of actions A . Learning in POMDPs with full-rank observations One popular learning approach in partiallyobservable systems is the tensor-decomposition method, which extracts the realization of modelparameters from third-order tensors [1, 7, 14].
However, the recovery of model parameters requirespeciﬁc 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 reachedwith non-negligible probabilities by any exploration policy for the parameter recovery.
Recent resultsin [17, 28] have shown that the uniform reachability assumption can be dropped with the optimismprinciple.
However, they still require the full-rankness of a test-observation matrix to keep thevolume of a conﬁdence set explicitly bounded.
Since LMAB instances do not necessarily satisfy thefull-rankness assumption, their results do not imply an upper bound for learning LMABs.
Preliminaries
We deﬁne the problem of episodic latent multi-armed bandits with time-horizon H  2 as follows:
Deﬁnition 2.1 (Latent Multi-Armed Bandit (LMAB)) Let
LMAB be a tuple
B =( A , { w m } Mm =1 , { µ m } Mm =1 ) , where A is a set of actions, { w m } Mm =1 are the mixing weights 101 such that a latent context m is randomly chosen with probability w m , and µ m is the model parameterthat describes a reward distribution, i.e., P µ m ( r | a ) :
= P ( r | m, a ) , according to an action a 2 Aconditioning on a latent context m .
(NB: each µ m represents a probability model, not necessarily amean reward vector).
We do not assume a priori knowledge of mixing weights.
The bulk of this paper considers discretereward realizations, when the support of the reward distribution is ﬁnite and bounded.
In Appendix D, we also show that our results can be adapted to Gaussian reward distributions.
Assumption 2.2 (Discrete Rewards)
The reward distribution has ﬁnite and bounded support.
Thereward attains a value in the set Z .
We assume that for all z 2 Z we have | z |  1 .
We denote the 110 cardinality of Z as Z and assume that Z = O (1) .
As an example, Bernoulli distribution satisﬁes Assumption 2.2 with Z = { 0 , 1 } and Z = 2 .
We denotethe probability of observing a reward value z by playing an action a as µ m ( a, z ) :
= P ( r = z | m, a ) ina context m .
We often use µ m as a reward-probability vector in R AZ indexed by a tuple ( a, z ) 2 A ⇥ Z .
At the beginning of every episode, a latent context m 2 [ M ] is sampled from a mixing distribution{ w m } Mm =1 and ﬁxed for H time steps, however we cannot observe m directly.
We consider a policyclass ⇧ which contains all history-dependent policies ⇡ : ( A ⇥ Z ) ⇤ ! A .
Our goal is to ﬁnd 117 a near optimal policy ⇡ 2 ⇧ that maximizes the expected cumulative reward V for each episode
V ? = max ⇡ 2 ⇧ V ( ⇡ ) :
= E ⇡ hP
Ht =1 r t i , where the expectation is taken over latent contexts andrewards generated by an LMAB instance, and actions following a policy ⇡ .
Deﬁnition 2.3 (Approximate Planning Oracle)
A planning oracle receives an LMAB instance Band returns an ✏ -approximate policy ⇡ such that V ?  V ( ⇡ )  ✏ .
Concretely, the point-based value-iteration (PBVI) algorithm [31] is an ✏ -approximate planningalgorithm which runs in time
O
( HMAZ ( H 2 / ✏ ) O
( M ) ) .
Additional notation.
For any quantity q with respect to the true LMAB, B , we use ˆ q to refer to itscorresponding empirical estimate.
We use w min :=
min m 2 [ M ] w m for the minimum mixing weight.
For any l th -order tensor T l 2 R n ⇥ ... ⇥ n ( n repeated l times), we denote k T l k 1 for the element-wiselargest absolute value.
For any vector in v 2 R n in dimension n 2 N + , we use v N l to denote adegree l tensorization of v .
We often use k µ m ( a, · )  ˆ µ m ( a, · )
k 1 to mean the l 1 statistical distance 129 between reward distributions: P z 2 Z | µ m ( a, z )  ˆ µ m ( a, z ) | .
Lastly, we denote by a t and r t as theaction and reward realizations observed at time step t 2 [ H ] .
Sample-Complexity of Learning LMABs
In this section, we develop our main algorithm that learns a near-optimal policy of an LMAB instancewith O (poly( A )+poly( M, H ) min( M,H ) ) samples.
Towards achieving this goal, we ﬁrst elaborate ona low rank representation of the LMAB problem.
Then, we show how experimental design techniquescan assist in utilizing this structure to improve over the naive upper bound to the problem.
3.1 Dimensionality Reduction via Experimental Design
When H = 1 , estimating the latent context is not possible, and playing a standard MAB strategy(e.g., UCB) is optimal.
However, to obtain a near-optimal history-dependent policy with longertime-horizons H > 1 , tracking the mean reward of an action is not enough, since rewards fromprevious time steps are correlated with current and future rewards. Moreover, when H > 1 wecan collect more information such as correlations between di  erent actions. This allows the use ofhigher-order statistics.
Speciﬁcally, we consider a model-based learning approach from the l th -order moments of rewardobservations.
Since there are A actions, this implies that we have O ( A l ) quantities to estimate, whichwould incur O ( A l ) sample complexity if we separately measure every correlation between l pairs ofactions.
On the other hand, when M ⌧ A , the distributions µ m (recall these are reward probabilitiesfor each action, in each context) occupy only a M -dimensional subspace of AZ -dimensional space:
U := span { µ 1 , µ 2 , . .
. , µ M } .
Thus if we could estimate U , and project all observations to U , wecould remove the unfavorable dependence on A .
Even when the subspace U is known a priori, dimensionality reduction with bandit feedback isnon-trivial.
We need to compute directly the projection onto U of estimates of the reward probabilityvectors, { µ m } Mm =1 .
To see the challenge, let {  j } Mj {+=
 ⇢+} R AZ be an orthonormal basis of U .
Tocompute the projected estimate µ > m  j , we need a sampling policy ⇡ such that ⇡ ( a ) / P z |  j ( a, z ) | .
The variance of this sampling policy can be as much as O ( k  j k 21 ) , which in general scales withk  j k
= O ( A ) .
Therefore, reliable estimation for any statistics in the reduced dimension wouldstill have a dependence on A .
In particular, since our approach is based on higher-order method-ofmoments, an estimation of l th -order statistics for l  2 would require ⌦ ( A l/ 2 ) samples.
The general idea of dimensionality reduction is not fatally ﬂawed.
Instead, we need to avoid thepitfalls of the approach outlined above.
First, the calculation above shows that we need to control thek · k 1 -norm of the vectors  j , ideally, k  j k 1 =
O (1) .
Second, we do not need to approximate our the 161 reward-probability vectors in l 2 -norm, but only in a k · k 1 , 1 -norm.
Indeed, to compute an ✏ -optimal 162 policy, we only need a good estimate for each µ m such that max a [Equation] A k µ m ( a
[Equation] µ m ( a [Equation] k
[Equation] ✏ . The key is to show that we can choose  j ’s to be a subset of the standard basis in R AZ (hencethey will have k  j k 1 = 1 ), in such a way that guarantees the approximation quality for µ m fromestimating µ > m  j ’s. In terms of the original problem, the existence of such a subset of the standardbasis is equivalent to the existence of a small set of informative action-value pairs, that are su  cientlycorrelated with all other action-value pairs.
This is called a core set in the experimental designliterature, and its existence in our context, is an important consequence of the Kiefer–Wolfowitztheorem.
Speciﬁcally, 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 R d where d  k and let u, ¯ u 2 U .
There exists an algorithm that runs in time ˜ O
( dk 2 ) and returns the following.
A core set of at most n = 4 k log log k + 16 coordinates { i j } nj =1 ✓ [ d ] such that
[Equation]
A linear transformation that maps [¯ u ( i 1 ) , . .
, ¯ u ( i n )] to its corresponding ¯ u 2 U .
Note that in our setting, d = AZ and k = M .
We prove Lemma 3.1 in Section C.1, essentially asa corollary of the Kiefer–Wolfowitz theorem (and its geometric interpretation) for (near)-optimalexperimental design [18, 37].
In the context of bandits and RL, the Kiefer–Wolfowitz theorem has been used to study how themisspeciﬁcation in linear representation changes the problem landscape ( e.g., [25, 30]).
However,experimental design has not been previously used for dimensionality reduction for problems withbandit feedback (as far as we know).
We believe it is a powerful tool.
We review the basics ofexperimental design in Appendix B.2.
A direct consequence of Lemma 3.1 is an algorithm to ﬁnd a set of coordinates of size ˜ O ( M ) that aresu  cient to reconstruct the latent reward model { µ m } M m =1 where µ m 2 R AZ .
We refer to this set of 185 coordinates as the set of core action-value pairs.
Corollary 3.2 Suppose the subspace U = span { µ 1 , . . .
, µ M } is given.
Then for any ˆ µ 2 U , thereexists an algorithm that runs in time ˜
O
( ZAM 2 ) and returns the following.
A set of core action-value pairs { ( a j , z j ) } nj = ✓ A ⇥ Z of size at most n = 4 M log log M + 16 , such that for all m 2 [ M ]
[Equation]
where ⌫ m , ˆ ⌫ m 2 R n such that ⌫ m ( j ) :
=
µ m ( a j , z j ) and ˆ ⌫ m ( j ) :
= ˆ µ m ( a j , z j ) .2.
A linear transformation that maps [ˆ ⌫ i 1 ) , . .
, ˆ ⌫ ( i n )] to its corresponding ˆ µ 2 U .
Assuming access to the subspace
U (we remove this assumption in Section 3.4), Corollary 3.2 impliesa strategy for estimating the latent model parameters { µ m } Mm =1 : obtain core action-value pairs{ ( a j , z j ) } nj =1 , estimate { ⌫ m } Mm =1 , and construct latent model parameters via a linear transformationof { ⌫ m } Mm =1 .
In the following sections, we show that by matching higher-order moments we canestimate { ⌫ m } Mm =1 to su  ciently good accuracy with only polynomial dependence in A .
3.2 H  2 M  1 : Identiﬁable Regime in Wasserstein Metric
In the regime where that H  2 M  1 we can measure moments up to order (2 M  1) th .
Then,we can leverage recent advances in learning parameters of mixture distributions from higher-ordermoments [38, 11].
That is, we can recover { ⌫ m } Mm =1 by estimating 2 M  1 higher-order momentsand ﬁnd { ˆ ⌫ m } Mm =1 that matches them.
This further emphasizes the importance of the dimensionality 200 reduction we take to obtain the set of core action-values pairs; otherwise, a moment-based approachto recover { µ m } Mm =1 would have an exponential dependence in A .
We now elaborate on the estimation procedure of { ⌫ m } Mm =1 from higher-order moments.
We deﬁnethe l -order tensor as T l := [Equation]
[-Mm =
-]m .
For an LMAB instance, we can access the tensor T l usingobservational data by simply estimating correlations between l core action-value pairs within eachepisode.
Speciﬁcally, for any I
= ( i 1 , i 2 , ..., i l ) 2 [ n ] l the tensor T l ( i 1 , i 2 , ..., i l ) is also given by
[Equation]
where ⇡ I is a policy that performs the sequence of actions ( a i 1 , a i 2 , ..., a i l ) for t = 1 , . .
, l .
Suppose we ﬁnd estimators { ( ˆ w m , ˆ ⌫ m ) }
Mm =1 such that moments match T l up to error  for all l 2[2 M  1] .
The results in [38, 11] imply that the closeness in moments of up to (2 M  1) th degreeimplies the closeness in model parameters:
Lemma 3.3 Suppose k T l  P
Mm =
ˆ w m ˆ ⌫ N l m {+k 
+} <  for all l = 1 , 2 , ..., 2 M  1 for some su  cientlysmall  > 0 .
Then,
[Equation]
where  is a joint distribution over ( m, m 0 )
[ M ] 2 satisfying
[Equation]
The form of guarantee given for { ˆ w m , ˆ ⌫ m } Mm =1 is in the Wasserstein distance between two latentmodel parameters.
A useful property of the Wasserstein metric is that the distance measure is invariantto permutation of individual components, and ﬂexible with arbitrarily small mixing probabilities orarbitrarily close components [38] (see also Appendix B.3 for the review on Wasserstein distance).
Once we obtain estimates of { ⌫ m } Mm =1 , we can estimate { ˆ µ m } Mm =1 as implied by Corollary 3.2 (see 218 Appendix C.8 for further details).
We then show that for any history-dependent policy ⇡ , the expectedcumulative rewards of ⇡ are approximately the same for any close LMAB instances.
Proposition 3.4 Let B =
( A , { w m } Mm =1 , { µ m } Mm =1 ) and ˆ B
=
( A , { ˆ w m } Mm =1 , { ˆ µ m } Mm =1 ) be any 221 two LMABs.
Then, for any history-dependent policy ⇡ : ( A ⇥ Z ) ⇤ !
A , we have
[Equation]
where the inﬁmum over  is taken over joint distributions over ( m, m 0 ) satisfying (3) .
From Corollary 3.2 we have max a 2 A k µ m ( a, · )  ˆ µ m 0 ( a, · ) k 1  2 Z p 2 M k ⌫ m  ˆ ⌫ m 0 k 1 for any 224 m, m 0 2 [ M ] .
Plugging this into Proposition 3.4, we have
[Equation]
Thus, using Lemma 3.3 with  < (poly( H, Z, M, n ) / ✏ ) 2 M  1 , we can conclude that any ✏ -optimalpolicy for ˆ B is O ( ✏ ) -optimal for the underlying LMAB B .
3.3 H < 2 M  1 : Unidentiﬁable Regime with Short Time-Horizon
A more interesting regime is the one in which the time-horizon H is smaller than the required degreeof moments 2 M  1 .
For such a setting we cannot measure moments of degree higher than H .
Therefore, if H < 2 M  1 , we cannot rely on the identiﬁablility of the underlying LMAB model.
Instead, we make the following claim: to compute an optimal policy only for time horizon H < 2 M  1 ,we only need to match the H th order moments.
This is formalized in the next lemma.
Lemma 3.5 Suppose that k
[Equation]
k
[Equation]
  , then for any history de- 234 pendent policy ⇡ [Equation] we have
[Equation]
Therefore, it is su  cient to ﬁnd estimates of mixing weights and latent models that match the measured
H th -order moment, and then compute an optimal policy for the estimated model.
Lemma 3.5 isnatural for discrete reward distributions with bounded support.
Interestingly, we extend this result to
Gaussian rewards, which are continuous and unbounded, in Appendix D.
Hence, a natural idea is to estimate parameters { ˆ w m } Mm =1 , { ˆ µ m } Mm =1 that satisfy the condition in 240 Lemma 3.5 without incurring O
( A H ) sample complexity.
This is possible if we have good estimatesof moment-matching parameters for a set of core action-values.
Proposition 3.6 For any given l  1 , if k P Mm =1 w m ⌫ N l m  P
{+Mm =
 ˆ w m ˆ ⌫ N l m k 1   , then M X m = w m µ N l m +}
M X m = ˆ w m ˆ µ N l m  (2 M ) l/ 2 ·  .
By Proposition 3.6, we can conclude that it is su  cient to estimate T H :=
P
Mm [-=
-]=1 w m ⌫ N H m and ﬁnd{ ( ˆ w m , ˆ ⌫ m ) } Mm =1 that matches T H element-wise up to accuracy  := ( ✏ /H ) /
( Z p 2 M ) 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-optimalpolicy of an LMAB instance.
We give the details in Algorithm 1.
Our algorithm is divided to threesteps as detailed below.
Step 1: Estimating U .
Algorithm 1 ﬁrst estimates the second-order moments M=
P
Mm =1 w m µ m µ > m .
Let b U be the top- M eigenvectors of an empirical estimate of M 2 , where ˆ M 2 isconstructed by collecting samples of reward correlations by taking random actions at ﬁrst two timesteps.
Note that b U may not be a good proxy for some µ m with small mixing probability w m ⇡ 0 ,and thus all elements in U are not necessarily close to b U .
We defer the details of subspace recoveryprocedure to the proof of the following lemma in Appendix C.7.
Lemma 3.7 Let b U be a subspace spanned by top- M eigenvectors of ˆ M 2 .
After we estimate ˆ M 2 using
N
= O
( A 4 Z 2 log( ZA/ ⌘ ) /  4sub ) episodes, with probability at least 1  ⌘ , for all m 2 [ M ] , thereexists  m : k
m k 1   sub /w 1 / 2 m such that µ m +  mb U .
Our choice of  sub di  ers in two regimes as the following:
[Equation]
where w min = min m {+
+}
[ M ] w m .
The main di  erence in two regimes is that for H {+ 
+}
M  1 , 261 when the parameters are identiﬁable, latent contexts with small mixing probabilities can be ignoredonce w m =
o ( ✏ /M )
to guarantee the closeness in distributions of observations.
However, in theparameter unidentiﬁable regime, total variation distance is bounded only through errors in the momentspace.
Therefore, the estimated subspace needs to be accurate even for contexts with small mixingprobabilities to keep the higher-order moments well approximated.
Note that for instances withwell-balanced mixing probabilities, i.e., if w min = ⌦ (1 /M ) , the order of  sub remains the same as inthe H  2 M  1 case.
Step 2: Moment Matching .
Given b U , the subspace spanned by top- M eigenvectors of ˆ M 2 , Algorithm 1, follows the procedure described in Sections 3.1-3.3.
It constructs the set of (approximate) coreaction-values pairs { ( a j , z j ) } nj =1 (see Appendix C.8 for the detailed algorithm).
Then, we constructhigher-order moments of the core action-value pairs.
For every multi-index ( i 1 , i 2 , ..., i l ) 2 [ n ] l , using
N 1 =
O (log( ln l / ⌘ ) /  2tsr ) episodes, we execute a kt = a i t for t = 1 , ..., l and estimate higher-ordermoments (as also described in equation (1))
[Equation]
Using standard concentration inequalities and applying the union bounds over all elements in tensors,we get k ˆ T l  T l k 1 <  tsr with probability at least 1  ⌘ .
Then we ﬁnd empirical parameters{ ( ˆ w m , ˆ ⌫ m ) }
Mm =1 that satisfy
[Equation]
[Table]
[Table]
[Table]
[Table]
We set  tsr
= O ( ✏ / ( ZH 2 M 3 .
5 n )) 2 M  1 when H  2 M  1 .
In the parameter unidentiﬁable regime
H < 2 M  1 , we set  tsr
= O ( ✏ /H ) /
( Z p 2 M ) H .
Step 3: Constructing Empirical LMAB .
Finally, Algorithm 1 uses the estimates { ( ˆ w m , ˆ ⌫ m ) } Mm =
1to construct an empirical model ˆ B =
( A , { ˆ w m } Mm =1 , { ˆ µ m ) }
Mm =1 ) after proper clipping and normal- 281 ization.
For this step to succeed, we require ˆ w m and ˆ ⌫ m to be valid parameters for the reconstructionof a valid empirical model.
We state details on the recovery procedure in Appendix C.8.1.
Once a valid empirical model ˆ B is obtained, the remaining step is to call the planning oracle thatgives an ✏ -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 1  ⌘ , Algorithm 1 returns an ✏ -optimal policy given total number of episodes of at mostpoly( H, M, Z, A, 1 / ✏ ) · log( AZ/ ⌘ )
+ poly( H, Z, M ) 2 M  1 · log( M/ ⌘ ) / ✏ 4 M  2 , if H  2 M  1 , poly( H, w  1min , Z, A, 1 / ✏ ) · log( AZ/ ⌘ ) + H 2 log( M/ ⌘ ) ·  2 MZ H / ✏ 2 , otherwise .
Note that the dependency on A is polynomial.
This polynomial term is needed to control the subspaceestimation error.
The exponential term is derived from the closeness in moments as discussed earlier.
Remark 3.9 (Continuous Rewards: Gaussian Case)
So far we have focused on rewards with ﬁnite 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 is at most poly( H, M, A, 1 / ✏ ) · log( A/ ⌘ ) + poly ( H, M, log(1 / ( ⌘✏ )))
min( H,M ) / ✏ min(2 H +2 , 4 M  2) , 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 a procedure that learns a near-optimal policy of an LMAB instance.
However, our procedure relies on matching higher-order moments, which, in general, is not suitablefor practical implementation.
In this section, we describe a maximum likelihood (MLE) method,motivated by our previous results.
Importantly, we can use the Expectation-Maximization (EM) [10]heuristic to ﬁnd an approximate solution to the MLE optimization problem.
We can start from the set of core action-value pairs { ( a j , z j ) } nj =1 given by Corollary 3.2.
Now forevery time step t , we choose i t randomly from a uniform distribution over [ n ] , play action a i t , andobserve b t := 1 { r t =
z i t } .
After repeating this for N episodes, we formulate the log-likelihoodfunction with parameterization ✓ = { ( w m , ⌫ m ) }
Mm =1 as follows:
[Equation]
To prevent the confusion with searching parameters, we use q ⇤ to denote any quantity q constructedwith ground truth parameters, e.g., ✓ ⇤ :
= { ( w ⇤ m , ⌫ ⇤ m ) }
Mm =1 .
Let ✓ N be the maximum likelihoodestimator, i.e., ✓ N = arg max ✓ 2 ⇥ l N ( ✓ ) in some valid parameter set ⇥ .
We can recover the sample-complexity guaranteed by moment-matching methods studied in theprevious section.
Speciﬁcally, we show that the maximum likelihood estimator with su  ciently manysamples have nearly matching moments:
Lemma 4.1 Consider the maximum likelihood estimator ✓ N = { ( ˆ w m , ˆ ⌫ m ) }
Mm =1 with N episodesfor large enough N .
If N  C · n min(2 H +1 , 4 M  1) log( N/ ⌘ ) /  2 for some su  ciently large constant
C > 0 , then with probability at least 1  ⌘ ,
[Equation]
Therefore, by setting the same accuracy parameter  used in Algorithm 1, we can obtain the required✓ N = { ( ˆ w m , ˆ ⌫ m ) }
Mm =1 for matching moments.
We only replace the moment-matching step (Step 2in Algorithm 1) with solving the MLE optimization problem (8).
Implementation Details While the log-likelihood formulation (8) is still non-convex and thusintractable, we can rely on powerful heuristics: we initialize the parameters with clustering methods[2] or spectral methods [1, 7], and run the EM algorithm to improve the accuracy [10] (regardless ofconditions or guarantees).
Therefore, while the sample-complexity upper bound for both methodsremains the same, we beneﬁt from the log-likelihood formulation despite the non-convexity.
Speciﬁcally, we follow the same steps in Algorithm 1.
We ﬁrst ﬁnd the core action-event pairs fromsecond-order moments as described ( Step 1 ).
Then we form the maximum log-likelihood objective(8) and ﬁnd an approximate MLE solution with the EM algorithm ( Step 2 ).
After obtaining ✓ N ,we recover an empirical model which we use as an input to an approximate planning oracle ( Step3 ).
We compare our MLE-based implementation with experimental design (ED + MLE) to otherbaselines applicable to learning in LMABs ( e.g., L-UCRL + EM [23], tensor-decomposition methods[1, 7, 39], naive-UCB [6]).
In 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.
Due tospace constraints, we defer the details to Appendix A.
Provable Beneﬁts of MLE In Appendix E.2, we also show that MLE solutions can automaticallyadapt to mild separation conditions.
Speciﬁcally, we show that if H =
˜ O
( M 2 /  2 ) for some separationparameter  > 0 (see Assumption E.1), then we can achieve the polynomial sample complexity forlearning a near-optimal policy with the MLE solution ✓ N .
Conclusion and Future Work
In this work, we designed an algorithm that learns a near-optimal policy for an LMAB instance thatrequires only poly( A ) + poly( M, H ) min( H, 2 M  1) number of samples.
We achieved this with theexperimental design and method-of-moments, and the maximum likelihood estimation which is moresuitable in practice.
We discuss a few limitations of this work and future directions below.
Lower Bounds
The question of a lower bound on the sample complexity of the LMAB problemremains unresolved.
For LMDP, an MDP extension of LMAB, a lower bound of ⌦ ( A M ) is known dueto [23].
While we conjecture that some exponential dependence in M is unavoidable, characterizingthe minimax dependence of the exponent is left as an interesting future research question.
Latent MDPs We believe that the moment-matching based approach we took in this work o  ers apromising way for designing RL algorithms in the presence of latent contexts.
Speciﬁcally, we believethese techniques can be used to design algorithm that ﬁnds a near-optimal policy of LMDPs [23] orreward-mixing MDPs [22] with M =
O (1) and without further separation assumptions.
Linear Bandits / Continuous Rewards
Our work has focused on the tabular setting, where allarms are independent.
It will be an interesting future work to consider the same objective in linearbandit settings.
Also, while we only considered Gaussian rewards, investigating a general parametricclasses of rewards is an important future research direction from practical perspective.
Checklist
For all authors...
• (a) Do the main claims made in the abstract and introduction accurately reﬂect 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] : Thiswork is primarily focused on the theoretical aspect and its purpose is too general for us
to speculate about broader societal impacts.
• (d) Did you describe the limitations of your work?
[Yes]
If you are 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
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 a URL)?
[Yes]
• (b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere 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 ofGPUs, internal cluster, or cloud provider)?
[N/A]
The remaining issues are not applicable to
us