On the Sample Complexity of Representation Learning in Multi-task Bandits with Global and Local structure

We investigate the sample complexity of learning the optimal arm for multi-task bandit problems. Arms consist of two components: one that is shared across tasks (that we call representation) and one that is task-specific (that we call predictor). The objective is to learn the optimal (representation, predictor)-pair for each task, under the assumption that the optimal representation is common to all tasks. Within this framework, efficient learning algorithms should transfer knowledge across tasks. We consider the best-arm identification problem for a fixed confidence, where, in each round, the learner actively selects both a task, and an arm, and observes the corresponding reward. We derive instance-specific sample complexity lower bounds satisfied by any $(\delta_G,\delta_H)$-PAC algorithm (such an algorithm identifies the best representation with probability at least $1-\delta_G$, and the best predictor for a task with probability at least $1-\delta_H$). We devise an algorithm OSRL-SC whose sample complexity approaches the lower bound, and scales at most as $H(G\log(1/\delta_G)+ X\log(1/\delta_H))$, with $X,G,H$ being, respectively, the number of tasks, representations and predictors. By comparison, this scaling is significantly better than the classical best-arm identification algorithm that scales as $HGX\log(1/\delta)$.


Introduction
Learning from previous tasks and transferring this knowledge may significantly improve the process of learning new tasks. This idea, at the core of transfer learning [Pan and Yang, 2009, Skinner, 1965, Woodworth and Thorndike, 1901, lifelong learning [Thrun, 1996] and multi-task learning [Baxter et al., 2000, Caruana, 1995, 1997, has recently triggered considerable research efforts with applications in both supervised and reinforcement learning. Previous work on transfer and multi-task learning has mostly focused on batch learning problems [Lazaric, 2012, Pan andYang, 2009], where when a task needs to be solved, a training dataset is directly provided. Online learning problems, where samples for a given task are presented to the learner sequentially, have been much less studied Stone, 2009, Zhan andTaylor, 2015]. In this paper, we consider a multi-task Multi-Armed Bandit (MAB) problem, where the objective is to find the optimal arm for each task using the fewest number of samples, while allowing to transfer knowledge across tasks. The problem is modelled as follows: in each round, the learner actively selects a task, and then selects an arm from a finite set of arms. Upon selecting an arm, the learner observes a random reward from an unknown distribution that represents the performance of her action in that particular task. To allow the transfer of knowledge across the various tasks, we study the problem for a simple, albeit useful model. We assume that the arms available to the learner consist of two components: one that is shared across tasks (that we call representation) and one that is task-specific (that we call predictor). Importantly, the optimal arms for the various tasks share the same representation. The benefit of using this model is that we can study the sample complexity of learning the best shared representation across tasks while learning the task-specific best action. Contribution-wise, in this work we derive instance-specific sample complexity lower bounds satisfied by any (δ G , δ H )-PAC algorithm (such an algorithm identifies the best representation with probability at least 1 − δ G , and the best predictors with probability at least 1 − δ H ). Again, our lower bounds can be decomposed into two components, one for learning the representation, and one for learning the predictors. We devise an algorithm, OSRL-SC, whose sample complexity approaches the lower bound, and scales at most as H(G log(1/δ G ) + X log(1/δ H )). Finally, we present numerical experiments to illustrate the gains in terms of sample complexity one may achieve by transferring knowledge across tasks.
Related work. Multi-task learning has been investigated under different assumptions on the way the learner interacts with tasks. One setting concerns batch learning (often referred to as learning-to-learn), where the training datasets for all tasks are available at the beginning [Baxter et al., 2000, Maurer, 2005, 2009, Maurer et al., 2013. The so-called batch-within-online setting considers that tasks arrive sequentially, but as soon as a task arrives, all its training samples are available [Balcan et al., 2015, Pentina and Ben-David, 2015, Pentina and Urner, 2016, Alquier et al., 2017. Next, in the online multi-task learning [Agarwal et al., 2008, Abernethy et al., 2007, Dekel et al., 2007, Cavallanti et al., 2010, Saha et al., 2011, Lugosi et al., 2009, Murugesan et al., 2016, Yang et al., 2020, in each round, the learner observes a new sample for each task, which, in some cases, this may not be possible. Our framework is different as in each round the learner can only select a single task. This framework has also been considered in [Lazaric et al., 2013, Soare et al., 2014, Soare, 2015, Alquier et al., 2017, Wu et al., 2019, but typically there, the learner faces the same task for numerous consecutive rounds, and she cannot select the sequence of tasks. Also, note that the structure tying the reward functions of the various tasks together is different from ours. The structure tying the rewards of actions for various contexts is typically linear, and it is commonly assumed that there exists a common low-dimensional representation, or latent structure, to be exploited [Soare et al., 2014, Soare, 2015, Deshmukh et al., 2017, Kveton et al., 2017, Hao et al., 2020, Lale et al., 2019, Yang et al., 2020, Lu et al., 2021, or that the reward is smooth across tasks and/or arms [Magureanu et al., 2014, Slivkins, 2014. The aforementioned papers address scenarios where the context sequence is not controlled, and investigate regret. Meta-learning is also closely connected to meta-learning [Cella et al., 2020, Kveton et al., 2021, Azizi et al., 2022. In [Azizi et al., 2022] the authors investigate the problem of simple regret minimization in a fixed-horizon when tasks are sampled i.i.d. from some unknown distribution. To our knowledge, this paper is the first to study how tasks should be scheduled towards a sample-optimal instance-specific best-arm identification algorithm.

Model and assumptions
In this section, we describe the analytical model of the multi-task MAB problem considered, and describe the framework of best-arm identification for this class of multi-task MAB models.
Model. We consider multi-task MAB problems with a finite set X of X tasks. In each round t ≥ 1, the learner chooses a task x ∈ X and an arm (g, h) ∈ G × H. The components g and h are, respectively, referred to as the representation and the predictor. When in round t, x(t) = x and the learner selects (g, h), she collects a binary reward Z t (x, g, h) of mean µ(x, g, h) (for the sake of the analysis we only analyse the binary case, although it can be easily extended to the Gaussian case as in [Garivier and Kaufmann, 2016]). The rewards are i.i.d. across rounds, tasks, and arms. Consequently, the system is characterized by µ = (µ(x, g, h)) x,g,h , which is unknown to the learner. The main assumption made throughout the paper is that tasks share a common optimal representation g : for any task x ∈ X , there is a predictor h x such that (g , h x ) yields an optimal reward. Formally, µ(x, g, h).
(1) Figure 1: Example of two symmetric tasks x 1 and x 2 , where learning the optimal representation g can be accelerated by considering both tasks, instead of focusing only on a single task. Task x 1 can be used to learn that g 3 is suboptimal, while task x 2 can be used to learn that g 2 is suboptimal.
Moreover, note that there is no assumption on the smoothness of µ with respect to (x, g, h). This type of model represents the case where a learner can actively choose the task to execute (as if a generative model is available to the learner), and in this way maximize the learning efficiency by accurately picking tasks to reduce the sample complexity. Since the model is quite generic, it can be applied to a variety of problems where a collection of tasks have a local component, and a shared global component: (i) influence mechanisms with global/local groups; (ii) hyperparameters learning across multiple tasks; (iii) for advanced clinical trials, where, depending on the group of patients (tasks, that vary according to factors such as age, severity of the disease, etc.), different drugs and dosages can be used for inoculation (g and h).
Sample complexity. The objective of the learner is to devise a policy that learns the best arms (g , h 1 , . . . , h X ) with the least number of samples. Here, a policy π is defined as follows. Let F π t denote the σ-algebra generated by the observations made under π up to and including round t. Then π consists of (i) a sampling rule: in each round t, π selects a F t−1 -measurable task x π (t) and an arm (g π (t), h π (t)); (ii) a stopping rule defined by τ , which is a stopping time w.r.t. the filtration (F t ) t≥1 ; (iii) a decision rule returning a F τ -measurable estimated best arm for each task (ĝ,ĥ 1 , . . . ,ĥ X ). Then, the performance of a policy π is assessed through its PAC guarantees and its expected sample complexity E[τ ]. PAC guarantees concern both learning g and (h 1 , . . . , h X ). Denote by M = {µ : ∃(g , h 1 , . . . , h X ) : Eq. (1) holds} the set of systems where tasks share a common optimal representation. Then, we say that π is (δ G , δ H )-PAC if for all µ ∈ M,

Sample complexity lower bound and the OSRL-SC algorithm
This section is devoted to the best-arm identification problem for the model considered in this work. We first derive a lower bound for the expected sample complexity of any (δ G , δ H )-PAC algorithm, and then present an algorithm approaching this limit. In what follows, we let δ = (δ G , δ H ).

Sample complexity lower bound
We begin by illustrating the intuition behind the sample complexity lower bound, and then state the lower bound theorem. To identify the optimal representation g in a task as quickly as possible, an algorithm should be able to perform information refactoring, i.e., to use all the available information across tasks to estimate g . To illustrate this concept, we use the model illustrated in Fig. 1. In this model, there are only 2 tasks x 1 , x 2 , and 3 arms in G (and only 1 arm in H, thus it can be ignored). For this model, to learn that g 3 is sub-optimal, we should mainly sample task x 1 , since the gap between the rewards of g and g 3 is the largest. Similarly, learning that g 2 is sub-optimal should be mainly obtained by sampling task x 2 . Using the same task, to infer that g 2 and g 3 are sub-optimal, would be much less efficient. This observation also motivates why it is inefficient to consider tasks separately, even in the case where µ is highly non-smooth with respect to (x, g, h), and also motivates the expression of the sample complexity lower bound that we now present.
Sample complexity lower bound. Computing the sample complexity lower bound amounts to finding the lower bound of a statistical hypothesis testing problem, which is usually done by finding what is the most confusing model. In this case the lower bound is given by the solution of the following optimization problem.
Theorem 1. The sample complexity τ δ of any δ-PAC algorithm satisfies: is the value of the optimization problem 1 : where in the first constraint ) accounts for the difficulty of learning the best predictor h x for each task x. The term I α (µ 1 , µ 2 ) = αkl(µ 1 , d α (µ 1 , µ 2 ))+ (1 − α)kl(µ 2 , d α (µ 1 , µ 2 )) is a generalization of the Jensen-Shannon divergence, with d α (µ 1 , µ 2 ) = αµ 1 + (1 − α)µ 2 , and α ∈ [0, 1]. Finally, the term α x,g,h := η(x, g , h x )/(η(x, g , h x ) + η(x, g, h)) represents the proportion of time (x, g , h x ) is picked over (x, g, h). The second constraint, µ (η,ḡ,h x ) accounts for the difficulty of learning the optimal g . To define it, let the average reward over C ⊆ G × H for some task x and allocation η to be defined as Then, µ (η,ḡ,h x ) is given by: where kl(a, b) is the KL divergence between two Bernoulli distributions of respective means a and b, and U µ x,ḡ,hx is the set of confusing arms for task x. Then m Ä x, η, U η,µ x,ḡ,hx ä represents the average reward of the confusing model for estimating g .
The set U η,µ x,ḡ,hx is defined by through N µ x,g,h;ḡ,hx , which is the set of arms whose mean is bigger than µ(x, g, h) and that also include (ḡ,h 1 Refer to the appendix for all the proofs. Figure 2: Example of the set U η,µ x,ḡ,hx . All points in the shadowed area are in the set. All arms with average reward above m Ä x, η, U η,µ x,ḡ,hx ä belong to U η,µ x,ḡ,hx (also (ḡ,h x ) ∈ U η,µ x,ḡ,hx ).
Then, the set of confusing arms U η,µ x,ḡ,hx is given by In Theorem 1 η(x, g, h) can be interpreted as the minimal expected number of times any δ-PAC algorithm should select (x, g, h). Eq. (5) corresponds to the constraints on η one has to impose so that the algorithm learns the optimal predictor h x for all x, while Eq. (6) is needed to ensure that the algorithm identifies the best representation g across all tasks. Both Eq. (5) and Eq. (6) define two convex sets in terms of η, and hence K (µ, δ) is the value of a convex program.
Observe that the constraints Eq. (5) and Eq. (6) share the components of η that concern g only. We believe that actually separating the problem into two problems, one for each constraint, as formulated in the proposition below yields a tight upper bound of K (µ, δ).
Note that K H (µ, δ H ) scales as HXkl(δ H , 1 − δ H ) (since the corresponding optimization problem is that obtained in a regular bandit problem for each task [Garivier and Kaufmann, 2016], which scales as Hkl(δ H , 1 − δ H ) for each task). Now, to know how K G (µ, δ G ) scales, we can further derive an upper bound of K G (µ, δ G ).
One can show that L G (µ, δ G ) scales as GHkl(δ G , 1 − δ G ) ( since, even with one task, we need to sample all GH arms to identify g ). To summarize, we have shown that the lower bound K (µ, δ) scales at most as H(Gkl(δ G , 1−δ G )+Xkl(δ H , 1−δ H )). The latter scaling indicates the gain in terms of sample complexity one can expect when exploiting the structure of M, i.e., a common optimal representation. Without exploiting this structure, identifying the best arm for each task would result in a scaling of GHXkl(δ, 1 − δ) for δ = δ G + δ H .
Differences with classical best-arm identification . To better understand the lower bound in Theorem 1 it is instructive to compare the with a classical MAB problem. Consider the best-arm identification problem in MAB with K arms. There the set of confusing problems is Λ(µ) = {λ ∈ [0, 1] K : a (λ) = a (µ)}, where a (µ) denotes the optimal arm under µ, i.e., a (µ) = arg max a∈K µ(a). The sample complexity lower bound derived for such classical MAB problems [Kaufmann et al., 2016] exploits the fact that the set Λ(µ) can be written as Unfortunately, this way of rewriting the set of confusing problems cannot be used in our problem setting. The reason is that the constraint in Λ a (µ) does not account for the model structure, i.e., the optimal representation g needs to be the same across all the tasks (that is equivalent to imposing that a is optimal for all tasks), which is also the reason the lower bound appears more complex. In Appendix A we show how to account for this kind of structure.
Because of this difference, with the model specified in Eq. (1) the confusing parameter λ differs from µ for more than 2 arms (i.e., we need to consider all arms in the the set U η x,ḡ,hx , see also Lemma 1 in Appendix A), whereas in classical MAB to learn that a is sub-optimal, the confusing parameter λ ∈ Λ a (µ) differs from µ only for arms a and a (µ). In fact, in our model to identify that (ḡ,h x ) is sub-optimal we need to consider an alternative model where only the average reward of the arms in the set U η,µ x,ḡ,hx changes. The set U η,µ x,ḡ,hx can be computed easily by recursively enlarging the set. We start with a set with only (g , h x ) and (ḡ,h x ). We compute the corresponding value of m Ä x, η, U η,µ x,ḡ,hx ä , and we add to the set U η,µ x,ḡ,hx the arm (g, h) with the highest mean not already in the set. We iterate until convergence. Fig. 2 provides an illustration of the set U η,µ x,ḡ,hx .

Algorithm
We now present OSRL-SC (Algorithm 1), a δ-PAC algorithm whose expected sample complexity is asymptotically upper bounded by K G (µ, δ G ) + K H (µ, δ H ). The algorithm proceeds in two phases: a first phase aimed at learning g , and a second phase devoted to learning the optimal predictor for each task. At the end of the first phase, we have an estimateĝ of the best representation. In the second phase, for each task x, we use the optimal track-and-stop algorithm [Garivier and Kaufmann, 2016] to identifyĥ x , the best predictor associated toĝ. In the remaining part of the section, we just describe the first phase.
A track-and-stop algorithm to learn g . The lower bound describes the minimal expected numbers η of observations of the various tasks needed to learn g . These numbers minimize x,g,h η(x, g, h) over η ≥ 0 satisfying (Eq. (6)). In other words, the sampling budget should be allocated according to the following distribution: q (µ) ∈ Σ, solving: max q∈Σ ρ(q, µ), where and Σ denotes the set of distributions over X × G × H. We design an algorithm tracking this optimal allocation: it consists of (i) a sampling rule, (ii) a stopping rule, and a (iii) decision rule, described below.
(i) Sampling rule. We adapt the D-tracking rule introduced in [Garivier and Kaufmann, 2016]. The idea is to track q (µ), the optimal proportion of times we should sample each (task, arm) pair. One important issue is that the solution to max q∈Σ ρ(q, µ) is not unique (this happens for example when two tasks are identical). To circumvent this problem, we track q σ (µ), the unique solution of (P σ ) : max q∈Σ ρ(q, µ) − 1 2σ q 2 2 , where σ > 0. When σ is large, Berge's Maximum theorem [Berge, 1963] implies that q σ (µ) approaches the set of solutions of max q∈Σ ρ(q, µ), and that the value C σ (µ) of (P σ ) converges to Select (x(t), g(t), h(t)), observe the corresponding reward and update statistics; t ← t + 1 12: end while 13: returnĝ = arg max gμτG (g) 14: Phase 2. Learning (h 1 , . . . , h X ) 15: For all task x,ĥ x ← [track-and-stop [Garivier and Kaufmann, 2016] with arms (ĝ, h) h∈H and confidence δ H ] 16: return (ĝ,ĥ 1 , . . . ,ĥ X ) More precisely, if the set of under-sampled tasks and arms is not empty, or whenμ t / ∈ M, we select the least sampled (task, arm) pair. Otherwise, we track q σ (μ t ), and select arg max (x,g,h) (ii) Stopping rule. We use Chernoff's stopping rule, which is formulated using a Generalized Likelihood Ratio Test [Chernoff, 1959]. The derivation of this stooping rule is detailed in appendix B. The stopping condition is maxg Ψ t (g) > β t (δ G ), where the exploration threshold β t (δ G ) will be chosen appropriately, and where Ψ t (g) = min (iii) Decision rule. The first phase ends at time τ G , andĝ is chosen as the best empirical representation: g = arg max gμt (g).

PAC and sample complexity analysis
We now present the sample complexity upper bound for Algorithm 1. First, we outline the stopping rule used by the algorithm. Following [Kaufmann and Koolen, 2018], we first need to define φ : Then, the following theorem states that with a carefully chosen exploration threshold, the first phase of OSRL-SC returns the optimal representation w.p. greater than 1 − δ G .
Theorem 2. Let δ G ∈ (0, 1): for any sampling rule, Chernoff's stopping rule with threshold From [Garivier and Kaufmann, 2016], the second phase of OSRL-SC also returns the optimal predictors for each task w.p. greater than 1 − δ H . Finally, in the next theorem, we show that OSRL-SC stops in finite time a.s., and that its expected sample complexity approaches K G (µ, δ G ) + K H (µ, δ H ) for sufficiently small values of the risks δ G , δ H , and sufficiently large σ.
Theorem 3. If the exploration threshold of the first phase of OSRL-SC is chosen as in Theorem Theorem 2, then we have: P µ [τ G < ∞] = 1 and P µ [τ H < ∞] = 1 (where τ H is the time at which the second phase ends). In addition, the sampling complexity of OSRL-SC satisfies: lim sup δ G ,δ H →0 Corollary 1. Additionally, due to Berge's theorem, since ρ(q, µ) − 1 2σ q 2 2 is continuous in q for each (σ, µ), we have:

Numerical results
We analyze the performance of OSRL-SC, and compare it directly with TRACK AND STOP (TAS) [Garivier and Kaufmann, 2016]. We are interested in answering the following question: is it easier to learn the best representation by just focusing on one task, or should we consider multiple tasks at the same time?
Simulation setup. We consider 2 tasks (x 1 , x 2 ), 3 representations (g 1 , g 2 , g 3 ) and 2 predictors (h 1 , h 2 ). This setting is rather simple, although not trivial. Note that as the number of (task, arm) pairs decreases, we expect the gap between the two algorithms to decrease and thus favour TAS. We set up tasks x 1 and x 2 so that they are very similar: actually, the only difference is that the 2nd and 3rd row of the above matrices are swapped. Therefore it should not matter which task TAS picks, but, on the other hand, OSRL-SC should benefit from this small difference. Hence, for each simulation of TAS, we picked one task uniformly at random. We average results over 1120 runs.
Algorithms. We test TAS and OSRL-SC with various risks δ ∈ {0.01, 0.05, 0.1} (with δ = δ G = δ H for OSRL-SC). For TAS, we use the following threshold β t (δ G ) = log Ä 2t(GH−1) δg ä . We tried the same threshold as in OSRL-SC, but it yielded worse results. For OSRL-SC, we set σ = 10 5 . For the example considered, one can see that arg max q∈Σ ρ(q, µ) has a unique maximizer. Therefore, σ will not influence the value of the lower bound ifμ t is approximately close to µ, in norm. However, whenμ t is visibly different from µ, then the value of C σ (μ t ) may be affected by the value of σ. We have not thoroughly explored different values of σ, but we may suggest that a value of σ > 10 3 is a safe choice.
We compute q σ (μ t ) every 12 rounds (which is equal to GHX) to reduce the computational time (this is better motivated in Appendix B). Despite that, one needs to keep in mind that tracking a sub-optimal, or wrong, reference vector q σ may sensibly affect the sample complexity. We can also motivate this period update by the fact that q σ (μ t ) in a small time interval does not change much, as shown numerically in what follows.
To compute q σ (μ t ), in round t + 1, we use as initial condition a convex combination of the previous solution and a uniform point in the probability simplex (with a factor 0.5). This is done to speed-up the algorithm (for more details refer to Appendix B).
Comparison of OSRL-SC and TAS. In Table 1, we report the sample complexity of the two algorithms. In bold, we highlighted results for the first phase of OSRL-SC. Even if the number of representations is higher than the number of predictors, somewhat surprisingly, the first phase OSRL-SC looks on average very stable, with a small confidence interval (when compared to Phase 2 or TAS). It is worth observing that with the smallest number possible of tasks (X = 2), OSRL-SC manages to reduce the required number of rounds to identify the optimal representation, and the predictors, compared to TAS. Furthermore, the first phase of OSRL-SC appears stable also when δ decreases. Between δ = 0.1 and δ = 0.05 there is a relative increase of average sample complexity of roughly 2% for OSRL-SC; between δ = 0.05 and δ = 0.01 we have that the average sample complexity for OSRL-SC has a relative increase of roughly 5%. Overall, these results indicate that OSRL-SC is able to re-factor information efficiently.
Analysis of OSRL-SC: First phase To analyze the convergence of OSRL-SC, we focus on its first phase, specifically on the following quantities:μ t , q σ (μ t ) and C −1 σ (μ t ). We will use the right y-axis of each plot to display the difference between the value in round t and round t − 1 of the quantities considered. Figure 3(b) shows how μ t −μ t−1 2 (normalized by √ XGH) and μ t 2 evolve over time. We clearly see that μ t 2 quickly converges to some fixed value. This convergence appears in all the plots. Figure 3(a) shows the value of C −1 σ (μ t ), the true value C −1 σ (µ), and the relative change of C −1 σ (μ t ) between two consecutive steps. We observe that the convergence rate ofμ t dictates also the convergence of C −1 σ (μ t ). This suggests that we do not need to solve the lower bound optimization problem too often to update the target allocation, which helps reduce the computational complexity.
show the average change of each component, and low-pass filtered using a 8-th order Butterworth filter with critical frequency ω 0 = 0.025. The shadowed areas represent 97.5% confidence interval.    (b) show 2 curves each: the left plot shows the 2-norm of q σ (μ t ) and q σ (μ t ) − q σ (μ t−1 ) (the latter normalized by √ GHX), whilst the right plot shows the same signals computed using the L ∞ -norm. In Figure 4(b), notice that the average absolute change in each component of the reference vector is very small, below 3% after few dozens of steps. Furthermore, we can see that this quantity has a convergence rate that is directly dictated by the convergence ofμ t (even if its convergence rate is smaller). In Figure 4(a), observe that the relative difference between q σ (μ t ) and q σ (µ) around t = 2500 is upper bounded by roughly 1/9.
Finally, and importantly, in Table 2, we show the average proportion of arm pulls under OSRL-SC at Table 2: Analysis of OSRL-SC. Comparison of the optimal allocation vector q σ (µ) and the average proportion of arm pulls N τ G /τ G at the stopping time.

Conclusion
In this work, we analyzed knowledge transfer in stochastic multi-task bandit problems with finite arms, using the framework of best-arm identification with fixed confidence. We proposed OSRL-SC, an algorithm that transfers knowledge across tasks while approaching the sample complexity lower bound. We believe that this paper constitutes a sound starting point to study the transfer of knowledge in more general online multi-task learning problems. The limitation of this work is that we only consider models with a finite number of tasks and arms, which could limit their application in real life. Furthermore, our algorithm converges to an upper bound of the lower bound. Whether it is possible to approach the sample complexity lower bound is an interesting area of research. Lastly, we think it would be interesting to study different structural assumptions (e.g. linearity) that tie reward functions across tasks together, or extend this work to multi-task reinforcement learning in MDPs.

A.1 Proof of the sample complexity lower bound of Theorem 1
The proof relies on classical change-fo-measure arguments, as those used in the classical MAB [Kaufmann et al., 2016].
Define the set of confusing problems As in the analysis of the regret lower bound, we will split the analysis by considering two subsets of Λ, defined as follows: We will now focus on the first subset Λ µ 1 , from which we will derive the first constraint (5) of Theorem 1. Then, we will focus on the second subset Λ µ 2 , from which derives the second constraint (6). First constraint (5). We restrict our attention to Λ µ 1 . Define the set of confusing problems for task x ∈ X as: Λ µ 1 (x) := {λ ∈ Λ µ 1 : h y (λ) = h y (µ), ∀y ∈ X \ {x}}, which is a handy definition that we will use in few steps. Now, consider a (δ G , δ H )-PAC algorithm, and for a specific task x ∈ X define the event E = {ĥ x = h x ,ĝ = g }, whereĥ x andĝ denote respectively the estimated predictor for task x and the estimated optimal representation at the stopping time τ . Let then λ ∈ Λ µ 1 (x), be an alternative bandit model: the expected log-likelihood ratio L τ at the stopping time τ of the observations under the two models µ and λ is given by E µ [L τ ] = (y,g,h)∈X ×G×H η(y, g, h)kl µ|λ (y, g, h), and in view of the transportation Lemma 1 in [Kaufmann et al., 2016] and the definition of (δ G , δ H )-PAC algorithm, we can lower bound the previous quantity at the stopping time τ : Then, we can get a tight lower bound by considering the worst possible model λ. To do so, first introduce the set Λ µ 1 (x, h) = {λ ∈ Λ µ 1 (x) : λ(x, g , h) > λ(x, g , h x )} which is the set of confusing problems where the predictor h x is not optimal in task x. Observe that one can write Λ µ 1 (x) = ∪ h =h x Λ µ 1 (x, h). This rewriting allows us to derive the first constraint as follows: where (a) follows from the fact that in λ we are changing the predictor of only one task x; (b) follows by solving the infimum problem as in Lemma 3 of [Garivier and Kaufmann, 2016], and from the definition of generalized Shannon divergence I α (x, y) with α x,g,h defined as .
The last equality holds because the quantity only depends on whichh x is chosen in task x. The proof is hence completed.
Lemma 1. For any η ∈ R X ×G×H ≥0 , for a fixed x ∈ X and an action (ḡ,h) ∈ (G \ {g }) × H X , , and U µ x,ḡ,hx (η) is defined as Proof. Bear in mind that Λ µ 2 (ḡ,h) is the set of models whereḡ is the optimal representation, andh is the set of optimal predictors. Consequently, for a specific task x, we can rewrite inf λ∈Λ µ 2 (ḡ,h) g,h η(x, g, h)kl µ|λ (x, g, h) as the solution of the following optimization problem: Define N µ x,g,h to be the set of arms whose mean is bigger than µ(x, g, h), and that also include (ḡ,h x ), i.e., Then, one can directly verify that the solution of the previous optimization problem is given by , and the set U µ x,ḡ,hx (η) is defined as Observe that the set U x,ḡ,hx (η) is nonempty, since it includes (ḡ,h x ) and (g , h x ). Therefore we have η(x, g, h)kl(µ(x, g, h), m(x,ḡ,h x ; η, µ)).
Therefore, using the previous result, together with the result from Theorem 1, we can derive that for allḡ = g the following inequality holds from which we get the statement of Proposition 2.

B Sample complexity of OSRL-SC
This section is devoted to the analysis of OSRL-SC. We first provide useful notations and then discuss the implementation of the algorithm in detail. In particular, we provide a full description of our stopping rule (not described in the main text due to space constraints). Finally, we analyse the sample complexity of OSRL-SC.

B.1 Preliminaries
We introduce some useful notation used in the remaining of the paper. Recall the following quantity from Theorem 1: where g = g (µ). If one lets q ∈ Σ be defined as q(x, g, h) = η(x, g, h)/ x,g,h η(x, g, h) then one obtains the definition of ρ(q, µ): :=ρ(q,µ) .

B.2 Algorithmic details of OSRL-SC
An essential part of the algorithm is computing the value of q σ (μ t ) and ρ σ (μ t ). To this aim, we make use of the following lemma, which allows us to compute the aforementioned value by considering the dual problem of ρ(q, µ).
Lemma 2. For any µ ∈ M, the optimization problem: is equivalent to with Σ being the probability simplex over X × G × H and ∆(k), k ∈ N, being the probability simplex of dimension k − 1, with A = G − 1 and B = HX(G − 1).
One can proceed by reformulating the original problem in the following manner: We can further expand the optimization problem and write The previous problem is a linear program with convex non-linear constraints. To notice this, remember from Lemma 1 that f x,ḡ,h can be equivalently written as Thus f x,ḡ,h (q, µ) can be seen as a minimization problem over a set of concave functions, thus for a fixed µ we have that f x,ḡ,h (q, µ) is concave in q, and therefore p x,ḡ − f x,ḡ,h (q, µ) is convex. Now, denote by L the Lagrangian function, and introduce the dual variables λ, ν: We know that at the optimum z = sup q,z,t inf (λ,ν)≥0 L(q, z, t, λ, ν). But, we can easily verify Slater's condition, therefore strong duality holds. Hence z = inf (λ,ν)≥0 sup q,z,p L(q, z, p, λ, ν). But if ḡ =g λ g = 1 then the problem is unbounded, therefore it needs to be ḡ =g λḡ = 1. Similarly, we need λḡ = h ν x,ḡ,h .
Introduce θ x,ḡ,h = ν x,ḡ,h /λḡ. It follows that the dual problem is We can directly solve the dual problem over the probability simplex by making use of standard optimization tools, such as Mirror Ascent. By using the KL-divergence as Bregman divergence this will result in the Entropic Descent algorithm [Beck and Teboulle, 2003].
In order to use Mirror Ascent, we will make use of the fact that for a fixed µ ∈ M then f x,ḡ,h (·, µ) is Lipschitz with respect to the L 1 norm, where the Lipschitz constant L satisfies L ≥ max i,j kl(µ i , µ j ). This allows us to choose an appropriate learning rate η k for the computation of q: Furthermore, we also have the following theoretical guarantee from Theorem 5.1 in [Beck and Teboulle, 2003]: Theorem 4 (Theorem 5.1 in [Beck and Teboulle, 2003]). Let {q t } t≥1 be a sequence of points generated by the mirror ascent algorithm, with a uniform starting point q 1 . Then, for all t ≥ 1 one has Now that ones knows which procedure can be used to compute q , what is left is: (i) how to numerically compute the problem and (ii) also how to compute the gradient of f x,ḡ,h (q, µ) with respect to q. The first problem (i) is taken care of by Lemma 1. To solve (ii), we resort to libraries that implement automatic differentiation, such as PyTorch [Paszke et al., 2017]. One needs to pay attention in doing so: rewrite f x,ḡ,h as follows: One can observe that the series depends on the set U x,ḡ,hx , which in turn depends on the variable q. Apart from this, one needs to pay attention to numerical errors when applying differentiable programming techniques. In the mirror ascent algorithm we clamp the gradient in [−6, 6] to prevent the exploding gradient effect.
Regarding the use of σ in the algorithm, what we can say is that it is mostly needed only for theoretical purposes. The algorithm given an initial condition will always converge to some point, and we found no major differences by not having σ, or having a large value of σ, such as σ ≥ 10 5 . Moreover, for a bandit problem µ that admits a unique maximizer, we can say that it rarely happens thatμ t has more than one maximizer.
Finally, denote by c t = (q, λ, θ) t the point found at the t-th step. We found that to speed-up the algorithm at the t + 1-th round one can use as initial condition x 1,t+1 the point c t = (q, λ, θ) t of the previous iteration. By doing so one may converge to a solution where some components of q, λ or θ may be 0. Furthermore, we do not have anymore the guarantees of Theorem 4. To avoid these two problems, we used as initial condition a convex combination between the previous point and a uniform point: where 1 is a vector of ones, and S = XGH + G − 1 + HX(G − 1) is the dimensionality of a point x. We found out that a large value of α, such as α ≥ 0.5, gave a computational speed-up improvement.
Last, but not least, notice that in the pseudocode of OSRL-SC 1 the computation of Ψ t (g) is undefined if µ t / ∈ M. In that case, we will continue running the loop, untilμ t ∈ M, which happens almost surely (for a proof, please refer to section E).

B.3 Analysis of OSRL-SC B.3.1 Analysis of the stopping rule
In this subsection, we show how to devise a stopping rule that guarantees the δ G -PAC property for any sampling strategy.
Taking inspiration from the classical work of Wald [Wald and Wolfowitz, 1948] and Chernoff [Chernoff, 1959] in sequential hypothesis testing, and the more recent work from Garivier and Kaufmann [Garivier and Kaufmann, 2016, 2019, Kaufmann and Koolen, 2018, we now describe a stopping rule based on hypothesis testing, by expressly making use of a Generalized Likelihood Ratio Test to decide between a set of symmetric simple hypotheses.
We shall now consider the decision problem of identifying the best representation g : we have G different symmetric simple hypotheses of which only one is true: These hypotheses play a symmetric role (in the sense that both type I and type II errors have the same importance), and one wants to decide for either one of the hypotheses. Now, we know that a certain decision g = g for a problem µ ∈ M is appropriate if g (µ) = g. Then, consider subsets Λ(g) = {µ ∈ M : g (µ) = g}: it follows that M = ∪ g∈G Λ(g), and that a decision g = g is appropriate if µ ∈ Λ(g). We are thus facing the equivalent problem of choosing amongst the following set of hypotheses {H j : (µ ∈ Λ(g j ))} G j=1 . In order to decide between a finite set of symmetric hypotheses, it is common to consider a test on two non-overlapping and asymmetric hypotheses: H 0 : (µ ∈ M \ Λ(g)) andH 1 : (µ ∈ Λ(g)).
This allows us to define our Generalized log Likelihood Ratio statistics Ψ t (g) for each arm g up to time t. Introduce Y t = {Y s } s≤t to be the collection of samples observed from all the arms up to time t, then the likelihood ratio at time t for arm g is where (Y t ; λ) represents the likelihood of having observed Y t given a model λ ∈ M. Notice that large values of L t (g) tend to rejectH 0 , and when the arms belong to a canonical regular exponential distribution then we can write the log likelihood ratio as follows Ψ t (g) := ln L t (g), where Define Λ 2 (ḡ,h) = {λ ∈ M : (g , h 1 , . . . , h X )(λ) = (ḡ,h)}. Then, forμ t ∈ M, we can rewrite Ψ t in the following way where the last equation is an application of Lemma 1. To ease the notation letm t (x) := m t (x,ḡ,h x ). Then, we can then define the stopping rule using a threshold function β t (δ G ), with δ G ∈ (0, 1), as follows Notice, furthermore, that for the type of model considered the maximum is achieved forg = g t whenever µ t ∈ M.
To define the thresholds, we will make use of Theorem 14 in [Kaufmann and Koolen, 2018], which is stated after the proof of Theorem 2.

B.3.2 Proofs of Theorem 2
Proof. The proof can be seen as multi-task version of the Proposition 12 in [Garivier and Kaufmann, 2016] with an improved bound from [Kaufmann and Koolen, 2018]. Let g t be the decision rule at time t, and for the sake of notation letm t (x) := m t (x,ḡ,h x ),μ(x) := µ(x,ḡ,h x ). The probability of error can be bounded as follows: Inequality (a) follows by applying a union bound over the set G \ {g }; (b) on the other hand is a consequence of the fact that m t minimizes the self-normalized sum. The last step (c) follows by applying Theorem 5 with x = ln Ä G−1 δ G ä and S = X × G × H.
Theorem 5 (Theorem 14 in [Kaufmann and Koolen, 2018]). Let φ : R + → R + be the function defined by where ζ(s) = n≥1 n −s , p(u) = u − ln(u) for u ≥ 1 and for any z ∈ [1, e] and x ≥ 0: Then, for S a subset of arms, P ∃t ∈ N : for some constants c, d. For Gaussian or Gamma distributions one can use c = 2, d = 4, while c = 3, d = 1 apply for one-dimensional exponential families.

B.3.3 Analysis of the tracking rule
We will first state the following lemma, that guarantees that our estimate q σ (μ t ) converges to the true maximizer q σ (µ) in the limit: Lemma 3. Given a sampling strategy that guarantees P µ (lim t→∞μt = µ) = 1, then, Proof. Given that ρ σ (µ), q σ (µ) are continuous in M, the proof is an application of the Continuous mapping theorem to the sequences {ρ σ (μ t )} and{q σ (μ t )}.
Because of the forced exploration, the event E = {lim t→∞μt → µ} is of probability 1 by the law of large numbers. From Lemma 3 it follows that for ε > 0 there exists t ε such that for t > t ε then q σ (μ t ) − q σ (µ) ∞ ≤ ε. It follows that {q σ (μ t )} t is a converging sequence, which allows us to make use of the following result from Lemma 17 in [Garivier and Kaufmann, 2016].
Lemma 4 (Tracking lemma in [Garivier and Kaufmann, 2016] where N t ∈ N X ×G×H denotes the vector of arm pulls. It should be noted that OSRL-SC tracks the optimal allocation q σ (μ t ) whenever U t = ∅ andμ t ∈ M. Such conditions guarantee that we end up sampling more than the lower bound provided in Lemma 4. A similar condition holds also for the Track and Stop algorithm proposed in [Garivier and Kaufmann, 2016], where they need to have a unique optimal arm to track the optimal allocation, although not explicitely stated in the paper.
Finally, the previous lemma allows us to state the following result Lemma 5. The sampling rule satisfies Proof. Because of the forced exploration, by the law of large numbers we have that P (lim t→∞ µ t = µ) = 1. Let ε > 0: then there exists t 0 such that for all t ≥ t 0 the distance q σ (μ t ) − q σ (µ) ∞ ≤ ε/3(XGH − 1). Hence, using Lemma 4 we can conclude that there exists t ε ≥ t 0 such that for all t ≥ t ε we have Additionally, as mentioned previously in the experiment section (Appendix 4), we will make use of the fact that q σ (μ t ) is a sequence that convergence almost surely: this allows us to periodically compute it, instead of computing it at every round.
Corollary 2. Given a sampling strategy that guaranteesμ t → µ for t → ∞ almost surely, then, for any integer k ∈ N >0 we have that Proof. The proof is trivial, and follows from Lemma 3. Due to the almost sure convergence of ρ σ (μ t ), then, the periodic subsequence {φ(kt)} t is convergent, with φ(t) = ρ σ (μ t ) and k ∈ N >0 .
Using this fact may save computational resources: furthermore, it will not influence the asymptotic analysis of sample complexity. But, despite that, it may affect the non-asymptotic sample complexity analysis.

B.3.4 Almost sure bound on the sample complexity of OSRL-SC
Given the previous analysis of the tracking rule, we are now able to state a theorem that guarantees that {τ G < ∞} happens with probability 1, and also an almost sure upper bound on the asymptotic sample complexity, that follows directly from the proof in [Garivier and Kaufmann, 2016].
First, we state the following bound on the thresholds that OSRL-SC uses: Lemma 6. Consider the threshold function defined in Theorem 2: Then there exist constants C, D > 0 such that Proof. We can easily see that β 1 (t) ≤ 3 x,g,h ln(1 + t) = 3XGH ln(1 + t), from which follows that β 1 (t) ≤ 3XGH ln(Bt), for t ≥ 1/(B − 1), and B > 1. β 2 (δ G ) on the other hand is just a constant term, and the function φ(x) ∼ x at infinity, and can be upper bounded by a linear function in x. Since x = ln((G − 1)/δ G )/K the claim follows.
We are now ready to state the almost sure bound on the sample complexity of OSRL-SC.
Theorem 6. For δ G ∈ (0, 1) let β t (δ G ) be a deterministic sequence of thresholds that is increasing in t and for which there exists constants C, D > 0 such that Let τ G be the sample complexity, with thresholds β t (δ G ), Then, the tracking rule of OSRL-SC ensures, for σ > 0: Proof sketch: the proof relies on the continuity of ρ σ , q σ and the fact thatμ t → µ and N t /t → q σ (µ) almost surely. This allows us to find a lower bound on max g Ψ t (g) that depends on the value of the problem, which in turns upper bounds the stopping time. Finally, a technical lemma is used to derive an upper bound.
Because of (i), for t ≥ t 0 we are allowed to write On the other hand, because of (ii) we have for t ≥ t 0 : where (a) follows from the definition of ρ σ , since we are subtracting a positive value to ρ; (b) follows by continuity of ρ σ and q σ and the fact that t ≥ t 0 . Consequently, we can write By applying the technical Lemma 8 with β = D/δ G , γ = (1 + ε)/ρ σ (q σ (µ), µ) we get Thus τ G is finite on E for every δ G ∈ (0, 1), and lim sup Where C σ (µ) = ρ σ (q σ (µ), µ) −1 . By letting ε → 0 we obtain the desired result.

B.3.5 Bound on the expected sample complexity of OSRL-SC
We prove our upper bound of the expected sample complexity of OSRL-SC. We first study the expected sample complexity of the first phase of OSRL-SC, and prove Theorem 7. Then, to upper bound the expected duration of the second phase of the algorithm, we will use the results from [Garivier and Kaufmann, 2016]. Finally, we show that we can easily put the upper bounds for the first and second phases together, and deal with the fact that the first phase could output the wrong representation.
Upper bound of the duration of the first phase.
Theorem 7. For δ G ∈ (0, 1) let β t (δ G ) be a deterministic sequence of thresholds that is increasing in t and for which there exists constants C, D > 0 such that Let τ G be the sample complexity with thresholds β t (δ G ) Then, the tracking rule of OSRL-SC with σ > 0 ensures lim sup Sketch of the proof: The proof follows along the same lines of the results proved in [Garivier and Kaufmann, 2016] and [Kaufmann and Koolen, 2018]. The novelty is to adapt the definition of "good tail event" (see later in the proof) to our model, and derive the upper bound given a fixed value of σ > 0.
Then, for every T ≥ max(T 0 (δ G ), T ε ) we have What is left is to derive an upper bound of T ε 0 (δ G ). Let ζ > 0 and introduce the following constant Remember that we use a deterministic sequence of thresholds that is increasing in t and for which there exist constants C, D > 0 such that ∀t ≥ C, ∀δ G ∈ (0, 1), β t (δ G ) ≤ ln Å Dt δ G ã .

C Experimental setup
Hardware and software setup. All experiments were executed on a stationary desktop computer, featuring an Intel Xeon Silver 4110 CPU, 48GB of RAM. Ubuntu 18.04 was installed on the computer. Ubuntu is an open-source Operating System using the Linux kernel and based on Debian. For more information, please check https://ubuntu.com/.