Tractable Optimality in Episodic Latent MABs
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 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 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.
Introduction
In Multi-Armed Bandits (MABs), an agent learns to act optimally by interacting with an unknown environment.
In many applications, interaction sessions are short, relative to the complexity of the overall task.
As a 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 to short horizon episodes, and have to learn how to process observations from different episodes 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 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 a 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 it is possible to find a near-optimal policy and to determine near-optimal actions for each episode, as if we knew the hidden context.
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.
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 latent parameters are learnable.
However, for short horizon, the latent parameters may not be identifiable and a different perspective is required.
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 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 ) H ) states.
Appealing to standard reinforcement learning algorithms ( e.g., [18, 6]) we can learn an ϵ -optimal policy with O ( AZ ) H /ϵ 2 samples.
A natural way to improve upon the naive approach 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 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 ) min
( H,M ) , and thus tractable when either H or M is small, even for very large A .
That is, we are particularly focused on the setting with 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 (see also 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 lengthier discussion 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 [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 a fixed Bayesian algorithm.
Latent MDPs 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 information-theoretic limit of learning in LMDPs with no assumptions on MDPs, i.e., an exponential number of sample episodes Ω( A M ) is necessary to 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 state-exploration 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 ≥ 3 without further assumptions.
Learning in POMDPs with full-rank observations One popular learning approach in partially observable systems is the tensor-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 a test-observation matrix to keep the volume of a confidence set explicitly bounded.
Since LMAB instances do not necessarily satisfy the full-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 weights such that a latent context m is randomly chosen with probability w m , and µ m is the model parameter that describes a reward distribution, i.e., P µ m ( r | a ) :
= P ( r | m, a ) , according to an action a ∈
A conditioning on a latent context m (each µ m parameterizes the probability model, and is not necessarily a mean reward vector).
We do not assume a priori knowledge of mixing weights.
The bulk of this paper considers discrete reward 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.
The reward attains a value in the set Z .
We assume that for all z ∈ Z we have | z | ≤ 1 .
We denote the cardinality 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 denote the probability of observing a reward value z by playing an action a as µ m ( a, z ) :
= P ( r = z | m, a ) in 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 class Π which contains all history-dependent policies π : ( A × Z ) ∗ → A .
Our goal is to find a 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 and rewards generated by an LMAB instance, and actions following a policy π .
Definition 2.3 (Approximate Planning Oracle)
A planning oracle receives an LMAB instance B and returns an ϵ -approximate policy π such that V ⋆ − V ( π ) ≤ ϵ .
Concretely, the point-based value-iteration (PBVI) algorithm
[33] is an ϵ -approximate planning algorithm 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 its corresponding 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-wise largest absolute value.
For any vector in v ∈ R n in dimension n ∈ N + , we use v (cid:78) l to denote a degree l tensorization of v .
We often use ∥ µ m ( a, · ) − ˆ µ m ( a, · )
∥ 1 to mean the l 1 statistical distance between reward distributions: (cid:80) z ∈Z | µ m ( a, z ) − ˆ µ m ( a, z ) | .
Lastly, we denote by a t and r t as the action 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 instance with O (poly( A )+poly( M, H ) min( M,H ) ) 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 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 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-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 for 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 , we could remove the unfavorable dependence on
A .
However, even when the subspace U is known a priori, dimensionality reduction with bandit feedback is non-trivial.
We need to compute directly the projection onto U of estimates of the reward probability vectors.
To see the challenge, let { β j } Mj =1 ⊂ R AZ be an orthonormal basis of U .
To compute the projected estimate µ ⊤ m β j , we need a sampling policy π such that π ( a ) ∝ (cid:80) z | β j ( a, z ) | .
The variance of 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 dependence on A .
In particular, since our approach is based on higher-order method-of-moments, 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 flawed.
Instead, we need to avoid the pitfalls 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 (hence they will have ∥ β j ∥
= 1 ), in such a way that guarantees the approximation quality for µ m from estimating µ ⊤ m β j ’s.
In terms of the original problem, the existence of such a subset of the standard basis is equivalent to the existence of a 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 Kiefer–Wolfowitz 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 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 as a corollary of the Kiefer–Wolfowitz theorem (and its geometric interpretation) for (near)-optimal experimental design [20, 39].
In the context of bandits and RL, 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 ˜ O ( M ) that are sufficient to reconstruct the latent reward model { µ m } Mm =1 where µ m ∈ R AZ .
We refer to this set of 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 , there exists an algorithm that runs in time ˜ O
( ZAM 2 ) and returns the following.
A set of core action-value pairs { ( a j , z j ) } n j = ⊆
A × Z of size at most n = 4 M log log M + 16 , such that for all m ∈
[ M ]
[Equation] [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 implies a 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 transformation of { ν m } Mm =1 .
In the following sections, we show that by matching higher-order moments we can 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-order moments [40, 12].
That is, we can recover { ν m } Mm =1 by estimating 2 M − 1 higher-order moments and find { ˆ ν m } Mm =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, a moment-based approach to 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 define the l -order tensor as T l := (cid:80) M m =
1 w m ν (cid:78) l m .
For an LMAB instance, we can access the tensor T l using observational data by simply estimating correlations between l core action-value pairs within each episode.
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 [40, 12] imply that the closeness in moments of up to (2 M − 1) th degree implies the closeness in model parameters:
Lemma 3.3 Suppose ∥ T l − (cid:80)
Mm =
ˆ w m ˆ ν (cid:78) l m ∥ ∞ < δ for all l = 1 , 2 , ..., 2 M − 1 for some sufficiently small δ > 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 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 { ν m } M m =1 , we can estimate { ˆ µ m } M 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.
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 two 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 any m, 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 ϵ -optimal policy 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 degree of 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 ∥ (cid:80)
[Equation]
=
[Equation]
m
[Equation]
m −
[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 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, a natural idea is to estimate parameters { ˆ w m } M m =1 , { ˆ µ m } M m =1 that satisfy the condition in Lemma 3.5 without incurring O
( A H ) sample complexity.
This is possible if we have good estimates of 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 =1 w m µ
[Equation]
l m −
[Equation]
=1 ˆ w m ˆ
[Equation]
l m
[Equation]
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-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=
[Equation]
Mm =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 constructed by collecting samples of reward correlations by taking random actions at first two time steps.
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 recovery procedure 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 ] , there exists ∆ 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 is that for H ≥ 2 M − 1 , when the parameters are identifiable, latent contexts with small mixing probabilities can be ignored
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 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-balanced mixing probabilities, i.e., if w min = Ω(1 /M ) , the order of δ sub remains the same as in 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- rithm 1, follows the procedure described in Sections 3.1-3.3.
It constructs the set of (approximate) core action-values pairs { ( a j , z j ) } nj =1 (see Appendix C.8.1 for the detailed algorithm).
Then, we construct higher-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 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 = to construct an empirical model ˆ B =
( A , { ˆ w m } Mm =1 , { ˆ µ m ) }
Mm =1 ) after proper clipping and normal- ization.
For this step to succeed, we require ˆ w m and
ˆ ν m to be valid parameters for the reconstruction 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 that gives 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 most poly( 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 subspace estimation 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 upper bound does not depend on w min .
However when H < 2 M − 1 with arbitrarily small w min , effectively we have w min ∝ ϵ/ ( Z √ M ) H as can be seen in equation (5) .
This is mainly because in the parameter unidentifiable regime, we give guarantees from moment-closeness, and thus the subspace has to be very accurate to approximate higher-order moments well.
Nevertheless, we are not aware whether this 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 continuous reward 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 suitable for 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) [11] 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 for every time step t , we choose i t randomly from a uniform distribution over [ n ] , play action a i t , and observe b t := 1 { r t =
z i t } .
After repeating this for N episodes, we formulate the log-likelihood function with parameterization
[Equation]
as follows:
[Equation]
To prevent the confusion with searching parameters, we use q ∗ to denote any quantity q constructed with ground truth parameters, e.g., θ ∗ :
= { ( w ∗ m , ν
∗ m ) }
Mm =1 .
Let θ
N
be the maximum likelihood 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 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 θ N = { ( ˆ w m , ˆ ν m ) }
Mm =1 with N episodes for 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 in Algorithm 1) with solving the MLE optimization problem (8).
Implementation Details While the log-likelihood formulation (8) is still non-convex and thus intractable, 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 [11] (regardless of conditions or guarantees).
Therefore, while the sample-complexity upper bound for both methods remains 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 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 θ N , we recover 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 baselines applicable to learning in LMABs ( e.g., L-UCRL + EM
[24], tensor-decomposition methods [1, 7, 41], 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 automatically adapt to mild separation conditions.
Specifically, we show that if H =
˜ O
( M 2 /γ 2 ) for some separation parameter γ > 0 (see Assumption E.1), then we can achieve the polynomial sample complexity for learning a near-optimal policy with the MLE solution θ N .
In Figure 1, we 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. , we have bigger separations, and thus MLE based approaches can benefit from 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 that requires only poly( A ) + poly( M, H ) min
( H, 2 M − 1) number of samples.
We achieved this with the experimental design and method-of-moments, and the maximum likelihood estimation which is more suitable 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 problem remains unresolved.
For LMDP, an MDP extension of LMAB, a lower bound of Ω( A M ) is known due to [24].
While we conjecture that some exponential dependence in M is unavoidable, characterizing the 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 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 a near-optimal policy of LMDPs [24] or reward-mixing MDPs [23] with M
= O (1) and without further separation assumptions.
Linear Bandits / Continuous Rewards Our work has focused on the tabular setting, where all arms are independent.
It will be an interesting future work to consider the same objective in linear bandit settings.
Also, while we only considered Gaussian rewards, we believe our approach can also be extended to a broader class of parametric distributions.
Investigating more general classes of 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).
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] 2.
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 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 a 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 to us.
Appendix A Experiments
In this section, we demonstrate our maximum likelihood estimation (8) on synthetic examples.
We generate random LMAB instances with Bernoulli rewards, where the mean-reward vectors { µ m } Mm = lie in a subspace of dimension roughly 4 , i.e., rank (cid:16)(cid:80)
Mm =1 w m µ m µ ⊤ m (cid:17) ≈ 4 .
We compare our MLE with experimental design method (8) (ED+MLE) to three benchmarks:
Naive UCB [4] without considering contexts (thus a returned policy is stationary).
Tensor-decomposition methods [1, 41].
L-UCRL with spectral initialization [24].
Even when the theoretical conditions required for the success of tensor-decomposition or spectral methods do not, in practice, standard tensor-decomposition technique by [1] serves as good initialization for the EM or L-UCRL algorithms.
After estimating the LMAB ˆ B , we compute a heuristic policy using Q-MDP [28] for ˆ B .
We refer to the computed policy by Q-MDP [28] using the true LMAB model B as the genie policy .
In the first experiment, we compare the performance of the aforementioned four alternatives.
We draw a random LMAB instance with M = 4 , A = 20 using N = 5 · 10 4 sampled episodes, and
M = 7 , A = 50 using N = 100000 episodes (Figure 1).
We compare the averaged per time-step rewards obtained with each policy with increasing length of episodes H .
When M = 4 , all methods (except naive UCB) exhibit the same pattern of improved performance as the algorithms access more data.
As we generated instances to satisfy the full-rank condition, i.e., rank (cid:16)(cid:80)
Mm =1 w m µ m µ ⊤ m (cid:17) ≈ 4 , even pure tensor-decomposition method works well in practice in this setting.
However, when M = 7 , pure tensor-decomposition method significantly under-performs
L-UCRL or ED+MLE.
This demonstrates that for LMAB instances with rank degeneracy, additional iteration steps with EM are necessary to get a good solution.
Furthermore, the performance of
ED+MLE and of L-UCRL does not significantly drop in M = 7 .
This demonstrated that practically it works much better than what can be guaranteed in theory in the worst case.
We conjecture that this is because the EM iteration converges to the MLE (even when converging to local optimums), and
MLE solutions in general show much better performance under some mild conditions.
( e.g., with
Assumption E.1, or if random perturbations are applied to the underlying LMAB model [8]).
In our second experiment, we test the performance of the different methods while scanning different control parameters.
We first fix A = 50 , H = 7 , and change the number of contexts (left of Figure 2).
Since we keep generating instances satisfying rank (cid:16)(cid:80)
Mm =1 w m µ m µ ⊤ m (cid:17) ≈ 4 , tensor-based method performs significantly worse as M increases.
We then fix M = 5 and observe the performance gain by increasing the sample episodes (right of Figure 2).
We see that the performance of both L-UCRL+EM and our method (ED+MLE) approaches to that of genie as we increase N .
From both experiments, we conclude that, in practice, MLE based methods perform much better than the worst case guarantee shown via method of moments.
We believe it will be an interesting future direction whether we can derive much better guarantees for MLE based solutions.
Appendix B Additional Preliminaries B.1 Additional Related Work
Comparison of Regimes Nature of learning a near-optimal policy in LMABs might change depending on the length of episodes H and the number of contexts M (recall Figure 3).
For instance, when H = 1 or M = 1 , the problem of learning in LMAB is essentially equivalent to the classical
Multi-Armed Bandits problem that has been extensively studied in literature ( e.g., see [26] and references therein).
In another well-studied literature of Bayesian learning, each episode is assumed to have a structured prior over contexts ( M → ∞ , e.g., all expected rewards of arms are independently sampled from beta conjugate priors).
Depending on the length of time-horizon H , the problem can be solved with Thompson Sampling (TS) [38] H → ∞ , or the problem is reduced to learning the prior [21].
When the time-horizon is sufficiently long but finite e.g., if H = Ω( A/γ 2 ) for some separation parameters between a finite number of contexts M < ∞ , then it is possible cluster observations from every episode.
This setting has been studied in the literature of multitask RL which we describe below.
In this work, we are particularly focused on the setting with a few contexts and relatively short episodes (in comparison to A ), i.e., M = O (1) and H
= O (1) , where the most natural objective is to learn a near-optimal history-dependent policy.
Reward-Mixing MDPs Another closely related work is to learn an LMDP (MDP extension of
LMAB) with common state transition probabilities, and thus only the reward function changes depending on latent contexts [23].
The authors in [23] have developed an efficient moment-matching based algorithm to learn a near-optimal policy without any assumptions on separations or system dynamics.
However, it can only handle the M = 2 case with balanced mixing weights w 1 = w 2 = 1 / 2 .
It is currently not obvious how to extend their result to M ≥ 3 cases without incurring O ( A M ) sample complexity in LMABs.
Regime switching bandits LMAB may be also seen as a special type of adversarial or non-stationary bandits ( e.g., [5, 3, 13, 36]) with time steps being specified for when the underlying reward distributions ( i.e., latent contexts) may change.
The standard objective in non-stationary bandits is to find the best stationary policy in hindsight with unlimited possible contexts.
Recently, [41] considered a non-stationary bandit with a finite number of contexts M
= O (1) and the objective of finding the optimal history-dependent policy.
Their setting and goal subsume our goal of learning the optimal policy in LMABs; however, results in [41] require linear independence between reward probability vectors, and thus their setting essentially falls into the category of tractable POMDPs with full-rank observations.
Multitask RL with explicit clustering We can cluster observations from each episode if we are given sufficiently long time horizons H = ˜Ω
( A/γ 2 ) in each episode [9, 14, 16].
Here, γ > 0 is the amount of ‘separation’ between contexts such that for all m
[Equation]
where µ m is a mean-reward vector of actions in the m th context.
We focus on significantly more general cases where there is no obvious way of clustering observations, e.g., when H ≪ A or µ m can be arbitrarily close to some other µ m ′ for m ̸
= m ′ .
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 ) (see also section E.2).
Note that this could still be in H ≪
A regime with large number of actions A .
Miscellaneous There are other modeling approaches where multiple episodes proceed in parallel without limits on the time-horizon [31, 14, 10, 17, 25].
In such problems, the goal is to quickly adapt policies for individual episodes assuming certain similarities between tasks.
In contrast, in episodic
LMAB settings, we assume that every episode starts in a sequential order, and our goal is to learn an optimal history-dependent policy that can maximize rewards for a single episode with limited time horizon.
B.2 Experimental Design
We now give a high-level overview on experimental design techniques used in this work.
Suppose we are given a matrix Φ ∈ R d × k where d ≫ k .
Define a distribution over the rows of Φ , ρ ∈ ∆ d an element in the d -dimensional simplex.
We want to select a small subset of rows of Φ which minimizes g ( ρ ) defined below:
[Equation]
where Φ i, :
∈ R k be the i th row of Φ and G ( ρ ) ∈ R k × k .
To achieve this task we use results from the experimental design literature.
The following theorem shows the existence of ρ that minimizes g ( ρ ) with a small support over the row indices of Φ :
Theorem B.1 (Theorem 4.4 in [27])
There exists a probability distribution ρ such that g ( ρ ) ≤ 2 k and | supp ( ρ ) | ≤ 4 k log log k + 16 .
Furthermore, we can compute such ρ in time
˜ O
( dk 2 ) .
As noted in [27], Theorem B.1 can be obtained from results of Chapter 3 in [39].
Using this fundamental theorem, [27] showed the following proposition:
Proposition B.2 (Proposition 4.5 in [27]) Let ρ be a distribution over the rows of Φ that satisfies the condition of Theorem B.1.
Suppose a vector µ ∈ R d can be represented as a sum µ = v + ∆ where v ∈ V , ∥ ∆ ∥ ∞ ≤ ϵ 0 .
Let η be any small noise with η ∈ [ − ϵ 1 , ϵ 1 ] d .
Then ∥ Φˆ θ − µ ∥ ∞ ≤ ϵ 0 + ( ϵ 0 + ϵ 1 ) √ 2 k where
[Equation]
Crucially, we use Proposition B.2 to reduce the sample complexity in A .
B.3 Wasserstein Distance
We now give a brief overview on the Wasserstein distance and its applications in latent mixture models.
Wasserstein distance is a convenient error metric to measure the parameter distance between two latent models { ( w m , ν m ) }
Mm =1 and { ( ˆ w m , ˆ ν m ) }
Mm =1 , where w m , ˆ w m are mixing probabilities and ν m , ˆ ν m are some parameters for individual contexts.
Wasserstein distance is defined as follows.
Let ν be a finite-support distribution over { ν m } Mm = with probabilities { w m } Mm =1 , i.e., ν = (cid:80) Mm
=1 w m δ ν m where δ v is a Direc-delta distribution with a single mass on v ∈ R n .
Similarly with parameters { ( ˆ w m , ˆ ν m )
} Mm =1 , define an atomic distribution
ˆ ν = (cid:80) m ˆ w m δ ˆ
ν m .
We define a Wasserstein distance between ν and ˆ ν with respect to l ∞ norm as the following:
[Equation]
where the infimum is taken over all couplings over joint distributions ν and ˆ ν which are marginally distributed as ν and ˆ ν respectively (identical to (3)), i.e.,
[Equation]
One nice property of Wasserstein metric is that the distance measure is invariant to permutation of individual components, and flexible with arbitrarily small mixing probabilities or close parameters for different contexts [40, 12].
Appendix C Proofs for SectionC.1 Proof of Lemma 3.
Let { β j } kj =1 ⊆
R d be orthonormal basis of a k -dimensional subspace U .
Let
Φ ∈ R d × k be a matrix of form [ ββ...
β k ] , i.e., the j th column of Φ is β j .
We need to show the existence of a small set of core coordinates, from which we can reconstruct µ ∈ U .
Proposition B.2 implies that we can find a set of coordinates { i j } nj =1 with cardinality at most n = 4 d log log d + 16 , such that if µ ∈ U and we can access the vector [ µ ( i 1 ) , . . .
, µ ( i n )] up to accuracy ϵ 1 , then we can reconstruct µ up to ϵ 1 √ 2 k error.
We can also infer that for any µ ∈ U , we have Φ θ = µ where θ is given by G ( ρ ) − 1 (cid:80) i ∈
[ d ] ρ ( i ) µ ( i ) a i .
To conclude Lemma 3.1, suppose we find ˆ µ ∈ U such that | µ ( i j ) − ˆ µ ( i j ) | ≤ ϵ 0 for all j ∈ [ n ] .
Let η be such that η ( i ) =
µ ( i ) − ˆ µ ( i ) if i ∈ { i j } nj =1 , and η ( i ) = 0 otherwise.
Applying Proposition B. with ∆ = 0 and η be as defined above, with ˆ µ
=
Φˆ θ , we have ∥ ˆ µ − µ ∥ ∞ ≤ ϵ 0 √ 2 k .
C.2 Proof of Corollary
Let { β m } Mm =1 ⊆
R AZ be orthonormal basis of a M -dimensional subspace that includes U = span ( { µ m } Mm =1 ) .
Then by Lemma 3.1, we have
[Equation]
with a set of core action-value pairs { ( a j , z j ) } nj =1 .
Plugging v m ( j ) =
µ m ( a j , z j ) and ˆ ν m ( j ) =
ˆ µ m ( a j , z j ) gives Corollary 3.2.
Remark C.1 (Eliminating factor Z for H ≥ M )
Instead of core action-value pairs, we can also define core action-event pairs to save a factor of Z .
For instance, define a basis { ˜ β m } Mm =1 in a lifted space R A × 2 Z defined as the following:
[Equation]
For each µ m , define ϕ m :
[Equation]
By definition, span ( { ϕ m } Mm
=1 ) ⊆ span ( { ˜ β m } Mm
=1 ) .
Since 12 ∥ µ m ( a, · ) − ˆ µ m ( a, · ) ∥= max S ⊆Z | ϕ m ( a, S ) − ˆ ϕ m ( a, S ) | , it follows from Lemma 3.1 that we can find a set of core action-events { ( a j , Z j ) } nj =1 of size at most n = 4 M log log M + 16 such that for any ϕ m , ˆ ϕ m ∈ V , we have
[Equation]
This approach enables to approximate reward distributions in l 1 statistical distance, and thus we can save a factor of Z in subsequent analysis.
However, it comes with exponentially more expensive (in
Z ) computations for finding the core set.
C.3 Proof of Lemma 3.
To prove this result, we use the recent results on converting closeness in higher-order moments to closeness in Wasserstein distances for atomic distributions [40, 12].
A key result in [40] states a connection between moments and Wasserstein distance in one-dimensional case:
Theorem C.2 (Proposition 1 in [40])
If n = 1 and (cid:12)(cid:12)(cid:12)(cid:80) Mm =
[Equation]
m
[Equation]
< δ for all d = 1 , 2 , ..., 2 M − 1 , then
[Equation]
Theorem C.2, which holds for 1-dimensional mixtures–can be generalized to the high-dimensional case, as the following theorem shows.
Theorem C.3 For n ≥ 2 , if
[Equation]
m ˆ
[Equation]
δ for all d = 1 , 2 , ..., 2 M − 1 , then
[Equation]
Proof.
The idea follows the proof of Lemma 3.1 in [12].
Suppose a standard Gaussian random variable θ ∼ N (0 , I ) in R n .
For any x ∈ R n , anti-concentration of Normal distribution says
[Equation]
, for any τ > 0 .
By union bound and the fact that ∥ x ∥ 2 ≥ ∥ x ∥ ∞ , if we define X = { v − ˆ v | v ∈ { ν m } Mm =1 , ˆ v ∈ { ˆ ν m } Mm =1 } , then we have
[Equation]
Also with high probability, we have P ( ∥ θ ∥ 1 ≥ 2 n ) ≤ n E t ∼N (0 , 1) [ | t | ] 2 n < 1 / 2 by Markov inequality.
Thus,
P
[Equation]
By setting τ = M 2 / 2 , this probabilistic argument implies that there exists θ ∈ R n with a unit Lnorm, ∥ θ ∥ 1 = 1 , such that for all x ∈ X , we have ∥ x ∥ ∞ ≤ 4 M 2 n | θ ⊤ x | .
Using this θ , we have
[Equation]
On the other hand, consider 1-dimensional M -atomic distributions γ θ := (cid:80)
Mm =1 w m δ θ ⊤ ν m and ˆ γ θ := (cid:80) Mm
=
ˆ w m δ θ ⊤ ˆ ν m .
Then by Cauchy-Schwartz inequality
[Equation]
for all d = 1 , 2 , ..., 2 M − 1 .
Since γ θ , ˆ γ θ are 1-dimensional atomic distributions, by Theorem C.2,
[Equation]
Note that any coupling over ν θ , ˆ ν θ can be converted to a coupling over ν, ˆ ν and vice versa.
Plugging the above result into (13), we have the theorem.
□
With the above theorems, if we set δ < ( ϵ/ ( M 3 n )) 2 M − 1 , we have W ( γ, ˆ γ ) ≤
O
( ϵ ) .
First note that the difference in expected values of a policy is bounded by total variation distance between trajectory distribution:
[Equation]
Now we only need to focusing on total variation distance bound.
Fron this point, when we sum over sequences, we sum over all possible sequences if not specified.
With a slight abuse in notation
[Equation]
we use a compact notation for probability of action sequences:
[Equation]
First
[Equation]
suppose any coupling Γ between contexts m and m ′ in two systems
[Equation]
and we can write (cid:88)
[Equation]
where the inequality holds by the triangle inequality.
Now we can proceed as (cid:88)
[Equation]
where we used triangle inequality at the t
= H time step.
Note that (cid:80) r H | µ m ( a H , r H ) − ˆ µ m ( a H , r H ) | =
∥ µ
[Equation]
a H , · ) − ˆ µ m
[Equation]
a
[Equation]
, · ) ∥ 1 .
Thus, we can bound the first term by
[Equation]
where we summed over all probabilities over sequences where we used the fact that
[Equation]
For the second term, we observe that
[Equation]
since (cid:80) r H ˆ µ m ′ ( a, r H ) = 1 for any a .
From here, we can apply the same decomposition which we used at t =
H level, and apply the same argument recursively until t = 1 .
That gives for any coupling Γ , (cid:88)
[Equation]
[Equation]
Since the inequality holds for all valid couplings Γ( m, m ′ )
such that
[Equation]
we can take the infimum over Γ , and conclude that
[Equation]
C.5 Proof of Lemma 3.
In order to bound difference between expected long-term rewards for a fixed history-dependent policy π , it is sufficient to bound the difference in distributions of observations following π .
We first explicitly write down total variation distance between observations from { ( w m , µ m ) }
Mm =1 and { ( ˆ w m , ˆ µ m ) }
Mm =1 : (cid:88) r 1: H
[Equation]
a 1: H
[Equation]
M
[Equation]
, where π ( a 1: H | r 1: H − 1 ) is as defined in (14), in the first inequality we used the condition ∥ (cid:80)
Mm =1 w m µ
[Equation]
m
[Equation]
∥ ∞ ≤ δ , and the last inequality follows from that:
[Equation]
Now the difference in expected rewards follows as
[Equation]
C.6 Proof of Proposition 3.
To show this result, we first express each coordinate of µ m in terms of ν m .
That is, for all ( a, z ) ∈
A × Z , using Proposition B.2, by setting η = 0 we can express µ m as µ
[Equation]
Tν m , for some linear mapping T ∈ R AZ × n (see Section C.8.1 for details on T ).
Furthermore, from the conclusion of Proposition B.2 which implies the robustness of µ m against the perturbation of ν m in l ∞ norm, we can infer that every row of T has l 1 norm bounded by √ 2 M , i.e., ∥ T ( a,z ) , : ∥ 1 ≤ √ 2 M for all ( a, z ) ∈
A × Z .
To bound ∥ (cid:80) M m =1 w m µ (cid:78) l m − (cid:80) M m =1 ˆ w m ˆ µ (cid:78) l m ∥ ∞ , we only need to check one entry of tensors at any position (( a 1 , z 1 ) , ..., ( a l , z l )) and all other entries are bounded in a similar fashion.
We first check that
[Equation] [Equation]
 
[Equation]
ν
[Equation]
 
[Equation]
Unfolding the product expression over i for the original parameter part ν ,
M
[Equation]
Plugging this expression, we conclude the proof:
[Equation]
(cid:12)(cid:12)(cid:12) (cid:12)(cid:12)
[Equation]
which is what we needed to show.
C.7 Proof of Lemma 3.
For each episode k where k ∈ [ N 0 ] , let the first and second actions be a k 1 , a k 2 ∼ Unif ( A ) , and let r k 1 , r k 2 be the corresponding reward feedback.
We construct an empirical second-order moments
ˆ Msuch that ˆ M 2 ( i, j ) is the mean of r 1 · r 2 when a k 1 = a i , a k 2 = a j .
Specifically, we construct ˆ M 2 as the following:
[Equation]
(15) where e ( a kt ,r kt ) is a standard basis vector in R AZ with 1 at position ( a kt , r kt ) .
The argument follows from a rather standard concentration argument for dimensionality reduction (e.g., Lemma 3.5 in [12]).
Let δ =
∥
M 2 − ˆ M 2 ∥ 2 .
With standard measure of concentration arguments, we can show that ∥ M 2 − ˆ M 2 ∥ ∞ < C · (cid:112)
A 2 log( AZ/η ) /N
with probability at least 1 − η , which is translated to δ ≤ C · (cid:113) A 4 Z 2 log( AZ/η ) Nfor some universal constant C > 0 .
Let P U be the orthogonal projector onto the top- M eigenspace of M 2 and P ⊥ U
=
I − P U .
We can also define similar quantities from the empirical estimate ˆ M 2 .
Let P (cid:98) U similarly be the orthogonal projector onto the top- M eigenspace of ˆ M 2 , and let P (cid:98) U ⊥ =
I − P (cid:98) U .
By Weyl’s theorem, we have λ M +
( ˆ M 2 ) ≤ δ since rank( M 2 ) is M .
Our goal is to bound
[Equation]
for any m ∈ [ M ] .
Let y be the vector maximizing the above.
Since w m µ m µ ⊤ m ⪯ M 2 and since y belongs to the eigenspace of ranks lower than M ,
[Equation]
We have shown that ∥ µ m − (cid:98) U µ m ∥ 2 ≤ (cid:112) 2 δ/w m .
Thus we need δ :
= O ( δ 2sub ) to bound the l 2 error by δ sub /w 1 / 2 m .
Thus N
=
O (cid:0)
Z 2 A 4 log( AZ/η ) /δ
4sub (cid:1) samples are sufficient for the estimation of the subspace U .
Note that we have not optimized for the polynomial factors which can be improved by more tightly bounding δ =
[Table]
C.8 Deferred Details of Theorem 3.
The proof of Theorem 3.8 follows by collecting the results from the preceding lemmas for subspace estimation (Lemma 3.7) and results from Section 3.2 and 3.3.
Before we get into the proof, let us describe some details about finding a set of core action-value pairs, and how to recover the original model parameters { µ m } Mm =1 from { ν m } Mm =1 .
C.8.1 Detailed Procedures for Experimental Design
Step 1.
Subspace estimation We first find the set of core action-value pairs { ( a j , z j ) } nj =1 following the same procedure in Corollary 3.2.
By Lemma 3.7, for every µ m , we have
[Equation]
where ∥ ∆ m ∥ ∞ ≤ δ sub /w 1 / 2 m .
Step 2.
Pick core action-event pairs Let { ˆ β j } Mj =1 be the orthonormal basis of (cid:98) U .
We can use this basis as input to Corollary 3.2 to get a set of (approximate) core action-event pairs.
Specifically, let ˆΦ ∈ R AZ
×
M be a feature matrix where each j th column ˆΦ : ,j is given as:
[Equation]
After invoking Theorem B.1 with supplying ˆΦ as input, we use the support of ρ as the set of core action-value pairs { ( a j , z j ) } nj =1 .
Let ˆ G ( ρ ) be defined as in equation (10).
Note that with too small mixing weights w m , we can instead use ∆ m = − µ m with ∥
∆ m ∥ ∞ ≤ 1 .
Step 3.
Search constraints for moment-matching After finding a core action-event pairs { ( a j , z j ) } nj =1 as in Corollary 3.2, we estimate ν m from higher-order tensors { ˆ T l } min( H, 2 M − 1) l
= .
When searching parameters for { ( ˆ w m , ˆ ν m ) } Mm =1 , we can put constraints to ensure that ˆ w m and
ˆ ν m belong to a set of valid parameters for all m ∈
[ M ] :
[Equation]
where ⪯ is an element-wise inequality.
That is, we want that ˆ v m =
ˆ T ˆ ν m is not too far from ˆ µ m after clipping and normalization.
Without loss of generality, we assume that a rough estimate of w min is known in advance (otherwise, we can repeat the same procedure with geometrically decreasing estimates of w min , e.g., 1 /M, 1 / 2 M, 1 / 4 M, ..., 1 /M
O (min( M,H )) , and pick the best returned policy).
A solution satisfying all constraints is guaranteed to exist since the true model { ( w m , ν m ) } Mm =1 also satisfies constraints.
Step 4.
Recovery of parameters Let ˆ µ m be computed by (11) using ˆ ν m and ˆΦ , i.e., let ˆ T ∈ R AZ × n be defined as
[Equation]
We let v m =
ˆ
Tν m , ˆ v m =
ˆ T ˆ
ν m .
We recover ˆ µ m from ˆ v m =
ˆ T ˆ
ν m as ˜ µ m :
[Equation]
clip (ˆ v m , 0 , 1) ,
[Equation]
A simple algebra can show that the normalized estimates of ˆ µ m are close to µ m .
Specifically, we show the following with Proposition B.2:
[Equation]
C.8.2 Proof of Theorem 3.8, Case I: H ≥ 2 M −
We search the empirical parameters { ( ˆ w m , ˆ ν m ) }
Mm =1 ) over the set (17).
After finding ˆ w m , ˆ ν m that satisfy the moment matching condition (7)
[Equation]
we recover ˆ µ m from ˆ v m =
ˆ T ˆ ν m after clipping and normalization (19).
Then we first observe that for any a ∈ A ,
[Equation]
where in the first inequality
[Equation]
we used the fact that clipping can only improve the l 1 error.
Then, errors from the normalization can be bounded as
[Equation]
The second line holds since
[Equation]
where the third relation holds since (cid:80) z ∈Z | ˜ µ m ( a, z ) | / (cid:12)(cid:12)(cid:80) z ∈Z ˜ µ m ( a, z ) (cid:12)(cid:12)
=
since ˜ µ m ( a, z ) ≥ 0 .
By the choice of δ sub , we have ∥ ∆ m ∥ ∞ ≤ δ sub /w 1 / 2 m ≤ ϵ/ (2 ZMH 2 w 1 / 2 m ) .
Now we can call
Proposition 3.4, and proceed as | V ( π ) − ˆ V ( π ) | ≤ H 2 · inf Γ
[Equation]
where in the last inequality, we used Cauchy-Schwartz inequality (cid:80)
Mm =1 √ w m ≤ √ M .
Hence if we have 2 ZH 2 √ 2 MW ( γ, ˆ γ ) ≤ ϵ , which is given by the choice of δ sub and Lemma 3.3, we have | V ( π ) − ˆ V ( π ) | ≤ O ( ϵ ) .
C.8.3 Proof of Theorem 3.8, Case II: H < 2 M −
If H < 2 M − 1 , we start by observing that for any 0 ⪯ ˆ ν m ⪯ 1 , using Proposition B.2 and setting µ = 0 , ϵ 0 = 0 , η = ν m and ϵ 1 = 1 , we have − √ 2 M ⪯ ˆ T ˆ ν m ⪯ √ 2 M .
To exploit the moment-closeness property, we define an auxiliary model { ( w m , v m ) }
Mm =1 where v m := ˆ Tν m .
Similarly to H ≥ 2 M − 1 case, for any a ∈ A , we have that
[Equation] [Equation]
We also have that
[Equation]
Let ˆ v m =
ˆ T ˆ ν m
and ˆ µ m be defined as in (19).
Recall our goal to bound
V
[Equation]
Define auxiliary value functions V aux ( π ) and
ˆ V aux
( π ) as
[Equation]
Then | V ( π ) − ˆ V ( π ) | ≤ | V ( π ) − V aux
( π ) | + | V
aux ( π ) − ˆ V aux ( π )
| + | ˆ V aux ( π ) − ˆ V ( π ) | .
We bound each term separately.
Term I. | V ( π ) − V aux
( π ) | : This is less than | V ( π ) −
[Table]
[Equation]
For the second term, we have
[Equation]
for all m ∈ [ M ] .
For the first term, we can recursively apply the same inequality until the time step reaches to t = 1 .
Applying this recursively,
[Equation]
To bound the above, we first consider the case when w m ≥ ϵ H 2 M
( Z √ 2 M ) H .
Note that if w min ≥ ϵ
H 2 M ( Z √ 2 M ) H , then this is always the case.
In this case, for every a ∈ A , we have (recall (20))
[Equation] [Equation]
where we use our choice of δ sub in (5).
By the second condition,
[Equation]
The first condition can be combined with Proposition B.2 similarly to the H ≥ 2 M − 1 case to get
H(cid:88)
[Equation]
On the other hand, if w m <
[Equation]
H , then we can directly bound as
[Equation]
Thus, we have | V ( π ) − V aux
( π ) | ≤ O ( ϵ ) .
Term II.
| V aux ( π ) − ˆ V aux ( π ) | : We use the moment-closeness properties between v m and ˆ v m given similarly to Lemma 3.6.
Lemma C.
For any given degree l ≥ 1 , if ∥ (cid:80) Mm =1 w m ν (cid:78) l m − (cid:80) Mm =1 ˆ w m ˆ ν (cid:78) l m ∥ ∞ ≤ δ , then v m and ˆ v m satisfy
[Equation]
Proof.
Note that ∥ (cid:80)
[Equation]
Mm =
ˆ w m ( ˆ T ˆ ν m ) (cid:78) l ∥ ∞ ≤ (2 M ) l/ 2 δ , following the same argument in Appendix C.6: the conclusion of Proposition B.2 also implies that the l 1 norm of every row in ˆ T is less than √ 2 M , i.e.,
[Equation]
Lemma follows since v m is a vector consisting of partial coordinates of ˆ Tν m .
□
Now we proceed as
[Equation]
a
[Equation]
−
[Equation]
≤ HZ H (2 M ) H/ 2 δ tsr .
Choice of δ tsr = ( ϵ/H ) / ( √ 2 MZ ) H for H <
2 M − 1 gives | V aux ( π ) − ˆ V aux ( π ) |
= O ( ϵ ) .
[Table]
[Equation]
For each m ∈ [ M ] , if ˆ w m ≥ ϵ H 2 M ( Z √ 2 M ) H , then
[Equation]
where we used the choice of δ sub in (5).
We also need to show that ∥ ˆ v m ( a, · ) − ˆ µ m ( a, · ) ∥ 1 is bounded.
Let ˜ µ m be the intermediate step after clipping ˆ v m before normalization as in 19.
Due to the third condition of (17), clipped amount can be at most
[Equation]
With the second condition of (17), we have
[Equation]
Similarly,
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Therefore, we can show that
[Equation]
where we used 1 − 3 Zδ sub / ˆ w 1 / 2 m ≥ 1 / 2 due to the first constraint and the choice of δ sub .
If ˆ w m <
ϵ H 2 M ( Z √ 2 M ) H , then we can use the fact that all l 1 -norm of rows of ˆ T are less than √ 2 M (equation (22)), and thus ∥ ˆ v m ( a, · ) ∥ 1 ≤ √ 2 M for all a ∈ A .
We also have that ∥ ˆ µ m ( a, · ) − ˆ v m ( a, · ) ∥ 1 ≤ 1 + √ 2 M .
Now we plug all things together, and proceed as |
[Equation]
|
[Equation]
Collecting all three terms, we have | V ( π ) − ˆ V ( π ) | ≤ O ( ϵ ) .
This concludes the proof of Theorem 3.8.
Appendix D LMAB with Gaussian Rewards
So far we have focused on rewards with finite support Z
= O (1) .
In this section we consider and
LMAB setting with Gaussian rewards and generalize Algorithm 1.
Indeed, some steps in the algorithm cannot be straightforwardly extended if Z = ∞ .
In this subsection, we consider a standard Gaussian reward distribution–a special case of continuous rewards–and generalize Algorithm 1 to this setting.
We make the following assumption.
Assumption D.1 (Gaussian Rewards)
The reward distribution conditioning on an action a ∈ A in a context m ∈
[ M ] is N ( µ m ( a ) , 1) for some | µ m ( a ) | ≤ 1 .
Even though the rewards have infinite support, we show that the same conclusion holds as in finite-support case, i.e., the sample-complexity is upper bounded by O (( MH/ϵ )
O (min( H,M ))
+ poly( A, H, M )) .
Algorithms for the Gaussian case differ significantly in identifiable H ≥ 2 M − 1 and unidentifiable regimes
H < 2 M − 1 .
When H ≥ 2 M − 1 , there are only minor changes in the algorithm design for defining core actions and how tensors are constructed.
We handle this case in Appendix D.1.
A more interesting case is the parameter unidentifiable regime, where we follow an alternative approach and discretize the support of rewards by O ( ϵ/H 2 ) -level.
This approach is described in Appendix D.2.
We can reach similar conclusions for the Gaussian rewards:
Theorem D.2 Consider any LMAB with M contexts under Gaussian reward Assumption D.1.
There exists an algorithm such that with probability at least 1 − η , it returns an ϵ -optimal policy using a number of episodes at most poly( H, M, A, 1 /ϵ, log( A/η ))
+ poly(log( M/η ) , H, M ) 2 M − 1 · ϵ − (4 M − 2) , if H ≥ 2 M − 1 , poly( H, w − 1min , A, 1 /ϵ, log( A/η ))
+ poly(log( MH/ ( ηϵ )) , H, M ) H · ϵ − (2 H +2) , otherwise .
Note that the dependency on ϵ is at most ϵ − (4 M − 2) and smaller
when H < 2 M − 2 .
In the parameter identifiable regime, similarly to discrete reward cases, near-optimality of returned policy comes from the closeness of latent model parameters in the Wasserstein metric.
In the parameter unidentifiable regime, we first discretize the support of rewards in O ( ϵ ) level.
Then we can apply the same procedures for handling discrete rewards as in Section 3.3.
We provide the full details in Appendix D.
Let ˆ M 2 ∈ R
A × A be the empirical second-order moments as
[Equation]
Then let (cid:98) U be the subspace spanned by top- M eigenvectors of ˆ M 2 .
A similar conclusion to Lemma 3.7 holds:
Lemma D.3 Let (cid:98) U be a subspace spanned by top- M eigenvectors of ˆ M 2 .
After we estimate ˆ Musing N
= O
( A 4 log( A/η ) /δ 4sub ) episodes, with probability at least 1 − η , for all m ∈
[ M ] , there exists ∆ m : ∥
∆ m ∥ ∞ ≤ δ sub /w 1 / 2 m such that µ m + ∆ m ∈ (cid:98) U .
Proof of Lemma D.3 is identical to the proof of Lemma 3.7.
Let δ sub = ϵ/ (2 MH 2 ) .
Similarly to finite-support reward distributions, let { ˆ β j } Mj =1 be the orthonormal basis of (cid:98) U and construct ˆΦ
∈ R A × M such that the j th column of ˆΦ is ˆ β j , i.e., ˆΦ : ,j = ˆ β j .
We invoke Theorem B.1 to get a set of core actions { a j } nj =1 .
The main difference to the finite-support case is that we do not need to specify a corresponding event of rewards.
Instead, we measure the correlation in terms of actual reward values.
Specifically, for every multi-index ( i 1 , i 2 , ..., i l ) ∈ [ n ] l , using N 1 episodes where
[Equation]
We play a kt = a i t for t = 1 , ..., l and k ∈ [ N b ] , and estimate higher-order moments:
[Equation]
=
[Equation]
We can easily verify that
[Equation]
We can apply the concentration of higher-order polynomials of sub-Gaussian random variables element-wise, given from the following lemma on hypercontractivity inequality:
Lemma D.4 (Hypercontractivity Inequality (Theorem 1.9 in [34])) Consider a degree- l polynomial f defined over a set of N independent samples of zero-mean unit-variance Gaussians, such that f ( X 1: N ) :
= f
( X 1 , ..., X N ) .
Then,
[Equation]
, for some absolute constant C > 0 .
To show the concentration of ˆ T l ( i 1 , ..., i l ) around T l ( i 1 , ..., i l ) , we can apply Lemma D.4 with plugging λ
= O ( δ tsr ) and var ( f ) ≤ 2 l · var ( X l ) /N 1 ,
where X ∼ N (0 , 1)
Here, f can be viewed as a degree- l polynomial of X kt := r kt − µ m k ( a kt ) for t = 1 , . . .
, l and k = 1 , ..., N 1 , where m k is a latent context for the episode k .
Since var ( X l ) ≤ O ( l l ) , we need N
=
O
( l l log l (cid:0) ln l /η (cid:1) /δ
2sub ) to make the exponent less than η/ ( ln l ) .
Take union bound over all elements in ˆ T l ensures that ∥
T l − ˆ T l ∥ ∞ ≤ δ tsr with probability at least 1 − η .
Now we find a set of parameters { ( ˆ w m , ˆ ν m ) }
Mm =1 with the only constraint: ˆ w m ∈ R + ,
M (cid:88) m
= ˆ w m = 1 .
We aim to find
[Equation]
Now we can construct an empirical model by recovering ˆ µ m =
ˆ T ˆ
ν m where ˆ T ∈ R
A × M is defined similarly to (18) as
[Equation]
where ρ is a distribution over rows of ˆΦ found by Theorem B.1 and ˆ G ( ρ ) is defined as in (10).
We do not need extra clipping and normalization steps here since any ˆ µ m is a valid model parameter.
Now with { ( ˆ w m , ˆ µ m ) } , we call the planning oracle 2.3 and obtain an ϵ -optimal policy.
D.2 Algorithm with Short Time-Horizon H < 2 M −
We first discretize possible reward values: let Z = { z 1 , z 2 , ..., z Z } where z i = − 4 (cid:112) log( H/ϵ )
+ ( i − 1) · ϵ/H 2 and Z =
⌊ 8 H 2 (cid:112) log( H/ϵ ) /ϵ ⌋ .
We define an auxiliary reward (pseudo) p.d.f p m ( a, · ) for each a and m as the following: for all s ∈ [ Z − 1] ,
[Equation]
and p m ( a, r ) =
for all r ∈ ( −∞ , z 1 ) ∪ [ z Z , ∞ ) .
Define ˜ V ( · ) be the policy evaluation function in a (pseudo) LMAB model ˜ B
=
( A , { w m } Mm =1 , { p m } Mm =1 ) :
˜ V ( π ) :=
[Equation]
We first show that p m ( a, · ) is good approximation of true reward distributions:
Lemma D.5 Let B and ˜ B defined as above.
Then for any history-dependent policy π , | V ( π ) − ˜ V
( π ) | ≤ 10 ϵ .
Given Lemma D.5, we will discretize reward values are so that we can leverage the result of Section 3.3 for short time-horizon.
Specifically, let B be a model { ( w m , q m ) }
Mm =1 with discrete reward distributions taking values in
Z ∪ { 0 } , and let q m ( a, z ) = ϵH 2 · p m ( a, z ) for z ∈ Z and q m ( a, 0) =
P m ( r / ∈ [ z 1 , z Z ) | a ) .
As if the underlying model is B , we run Algorithm 1 with manually modifying the observed rewards r t → r t : r t = 0 , if r t < z 1 or r t ≥ z Z , r t =
z s , for some s ∈
[ Z − 1] , s.t. r t ∈ [ z s , z s +1 ) .
From actions ( a 1 , ..., a H ) and reward observations ( r 1 , ..., r H ) , we can now apply Algorithm 1 for the parameter unidentifiable case
H < 2 M − 1 .
D.3 Proof of Theorem D.
Now we are ready to prove the Theorem D.2 for Gaussian rewards.
D.3.1 Identifiable Regime H ≥ 2 M − 1 :
By Lemma 3.3, we know that W (
[Equation]
(cid:17) where ν
= (cid:80) M m =1 w m δ ν m
and ˆ ν = (cid:80)
Mm =
ˆ w m δ ˆ ν m .
With experimental design, by Corollary 3.2, we can ensure that
[Equation]
Observe that for any a ∈ A , total variation distance between standard Gaussians is bounded by the distance between centers of Gaussians, i
[Equation]
e.
[Equation]
Using Proposition 3.
[Equation]
we can show that
[Equation]
Plugging the choice of δ tsr
= O ( ϵ/ ( H 2 M 3 .
5 n )) 2 M − 1 , this is less than ϵ .
D.3.2 Unidentifiable Regime H < 2 M − 1 :
Let us first compare the expected rewards from ˜ B and B with any fixed policy π .
˜
[Equation]
With slight abuse in notation
[Equation]
let r t be a quantized value of r t .
Then we can show that
[Equation]
+
[Equation]
t =
[Equation]
where in the first inequality, we used p m ( a, r ) = 0 for r / ∈ [ z 1 , z Z ) .
Note that P X ∼N (0 , 1) ( | X | ≥ z Z − 1) ≤ ( H/ϵ ) 4 with z Z = 4 (cid:112) log( H/ϵ ) .
Note that a system with manually discretized rewards can be described by the model B .
Assuming the returned policy ˆ π from Algorithm 1 is O ( ϵ ) -optimal for B , by triangle inequality for the policy evaluation for any policy π ,
[Equation]
we conclude that ˆ π is O ( ϵ ) -optimal for B with Gaussian rewards.
D.4 Proof of Lemma D.
We start by unfolding the expression for policy value differences.
[Equation]
where g m ( a t , r t ) :
=
[Equation]
exp (cid:0) −
[Equation]
r
[Equation]
− µ m
[Equation]
a t )) 2 / 2 (cid:1) .
We first rule out reward values greater than z Z = 4 (cid:112) log( H/ϵ ) .
Define a set of bad reward sequences E b
= { r 1: H |∃ t ∈ [ H ] , s.t., | r t | > z Z } .
Then for any t
[Equation]
∈ [ H ] ,
[Equation]
where the last inequality results from integrating over probabilities for time steps t 0 + 1 , ..., H .
The last summation term can be further bounded by integrating out t th0 time step since (cid:90) {
[Equation]
where in the first inequality we used | µ m ( a t ) | ≤ 1 ≤ (cid:112) log( H/ϵ ) .
In last inequality we used
E [ | X | · 1 {| X | ≥ t } ] ≤ (cid:112) P ( | X | ≥ t ) by Cauchy-Schwartz inequality, and then used the Gaussian tail bound P ( | X | ≥ t ) ≤ exp( − t 2 / 2) for t = 3 (cid:112) log( H/ϵ ) .
Therefore, we now get (cid:88)
[Equation]
where we used P m ( E b ) ≤ H · P ( | X | ≥ z Z − 1) ≤ ϵ 4 /H 3 with sufficiently small ϵ > 0 .
This can be similarly done for all t 0 , and thus (cid:88) a 1: H
[Equation]
r
[Equation]
∈E b (cid:32)
[Equation]
t = | r t | (cid:33)
[Equation]
Π Ht =1 p m ( a t , r t ) − Π
Ht
=
g m ( a t , r t ) (cid:12)(cid:12) π
( a 1: H | r 1: H − 1 ) d ( r 1: H ) ≤ 4 ϵ 2 .
We remain to bound (cid:88)
a
[Equation]
−
[Equation]
d
[Equation]
r
[Equation]
Note that | r t | <
z Z for all t ∈
[ H ] when r 1: H ∈ E cb .
For each t ∈ [ H ] , we aim to bound (cid:88) a 1: H (cid:90) r 1: H ∈E cb | r t | (cid:12)(cid:12)
Π
Ht =1 p m ( a t , r t )
− Π Ht =1 g m ( a t , r t ) (cid:12)(cid:12)
π ( a 1: H | r 1: H − 1 ) d ( r 1: H ) .
Next, for any r t ∈ [ z s , z s +1 ] for s ∈
[ L − 1] , we observe that | p m ( a t , r t ) − g m ( a t , r t ) | =
z
[Equation]
−
[Equation]
s (cid:90)
[Equation]
+
[Equation]
| g
[Equation]
( a t , r t ) | dx
≤ z s +1 − z s (cid:90)
z s + z s | g m ( a t , x ) | + z s +
− z s (cid:90) z
s + z s | g m ( a t , x ) − g m ( a t , r t ) | dx.
Then we observe that
[Equation]
for some x ′ ∈ [ z s , z s +1 ] where we used the mean-value theorem.
A simple algebra shows that for any x, x ′ ∈ [ z s , z s +1 ] ,
[Equation]
where we used the second derivative of x exp( − x 2 / 2) , which is ( x 2 − 1) exp( − x 2 / 2) , is always less than 1 in absolute value.
Plugging above relations into bounding the difference between p m and g m yields
[Equation]
Using this the above
[Equation]
we bound (cid:88)
[Equation] [Equation]
If t = H , then
[Equation] [Equation]
≤
[Equation]
For the first term, (cid:90)
[Equation]
where we used (cid:82) r H 1 √ 2 π r 2 H exp( − r 2 H / 2) ≤ 1 and 32( ϵ/H ) · log( H/ϵ ) < 1 for sufficiently small ϵ .
Note that we also have
[Equation]
π ( a
[Equation]
H | r
[Equation]
H
[Equation]
[Equation] [Equation]
H −
[Equation]
≤ 1 .
For the second term, we first have
[Equation]
Furthermore, we can show that (cid:88)
[Equation]
from which we recursively apply similar arguments.
Thus we can conclude that
[Equation]
| r t | with other time steps can also be similarly bounded.
Thus, we can conclude that
[Equation]
and therefore | V ( π ) − ˜ V ( π ) | ≤ O ( ϵ ) since (cid:80) Mm =1 w m = 1 .
Appendix E Deferred Details in SectionE.1 Additional Definitions
Let us define a few notation and interaction protocol.
Suppose at the beginning of episode, a latent context m 0 ∈
[ M ] is chosen, and and at each time step t ∈
[ H ] , we play a i t where i t is sampled from Unif ([ n ]) .
Let i =
( i 1 , i 2 , ..., i H ) be the sequence of indices of played core actions, and b = ( b 1 , b 2 , ...
b H ) be the event-observation sequence in the episode, where b t := 1 { r t =
Z i t } .
Let the parameter space
Θ be the set of valid parameters:
[Equation]
We use superscript k to denote quantities observed in the k th episode.
The probability of a trajectory under a model θ ∈ Θ in the k th episode is defined by
[Equation]
E.2 Polynomial Upper Bounds with Separation
In this subsection, we specify the details on separation conditions that make the polynomial sample complexity possible with MLE solutions.
Suppose that there exists a context revealing action for any m ̸ = m ′ ∈
[ M ] , i.e., we are given the following assumption:
Assumption E.1 (Separated Bandit Instances)
For any m ̸ = m ′ ∈
[ M ] , there exists some (unknown) a ∈
A such that ∥ µ ∗ m ( a, · ) − µ ∗ m ′ ( a, · ) ∥ 1 ≥ γ for some known γ > 0 .
Under Assumption E.1, if the time horizon H =
˜ O ( Z 2 M 2 /γ 2 ) is given enough to identify the context within each episode, then we can significantly improve the sample complexity for learning LMAB, from exponential to polynomial.
Note that the time-horizon H can be still much smaller than A and thus we cannot explore all actions within a single episode, which is in contrast to explicit clustering based approaches studied in [9, 16].
Maximum likelihood estimator for LMABs with separation can guarantee the following:
Lemma E.2 Consider the maximum likelihood estimator θ N = { ( ˆ w m , ˆ ν m ) }
Mm =1 under Assumption E.1 with time-horizon H ≥ C 1 · nMZ 2 log(1 / ( ϵw min ))
/γ
for some universal constant C 1 > 0 .
If N = C 2 · w − 2min n · log( N/η ) /ϵ 2 for some large constant C 2 > 0 , then with probability at least 1 − η , we have (up to some permutations in θ N )
[Equation]
To get the above result, we first observe a consequence due to experimental design: the converse of
Corollary 3.2 implies that if Assumption E.1 holds for all m ̸
= m ′ , then we have ∥ ν ∗ m − ν ∗ m ′
∥ ∞ ≥ γ/ ( Z √ 2 M ) .
Thus, if H/n = ˜ O ( MZ 2 /γ 2 ) , then we can play each core action O ( H/n ) -times and get an O ( γ/ ( Z √ M ))
-accurate estimator ˆ ν m for one of { ν ∗ m ′ }
Mm ′ =
Since we have good separation between samples from different contexts, by proper clustering arguments, the sample complexity of recovering θ ∗ can be polynomial.
With Lemma E.2, we can use equation (4) to connect the near-optimality of returned policy computed with { ( ˆ w m , ˆ µ m ) } Mm =1 and the closeness in Wasserstein metric.
We mention that for H , the dependence on Z can be removed with more computationally expensive experimental design (see also Remark C.1).
E.3 Proof of Lemma 4.
We connect the maximum likelihood estimator to total variation distance between observations from θ ∗ and θ N .
The connection between MLE θ N and closeness in distributions of observations b can be established by the following lemma.
Lemma E.3 There exists a universal constant C > 0 such that with probability at least 1 − η ,
[Equation]
N .
That is, total variation distance between two observation distributions is bounded by ˜ O (cid:16)(cid:112) n/N (cid:17) .
On the other hand, for any l ∈ [min( H, 2 M − 1)] and any multi-index ( i 1 , i 2 , ..., i l ) ∈ [ n ] l , we have
[Equation]
This implies that
[Equation]
Applying this to all l ∈ [min( H, 2 M − 1)] , we get Lemma 4.1.
E.4 Proof of Lemma E.
From Lemma E.3, without loss of generality, we assume that the total variation distance between observations from θ N and θ ∗ is bounded by ϵw min / 2 since N = ˜
O ( w − 2min n/ϵ 2 ) :
[Equation]
We will verify that for every m ∈ [ M ] , there exists m ′ ∈
[ M ] such that ∥ ˆ ν m − ν ∗ m ′ ∥ ∞ ≤ 2 ϵ and | ˆ w m − w ∗ m ′
| ≤ ϵw min .
First note that for all m ̸ = m ′ ∈
[ M ] , we have ∥ ν ∗ m − ν ∗ m ′
∥ ∞ ≥ λ :=
γ/ ( Z √ 2 M ) .
If it is not, then the model does not satisfy Assumption E.1.
We start with the following lemma:
Lemma E.4 Suppose that there exists m ∈
[ M ] such that ∥ ν ∗ m − ˆ
ν m ′ ∥ ∞ ≥ λ/ 4 for all
m ′ ∈
[ M ] .
For every j ∈ [ n ] , define E m,j an event defined as:
[Equation]
and let E m =
∩ nj =1 E m,j .
Then
[Equation]
Proof.
Let us first check that P θ ∗ ( E ) ≥ w m / 2 .
Let n j = (cid:80)
Ht =1 1 { i t = j } .
Since,
[Equation]
it suffices to show that
[Equation]
where in the last inequality we applied Hoeffeding’s concentration inequality for first term
[Equation]
and then used Bernstein’s inequality.
Plugging H > C · nMZ 2 log(1 / ( ϵw min ))
/γ
for some sufficiently large constant
C > 0 , we get
[Equation]
By union bound
[Equation]
we have P θ ∗ ( E ) ≥ w m
[Equation]
(1 − nϵ 2 w 2min ) ≥ w m / 2 .
Now we check that P θ N ( E ) ≤ ϵw min .
Starting from
[Equation]
, it suffices to show that P θ N ( E | m 0 = m ′ )
≤ ϵw min for all m ′ ∈
[ M ] .
Let us fix m ′ and define
[Equation]
and let E m ′
=
∩ nj =1 E m ′ ,j .
Following the same argument for P θ ∗ ( E j | m
= m ) , we can show that
[Equation]
If this happens, then it implies E c since ∥ ν ∗ m − ˆ ν m ′
∥ ∞ ≥ λ/ 2 .
Thus,
M (cid:88)
[Equation]
In other words, we have
[Equation]
□
The conclusion of Lemma E.4 contradict that | P θ ∗ ( E m ) − P θ N ( E m ) | ≤ ϵw min / 2 due to the total variation distance bound from Lemma E.3.
Thus, we ensure that for all m ∈
[ M ] , there exists m ′ ∈
[ M ] such that ∥ ν ∗ m − ˆ ν m ′
∥ ∞ ≤ λ/ 4 .
Now without loss of generality, we can ignore the permutation invariance of models and assume that
[Equation]
Then we now show that for all m ∈
[ M ] , it holds that
[Equation]
For every m , let us define E m similarly to Lemma E.4:
[Equation]
and E m :=
∩ nj =1 E m,j .
Then we proceed as the following:
[Equation]
=
[Equation]
Recall that for any m ′ ̸
= m , the event E m implies E cm ′ since every ν ∗ m and ν ∗ m ′ are separated by at least λ =
γ/ ( Z √ 2 M ) , which implies that ν ∗ m and ˆ ν m ′ are separated by (3 / 4) λ .
Following the same argument as in the proof for Lemma E.4,
[Equation]
Similarly, we can also check that
[Equation]
Thus, we have (cid:88)
i
[Equation]
We remain to lower bound (cid:80) b , i |
[Equation]
b ,
[Equation]
− w
[Equation]
m P θ ∗
[Equation]
b ,
[Equation]
| m
[Equation]
= m ) | .
First note that
[Equation]
Thus we have | ˆ w m − w ∗ m | ≤ ϵw min .
Now given this, we can proceed as
[Equation]
Now let i m := arg max j ∈ [ n ] | ν ∗ m ( j ) − ˆ ν m ( j ) | .
Define an event b t being 1 at the first time a i m is played:
[Equation]
Note that since H ≫ n log(1 / ( ϵw min )) , i m is played at least once with probability at least 1 − ( ϵw min ) 2 .
Then we can lower bound the total variation distance as (cid:88)
[Equation]
Hence we can conclude that
[Equation]
which implies ∥ ν ∗ m − ˆ ν m ∥ ∞ ≤ 2 ϵ .
Note that the Wasserstein distance can thus be bounded as
[Equation]
Remark E.
In our guarantee, we assumed that H ≫ log(1 /ϵ ) , i.e., H should be increased logarithmically with the final accuracy.
As one might imagine, this is not the optimal condition for separations between individual models.
A more delicate and technically involved analysis might reveal that the sub-optimal dependency on log(1 /ϵ ) can be dropped with the EM algorithm as in learning Gaussian mixture models (e.g., [22]).
Since it is technically much more complicated, we leave the task of verifying more tight separation conditions as future work.
This is a rather standard consequence of MLE for parameterized distributions.
Let us define wellconditioned parameters Θ ′ ⊆
Θ defined as follows:
[Equation]
Let Θ ϵ be ϵ 2 -covering of Θ ′ .
Since there are n + M free parameters, log-cardinality of Θ ϵ is at most
O (1 /ϵ ) 2( n + M ) .
Note that Θ ϵ is a (not necessarily minimal)
ϵ -cover for Θ as well.
To simplify the notation, we often use X :=
( b , i ) (and X k = ( b k , i k ) ) to replace a sample trajectory.
Our goal is to bound
[Equation]
where TV ( θ 1 , θ 2 ) is a total variation distance between P θ 1 and P θ 2 for any θ 1 , θ 2 ∈ Θ .
Let θ = { ( w m , ν m } Mm =1 ∈ Θ ϵ such that
[Equation]
Such θ is guaranteed to exist in Θ ϵ by construction.
A simple algebra shows that for any trajectory X , we have
[Equation]
As long as ϵ < 1 /H 2 , this is bounded by constant.
Now, fix any θ ∈ Θ ϵ and let l ( X ) :
= 12 log (cid:16) P θ ( X ) P θ ( X ) (cid:17) .
Then using Chernoff’s method, we get
[Equation]
where we used E [exp( S − log( E [exp( S )]))]
=
E [exp( S )] E [exp( S )]
= 1 for any S and distribution.
Taking union bound over Θ ϵ , we can conclude that with probability at least 1 − η ,
[Equation]
Now, since X 1 , ..., X N are independent, we have − log
[Equation]
log
[Equation]
Using − ln( x ) ≥ 1 − x , we get −
N (cid:88) k
= log
[Equation]
∗
[Equation]
P
[Equation]
X
[Equation]
≥
[Equation]
=
[Equation]
k
[Equation]
(cid:32)
[Equation]
−
[Equation]
where H ( θ 1 , θ 2 ) is a Hellinger distance between P θ 1 and P θ 2 .
Also note that TV ( θ ∗ , θ ) ≤ H ( θ ∗ , θ ) .
Collecting all, we now can say that
[Equation]
with probability at least 1 − η .
Finally, let θ N ∈
Θ
ϵ be the one close to θ N such that
[Equation]
similarly to when defining θ .
Then,
[Equation]
where we used the fact that θ N is the maximum likelihood estimator.
With a proper scaling of ϵ ≪ 1 / ( nHN ) 4 , we get
[Equation]
Finally, it is not hard to show that TV ( θ N , θ N ) ≤ 2 ϵH ≤ 1 /N .
We can conclude that
[Equation]