Fast Treatment Personalization with Latent Bandits in Fixed-Confidence Pure Exploration
Authors: Newton Mwai Kinyanjui
Emil Carlsson
Fredrik D. Johansson
Venue: Accepted by TMLR ()

Abstract: Personalizing treatments for patients  often involves a period of trial-and-error search until an optimal choice is found. To minimize suffering and other costs, it is critical to make this process as short as possible. When treatments have primarily short-term effects, search can be performed with multi-armed bandits (MAB), but these typically require long exploration periods to guarantee optimality. In this work, we design MAB algorithms which provably identify optimal treatments quickly by leveraging prior knowledge of the types of decision processes (patients) we can encounter, in the form of a latent variable model.  We present two algorithms, the Latent LP-based Track and Stop (LLPT) explorer and the Divergence Explorer for this setting: fixed-confidence pure-exploration latent bandits. We give a lower bound on the stopping time of any algorithm which is correct at a given certainty level, and prove that the expected stopping time of the LLPT Explorer matches the lower bound in the high-certainty limit. Finally, we present results from an experimental study based on realistic simulation data for Alzheimer's disease, demonstrating that our formulation and algorithms lead to a significantly reduced stopping time. 

Published in Transactions on Machine Learning Research (04/2023)
Fast Treatment Personalization with Latent Bandits in Fixed-Confidence Pure Exploration
Newton Mwai
mwai@Ochalmess.se
Deportment of Computer Science and Engineering
Chalmers Universily of Technology
Emil Carlsson
caremal@@chbmmers.se
Department of Computer Science and Engineering
Chalmers University of Technology
Fredrik D. Johansson
fradkk:jphannoon@hhhhnees..
Department of Computer Science and Engineering Chalmers University of Technology
Reviewed on OpenReview: https: openreview. net/ net/forunt forum? id- NNRIGE8buF
Abstract
Personalizing treatments for patients often involves a period of trial-and-error search until an optimal choice is found. To minimize suffering and other costs, it is critical to make this process as short as possible. When treatments have primarily short-term effects, search can be performed with multi-armed bandits (MAB), but these typically require long exploration periods guarantee optimality. In this work, we design MAB algorithuns which provably identify optimal treatments quickly by leveraging prior knowledge of the types of decision processes (patients)
encounter, in the form of latent variable model. We present two algorithms, the Latent LP-based Track and Stop (LLPT) explorer and the Divergence Explorer for this setting: fixed-confidemt pure-exploration latent bandits. We give lower bound on the stopping time of any algorithm which is correct at given certainty level, and prove that the expected stopping time of the LLPT Explorer matches the lower bound in the high-certainty limit. Finally, we present results from an experimental study based on realistic simulation data for Alzheimer's 's disease, demonstrating that our formulation and algorithmss lead to a significantly reduced stopping time.
1 Introduction
There is growing interest in using machine learning for personalizing medical treatments to account for heterogeneity in patients' responses. Finding suitable choice for an individual often involves phase of trial and error before settling on a therapy that works for them, especially in the treatment of chronic diseases (Fraenkel et al., 2021; Stern, 2009). In rheumatoid arthritis, for example, when first and second-line treatment fails, there large variability the choice of next therapy, and several drugs may be considered oqually good choices a priori (Zink al., 2001). Further, switching therapics has associated costs: every time a therapy changed, the patient. has to get used to the new therapy and its potential side elfects. It is therefore desirable minimize such switches, even if changes to other equally good treatments after treatment has been identified in the search phase. Learning algorithms could improve the efficiency of this search, reducing the number of avoidable trials (Chakraborty and Moodie, 2013)
A classical framework for exploring alternative treatments is Multi-armed Bandits (MAB) (Gittens and Dempster, 1979; Lai and Robbins, 1985), originally motivated by reducing suffering drug testing (Thompson, 1933). However, MABs tend to be sample-hungry to the point of being unsuitable for finding personalized treatments in real-world clinical settings. Because long search phase can prolong unnecessary suffering, it

---

Published in Transactions on Machine Learning Research (04/2023)
Population
Estimated model
State
Context
Treatments
a1
a
at
p(5)
X
+
P(XIS)
p(R X,S,
71
T2.
Tt
(S,
v
New patient
Exploration period
Historical data
Figure 1: Illustration of the pure-exploration latent bandit problem and the example of treatment personal- ization. A population of patients have been observed in historical data to learn the distribution of latent states P(S), P(X SS and the conditional reward the distribution P( S, A). A new patient, represented X,S,
by the instance v (ar, s) is treated with actions at, observing rewards Tt until the stopping time т.
must be avoided and minimized whenever possible. Existing methods for the fixed-confidence pure exploration setting MABs, which aim to minimize the time it takes find an optimal treatment at given certainty level (Even-Dar Dar et al., 2006; Garivier and Kaufmann, 2016; Russo, 2016; Shang al., 2020) also yield long exploration phases.
One reason for the long exploration of bandit algorithms is that each instance each patient, in example is treated as independent, learning parameters from scratch each time. This allows for complete personalization, often incorporating contextual or side information (Ti et al., 2010; Chu et al., 2011), but disregards any similarities between instances. For many conditions, differences responses (rewards) to treatment between patients are believed be explained by a small number of discase subtypes (Devi and Scheltens, 2018; Borish and Culp, 2008). Thus, for patient with a known subtype, an optimal treatment could be identified from the treatment responses of previous patients with the same subtype.
The subtype of a palient may be viewed as latent state, as it is unobserved al the start of treatment, but manifests in a patient's responses to different therapies. With access to data on the treatment of previous patients, it is possible to fit a model of the distribution latent states and their association with actions and rewards, for instance with variational inference methods (Kingma and Welling, 2013; Jang et al., 2016). Given such a model, for a new patient. (bandit instance), our task becomes identify which latent state they belong to, see Figure 1. Latent Bandits and recent iterations formalize this idea but are limited to regret minimization, aiming to minimize the regret compared to optimal actions over possibly infinite period (Maillard and Mannor, 2014; Zhou and Brunskill, 2016; Hong et al., 2020a;b; Kwon et al., 2021). This differs from our goal of finding the optimal treatments within a desirably short exploration period, while also ensuring that the algorithm commits to good treatment after exploration, without treatment switches.
In this work, we derive fixed-confiddncee pure-exploration bandit algorithms which aim minimize the number of trials required find an individual-optimal treatment by incorporating existing knowledge of latent structure.
We propose formulation of the personalized treatment search problem with known latent structure in the fixed-confiencc pure-cxploration setting (Section 2). 2) We prove lower bound for the search time of any algorithu in our latent bandit setting and prove a matching upper bound for the Latent LP-based Track and Stop (LLPT) Explorer (Section 3, 5). 3) We present two algorithms, the LLPT Explorer and the Divergence Explorer (Soction 4). 4) We perform an extensive empirical evaluation on simulator of Alzheimer's discase and illustrate that our formulation and algorithms lead to significantly reduced stopping time compared to classical pure-exploration algorithms in the MAB framework (Section 6).
Main contributions. 1)
2

---

Published in Transactions on Machine Learning Research (04/2023)
2 Problem formulation
We think of treatment personalization strategy as an agent which interacts with a patient over E N rounds, aiming to try as few treatments as possible before the best possible treatment has been identified with a confidence level of at least 1 1-6 8, for a pre-specified 5 > O. At the start of the sequence, the agent observes patient's ' context covariates (e.g,, lab measurements) as draw of a random variable X E Rd Then, at each step t = 1,2,... ..., the agent takes an action A, E A = {1....,K] (trial treatment) and gets a reward Rs E IR (treatment outcomes). When an optimal treatment has been found, exploration stops and the agent rocommends this treatment. The setting is illustrated in Figure 1. In the multi-armed bandit litcrature this setting called fixe-ccofiidence pure exploration (Garivier and Kaulmann, 2016; Shang el al., 2020).
A fixed-confidence pure-exploration strategy é comprises sampling rule for exploring actions A at each step 1, a stopping rule to decide the time at which the exploration is over, and recommendation rule which returns the best action 0 at the stopping time т. Our goal is to design a strategy ø to minimize the expected stopping time E[r]. In our healthcare example, this serves to minimize the search for optimal treatments, and thus minimize patient suffering in the treatment search phase while also ensuring that the algorithm commits to a good treatment after exploration, without treatment switches.
Even for state-of-thh-art pure exploration algorithms, the necessary exploration tends to be long in realistic settings (see Figure 2). To overcome this, will make structural assumptions about contexts, actions and rewards regarding patient similarity. In our healthcare example, it is plausible that a new patient (bandit problem instance) shares significant. similarity with historical palients (logged bandit data), and that the optimal treatment for them is the same as for similar patients. However, in many domains, the context X is not sufficient to identify optimal treatment since it does not account for all individual variation (Hákansson et al., 2020). To account for remaining individual variation between patients with the same X, we will assume that there is finite number of latent states S E S {1, M), e-g.; patient types, which cannot be directly observed. Thus, the optimal treatment is determined by the context X and the latent state S two instances (e.g., two paticnts) are similar if they have the same context and latent state (c.g., discase subtype).
Identifying the true latent state 5 is sufficient but not strictly necessary to solve our problem. For successful treatment, we are only interested to identify the optimal treatment at exploration stop, â,. Therefore, it is not necessary to estimate the correct latent state, but the set latent states that have the same optimal arm. Having context X is desirable as it helps reduce the number of trials if is informative of the underlying latent state S, with unexplained variation further discoverable by trying different treatments.
A latent variable model (LVM) of the distribution of latent stales S, contexts X, actions A and rewards R can be estimated from historical data and used to speed up exploration for a new subject. Maillard and Mannor (2014) and Hong et al. (2020a) made use of LVMs for "Latent Bandits" " in the related sctting of regret minimization. As these algorithms do not come with stopping and/or recommendation rules, they are not applicable to the fixed-confidence setting where the goal is to terminate search as quickly as possible.
In the MAB formalism, our problem can be defined fixed-confidence purc-exploration latent bandits with single initial context. In doing so, assume that the latent subtype and the distributions of rewards is unaffected by time and previous actions. This is plausible for conditions treated with symptomatic therapies, such as for chronic degenerative disease like AD or Rheumatoid Arthritis (RA), where treatments typically target the symptoms and not the underlying disease pathology (Fish et al., 2019). Under these assumptions, the imal choice of treatment remains fixed through exploration.
2.1 Fixed-confidence pupe-exploration latent bandits
Given a state s, a context and an action a, let
Pa := E A = 0, X x S= s]
By convention, we capital letters for random variables and lowercase for observed variables
3

---

Published in Transactions on Machine Learning Research (04/2023)
denote the expected reward for that action, and let
Ma.s max Ma,r, and 0z = arg max Ma.z,a
denole, respectively, the optimal expected reward and arm in latent state and observed context x. We assume that the maximizer dx is single action for each state-context pair (æ,, s), but arguments can be generalized to the case with multiple optimal actions. Further, let H2 = (X, A1, Ri...., A,. Re) denote the to time i, letting Ho (X). The utility of the context X is in
history of context, actions and rewards, computation of the likelihood H(s)H, H,) and this agnostic either finite or infinite context assuming that a good model of the likelihood is known.
Our goal is to design a search strategy o to mènimize the expected number of trials I required identify an optimal action, with confidence at least 6, for new subjects with contert X and unknown latent state S.
minimize
Eq..SH [T]
(1)
subject to
P(pa, w. P2, X = x, = =
S Vx,
We say that a search strategy 8-PAC if the error probability is bounded by Here, this is captured by our constraint, Vx,s: : P(p pr,s X x, s) S 6, as long as the probability model correct.
In equation 1, minimize the expected stopping time (e.g., over population of patients) while satisfying instance-dependent constraints (per patient). We justify this formalization by noting that, in our rumning example, any single patient will have single random stopping time, which we can estimate and analyze only in expectation. However, it is desirable and possible to guarantee, per patient, that our confidence exceeds 1 - whenever we stop.
We assume that a model Ma (Po(S), Po S). A. of the marginal state probability p(S) and the likelihood of observed variables under S, including the set of reward means Pa, is avoilable when search begins, akin to Hong et al. (2020a). This means that once is known, so is the optimal arm in s, and no further exploration is necessary. Such model can be learned from logged bandit instances, for example, using a variational autoencoder (Kingma and Welling, 2013), but this is outside the scope of this work.
For simplicity, we will assume that all reward distributions are stationary in time and Gaussian with equal variance o², that is, given A = a, X = x., S = 8, for all
R2 ~ N(po......
The algorithms presented in Section 4 are applicable in the non non-Gaussian case as well, assuming that the reward distribution is known through Ma, but our analysis in Section 5 is limited to Gaussian rewards for now. Our analysis makes heavy use of the Kullback-Leibler (KL) divergence, and we will adopt the notation KL(Ha, ) = (L(p(R)) x, s) p(R a, a,z,s')) for the KL-divergence between the two Gaussian rewards for arm a under states 8, s' with equal variance o² and means as indicated.
3 Lower bound on stopping time
To serve as benchmark for our algorithms, we derive a lower bound on the worst-case solution to objec- tive equation 1 for any algorithm which satisfies its constraints.
The seminal work of Kaufmann et al. (2016) presented general inequality from which one can derive lower bounds for 8-PAC algorithms in the best-arm identification framework. In lemma 1, we present variant of their key result, adapted to our latent bandit setting. For brevity, we let
pl s, s') = log[p(æ s)/p(x s')]
denote the log- likeliioddood ratio of the observed context x under latent states and 8', and use the shorthand
KT.
R.O.z KT, (a,,,)) g.8
4

---

Published in Transactions on Machine Learning Research (04/2023)
for the KL-divergence between rewards under states 8, s, Our bounds and algorithms use state reference point for the set of alternative states s' with different optimal arms,
(s) = {s' 2z + az
We can now derive the following result.
Lemma 1. Given problem instance with latent state and observed context æ, any 8-PAC algorithm ø must satisfy for any alternative state s' E (s),
EolAN [,,8 s|KL
R.a.a 5,5"
p( kl -5),
(2)
where Na is the namber of plays of arm drawm under & and kl - 8) is the Kl-ddinergence between two Bernoulli random variables with porameters and 1 1- 5.
Proof summary. The proof follows the argument of the original Lemma in Kaufmann et al. (2016). We start from the KL-diveerencc between the distribution of histories H, under and s' and expand this using the chain-rule of the KLL-ivvegeence We then apply the cetommation-proesssiii by kl - 8). The difference from Kaufmann et al. (2016) is that we get an additive term which depends the contert distribution under different latent models. For a full proof, see Appendix A.1.
inequality to lower bound this
From lemma 1, we can derive a lower bound on the expected stopping time. Here, we assume that the optimal arm unique for each state-context pair (s, z), that is, (s) S {s}. This assumption is not necessary
rum our proposed algorithms.
Proposition 1. For any 8-PAC learner & with 8 (0, 1/2) and any latent state and contert x, the expected stopping time satisfies
1
EglT s, xz(s
kl -8)
where 1/C Eo (s) with (s) the minimizers of the following linear program,
minimize
Σ 7/eaa
(3)
20.220
p(1;,,88) kl -) 2 1, AA E Al(
subject to Σ Tia,a KKll akll s.ss
R.a.z
Proof summary. By lemma 1, we have a constraint on the sum of the expected number of times each arm is played by any 8-PAC algorithm Φ. By dividing each side of equation 2 by kl(ōl[1 - -) and minimizing the the stopping time under the resulting constraint, we obtain the linear program (LP) in equation 3. For a new bandit instonce, x observed before search begins. Thus, given model Mo, the only unknowns in equation 3 construct fonite set of linear constraints and
are Ya.a. As we have finite set of latent states s ,, solve for the minimall slopping time. A full proof is given in appendir A.1.
Remark 1. As sanily check, we verify that the contexiual information makes the pure-exploration problem fundamentally easier. Indeed, when an observation x clearly separates the true latent state from é P increases, the constraint in equation is satisfied by a larger set of parameters "s,,, and the lower bound attains a smaller value. However, as we require increasing certainly and + the influence from contextual information X C (s,2) vanishes. This is expected since don't collect more information through x requirement on certainty increases it remains constant.
As a consequence of proposition 1, we can obtain a bound for the population (marginal) search time. If we assume that G; = =Exss Ee Vi.a exists, with 2e.a the minimizers as in proposition 1, we have
1 C
Esxx,s|T| 2 k1
The lower bound indicates that the optimal worst-case solution to equation limited by the hardest-to- separate states s, s'. We make of this insight next to develop algorithms.
5

---

Published in Transactions on Machine Learning Research (04/2023)
Algorithm LLPT Explorer and Divergence Explorer
Input K, Output 7,
1: Observe h1 (x)
LLPT Explorer then Compute (s) for all 8,8
See equation 4, equation 3
end if
6: while Ze <1- and t T do
7:
LLPT Explorer then
8:
St = arg max,es p(s[ht)
4++1 = arg max L. w (8+) Na, (t)
10:
else if Divergence Explorer then
11:
8t p(s)
12:
fe(a) = E.' (o(s"))))TTT(/.... (
13:
a++1 = arg max Se(a)
14:
end if
15:
Choose ac+1: and Observe r4+1
16:
Update h4 h4-1 (1t-1,
17:
Update (1) N No +1
18:
19:
Update 3, arg m po(s h4)
20:
Update å = arg aaX, Ha,
21:
Update Z, = E. (o (11) [ā; = [%
22:
end while
23:
24:
Return à
4
Algorithms
We present two best-arm identification strategics, each comprising sampling rule for selecting arms A, stopping rule for determining and recommendation rule for selecting å,. Both algorithms, defined in Algorithm 1, are given access to an already estimated latent variable model Mg including all reward means Ha,z, 4 8 E S, a E A given a context and differ only in their sampling rules; the stopping and recommendation rules are equivalent. Either algorithm starts by observing the random context X, and proceeds from there.
4.1 Sampling rule 1: Latent LP-based Track and Stop (LLPT) explorer
Our first sampling rule is based on the Track-and-Stop method (Garivier and Kaufmann, 2016), where arm allocations are determined by tracking proportions w", obtained by solving the lower bound optimization problem in equation 3. Since we have finite sets of states and actions, and x is observed at the start of the search, we can compute (s) for all s E S,a E A directly. Then, we define playing proportions ws(s), for each possible state S, as
(s)(s) (s) ( ((.a (s))
(4)
At each time step 4, the algorithm picks latent state 8 arg max, p(s)h)) from the (known) posterior given the current history h4, and plays the arm which most closely tracks w (st). Let N.(1) be the number of times arm has been played up until and including t. Then, the LLPT Explorer sampling rule is defined by
A++1 arg max t- (s Natt).
aE[s]

---

Published in Transactions on Machine Learning Research (04/2023)
The LLPT Explorer aims to play the minimum total number of trials using arms which distinguish latent states the most, as given by the KT, term in the constraint of equation 3. It aims only to distinguish latent. states with different optimal arms, as the goal is to identify the best action, not the state.
4.2 Sampling rule 2: Divergence explorer
The LLPT Explorer plays according to the optimal proportions for the worst-case alternative state given the current estimate. This is because the constraint in equation 3 will be hardest to satisfy (require largest Ta,a) for states which are the most similar to 8. drawback of this idea is that it ignores the likelihood of said alternative state under the posterior. If there strong evidence that is unlikely collecting more evidence rule it out may be suboptimal. In the extreme case, state s' with posterior probability p(S= s' h2) R O may still (unnecessarily) inform the sampling rule for the LLPT Explorer.
be the true state,
As an alternative, we define the Divergence Explorer sampling rule. This algorithm aims to play arms according to how much information is gained by playing an arm in expectation given the current postcrior probability of states in (s). At each time t, a latent state St ~ P((s(he) is sampled as reference. Then, the sampling rule uses the expected divergence between 8t and alternative states 8/
a) = Σ
P(st)le)KT.((aa (...,s))
The arm A - arg max Je(Se, a) is played next.
Because KT,(Ha,,,,z 2,, measures the information distance between the reward distribution of arm a under the two latent models & and 8/ f(t, a) does a one-to-many test assuming 8t is the true model and st is another latent model with probability Phs))
4.3 Recommendation rule
Both algorithms recommend the best arm in the state most believed to be correct in a given instance, so the recommendation rule is à, = arg max A Ha,,, where 8, the most probable state under the posterior as defined in Algorithm
4.4 Stopping rule
It is natural to stop search at & when we are confident enough that the recommended arm â is optimal under the posterior over latent states. Since we assume to have access to the full posterior over S, we use the simple stopping rule
= min(t: Z, 2 8} 1 where Z, ===>pm P(sh))) = d",o]
(5)
and the threshold is the desired confidence level. Whenever this rule satisfied, so is Chernoff's stopping rule based on a threshold ( on the log- likelihoooodd ratio between states, as used by Garivier and Kaufmann (2016). See the proof of proposition 2 in appendix A.2 for derivation.
In many applications, it us sufficient to identify a action which is e-optimal with respect to the best possible action in the true latent state. We can accommodate this in our algorithm by redefining the set of alternative states s' to include only those for which the optimal arm in s is more than e worse than the optimal arm in s',
(s) {s' a pa. -)
This change involves only a minor modification to the stopping criterion in equation 5 and could also be used in the Divergence explorer sampling rule.
5 Upper bound on the expected stopping time of LLPT explorer
Next, we show that the lower bound derived in Section is matched by an upper bound on the stopping time for the LLPT Explorer algorithm in the high-confidence limit, 5 + 0. Similar to the lower bound, we make
7

---

Published in Transactions on Machine Learning Research (04/2023)
the simplifying assumption that each latent state has unique optimal arm, shared with no other states, (s) = S {s}. As a consequence, finding the optimal arm equates to finding the true underlying state. We have the following result.
Proposition 2. Let t be the stopping time of LLPT Explorer o, as defined in Algorithm 1. With s the true stole and (8, the oplimum in equation 3 with the p-term removed, there is consland a > 0 such that
BØlT s 6-20 1og(1/5)
lim
S (s,
(6)
Proof summary. The proof combines and expands arguments from Gorivier and Kaufmann (2016) and Chernoff (1959) to show that after sufficently many samples, a) the true latent state is identified, b) the tracked proporlions are near oplåmal for the identified stale, e) the probabilily that the slopping criterion is not satisfied decays exponentially quickly. As a result, the expected stopping time can be bounded using concentration arguments. For a proof, see appendiz A.2.
As stated, proposition 2 applies to the LLPT Explorer, as defined in Algorithm 1, in which the MAP state 8, is used for tracking. We have also implemented a slight variation of LLPT with a sampled state & ~ p(s)h)) and found that the latter worked slightly better empirically. We report only results for the version in Algorithm
Similarly to the lower bound, we obtain an upper bound on the population search time by taking the expectation of equation 6 with respect to S and X.
Remark 2. Comparing the result in equation 6 to bounds for pure-ccploration without latent variable models; e.g., Russo (2016); Garivier and Kaufmann (2016), superficially, they appear very similar. However, the critical quantity in the classical setting is the smallest separation of reward means for alternative, free vectors of arm paramelers. Here, the equivalent quantily is the set of paramelers of the discrele lalent states, which is generally much smaller than the set of free parameters, leading to a tighter bound.
More precisely, the sample complezity term (*(s shrinks when we have knowledge of the latent state structure because the set of plausible alternative parometers (s) is smaller compored to the case with no structure in, for ezample, Garivier and Kaufmann (2016). In our case, Ast, comprises finâte set of paramelers, whereas the case where paramelers are estimated online without latent structure corresponds to an infinite set of alternative parameters. As result, the worst-case (supremum) over alternative parameter sets shrinks, as do the lower and upper bounds on the stopping time.
6 Experimental study
We evaluate our proposed algorithms in a series of experiments, comparing them to baseline algorithms for fixed-confieece pure exploration.
6.1 Baseline algorithms
Previous work incorporating latent states in pure exploration was not available at the time of writing, so to get comparable baselines, we adapted the Top-Two Thompson Sampling (TTTS) rule (Russo, 2016) to compare to our algorithms.
Top-Two Thompson Sampling (TTTS) TTTS operates with the goal of estimating parameters II, (c.g., mean vectors of arms with Gaussian distribution) that yicld the best arm for given confidence level 1- 6. It proceeds as follows; al each lime step either; (i) with probability p, sample a parameter vector 0: ~ IT, and play the arm a. = arg max 0g or (ii) with probability 1 1- resample 0,' ~ πT until it gets and subsequently plays arm a, (2) of a (1), We implemented the T3C (Shang et al., 2020) variant of TTTS which finds a, (2) + ag (I) faster. TTTS does not make use of a latent variable model.
TTTS-Latent Explorer This is an adaptation of TTTS to our sctting where, instead of estimating arm parameters, the goal purely to identify the latent state. does not account for the case where there is shared optimal arm over different states which is accounted for in the LLPT and Divergence Explorer.
8

---

Published in Transactions on Machine Learning Research (04/2023)
At each time step I, the sampling rule samples latent state 81 (s(s)h)) and either (i) with a Bernoulli parameter p evaluates the latent state, 8t = = (
or (ii) with a Bernoulli parameter 1 1-p resamples Pr(s(he)
(1) It then plays the arm A, arg max Ma, until it gets a latent state 8, (2) + st
Greedy Explorer This is a naive sampling rule which plays the reward-optimal arm in a state sampled from the current postcrior, akin to TTTS-Latent but without the ro-sampling step. At cach time t, it picks 9t = arg max, p(s(h, and then plays the locally reward-maximin naïve in the sense that it only considers the rewards from state, but this is not always informative for distinguishing alternative states. It also corresponds to standard Thompson Sampling (Thompson, 1933) which has been shown to perform poorly for pure exploration tasks, hence the motivation for TTTS.
arm A, arg aag Ma, It is
6.2 Experimental environment
As treatment personalization task, we use the Alzheimer's Disease Causal estimation Benchmark (ADCB) enviromment (Kinyanjui and Johansson, 2022). In this environment, simulated subjects go through cognitive decline, eventually progressing into Alzheimer's disease. Outcomes Yú represent their cognitive abilities and treatments A, are symptomatic, affecting only immediate outcomes. Both treatment responses and an initial 33-dimensional observed context X I& are affected by a latent state S, representing the disease subtype.
In the ADCB environment, the number of actions is K = 8 and the number of latent states, S S= 6. The outcome Y, at time & is generated as (1(A, X) == S) + A(A, S) + &. where & ~~ 5~N(0,00) and + is an non-linear function fit to real data to model the cognitive function of subjects when not treated. For the environment we are using, + is a Random Forest. Regressor fit to observed outcomes of untreated patients. A is a function that is defined to moderate the heterogeneity of simulated treatment effects over the latent dimensions. Here, A == v + 1 (1n33") where v E RRKK y is heteregoneity scaling parameter, and B E RR is a factor matrix whose rows sum to 0.
K is the average trcatment effect of the trcatments,
We define two alternative reward settings (see below), both with Gaussian rewards, based on the ADCB outcomes of treatments, Y. We give algorithms which make use of latent variables perfect knowledge of the true latent variable model, as defined by the simulator. Hence, for each context x E IRd, latent state E [S] and action a E [K], the corresponding posterior p(s h4) and reward means, Pa,z, are known.
Reward setting 1: Non-contextual rewards Here, for each latent state we define the reward R -(Y(A,,,)) - (0, S)). From the definition of the outcome Y above, this removes the effect of context from the reward, by cancelling ((),s), and takes us closer to a typical best arm identification setting with additional latent state structure, where the structure is given by A. In appendix B.1, Figure 5(a) shows the structure of the mean rewards Pa,a, under the different latent states 8 E 8, a E for this setting.
Reward setting 2: Contextual rewards Here, we define the reward R YY,, thus preserving the effect of context in the reward. As seen from appendix B.1, Figure 5(b), which is an example of the mean rewards structure Ma,r,s s E S, a E K for some given context I, the reward structure stays the same as in the previous setting, but the scale shifted depending on the context. The similarity a property of the environment. The results presented in the results section below are for this setting, and those of setting above are appended in the supplementary materials.
Repeated experiments Each experiment proceeds as follows; A new patient is sampled from the envi- roument (sampled patients have potentially different latent states and contexts). The algorithms do not observe the latent state and they proceed as described in Section 4 and Section 6.1. For run, all algorithms are provided with the same context. All results are presented for 100 different patients and averages are computed for the different quantities compared. Errorbars represent the standard deviation across patients.
Evaluation metrics We compare empirical estimates of the expected stopping time E[r], convergence of the posterior probability p(&, | he)) with t, and the average correctness level, [[ a a"]] of the different algorithms for i) different levels of confidence d E (0, 1/2) under fixed noise level 0 and ii) different levels
9

---

Published in Transactions on Machine Learning Research (04/2023)
of noise for fixed 8. Results for correctness presented in Figure 6 in the Appendix, and correspond closely with the parameter 8.
6.3 Results
In Figure 2, we see an example of the drastic effect that incorporating latent structure can have on the stopping time of pure-exploration algorithms.. All latent-variable methods outperform the non-latent baseline TTTS by a substantial margin.
LLPT (With latent structure)
200
TTTS (No latent structure))
150
100
0.9
0.95
1-6
Using latent state structural information significantly reduces the expected number of trials E[r]] required to identify optimal treatment with confidence at least 1 in simulator of Alzheimer's discase progression.
Figure 2:
Moreover, Figure 3a, we see that, even for the worst-case instances, the LLPT algorithm is faster than the average for standard TTTS observed in Figure 2. This supports our hypothesis that exploiting latent structure between instances (paticnts), which could be estimated from historical data, contexts, is useful design sample-efficient pure-exploration algorithms.
In the graph of latent state posterior convergence, Figure 3b, Explorer converge quicker in their belief of the inferred latent state. We also observe less variance across bandit instances (shaded area) compared to the Greedy and TTTS-Latent baselines. The implication for this is that these algorithms stop exploration earlier thus attaining our goals outlined Section 2.
that LLPT Explorer and Divergence
In Figure 3c, we study the average stopping time, B[r] for all algorithms with access to the same latent variable model, under changing certainty level 6. LLPT Explorer and Divergence Explorer consistently more efficient than baselines, demonstrating benefit of the insights derived from the lower bound in proposition 1. The difference is especially pronounced in the high-certainly regime, d ~ 0, which is the regime that would be ideal for safety-critical healthcare applications. Interestingly, we find that the Divergence Explorer performs consistently better than the LLPT Explorer and its average stopping time approaches the lower bound as + O. We believe this due selecting actions based on comparison with alternative states on average under the current posterior, rather than the worst-case alternative state some latent states are ruled out by the posterior and no longer affect the action selection of the divergence explorer.
Studying our algorithms with respect to noise in the rewards, Figure 3d, shows that our proposed methods are also more robust to noise compared to the baseline algorithms. At a = 10, which is comparable to the marginal standard deviation of rewards due to X and S, we see that our algorithms perform better. We also observe that they are also more robust to over- and under-cstimation of the noise level in the rewards as shown by [E]]] at other noise levels.
7 Related work
The problem of finding optimal decisions under uncertainty has long history (Thompson, 1933; Chernoff, 1959; Gittens and Dempster, 1979; Jennison el al., 1982; Lai and Robbins, 1985; Glynn and Juneja, 2004) and has recently been studied as pure exploration problem in the multi-armed bandit framework under
The small discrepancy the where due the exclusion of the p term in the computed bound.
10

---

Published in Transactions on Machine Learning Research (04/2023)
Posterton
0.6
0.25
T. Lower Bound
0.20
LLPT,
El LLPT, worst-case
0.15
LLPT, all states
PT Explorer
0.10
case
0.2
Divergence Explorer
TTTS Latent Explorer
0.05
Greedy Explorer Eld] LB 6-0.01
0.00
10
20 30
40
50
0.0
10
20
50
60
((t)
(a) Density of stopping times under LLPT(ouus) showing worst-case latent state (6 0.01, Number of patients, N 10, 000). The variance of the stopping time under all the latent states is reasonably low. The higher stopping times be attributed the worst-case latent states, though they still reasonably low.
(b) Comparison of posterior convergence of the different al- gorithms [8 0.01, Number of patients, N 100). The posteriors for our algorithms, LLPT Explorer and Divergence Explorer, concentrate quickly.
100
LLPT Explorer
Explorer
35
Divergence Explorer
Divergence Explorer
TTTS-Latent Explorer
TITS-Latent Explorer Greedy Explorer
30
Greedy Explorer
25
E
20
5
0
0.7
0.8
0.9
0.95
0.99
0-
3
10
15
1-5
(c) Comparison of stopping time vs confidence (1 - 8) the algorithms. Our algorithms, LLPT Explorer and Divergence Explorer, have stopping times that are consistently lower. The dashed line shows the lower bound from Proposition 1.
(d) Comparison stopping time vs noise for algorithms Our algorithms, LLPT Explorer and Divergence Explorer, are consistently more robust to noisy rewards compared to the basclines. The dashed line shows the lower bound from Proposition 1.²
Figure 3: Selected results from our experimental study.
various assumptionssBveen-Da et al., 2006; Bubeck et al., 2009; Jamieson et al., 2013; Kaufmann et al., 2016; Garivier and Kaufmann, 2016; Jedra and Proutiere, 2020; Wang et al., 2021; Agrawal et al., 2021; Tirinzoni and Degenne, 2022).
The work of Garivier and Kaufmann (2016) the first to introduce an optimal algorithm, Track and Stop, the fixed confidence sctting for classical multi-armed bandits and our LLPT Explorer takes inspiration from their algorithm, adapting it the latent bandit setting. Russo (2016) introduces class of top-two sampling strategies for the pure-exploration problem, which here originally analyzed using a different performance measure but have recently been theoretically analyzed in the fixed-confidence setting by Jourdan et al. (2022). Our work is also related to (Maillard and Mannor, 2014; Zhou and Brunskill, 2016; Hong et al., 2020a;b), who study regret minimization in latent. bandits, contrast. to our work which studies the pure-exploration problem in latent bandits.
baselines. These top-two algorithms were
Kato and Ariu (2021) studied pure exploration in contextual bandits, where a new context observed each time point, and found that contextual information improves the speed at which the average treatment.
11

---

Published in Transactions on Machine Learning Research (04/2023)
effects (Imbens and Rubin, 2015) of actions across contexts can be estimated. Our problem is related to this setting but differs in that we see only single context per bandit instance, and are interested in the effects of actions for this specific x, not on average. Hákansson et al. (2020) studied fast search for near-optimal treatments, based on a model Icarned from historical trajectorics, but did not consider online learning. In their setting, an optimal search stralegy can be found by solving a dynamic programming problem in an estimated discrete state space. This not feasible here due the high dimensionality of our history, H.
8 Discussion & conclusion
In this work we have studied the problem of finding the optimal in latent bandits using few trials possible. We have empirically and theoretically shown that our proposed algorithms are able to leverage the latent structure in a near-optimal way to substantially reduce the expected stopping time compared to available baselines. Our empirical evaluation in a simulator of Alzheimer's 's disease derived from real-world data, demonstrated that our algorithms are able to find the optimal treatment in just a few trials.
Our analysis is limited to the case in which the latent variable model given and exact. When forced estimate the model from historical data, sensitivity to misspecification misestimation becomes concern. Hong al. (2020a) analyzed latent bandits in regret minimization when the reward model is misspecifted but the resulting bound suffers linear regret scaled by the error, and Hong et al. (2022) provided an improved sub-linear regret bound for this with additional assumptions on the reward structure. In the pure-exploration setting, recovering quickly from misspecification is even more critical since the time scale shorter. We conjecture that an informative guarantee in the misspecified case will similarly require additional assumptions on the reward structure or additional sources data. We believe the setting where learner needs to recover the true model up to some pre-specified precision is an interesting direction for future work. Another useful generalization would be to go beyond the analysis of expected rewards. In high-stakes applications, it is desirable to manage also the risk of worst-case low-probability events, see e.g., Tamkin et al. (2019). This would further increase the suitability of our approach for the medical domain.
References
Shubhada Agrawal, Wouter M Koolen, and Sandeep Juneja. Optimal best-arm identification methods for tail-risk measures. In Advonces in Newral Information Processing Systems, volume 34, pages 2021.
Larry Borish and Jeffrey A Culp. Asthma: a syndrome composed of heterogeneous diseases. Annals of Allergy, Asthma & Immunology, 101(1):1-9, 2008
Sébastien Bubeck, Rémi Munos, and Gilles Stoltz. Pure exploration in multi-armed bandits problems. In Algorithmic Learning Theory: 20th International Conference, ALT 2009, Porto, Portugal, October 3-5, 2009. Proceedings 20, pages 23-37. Springer, 2009.
Bibhas Chakraborty and EE Moodie. Statistical methods for dynamic treatment regimes. Springer, 2013.
Herman Chernoff. Sequential design of experiments. The Annals of Mathematical Statistics, 30(3):755-770, 1959. ISSN 00034851.
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff functions. In Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, pages 208 214. JMLR Workshop and Conference Proceedings, 2011.
Gayatri Devi and Philip Scheltens. Heterogeneity of alzheimer's disease: consequence for drug trials? Alzheimer's 's Research & Therapy, 10(1):1 3, 2018.
M Dragomirescu and C Bergthaller. On the continuity of the optimum of a linear program. Studii Cercetari Mathemalice, 18:1 1200, 1966.
12

---

Published in Transactions on Machine Learning Research (04/2023)
Eyal Even-Dar, Shie Mannor, and Yishay Mansour. Action elimination and stopping conditions for the multi-armed bandit and reinforcement learning problems. Journal of Machine Learning Research, 7(39): 1079- 1105, 2006.
Paul V Fish, David Steadman, Ellioot D Bayle, and Paul Whiting. New approaches for the treatment of alzheimer's disease. Bioorganic & medicinal chemistry letters, 29(2):125 133, 2019.
Liana Fraenkel, Joan M Bathon, Bryant R England, William St. Clair, Thurayya Arayssi, Kristine Carandang, Kevin D Deane, Mark Genovese, Kent Kwas Huston, Gail Kerr, et al. 2021 american college of rheumatology guideline for the treatment of rheumatoid arthritis. Arthritis & Rheumatology, 73(7): 1108-1123, 2021.
Aurélien Garivier and Emilie Kaufmann. Optimal best arm identification with fixed confidence. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir, editors, 29th Annual Conference on Learning Theory, volune 49 of Proceedings of Machine Learning Research, pages 998 1027, Columbia University, New York, New York, USA, 23 26 Jun 2016. PMLR.
J. Gittens and Michael Dempster. Bandit processes and dynamic allocation indices. Journal of the Royal Statistical Sociely. Series B: Methodological, 41:148-177 02 1979.
Peter Glynn and Sandeep Juneja. large deviations perspective on ordinal optimization. In In Proceedings of the 2004 winter simulation conference, volume 1, pages 577-585, 01 2004.
Samuel Håkansson, Viktor Lindblom, Omer Gottesman, and Fredrik D Johansson. Learning to search efficiently for causally near-optimal treatments. Advances in Neural Information Processing Systems, 33: 1333 1344, 2020.
Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, and Craig Boutilier. Latent bandits revisited. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 13423-13433. Curran Associates, Inc., 2020a.
Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, Mohammad Ghavanzadeh, and Craig Boutilier. Non-stationary latent bandits. CoRR, 2bs/2012.00386, 2020b.
Joey Hong, Branislav Kveton, Manzil Zaheer, Mohammad Ghavamzadeh, and Craig Boutilier. Thompson sampling with mixture prior. In International Conference Artificial Intelligence and Statistics, pages 7565-7586. PMLR, 2022.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge University Press, 2015.
Kevin Jamieson, Matthew Malloy, Robert Nowak, and Sébastien Bubeck. lil' ucb : An optimal exploration algorithm for armt bandits. Journal of Machine Learning Research, 35, 12 2013.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016.
Yassir Jedra and Alexandre Proutiere. Optimal best-arm identification in linear bandits. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 10007-10017. Curran Associates, Inc., 2020.
Christopher Jennison, lain M Johnstone, and Bruce W Turnbull. Asymptotically optimal procedures for sequential adaptive selection of the best of several normal mcans. In Statistical decision theory and related topics III, pages 55 86. Elsevier, 1982.
Mare Jourdan, Rémy Degenne, Dorian Baudry, Rianne de Heide, and Emilic Kaufmann. Top two algorithms revisited. Xiv preprint ar arXiv:2006.00 05979, 2022.
Masahiro Kato and Kaito Ariu. The role of contextual information in best arm identification. arXiv preprint arXiv:2106.1400 2021.
13

---

Published in Transactions on Machine Learning Research (04/2023)
Emilie Kaufmann, Olivier Cappé, and Aurélien Garivier. On the complexity of best-arm identification in multi-armed bandit models. J. Mach. Learn. Res., 17(1):1-42, jan 2016. ISSN 1532-4435.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Newton Mwai Kinyanjui and Fredrik D Johansson. Adeb: An alzheimer's s disease simulator for benchmarking observational estimators of causal effects. In Conference on Health, Inference, and Learning, pages 103 118. PMLR, 2022.
Jeongyeol Kwon, Yonathan Efroni, Constantine Caramanis, and Shie Mannor. RI for latent mdps: Regret guarantees and a lower bound. Advances in Newral Information Processing Systems, 34:24523 24534, 2021.
T.L. Lai and H Robbins. Asymptotically efficient adaptive allocation rules. Advances in Applied Mathematics, 6:4-22, 22, 1985.
Tor Lattimore and Csaba Szepesvári. Bandit Algorithoms. Cambridge Universily Press, 2020. doi: 10.1017/ 9781108571401.
Lihong Li, Wei Chu, John Langford, and Robert E. Schapire. A contextual-bandit approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pages 661 670, 2010.
Odalric-Ambrym Maillard and Shic Mannor. Latent bandits. 31st International Conference on Machine Learning, ICML 2014, 05 2014.
Daniel Russo. Simple bayesian algorithms for best arm identification. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir, editors, 29th Annual Conference Learning Theory, volume 49 of Proceedings of Machine Learning Research, pages 1417-1418, Columbia University, New York, New York, USA, 23-26 Jun 2016. PMLR.
Xuedong Shang, Rianne Heide, Pierre Menard, Emilic Kaufmann, and Michal Valko. Fixed-confidence guarantees for bayesian best-arm identification. In International Conference Artificial Intelligence and Statistics, pages 1823 1832. PMLR, 2020.
John M Stern. Overview of evaluation and treatment guidelines for epilepsy. Current treatment options in neurology, (4):222 2244,, 2009.
Alex Tamkin, Ramtin Keramati, Christoph Dann, and Emma Brunskill. Distributionally-wware exploration for cvar bandits. In NewrIPS 2019 Workshop on Safety and Robustness on Decision Making, 2019.
Joy A. Thomas Thomas M. Cover. Entropy, Relative Entropy, and Mutual Information, pages 13-55. John Wilcy & Sons, Ltd, 2005.
William R. Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3/4):285-25) 1933. ISSN 00063444.
Andrea Tirinzoni and Rémy Degenne. On elimination strategies for bandit fixed-confidence identification. Xrx e-prinls, pages arXiv 2205, 2022.
Po-An Wang, Ruo-Chun Tzeng, and Alexandre Proutiere. Fast pure exploration via frank-wolfe. Advances in Neural Information Processing Systems, 34:5810-588 5021, 2021.
Li Zhou and Emma Brunskill. Latent contextual bandits and their application to personalized recommendations for new users. arXiv preprint arXiv:1604.0674, 2016.
ANGELA Zink, JOACHIM Listing, SABINE Ziemer, HENNING Zeidler, German Collaborative Arthritis Centres, et al. Practice variation in the treatment of rheumatoid arthritis among german rheumatologists. The Journal of rheumatology, 28(10):2001-2200 2001.
14


Reviews:
[{"review_id": "61syomqMFM", "reviewer": "Reviewer_k5Me", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Formulating personalized treatment search as pure exploration in latent MAB sounds interesting. Latent MABs appear to be the most reasonable candidate in modeling personalized treatment as a sequential decision making. Establishing this link through a solid mathematical formulation makes a nice contribution. Another strong point is to include both lower bound and algorithms. Finally, I found the setting of the numerical experiments quite relevant to the studied problem.  \n\nThe paper restricts to an intermediate setting where the latent structure is assumed “fully known”, an assumption which may render unrealistic in many practical situations. This assumption implies that the distribution of latent states, the distribution of contexts (conditioned on the latent state), and even conditional distribution of rewards are perfectly known –note that this is stronger than assuming known mean rewards. As a result, once the latent state is accurately identified, the optimal action is deduced. Although it is argued that the said distributions could be accurately estimated using state-of-the-art techniques, I personally found the assumption rather strong, thus limiting the contribution. I agree that some components of the latent structure could be accurately estimated. However, this may no longer be true for all the components, for example, reward distributions.   \n\nIt is nice to have included a lower bound for this setting. However, because of the perfectly known latent structure, its derivation straightforwardly follows from that of classical pure exploration (e.g., in (Kaufmann et al., 2016)). \n\nThe opening part of the article is written well, and the choice of problem formulation is reasonably justified. However, the rest, which constitutes the bigger part of the article, is written carelessly and perhaps in rush, which is strange for a journal submission with strict no deadline. Even the reference section contains strange mistakes: For example, see (Chernoff, 1959), where the paper title is missing. Another example is (Gittins and Dempster, 1979). The related work section discusses most relevant paper that I am aware of on the MAB side. There has been some recent work on RL in latent MDPs. Perhaps discussing this line of work and contrasting them to the considered problem could strength the literature review. \n\nNotation mistakes abound throughout the paper. For example, in one place there is $\\gamma^*$ depends on $s$, whereas this dependence disappears elsewhere without being justified. Another example: both $w^*_x(s)$\n \nand $w^*_{x,a}(s)$ are used (Section 4.1). I listed further at the of this section.  \n \nAlso, there are inconsistencies between the appendix and the main text in terms of notations and definitions. The mean reward in the main text correctly depends on $a,x,s$ whereas in the appendix it is a function of $s,a$ ---see Eq. 8, Appendix A.1. Also, $\\mathrm{Alt}(s)$ should be $\\mathrm{Alt}_x(s)$. Also, the expected value of $N_a$ should be conditioned on $x,a$, and not $X$. A similar comment applies to Eq. 12. Although all these appear fixable in one-shot, they evidently indicate that the authors did not do a careful proofreading. \n\n\nThe LLPT Explorer sampling rule is the same as the one in the conventional MABs. From the presented sampling rules, it is evident that the known latent structure is instrumental to straightforwardly tailor the sampling rules of classical pure exploration to this setting. Overall, my biggest concern is the imposed assumptions makes the contribution very marginal. On the theoretical side, the contribution is comparable to some light technical exercise. Moreover, the current execution suggests that the paper can still benefit from a careful rewriting. \n\n\n---- Further detailed comments:\n\n- In the constraint of the optimization problem (1), one must have strict inequality in the argument of $P$, i.e., $\\mu_{\\hat a_\\tau, x,s} < \\mu^*_{x,s}$. \n\n- The notation in the algorithm is not consistent with the text. For example, history is denoted by $h_t$, whereas $H_t$ is used to denote it in the text. As a matter of fact, there are many others. \n\n- The presentation of TTTS appears confusing to me.\n\n---- Minor comments:\n\n- Proposition 1: “…. the minimizers of the following linear program $x$” => Is $x$ redundant? \n\n- Section 5: sufficently => sufficiently \n\n- p. 8: an non-linear => a non-linear\n\n- Algorithm, throughout: $h$ => $H$\n\n- Algorithm, line 13: $N_{a_t}$ => $N_{a_t}(t)$ ---to be consistent with line 17\n \n- Algorithm, line 20: $[k]$ => $[K]$ \n\n- In the proof of Lemma 1, line 1: remove “for”. \n\n- In Appendix A.2, $L_n(s,s’)$ should be $L_t(s,s’)$. \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the problem of personalized treatment of patients. The problem is framed as a sequential as a sequential decision-making problem, while assuming that some historical data of a population of patients are available. The decision-making part is mathematically formulated as a latent MAB problem in the fixed-confidence pure exploration setting. The paper considers an intermediate setting where the latent structure and mean rewards are known, and the challenge remains in identifying the latent state (of the patient in question). The first contribution is an extension of the existing lower bound (for pure exploration) to this setting. Then, two algorithms are presented. The first one is LLPT, which is an adaptation of the Track-and-Stop algorithm for conventional MABs. The second one is Divergence Explorer. The two algorithms have the same recommendation and stopping rules but differ in the choice of sampling rules. LLPT is shown to achieve the presented lower bound in the high-confidence regime. Numerical experiments are performed to justify the superiority of the presented algorithms over the relevant baselines from the MAB literature. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above. ", "broader_impact_concerns": "Not necessary.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qGuQSZWe89", "reviewer": "Reviewer_f3zP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is very well written\n- The results seem novel, and the medical application well explored, both in terms of modelling as well as in terms of numerical experiments. The authors provide a complete set of results: lower bound, algorithms with a matching upper bound, and numerical experiments.\n\nWeaknesses\n- The main motivation for the model is the treatment of disease, which the authors use to motivate the model of best arm identification. Now, while this probably depends on practical considerations, I do not see why, in the treatment of a disease for an individual patient, one would like to perform best arm identification. Rather an individual patient cares much more about being allocated good treatments, which  in that context would mean minimizing regret. Imagine that there are 3 treatments, where both 1 and 2 are very good (say within a factor +- epsilon of each other, and 3 is extremely bad. Performing best arm identification is mostly overkill here, as it would require to differentiate between 1 and 2. Rather, minimizing regret would be more than sufficient, as it would guarantee that the bad treatment is seldom used. Identifying the best treatment seems useful at the population level (to make recommendations to the general public once the best treatment has been identified), but at the individual level identifying the best treatment seems much less useful. Also, if the disease is chronic, then the allocation of treatment never stops, rather it is a lifelong process, so that stopping at a given time to output the best treatment (like in best arm identification) is not necessary.\n- Some of the numerical results seem a bit too good to be true, in the sense that some of the figures give the impression that, when using the structure, one can sample 5 arms or less and still find the optimal arm with high confidence. This would suggest that the model is so highly structured that a simple greedy approach works just fine, and exploration is not that necessary (essentially playing any action $a$ gives so much information about the state $s$ that we quickly figure $s$ out and play the corresponding best action). And indeed, based on the error bars in Figure 3, it is not even clear that the proposed algorithms do much better than greedy exploration. It would be much more interesting to look at experiments where greedy strategies fail, so that the bandit setting is actually justified. \n- (Minor) The names used for $x$ \"context\" and $s$ \"state\" are somehow counter intuitive since none of them change during the learning process  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors consider a best arm identification problem for contextual bandits in the fixed confidence setting, with a latent state. Namely, an agent observes the context $x$, takes an action $a$, and obtains a reward with mean $\\mu_{a,x,s}$ where $s$ is a state unknown to the learner, and where the mapping $(a,x,s) \\mapsto \\mu_{a,x,s}$ is known to the learner. This is repeated until the learner has identified the best action maximizing $a \\mapsto \\mu_{a,x,s}$ with probability greater than $1-\\delta$ where $\\delta$ is a confidence threshold known to the learner. It is noted that both state $s$ and context $x$ are static, and do not change during the exploration process.\n\nIn fact the problem is similar to clustered bandit, where the notion equivalent to the unobserved state $s$ is the cluster to which the bandit problem belongs. \n\nThe authors prove an information theoretic lower bound on the amount of samples required in order to actually identify the best arm reliably, as well as algorithms matching those bounds. Their approach is mainly an adaptation of the techniques developped in the stateless context by Kauffman and Garivier. The main originality from the mathematical side is that the set of parameters $\\mu_{a,x,s}$ is not continuous like in the classical context, rather it is discrete of size $|S||X||A|$ which allows much faster learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Give motivation for choosing best arm identification over regret maximization in the context of treating diseases.\n- Show practically relevant numerical experiments where greedy strategies fail so that the problem is an actual bandit problem or numerical experiments where the greedy strategies are vastly outperformed by using bandit algorithms.\n", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OUC6DGlhGF", "reviewer": "Reviewer_F7AY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The problem is well-motivated, and the paper well-written.\n\nQuestions:\n\n1. How rich is the context set? The formulation would make sense only if the set of possible contexts is infinite. If it is finite, the optimization problem should be solved for each possible context. Essentially, contexts do not add to the complexity of the problem in this case. In light of this, I have a question from the first paragraph on page 3 (line 7) -- \"To account for remaining individual variation between patients with the same X, we will assume.....\" It seems from this line that the context set X is assumed to be finite (otherwise patients with same X would occur with probability 0 in general). Could the authors please clarify upon this?\n\n2. In continuation of the above, if X is indeed assumed finite, wouldn't it simplify exposition (without loss of any generality) if the context is completely removed from the model altogether? I think the paper fundamentally is about latent state estimation. Presence of context only adds unnecessary detail. The way I would motivate this setting (at least in my mind) would be to consider a classical 2-armed problem with a fixed confidence BAI objective. Where the problem would differ from traditional BAI would be in the presence of auxiliary information. For example, the bandit could have 2 states: one with $\\mathcal{N}(0.5,1)$ and $\\mathcal{N}(1,1)$ rewards for arms 1 and 2 respectively; and the other with reward distributions exchanged. The decision maker knows, for each possible state, the distribution of rewards for the two arms, but the state is latent, and therefore the task reduces to figuring out which of the two arms has distribution $\\mathcal{N}(1,1)$. Would this example fit into the framework?\n\n3. Continuing the above, I would like to see a more comprehensive discussion of how ex ante knowledge of the tuple of reward distributions for each possible state leads to a reduction in sample complexity (the $C_\\delta$ factor).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the fixed confidence pure exploration problem for a $K$-armed bandit with a latent state. There are $M$ possible states of the bandit, and the identity of the optimal arm for each of these $M$ states is known (together with the distribution of arm-rewards in each possible state), however, the state itself is unknown a priori. The problem thus reduces to that of latent state estimation using rewards as observations.\n\nA sample complexity lower bound is derived for this problem; the authors also propose two algorithms one of which is shown to attain the lower bound up to constant multiplicative factors. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Few typos:\n\n1. ALG1 line 13 should be: $a_{t+1} = \\arg\\max_{a\\in[K]} t w_{x,a}^\\ast(s_t) - N_a(t)$.\n\n2. The constant in equation (6) should be $C_0^\\ast(s,x)$, right? If this is not the case, then the bounds do not match, correct?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "UjtbJKRUYz", "reviewer": "Reviewer_nfdE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nStrengths:\n1. The problem of latent bandits in pure exploration is practical and well-motivated. \n\n2. LLPT Explorer solves minimal stopping time that distinguishes latent state by the LP, which is an intuitive idea. Theoretical analysis of the lower bound and the matching upper bound based on LP is a good contribution of the paper. Compared to the lower bound and the upper bound of Track-and-Stop strategy in Garivier and Kaufmann (2016), here the bounds of latent bandits depend on the identification of latent state that is described by the LP.\n\n3. Extensive experiments showed the proposed two methods consistently outperformed baselines under different confidence and noise level. \n\nWeaknesses and questions:\n\n1. It is unclear why Divergence Explorer consistently performed better than LLPT Explorer especially since Divergence Explorer does not have a stopping time upper bound analysis. The authors are encouraged to at least mention the challenge of analyzing of Divergence Explorer (e.g., maybe a worst-case guarantee cannot be analyzed for Divergence Explorer). If such worst-case bound does not exist, is it possible to show that there exist some hard instances that LLPT performs similarly or better than Divergence in experiments?\n\n2. TTTS-Latent Explorer performed similarly to LLPT Explorer in experiments. Can we adapt the theoretical guarantee of TTTS and be able to obtain an upper bound of TTTS-Latent Explorer for latent bandits? Why LLPT has a similar performance as TTTS-Latent? Is it possible to show that the performance gap could be larger with different latent state settings in experiments? It would be better to also explain the similarity/gap from a theoretical perspective.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper considered latent bandits in fixed-confidence pure exploration setting, with application to personalized treatments for patients. The authors considered a contextual bandit setting, where context information is not sufficient to identify the reward and optimal action. Latent state is considered as unobserved information that determines the reward jointly with context. The authors proposed two methods, Latent LP-based Track and Stop (LLPT) Explorer and Divergence Explorer. The authors also proved a lower bound on stopping time of latent bandits that matched the upper bound of LLPT. Experiments on simulation of Alzheimer’s disease validated the advantage of the proposed methods. \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses and questions for requested changes.\n", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies personalized treatment of patients. The problem is formulated as a latent bandit in the fixed-confidence pure exploration setting. Each patient type is represented by a latent state. To solve the problem, the state needs to be identified. The authors prove a lower bound on the stopping time of any algorithm for solving the problem. Two algorithms are proposed. One of them is analyzed and an upper bound on its stopping time matches the lower bound. Both algorithms are evaluated empirically.

I like this paper and the reviewers generally support it as well. The leaning reject decision points out that the analyses are on the simpler side. I agree. On the other hand, this is the first paper on latent bandits in pure exploration. The paper is well written and the experiments are also well done. I also wanted to bring

* [Thompson Sampling with a Mixture Prior](https://proceedings.mlr.press/v151/hong22b.html)

to the attention of the authors. This paper generalizes Hong et al. (2020) to imperfect models. It is essentially Thompson sampling with a mixture prior.
Decision: Accept (poster)