Tractable Optimality in Episodic Latent MABs
Anonymous Author(s)
AffiliationAddress email
Abstract
We consider a multi-armed bandit problem with A actions and 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 estimateaccurately the latent context.
The resulting partial observation of the environmentmakes the learning task significantly more challenging.
Without any additionalstructural assumptions, existing techniques to tackle partially observed settingsimply 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 polynomialsamples in A is possible.
We achieve this by using techniques from experimentdesign.
Then, through a method-of-moments approach, we design a procedure thatprovably 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 viamaximum likelihood estimation.
In our experiments, this significantly outperformsthe worst-case guarantees, as well 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 differentepisodes 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 Definition 2.1 for a formal description).
The LMAB frameworkis different from the setting considered in [28] where multiple episodes proceed in parallel withoutlimiting the horizon of an episode H .
For long horizons H , we show that it is possible to find 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 study the problem of learning a near optimal policy of an LMABs for both long andshort horizon H .
In the long horizon case, we show the problem’s latent parameters are learnable However, for short horizon, the latent parameters may not be identifiable and a different perspective isrequired.
A naive approach to learn a near optimal policy of an LMAB starts by casting the problem as a Markov Decision Process (MDP). Assume that the reward values are discrete and defined over a finite alphabetof 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 ) H ) states. Appealing to standard reinforcementlearning algorithms ( e.g., [16, 6]) we can learn an ϵ -optimal policy with O ( AZ ) H /ϵ 2 samples.
A natural way to sidestep these bottlenecks is through using the unique structure of the LMAB setting.
In this work we focus on the following question: can we learn a near-optimal policy with fewer than
O ( AZ )
H /ϵ 2 samples, in the absence of further assumptions ?
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 ) 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) , 50 where a natural objective is to learn a near-optimal history-dependent policy for H time steps (seealso Figure 3 in 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 56 Bayesian learning framework with short time-horizon [26, 32, 19].
The focus in this line of work ison 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 specialcase of LMABs, the baseline policy we compare ourselves to is the optimal H -step policy, which is aharder baseline than considering a fixed Bayesian algorithm.
Latent MDPs Some prior work considers the framework of Latent MDPs (LMDPs), which is the
MDP generalization of LMAB [15, 34, 22, 21].
In particular, [22] 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 [22], since their lower bound construction comes from the challenge in state-explorationwith 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 [22].
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.
Learning in POMDPs with full-rank observations One popular learning approach in partially 73 observable 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 requirespecific 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, 27] 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 confidence 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 define the problem of episodic latent multi-armed bandits with time-horizon H ≥ 2 as follows:
Definition 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 weightssuch 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 actiona ∈
A conditioning on a latent context m (each µ m parameterizes the probability model, and is notnecessarily a mean 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 finite 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 finite and bounded support.
Thereward attains a value in the set Z .
We assume that for all z ∈ Z we have | z | ≤ 1 .
We denote thecardinality of Z as Z and assume that Z = O (1) .
As an example, Bernoulli distribution satisfies 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 ) in 98 a context m .
We often use µ m as a reward-probability vector in R AZ indexed by a tuple ( a, z ) ∈ A×Z .
At the beginning of every episode, a latent context m ∈ [ M ] is sampled from a mixing distribution{ w m } Mm =1 and fixed for H time steps, however we cannot observe m directly.
We consider a policy 101 class Π which contains all history-dependent policies π : ( A × Z ) ∗ → A .
Our goal is to finda near optimal policy π ∈ Π that maximizes the expected cumulative reward V for each episode
V ⋆ = max π ∈ Π V ( π ) :
= E π (cid:104)(cid:80)
Ht =1 r t (cid:105) , where the expectation is taken over latent contexts andrewards generated by an LMAB instance, and actions following a policy π .
Definition 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 [30] 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 ∈ [ M ] w m for the minimum mixing weight.
For any l th -order tensor T l ∈ R n × ... × n ( n repeated l times), we denote ∥ T l ∥ ∞ for the element-wiselargest absolute value.
For any vector in v ∈ R n in dimension n ∈ N + , we use v (cid:78) l to denote adegree l tensorization of v .
We often use ∥ µ m ( a, · ) − ˆ µ m ( a, · )
∥ 1 to mean the l 1 statistical distance 114 between reward distributions: (cid:80) z ∈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 ∈ [ 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 first 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 executing 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.
We consider a model-based learning approach from the l th -order moments of reward observations.
Since there are A actions, this implies that we have O ( A l ) quantities to estimate, which would incur O ( A l ) sample complexity if we separately measure every correlation between l pairs of actions.
On the other hand, when M ≪ A , the distributions µ m (recall these are reward probabilities foreach 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 , we 132 could remove the unfavorable dependence on A .
However, even when the subspace U is known a priori, dimensionality reduction with bandit feedbackis non-trivial.
We need to compute directly the projection onto U of estimates of the reward probabilityvectors.
To see the challenge, let { β j } Mj [-=
 ⇢-]=1 ⊂ R AZ be an orthonormal basis of U .
To compute theprojected estimate µ ⊤ m β j , we need a sampling policy π such that π ( a ) ∝ (cid:80) z | β j ( a, z ) | .
The varianceof this sampling policy can be as much as O ( ∥ β j ∥ 21 ) , which in general scales with ∥ β j ∥
= O ( A ) .
Therefore, reliable estimation for any statistics in the reduced dimension would still have a dependenceon A .
In particular, since our approach is based on higher-order method-of-moments, an estimationof l th -order statistics for l ≥ 2 would require Ω( A l/ 2 ) samples.
The general idea of dimensionality reduction is not fatally flawed.
Instead, we need to avoid thepitfalls of the approach outlined above.
First, the calculation above shows that we need to control the∥ · ∥ 1 -norm of the vectors β j , ideally, ∥ β j ∥ 1 =
O (1) .
Second, we only need a good estimate for eachµ m such that max a ∈A ∥ µ m ( a, · ) − ˆ µ m ( a, · )
∥ 1 ≤ ϵ to compute an ϵ -optimal policy.
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 ∥ β j ∥ 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 sufficientlycorrelated 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.
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 R d where d ≫ k and let u, ¯ u ∈ 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
∥ u − ¯ u ∥ ∞ ≤ √ 2 k max j ∈ [ n ] | u ( i j ) − ¯ u
( i j ) | .
A linear transformation that maps [¯ u ( i 1 ) , . .
, ¯ u ( i n )] to its corresponding ¯ u ∈ 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, 36].
In the context of bandits and RL, the Kiefer–Wolfowitz theorem has been used to study how themisspecification in linear representation changes the problem landscape ( e.g., [24, 29]).
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 find a set of coordinates of size ˜ O ( M ) that aresufficient to reconstruct the latent reward model { µ m } Mm =1 where µ m ∈ R AZ .
We refer to this set of 167 coordinates as the set of core action-value pairs.
Corollary 3.2 Suppose the subspace U = span { µ 1 , . . .
, µ M } is given.
Then for any ˆ µ ∈ 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 ∈ [ M ]
[Equation]
where ν m , ˆ ν m ∈ R n such that ν m ( j ) :
=
µ m ( a j , z j ) and ˆ ν m ( j ) :
= ˆ µ m ( a j , z j ) .
A linear transformation that maps [ˆ ν ( i 1 ) , . .
, ˆ ν ( i n )] to its corresponding ˆ µ ∈ 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 can 176 estimate { ν m } Mm =1 to sufficiently good accuracy with only polynomial dependence in A .
3.2 H ≥ 2 M − 1 : Identifiable 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 [37, 11].
That is, we can recover { ν m } Mm =1 by estimating 2 M − 1 higher-order moments 181 and find { ˆ ν m } Mm =1 that matches them.
This further emphasizes the importance of the dimensionality 182 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 definethe l -order tensor as T l := (cid:80)
{+Mm =
+}
w m ν (cid:78) l m . For an LMAB instance, we can access the tensor T l using 186 observational data by simply estimating correlations between l core action-value pairs within eachepisode.
Specifically, for any I
= ( i 1 , i 2 , ..., i l ) ∈ [ 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 find estimators { ( ˆ w m , ˆ ν m ) }
Mm =1 such that moments match T l up to error δ for all l ∈[2 M − 1] .
The results in [37, 11] imply that the closeness in moments of up to (2 M − 1) th degreeimplies the closeness in model parameters:
Lemma 3.3 Suppose ∥ T l − (cid:80)
Mm =
ˆ w m ˆ ν (cid:78) l m [-k 
-]∥ ∞ < δ for all l = 1 , 2 , ..., 2 M − 1 for some sufficientlysmall δ > 0 .
Then,
[Equation]
where Γ is a joint distribution over ( m, m ′ ) ∈
[ M ] 2 satisfying
[Equation]
The form of guarantee given for { ˆ w m , ˆ ν m } Mm =1 is in the Wasserstein distance between two latent 196 model parameters.
A useful property of the Wasserstein metric is that the distance measure is invariantto permutation of individual components, and flexible with arbitrarily small mixing probabilities orarbitrarily close components [37] (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 200 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 anytwo LMABs.
Then, for any history-dependent policy π : ( A × Z ) ∗ →
A , we have
[Equation]
where the infimum over Γ is taken over joint distributions over ( m, m ′ ) satisfying (3) .
From Corollary 3.2 we have max a ∈A ∥ µ m ( a, · ) − ˆ µ m ′ ( a, · ) ∥ 1 ≤ 2 Z √ 2 M ∥ ν m − ˆ ν m ′ ∥ ∞ for anym, m ′ ∈ [ 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 : Unidentifiable 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 identifiablility 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 ∥
[Equation]
=
[Equation]
∥ ∞ ≤ δ , then for any history dependent policy π ∈ Π , we have
[Equation]
Therefore, it is sufficient to find 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 } M m =1 , { ˆ µ m } M m =1 that satisfy the condition in 222 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 ∥ (cid:80) Mm =1 w m ν (cid:78) l m − (cid:80) Mm =1 ˆ w m ˆ ν (cid:78) l m ∥ ∞ ≤ δ , then(cid:13)(cid:13)(cid:13)(cid:80)
[-Mm =
 ˆ w m ˆ ⌫ N l m k 1   , then M X m = w m µ N l m -][Equation]
l m (cid:13)(cid:13)(cid:13) ∞ ≤ (2 M ) l/ 2 · δ.
By Proposition 3.6, we can conclude that it is sufficient to estimate T H :
= (cid:80)
Mm {+=
+} w m ν (cid:78) H m and find{ ( ˆ w m , ˆ ν m ) } Mm =1 that matches T H element-wise up to accuracy δ := ( ϵ/H ) /
( Z √ 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 first estimates the second-order moments M=
[Equation]
=1 w m µ m µ ⊤ m .
Let (cid:98) U be the top- M eigenvectors of an empirical estimate of M 2 , where ˆ M 2 is 234 constructed by collecting samples of reward correlations by taking random actions at first two timesteps.
Note that (cid:98) 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 (cid:98) U .
We defer the details of subspace recoveryprocedure to the proof of the following lemma in Appendix C.7.
Lemma 3.7 Let (cid:98) 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 ∈ [ M ] , thereexists ∆ m : ∥
∆ m ∥ ∞ ≤ δ sub /w 1 / 2 m such that µ m + ∆ m ∈ (cid:98) U .
Our choice of δ sub differs in two regimes as the following:
[Equation]
where w min = min m [-
-]∈
[ M ] w m .
The main difference in two regimes for H [- 
-]≥ 2 M −
Algorithm
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table] [Table] [Table] once w m = o ( ϵ/M ) to guarantee the closeness in distributions of observations.
However, in the 245 parameter unidentifiable 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 in 249 the H ≥ 2 M − 1 case.
Step 2: Moment Matching .
Given (cid:98) U , the subspace spanned by top- M eigenvectors of ˆ M 2 , Algo- 251 rithm 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.1 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 ) ∈ [ 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-order 255 moments (as also described in equation (1))
[Equation]
Using standard concentration inequalities and applying the union bounds over all elements in tensors,we get ∥ ˆ T l − T l ∥ ∞ < δ tsr with probability at least 1 − η .
Then we find empirical parameters{ ( ˆ w m , ˆ ν m ) }
Mm =1 that satisfy
[Equation]
We set δ tsr
= O ( ϵ/ ( ZH 2 M 3 .
5 n )) 2 M − 1 when H ≥ 2 M − 1 .
In the parameter unidentifiable regime
H < 2 M − 1 , we set δ tsr
= O ( ϵ/H ) /
( Z √ 2 M ) H .
Step 3: Constructing Empirical LMAB .
Finally, Algorithm 1 uses the estimates { ( ˆ w m , ˆ ν m ) } M m =
1to construct an empirical model ˆ B =
( A , { ˆ w m } Mm =1 , { ˆ µ m ) }
Mm =1 ) after proper clipping and normal- 263 ization.
For this step to succeed, we require ˆ w m and ˆ ν m to be valid parameters for the reconstruction 264 of 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/η ) · (cid:0) 2 MZ 2 (cid:1) 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 (Note on w min ) Ideally, mixing weights are well balanced such that w min = Ω(1 /M ) . However, in general we may have arbitrarily small w min . Note that when H ≥ 2 M − 1 , the upperbound does not depend on w min .
However when H < 2 M − 1 with arbitrarily small w min , effectivelywe have w min ∝ ϵ/
( Z √ M ) H as can be seen in equation (5) .
This is mainly because in the parameterunidentifiable regime, we give guarantees from moment-closeness, and thus the subspace has to bevery accurate to approximate higher-order moments well. Nevertheless, we are not aware whetherthis is tight, and believe this to be an interesting question for future research. 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 continuousreward distributions.
In Appendix D, we show a similar upper bound that is at most
[Equation] 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 find 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 , and 290 observe [Equation] := 1 { r t =
z i t } .
After repeating this for N episodes, we formulate the log-likelihood 291 function with parameterization θ = { [Equation] 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 likelihood 294 estimator, i.e., θ N = arg max θ ∈ Θ l N ( θ ) in some valid parameter set Θ .
We can recover the sample-complexity guaranteed by moment-matching methods studied in theprevious section.
Specifically, we show that the maximum likelihood estimator with sufficiently 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 sufficiently 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 2 303 in 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 benefit from the log-likelihood formulation despite the non-convexity.
Specifically, we follow the same steps in Algorithm 1.
We first find the core action-event pairs fromsecond-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 θ N , werecover an empirical model which we use as an input to an approximate planning oracle ( Step 3 ).
We compare our MLE-based implementation with experimental design (ED + MLE) to other baselinesapplicable to learning in LMABs ( e.g., L-UCRL + EM [22], tensor-decomposition methods [1, 7, 38],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.
Provable Benefits of MLE In Appendix E.2, we also show that MLE solutions can automaticallyadapt to mild separation conditions.
Specifically, 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 .
In Figure 1, we compare the averaged per time-step rewards obtained with returned policies fromdifferent algorithms after some exploration episodes.
Note that with increasing H , it becomes easierto identify contexts, i.e. , we have bigger separations, and thus MLE based approaches can benefitfrom larger separations.
Due to space constraints, we defer the remaining details to Appendix A.
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 [22].
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 offers apromising way for designing RL algorithms in the presence of latent contexts.
Specifically, we believethese techniques can be used to design algorithm that finds a near-optimal policy of LMDPs [22] orreward-mixing MDPs [21] with M
= O (1) and without further separation assumptions.
Linear
Bandits / Continuous Rewards Our work has focused on the tabular setting, where all armsare independent.
It will be an interesting future work to consider the same objective in linear banditsettings.
Also, while we only considered Gaussian rewards, we believe our approach can also beextended to a broader class of parametric distributions. Investigating more general classes of rewardsis an important future research direction from practical perspective.
Checklist
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
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.