RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluation
Authors: Jeongyeol Kwon
Shie Mannor
Constantine Caramanis
Yonathan Efroni
Venue: NeurIPS 2024 poster (2024)
Keywords: Latent Markov Decision Process
Online Exploration
Off-Policy Evaluation
Partially Observed Systems
Sample Complexity
Information Theory
Primary area: reinforcement_learning
TL;DR: We introduce the first sample-efficient algorithm for Latent MDPs under no additional structural assumptions.

Abstract: In many real-world decision problems there is partially observed, hidden or latent information that remains fixed throughout an interaction. 
Such decision problems can be modeled as Latent Markov Decision Processes (LMDPs), where a latent variable is selected at the beginning of an interaction and is not disclosed to the agent initially. 
In last decade, there has been significant progress in designing learning algorithms for solving LMDPs under different structural assumptions. However, for general LMDPs, there is no known learning algorithm that provably matches the existing lower bound. We effectively resolve this open question, introducing the first sample-efficient algorithm for LMDPs without *any additional structural assumptions*. 
Our result builds off a new perspective on the role off-policy evaluation guarantees and coverage coefficient in LMDPs, a perspective, which has been overlooked in the context of exploration in partially observed environments. Specifically, we establish a novel off-policy evaluation lemma and introduce a new coverage coefficient for LMDPs. Then, we show how these can be used to derive near-optimal guarantees of an optimistic exploration algorithm. 
These results, we believe, can be valuable for a wide range of interactive learning problems beyond the LMDP class, and especially, for partially observed environments.

RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluation
Jeongyeol Kwon
Shie Mannor Technion / NVIDIA AI shie@ee. techni on ac.
University of Wisconsin- Madison jeongyeol. kwon@wisc. edu
Constantine Caramanis University of Texas at Austin constantine@utexas edu
Yonathan Efroni
Meta AI
jonathan efroni Samaiigmmal.... com
Abstract
In many real-world decision problems there is partially observed, hidden or latent information that remains fixed throughout an interaction. Such decision problems can be modeled as Latent Markov Decision Processes (LMDPs), where latent variable selected at the beginning an interaction and is not disclosed to the agent. In the last decade, there has been significant progress in solving LMDPs under different structural assumptions. However, for general LMDPs, even the tabular case, no algorithm is known provably match the existing lower bound [41). We introduce the first sample-efficiet algorithm for LMDPs without any additional distributional assumptions. Our result builds off perspective the role of off- policy evaluation guarantees and coverage coefficients in LMDPs, a perspective, that has been overlooked in the context of exploration partially observed environments. Specifically, we establish novel off-policy evaluation lemma and introduce a new coverage coefficient for LMDPs. Then, show how these can be used derive near-optimal guarantees of an optimistic exploration algorithm. These results, we believe, can be valuable for wide range of interactive learning problems beyond LMDPs, and especially, for partially observed environments.
Introduction
In Reinforcement Learning (RL) [54] an agent aims to maximize the long-term cumulative rewards through interactions within unknown environment. Markov Decision Processes (MDPs) are perhaps the most well-studied and popular framework for this goal. As the name suggests, MDPs heavily rely the Markovian assumption that requires the state be fully observable. However, many real-world decision problems involve critical partially observed or latent information, such as sensitive unknown preference information of users in recommendation systems
(28), undiagnosed illness in medical treatments [62 [53], and adaptation to uninformed tasks in robotics [67 48) Even when such latent factors remain fixed throughout period of interactions the fundamental Markovian property of MDPs is no longer valid.
A line of work has proposed efficient RL algorithms in the presence of latent contexts 12
24
377 within the framework that we here collectively refer to as Latent Markov Decision Processes [1]. In LMDPs, nature selects MDP from finite set of M candidate MDP models at the beginning of a period of interactions (a.k.a. episode), and an agent interacts with the chosen MDP for H time steps of episode (the horizon). However, the identity of the chosen MDP not given to the agent. We call this unknown identity the latent context.
(LMDP) following
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

Most prior work on LMDPs has relied on strict separation assumptions (e.g.. I 24 37). The applicability of these approaches is limited scenarios where the horizon is sufficiently large and identification of the latent model can be guaranteed, i.e,, H > (2(SA) [24] where S and A are the state and action spaces size. Without these explicit horizon requirements, as far we know, all existing algorithms suffer the curse of horizon, requiring sample complexity 9 AA which frequently arises
in the more general framework of Partially Observed MDPs (POMDPs) [52] [39]. Without the ability to identify the underlying latent model, it remains unclear how to address the curse of horizon inherent in partially observed systems
[39
4 proposed sample efficient algorithmss without separation assumptions when M (1). assuming the transition dynamics of models with different latent context is similar. While this is still a substantial contribution, their results cannot be easily extended to the general L.DPP setting with different transition dynamics (see Section T.1 Consequently, to date, the following question has remained open:
Recently, series of works 40
Can we break the curse of horizon in LMDPs if M = 0(1) without any assumptions?
In this work we provide the first sample-efficient exploration algorithm for LMDPs without any assumptions. Throughout the paper, we assume that H > 2M, and focus on whether we can improve the trivial upper bound that incurs complexity S2(A Since (2SAA) lower bound for LMDPs has been established [41] our goal is to achieve an upper bound of poly(S, Aoly(S,A) A)M without any assumptions, namely, to get a matching upper bound up to polynomial factors.
1.1 Technical Challenges
Many online RL algorithms follow a similar pattern. They make use of confidence set set of candidate models (hypothesis) that can explain the observed data with high probability and execute
policy that will shrink the volume of the confidence set is produced and executed (6 (8 27 7 33 36).
The entirety of the statistical problem is to analyze the decaying rate of confidence sets under proper model class assumptions
19
Challenge 1: Limitation of Existing POMDP Algorithms. Existing approaches for online explo- ration in partially observed systems largely fall into the category of Optimistic Maximum Likelihood Estimation (OMLE) 45 46 This class of algorithms often requires an assumption that allows the
construction of shrinking confidence sets. These algorithms also assume access to a set of special policies called core-tests - to be executed to generate trajectories [5] 16 17 45 57 [3] 22 46 26).
Without specifying the proper core-tests, the volume of confidence sets may not decay in desired rate, leading to the curse of horizon 92 (?(A") (3 15 Further, existing POMDP approaches require an ability to recover the belief of the underlying model from observations, e.g., by assuming the distribution of observations when executing the core-tests is invertible to the belief over hidden states.
Consequently, existing literature on POMDPs has two limitations: (i) it requires to specify a priori a set of core-tests policies, and (fi) it assumes the full-rankness of the state-observation emission matrix when the core-tests are being executed.
While LMDPs are special class of POMDPs, neither the existence of a set of core-tests is known priori, nor it is possible to recover the belief over latent contexts from distribution of trajectories (see Section B for details). This creates fundamental challenge for existing approaches when applied to LMDPs, Further, little is understood on learning a near-optimal policy among "doubly-exponentai number of candidate distory-dependee policies without either the visibility of contexts or core-tests. This calls for a new perspective on the question of efficient exploration in LMDPs.
Challenge 2: Limitation of Existing LMDP Algorithms. The work of 43 suggested an alternative strategy to learn a near-optimal policy in LMDPs: the moment-mann matching approach for exploration in LMDPs. When all contexts share the same state-transition dynamics, the notion of moments can be defined as the joint distribution of rewards under a fixed prior at a tuple of at most d' := 2M - 1 acaten pairs must learn how to visit these length-d action tuples simultaneously, i.e., find policy that ensures that æ appears as subsequence of the entire trajectory with high enough probability.
((s[))- anj);... ($[4). a(d)) This in turn suggests that the exploration algorithm
When the transition dynamics of different latent contexts similar, reaching optimally to state-action pairs is minor challenge; e.g., we can first learn the shared transition kernel with any reward-free

---

Vees +- Add a new test policy
Genarate O(E 3) samples with new
30 C, (.. PD)
segmented policies v E aay (Yess))
Update the confidence (loss constraints dataset)
Add new sample trajectories dataset
Figure 1: Highlevel description of LMDP-OMLE. In the online phase, we find a new test policy under which
models in the confidence set do not agree. Then the exploration policy is constructed with our new notion of segmentation of policies within Leas that executed throughout. In the offline phase, add the batched sample trajectories to dataset and update the confidence set of models.
exploration scheme for MDPs [33] and then execute the policy that maximizes the probability of reaching the d state-action pairs. However, for general LMDP, when the transition dynamics of different latent contexts may differ, this approach is no longer available since the latent transition dynamics may not be learnable general. Furthermore, to follow the notion of moments suggested in 43.. the data collected for estimating the correlation tensor must be collected under the same prior (belief) over all latent contexts. Unfortunately, ensuring this for general LMDPs, when the transition dynamics of different latent contexts not equal, is impossible, since even if we obtain the samples of correlations, different policies may result different and unknown priors over contexts. These challenges hint we need an alternative approach to solve general LMDPs, when the transition dynamics vary between latent contexts.
Challenge 3: Limitations of Existing Complexity Measures in RL. Numerous studies have examined complexity measures for RL with function approximation or in the rich-observation set- tings [30 34 [9]. These studies are based on the Markovian assumption, which does not hold in the LMDP setting where the entire history may be needed to decode the latent state. When defining the effective state as the entire history at each time step, it is unclear how analyze the complexity measures from these studies without resorting to exponential guarantees in the horizon.
1.2 Overview of Our Contribution
Recent studies have found some fundamental connections between off-policy evaluation (OPE) and online exploration in RL [59
In this work, we offer a fresh viewpoint, which deviates from existing works, on the connection between OPE and online exploration. This perspective, together with new analysis tools, allows us to provide sample-efficient algorithm for the LMDP setting. This further showcases the usefulness of OPE for online exploration in POMDPs.
55 3
Arguably, the fundamental question OPE the following: how much does behavioral policy e tell about a target policy π?? The simplest form of the OPE guarantee in MDPs relies on the notion of coverage coefficient given by:
I (s: s, -,
C( max
s,a, P* = 8, 5,0==0
How would this quantity be related to online exploration? A key observation to start developing intuition is the following: an unbounded coverage coefficient, i.e., C(v); = oo implies there exists state-action pair, at some time-steps, that cannot be reached under but can be reached with
The algorithmic framework develop in this work builds off OMLE 46 In Section 3 we consider the MDP setting to provide intuition of our analysis. There, OMLE iteratively tests new policies on models from the confidence set which predict different outcomes, until the trajectory distribution of all policies is reliably estimated. Since the number of new state-action pairs is bounded for MDPs, the number of times the coverage coefficient can be large must be bounded during an interaction. We provide new analysis for the MDP setting based on OPE tools.
To apply this approach for LMDPs, and new OPE tools. We propose a coverage coefficient that can be informally described as follows:
notion of coverage coefficient required to develop
P (T E m)
Clu) T) = max max (E,I) Tn
PT m, do Z)

---

where m is the unobserved latent context, T = (st, aq, Pe))e[m] is sampled trajectory, & is an event of interest, e.g., visiting length of at most d tuples of states and actions within an episode, and I is an intervention of interest, e.g., force an action at the 1th time step regardless of t (for the formal definition, see Definition [4.]] Note that the coverage coefficient cannot measured explicitly, since m is a latent variable; nevertheless, this concept is central to our analysis and our ability to analyze the sample complexity of the proposed algorithm. Its usefulness lies in an OPE guarantee we develop (see Lemma
TV (T)(T) & C():)) TVP. (T do I),
where TV(P1, is the total-variation (TV) distance between two probability measures P1, P2.
With these tools at hand, we design an iterative online exploration algorithm for the LMDP setting, and prove its sample complexity matches the lower bound, up to polynomial factors. The algorithm. we refer as LMDP-OMLE (see Figure I for highlevel illustration), repeats the following: (i) find policy for which the trajectory distributions between models in the confidence set is large or terminate, or (ii) collect new data with exploration policies constructed with set of (obtained) test policies and interventions, an exploration strategy for LMDPs that we introduce.
Preliminaries
We consider episodic RL with time-horizon H in LMDPs defined as follows:
Definition 2.1 (Latent Markov Decision Process (LMDP)) An LMDP M consists of a tuple (S,A, R, 0, H) with a state space S; action space A; reward space R, and a finite-time horizon H. 0 is model parameter consisting of multiple MDPs in the model 0 (fwm Tm, Rm 3 m==1- In each mhth MDP, Tm S x A x S - [0, 1] maps a state-action pair and a next state probability; Rm S x A R - [0, 1] is a probability of rewards; {wm 1M are the mixing weights such that at the beginning of every episode the th model is chosen with probubility Wm.
Without loss of generality, we assume that there exists null state that represents the starting and terminal state 8o = $H+1 0, and a null action at the beginning of an episode ap 0, even though actual policies do not take any action at the beginning. Ti Tm (-|s0, ag) is the initial state distribution of the th MDP. We assume that the number of latent contexts is constant M O(1), and the time-horizon is larger than the number of contexts > 2M. Further, we assume the reward values finite and bounded:
Assumption 2.2 (Finite and Bounded Reward) The reward distribution has finite support with (ar- bitrarily large) cardinality, and each reward is bounded: |r| < 1 for all E R.
We also note that this concept can be easily generalized to instantancous observations that include rewards, and thus, we do not lose much generality due to Assumption 2.2 We consider policy class 1l which contains all history-dependent policies : E x (S,A, R) x (S x [H]) + A(A), where E is the space of independent variables decided the beginning of execution. As special case, we consider the class of memoryless policies: Im (S x H)) + ((A) We are interested in finding an optimal history dependent policy II that maximizes the expected reward: Vga :=
max Ez-
H | where 6* E ④ is the true model parameter and E7- [-] is expectation taken over the true LMDP model M when policy æ is executed.
Notation We use [n]:= := (1. (I......]]) ... and [n] == (0) U [n]. We define d == 2MM and assume H > 2M. Let SubSeq d) be the set of subsequences of 21,2,.., .., H) with length less than or equal to d, Le., SubSeq H, d) == ((T), T2, Ty)lq [d. < <.. <T, S H). We often denote a state-action pair (s,a) a) as one symbol x (s.a) E X (S x A), and an reward-next state pair sr) as one symbol y = (r, s') y = (R x S). We often express the next state at time step as either 8 s/, and the pair of instantaneous observation and next state yr- (Tt, 8t-1) = (rr, s(). For any segment of a sequence (ž), 72, ..., zH) from L to 12, we often simplify the notation as ztyste- We denote the entire trajectory as T = (8, a, T)H and T ((s, a, r)..1,,))))))) for history of length t. For any set S, we define && k short-hand for the k-times Cartesian power of S.

---

Algorithm MDP-OMLE
1: Input: neteet N, CT,, 0, Co +
2: Initialize k =
3: while there exists x² E mmms. and ,,, 82 E c such that TV (PTT.
0, Po (T) > 4cry do
4: Generate data )7 J== by executing x", update D + D , )
5: Refine the confidence set with the dataset:
(0 E 0 | (TT,)) log PT(T) 2 arg max E(T log P7 (T) - (1)
+ +1
6: end while
7: Pick any C² and return the optimal policy of M == (S,A,0,0).
We define SubTraj(T, +) S (X' x 2/ as a valid subsequence of trajectories at time-steps SubSeq d), ie., if 0+) SubTraj(T, for any such that T; Ti+1 Vr, = (Tr,: s,) and ZT4T = (S ar.41) must have s'4 87,11*
For a tuple of state-action pairs (or states) of length q. we denote x: = (I (il (or 8
[$[1), .... ((r)) with bracketed indices for each element to distinguish from time steps. We use x for the length of sequence æ. We denote the cardinality of the state and action space as S == S and A = [A. For any two models 01, 02. we often denote P1(-) Po, (.) and (P(()) = Po ())
whenever the context is clear. We denote Pm( (.) for a probability measured conditioned on the context compare 01 and 02). We denote Unif (AA) the uniform distribution over a set A. Let TV(P,, P22 (X) be the total-variation distance between two probability measures P1(.), P-) over random variable X.
e [M]] over the ground-trun truth model (01 when
3 New Perspective on OMLE: Online Guarantees via Off-Policy Evaluation
In this section, we present our new approach for analyzing the OMLE algorithm, and, for establishing intuition in the Markovian setting. Differently than prior analysis 45 46 which is based on the generalized eluder-type condition assumption (see [46], Condition 3.2), we show that a certain type of an OPE guarantee can be used to study the performance of OMLE. This alternative perspective is instrumental in designing a sample-efficient algorithm for the LMDP class.
Consider MDP-OMLL depicted in Algorithm T HDP-OMLE is an adaptation of OMLE for the MDP setting with the goal of learning a near-optimal policy. The algorithm iteratively refines the confidence set, i.e. the set of statistically valid models, until it terminates. Specifically, it iteratively repeats the two steps: (i) find policy for which the TV distance between trajectory distributions of models the confidence set is sufficiently large, and (ii) collect data with that policy, and use the data to refine the confidence set. To bound the sample complexity of the algorithm we attempt to upper bound the number of iterations, namely, to bound the number of times the TV distance between trajectory distributions can be sufficiently large.
The following OPE lemma is a tool that allows us to bound the number of iterations of MDP-OMLE. Before discussing application, we present the result.
Lemma 3.1 (TV Bound via OPE for MDPs) For any behavioral and targel policies vb TI, let the coverage coefficient be defined by:
P (I+ = x)) *E[M] #EX p (x+ = x.)
C( m) = max max
(2)
For any two models 0 E O, the TV distance between trajectory folloin a target policy π. E II is bounded as follows:
TV(P3-, (T) 2CC(U1 Eieunn TVPP o.: P))
(I yt).
How can we use this result to bound the number of iterations of MDP-OMLE? Consider the infinite sample regime, when MDP-OMLE collects infinite data each iteration by executing policy n² on the
5

---

kth iteration, ie., Retest = oc. Further, assume the algorithm is at the beginning of its k + 1 iteration.
In the infinite sample regime all models the confidence set must have matching event distribution relatively to the underlying model measured when policy x² is tested. Specifically, for all 0 E C*
and E [#] holds that TV(PT-, P7 "zze) 1/) = 0. Then Lemma 3.1 implies the following: for all policies for which C(= x) < it also holds that TV(P2 P p0)(7
0. Conversely, assume the condition of the while loop at the beginning of the k + iteration holds true, namely, there exists a
policy म for which TV(P)- (T) O. Then Lemma 3.1 also implies that C(+"; л) 00, namely,
there exists an x E X and E [H] such that P (2e = x) > O whereas P (x+ = x) = 0. Next, recall that -MPE sets the data collection policy at the A + iteration to be π g+1 = स. Hence, the iteration will visit
data collection policy at the not visit. Hence, in the infinite sample regime, MDP-OMLE halts after at most HSA iterations, as there most H A different state-action pairs in different time steps.
state-action pair at time step л did
The intuition presented above is robust to sampling error, i.e, when Theest < xc. To simplify the discussion, let temporarily assume that P-- (2: z) for all II and S x A (we do not require this assumption in our final result by analyzing a perturbed MDP). The key intuition on which the finite sample analysis builds upon is formalized in the following lemma:
Lemma 3.2 (Coverage Multiplicative Increase) For all & > 0 in Algorithm 7 there exists at least one E [H] and x E X' such that
Po* (It = x)) 2 C
TY
Theast (HSA)
max Pg- (xt = x). jck
H
with some absolute constant
> O.
Therefore, by setting the number of samples be Theast 2 (4If2 we ensure that in every iteration MDP -MME doubles the coverage of at least one state-action pair at certain time step. Therefore, the algorithm terminates within most K O(HSA A (1/)) iterations with high probability After termination, we are guaranteed that any two models in the confidence set are étv-close TV-distance for any policy, hence we can obtain = (ffetv)-opting policy. To summarize, we state the following theorem:
Theorem 3.3 Let K O(HSA) log (HSA/c) and B log( K(O/))) least, 1 - 5, MDP- -OLLE terminates after K iterations with at most N episodes being generated, where
/n) Then, with probability at
N Z O S A²) llg S// (og(K))nn)/
and outputs an c-optimal policy with probability at least 1 17.
In typical tabular MDP setting, we take O(log (0) OSS by discretizing the class of MDPs.
Hence the sample complexity of MDP-OMLE is N 6 A²)e)) While this upper bound is suboptimal compared to the minimax rate (7). the appeal of this type of analysis is its ability to bypass the need for analyzing the decaying volume of the constructed confidence sets (Section I.I Challenge 1).
4 Efficient Exploration in LMDPs
In previous section we presented new approach to analyze the OMLE algorithm for MDPs. Next, we develop an analogous technique for the LMDP setting and design the SNDD-OMLE algorithm. Central to its design and analysis an OPE lemma and a new coverage coefficient which we now present.
Intuition from moment-exploration algorithm in 43 provide our intuition on how we construct the OPE lemma for LMDPs. Our construction is inspired by the moment- moment-exploration algorithm proposed in 43 when state-transition dynamics are identical across latent contexts, i.e., Ti = T2 = = Tar, we can first learn the transition dynamics with any reward-free type exploration scheme for MDPs [K]. and then set the exploration policy that sufficiently visits some tuples of state-actions
Before we dive into our key results, let us
length at most d. Specifically, they set a memorlyess
exploration policy vi E Ilma which sets ( x) sufficiently large for some SubSeq d)

---

Algorithm 2 LMDP-OMLE
1: Input: K, d. Rrest E N, Cteat, 7) > 0, B Log(K]ol/ng) V (Unif (A)). D
{(Tj. (Tf,))) Pon where each 7j is generated by executing Unif(A), C as defined in (I
2: Initialize kt =
3: while there exists x* E mmla and 21, 02 E C* such that TV (poz Poz) (7) Aereet do
4: Weest + Wtset U {,² initialize D = Dk-1
5: for all 2b (Vo, 11,,.., (1l,)..;a))) tha) E W x... x W x {Unif(A)) with at least one i E [d 1]+ being ti - x², and T SH, d), E {0, (0,1)17l 1) Generate data {Tz Mire by executing TT)
6:
7: Update D < D ((T),((>): 1311 Picess
8:
end for
9: Update the confidence set C++11 according to equation
10: k + k +1
11: end while
12:
Pick any E C and return the optimal policy of M == (S,A,
and x E X We note that the same moment-exploration strategy cannot applied to general LMDPs with different state-transition dynamics since learning the transition dynamics itself involves latent contexts. Nevertheless, the intuition from [3) suggests that our key statistics are this visitation probabilities to all tuples of state-actions within trajectory.
4.1 Off-Policy Evaluation in LMDPs
The OPE lemma we derive in this section makes use of a behavior policy special form which we refer to as segmented policy, inspired by the notion of moment-exploration in [43]. Let us formally define the key quantities to establish our OPE lemma. A segmented policy, which we denote by
u(zb; T, 2), takes as an input sequence of distory-dependent policies, &p (4/, .... tha). sequence
of time steps, we call checkpoints, T = (Ti, (.., 7]r) SubSeq (H, d), and sequence of binary numbers = (21. 2/1) E 10,1)17 where | S d, and returns history-dependee
policy.
The segmented policy u(bb; T, switches sequentially between different policies in sb The time steps in which the switch occurs are determined by starting from time step T; + policy t/s will be executed. Finally, the sequence determines whether an intervention with random action will occur at the time-step. If I the executed action at time step is the uniform action, Unif(A) and, otherwise, the policy -h;-1 executed. The segmented policy is also denoted by
p(b T, z)
40
1/
(11.21)
(12192)
(1)))))))
where "Ta o Wh" means switch to policy ne at starting from time step + and at time step take (t.=)
random action if and otherwise execute TT.
We are now ready define a coverage coefficient for the LMDP class models. This new coverage coefficient is central to the analysis and design of LMDP-OME -OME..
Definition 4.1 (LMDP Coverage Coefficient) The LMDP coverage coefficient sequence of poli- cies &b E [d++)) with respect to a target policy E 11 in is given by:
P = = = x,
C(bb;))=
max TESsbSeq(H.d) (0,1) inl (a (e .T) mE[M] (
max
mAx
max
(x x7 Vr y)
The LMDP coverage coefficient C(1b; x) between sequence policies, v/. and a target, history- dependent, policy depends on the worst-case way to generate segmented policy, v 7, 2) from vb. Further,
worst-case ratio of the probability of a sequence of observations within |-| d different time steps, namely, Fr. Vr- This is different than the standard coverage coefficient (see equation ( that depends on observation from a single time. Fortunately, C(a); ω) requires only partial set of observations, instead of using full trajectories. This is crucial towards developing
7

---

sample complexity guarantees that are not exponential in H. Lastly, observe that the LMDP coverage coefficient depends on the latent context m, and thus, we cannot measure C(ab; TT) from samples.
We are now ready to provide the key OPE lemma, which makes use of the LMDP coverage coefficient.
Lemma 4.2 (TV Bound via OPE for LMDPs) Let = 2M 1. For any models 0, 0 O and for any π E II and ab E (8+1)+), let C(b;; ) be defined as (4) over 0". Then the following holds:
TV(PS- PT)(T) & M -((b;))
Σ
Σ
TV P py( (Ir, Mr)
TuBubeg H.4) (0.1)))
This result is analgous to the OPE result for MDPs (see Lemma 3.1 It is tool that allows us to bound the TV distance between trajectory distributions of history-dependent policy т by term that depends on a segmented policy (< =) and an LMDP coverage coefficient. Importantly, the term on the RHS that depends on the segmented policy, u(sb: T, =), is sum of distributions of partial trajectories of size [T| S d, which is independent of the horizon length, H.
Remark 4.3 (Why is single latent-state coverability coefficient not enough?) One may wonder why it is not sufficient to consider a single latent-state coverability analogous to Lemma 3.7 namely an analogous to 2 defined as:
P (2² =z)
max max max H REX MEIM p (I, x)
In Appendix D.3 we provide counter-example where such single latent-state coverage coefficient is finite, and yet, off-policy evaluation guarantee cannot be established.
4.2 Coverage Doubling via Sufficiency of Memoryless Polices
To convert the OPE guarantee to an online exploration algorithm, we aim to use the coverage-doubling argument. Ideally. we could apply the coverage-doubling argument with the general policy class similarly to the MDP case as presented in Section 3 However, in its current form, Lemma 4.2 requires the behavior policy be a segmented policy, and is not valid for any general behavioral policy. Hence, it is not obvious on which probabilistic events we can apply the coverage doubling argument. We leave it as future work whether we can obtain an off-policy evaluation lemma with general history-dependent behavioral policies, and its clearer conversion to online guarantees.
In this work, we present an alternative plan the above issue: we reduce the search space from distory-dependent policies to memoryless policies. This allows to track quantities on a segmentwise coverage. Specifically, we first note that the LMDP coverage coefficient can be bounded (after maximizing over the sequence z) by:
maxTi::; Pa ( = 8/141]1s's = 5/19 T1:T))
C(P)T) EBBubSeq H.d)
max
max BE E Ir|-1
1l
max ME[M]
(1/A) (r((sit)) (ST.11 = 5(1+1]115 861))
(5)
where v(ti;; Ti) denotes segmented policy executing vh after the T! time-step with memory reset, hence ignoring the history up to Ti (the conditioning event 5% s/o at can be ignored), Inspired by the form denominator, we aim to double the following probability defined over context-segment pair:
max pu(lits) = ss = s'), NEW
m
for at least one m E [M], 8, E S and 4 < (2. However, another challenge remains: the RHS in equation (5) consists of the maximum over all possible histories in the numerator, whereas in the denominator we force the data collection policy to reset the memory at checkpoints. We still have to side-step this discrepancy apply the coverage doubling argument.
The restriction to the class of memoryless policies allows us to resolve these issues since
max PM m (8ta = s) 1/ s', ",Tr,;) = Pm (81g = s/86, =s),
Tiig
8

---

and ( s/ss, = = = Pm (st, = s|s;, =s),
if T. dh E 1mll since Pm represents the latent Markovian transition dynamics. Thus, we can aim to double up the quantity in equation ( To apply this argument, we establish our second key lemma, crucial building block for the coverage doubling argument:
Lemma 4.4 (Sufficiency of Memoryless Polices for LMDPs) Suppose the following holds:
max Elmie
TV(P ( S test
Then for any history-dependent policies II, the following holds:
& M(2 (MSA)* Étest-
Therefore, we reduced our goal of learning an optimal policy to finding a set of models which satisfies equation 7 with respect to all memoryless policies. Importantly, upon estimating the trajectory distribution up to accuracy Etest > 0 for memoryless policies, we have a bounded TV distance between the trajectory distribution of all history-dependent policies, includes the optimal policy.
4.3 The LMDP-OMLE Algorithm
Once the search space is reduced to memoryless policies, we aim to match trajectory distributions for all memoryless policies. At high-level, LMDP-OMLE follows similar recipe MDP-OMLE, and is summarized in Algorithm 2 It can be described as follows:
1. Find memoryless policy x² ml whose prediction on trajectory distributions does not match between two models in the confidence set CK. Add x to the collection of test policies Veaatt that forms each segment (segmented) exploration policies.
2. Collect new sample trajectories following the new set of segmented policies for exploration, generated by different combinations of collected test policies and switching operations.
3. Update the confidence set C* with Maximum Likelihood Estimation (MLE) on the updated dataset D* by equation (
The data collection policy of LMDP-OMLE (second step above) is inspired and leverages Lemma 4.2 to give upper bounds on the TV distance of untested policies. Next, by Lemma [444 we know that when the while loop terminates, any optimal policy of model contained in the confidence set is near-optimal policy of the underlying LMDP. We conclude this section with our main theorem on the sample complexity of learning the optimal policy in Latent MDPs:
Theorem 4.5 (Sample Complexity of LMDP-OMLE) Lei d = 2M - and assume H > 2M. After at most K = S* S2 H). log( M AHH/e) iterations, LMDP-OMLE (Algorithm) 2 terminates with at most N episodes being generated where
N SM A log( MSAH)))) M² Mog(((0)/n)//
and outputs an c-optimal policy with probability at least 1 17.
Note that in tabular LMDPs with finite support rewards, we have log((Ol) O(S2 log(1/c)) The appeal of lo]( dependence is an flexible extension the same result to parameterized reward distributions. In Section D.I and D.4 we provide proof overview and full proof of Theorem 4.5)
5 Conclusion and Future Work
In this work, we presented the first sample-efficient algorithm for LMDPs, resolving open question of efficient exploration with latent contexts. While our result is specialized to LMDPs, we believe our new perspectives and techniques on deriving online guarantees through the lens of OPE can be useful for a broader range of interactive learning, and, especially, partially observed problems. While resolving the open problem, there a few remaining questions for the LMDP setting.

---

Tightness of the Result. The upper bound in Theorem 4.5 scales with Ô (MSAH 1og(1/e)) ( while the existing lower bound ( A)M. Closing this polynomial gap in the exponent, and having a matching upper and lower bounds can be valuable for deeper understanding of LMDPs and possibly for POMDPs in general.
General OPE lemma and Regret Guarantees for LMDPs. The OPE lemma derived in this work (Lemma 4.2 assumes the behavior policy segmented policy with intervention different checkpoints. While this result allows us to provide guarantees LMDP-OMLE and prove it learns near-optimal policy, this result is restrictive, that does not provide general guarantees for OPE nor makes it possible to derive regret guarantees. In particular, can we evaluate π IT without policy-switching or intervention when the behavioral policy is a generic history-dependent policy di E II? Further, is there an algorithm with provable poly(S, A)M VT regret for the general LMDP setting?
Towards Practical Settings. Our result gives a worst-case guarantee. Yet, practical instances may be much simpler under different set of assumptions e.g., with provided side-information [66] [4] or Deriving new conditions can of great importance for real-world applications, e.g., there could be more practical notion of separation, instances that allows the notion of coverage-coefficiet with a (significantly) shorter length d = o(M) of state-action tuples. Further, developing practical RL methodologies for the LMDP setting remains an unexplored challenge with significant importance for numerous applications. These are remained to be explored in future works.
additional structural assumptions 41 63
the set of
Acknowledgment
This research was partially funded by AFOSR/AFRL grant no. FA9550-18-1-0166, NSF Grants 2019844 and 21 22147,,, and Israel Science Foundation Grant No. 2199/20.
References
[1] A. Agarwal, S. Kakade, A. Krishnamurthy, and W. Sun. FLAMBE: Structural complexity and representation learning of low rank MDPs. Advances in neural information processing systems, 33:20095-20107, 2020.
[2] A. Al-Marjani. Tirinzoni, and E. Kaufmann. Active coverage for PAC reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory, pages 5044-5109. PMLR, 2023.
[3] P. Amortila, D. J. Foster, N. Jiang, A. Sekhari, and T. Xie. Harnessing density ratios for online reinforcement learning. In The Twelfth International Conference on Learning Representations, 2023.
[4] P. Amortila, D. J. Foster, and A. Krishnamurthy. Scalable online exploration via coverability. arXiv preprint w:Xiv:2403.051) 2024.
[5] A. Anandkumar, R. Ge, D. Hsu, S. M. Kakade, and M. Telgarsky. Tensor dccompositions for learning latent variable models. The Journal of Machine Learning Research, 15(1):2773-2822 2014.
[6] P. Auer, N. Cesa-Bianchi, Y. Freund, and R. E. Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-77, 2002.
[7] M. G. Azar, Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning, pages 263-272 PMLR, 2017.
[8] K. Azizzadenesheli, A. Lazaric, and A. Anandkumar. Reinforcement learning of POMDPs using spectral methods. In Conference Learning Theory, pages 193-256, 2016.
[9] P. J. Ball, L. Smith, I. Kostrikov, and S. Levine. Efficient online reinforcement learning with offline data. In International Conference on Machine Learning, pages 1577-1594. PMLR, 2023.
[10] A. Bennett and N. Kallus. Proximal reinforcement learning: Efficient off off-policy evaluation in partially observed markov decision processes. Operations R0search,2023.
10

---

[11] E. Brunskill and L. Li. Sample complexity of multi-task reinforcement learning. In Uncertainty Artificial Intelligence, page 122. Citeseer, 2013.
[12] I. Chadès, J. Carwardine, T. Martin, S. Nicol, R. Sabbadin, and O. Buffet. MOMDPs: a solution for modelling adaptive management problems. In Twenty-Sixth AAAI Conference Artificial Intelligence (AAAI-12), 2012.
[13] F. Chen, Y. Bai, and S. Mei. Partially observable RL with B-stability: Unified structural condition and sharp sample-efficient algorithms. In The Eleventh International Conference on Learning Reprexentations, 2022.
[14] F. Chen, C. Daskalakis, N. Golowich, and A. Rakhlin, Near-optimal learning and planning in separated latent MDPs. arXiv preprint (rXxx24006009220 2024.
[15] Chen, H. Wang, C. Xiong, S. Mei, and Y. Bai. Lower bounds for learning in revealing POMDPs. In International Conference on Machine Learning, pages 5104-5161. PMLR, 2023.
[16] C. Dann, N. Jiang, Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. On oracle- efficient PAC RL with rich observations. In Advances in neural information processing systems, pages 1422-1432, 2018.
[17] Y. Efroni, C. Jin, A. Krishnamurthy, and S. Miryoosefi. Provable reinforcement learning with short-term memory. In International Conference Machine Learning, pages 5832-5850. PMLR, 2022.
[18] Y. Efroni, D. Misra, A. Krishnamurthy, A. Agarwal, and J. Langford. Provably filtering exoge- nous distractors using multistep inverse dynamics. In International Conference on Learning Representations, 2021.
[19] D. J. Foster, S. M. Kakade, J. Qian, and A. Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487 2021.
[20] A. Geer. Empirical Processes in M-estimation, volume 6. Cambridge university press, 2000.
[21] C. Gentile, S. Li, P. Kar, A. Karatzoglou, G. Zappella, and E. Etrue. On contex--dependent clustering of bandits. In International Conference on Machine Learning, pages 1253 1262. PMLR, 2017.
[22] N. Golowich, A. Moitra, and D. Rohatgi. Learning in observable POMDPs, without computa- tionally intractable oracles. Advances in neural information processing systems, 35:1458-1473, 2022.
[23] Z. D. Guo, S. Doroudi, and E. Brunskill. A PAC RL algorithm for episodic POMDPs. In Artificial Intelligence and Statistics, pages 510-518, 2016.
[24] A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv preprint 2015.
[25] D. Hsu, S. M. Kakade, and T. Zhang. spectral algorithm for learning hidden markov models. Journal of Computer and System Sciences, 78(5): 1453):146014888 2012.
[26] R. Huang. Y. Liang, and Yang. Provably efficient UCB-type algorithms for learning predictive state representations. In The Twelfth International Conference on Learning Representations, 2023.
[27] T. Jaksch, R. Ortner, and P. Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11:1563-1600, 2010.
[28] A. J. Jeckmans, M. Beye, Z. Erkin, P. Hartel, R. L. Lagendijk, and Q. Tang. Privacy in recommender systems. Social media retrieval, pages 263-281, 2013.
[29] 7. Jia, G. Li, A. Rakhlin, A. Sekhari, and N. Srebro. When agnostic reinforcement learning statistically tractable? Advances in Neural Information Processing Systems, 36, 2023.

---

[30] N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford. and R. E. Schapire. Contextual decision processes with low bellman rank are PAC-learnable. International Conference on Machine Learning, pages 1704-1713. PMLR, 2017.
[31] N. Jiang and L. Li. Doubly robust off-policy value evaluation for reinforcement learning. In International conference on machine learning. pages 652-661. PMLR, 2016.
[32] C. Jin, S. M. Kakade, A. Krishnamurthy, and Q. Liu. Sample-efficient reinforcement learning of undercomplete POMDPs. arXiv preprint arXiv:2006. /2484, 2020.
[33] C. Jin, A. Krishnamurthy, M. Simchowitz, and T. Yu. Reward- exploration for reinforcement learning. In International Conference on Machine Learning, pages 18870-8799 PMLR, 2020.
[34] C. Jin, Q. Liu, and S. Miryoosefi. Bellman eluder dimension: New rich classes of RL problems, and sample-efficiett algorithms. Advances in neural information processing systems, 34:1 13406- 13418, 2021.
[35] S. Kakade and Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning, pages 267-274, 2002.
[36] E. Kaufmann, P. Ménard, O. D. Domingues, A. Jonsson, E. Leurent, and M. Valko. Adaptive reward-free exploration. In Algorithmic Learning Theory, pages 865-891. PMLR, 2021.
[37] C. Kausik, K. Tan, and A. Tewari. Learning mixtures of markov chains and MDPs. In Interna- tional Conference Machine Learning, pages 16077-16017. PMLR, 2023.
[38] M. Kearns and S. Singh. Near-optimal reinforcement learning in polynomial time. Machine learning, 49:209 232,, 2002.
[39] A. Krishnamurthy, A. Agarwal, and J. Langford. PAC reinforcement learning with rich observa- tions. In Advances in Neural Information Processing Systems, pages 1840-1848, 2016.
[40] Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reinforcement learning in mixird MDPs. Advances in Neural Information Processing Systems, 34, 2021.
[41] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor, RL for latent MDPs: Regret guarantees and a lower bound. Advances in Neural Information Processing Systems, 34, 2021.
[42] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Tractable optimality in episodic latent MABs. Advances in Neural Information Processing Systems, 35:23634-23645 2022.
[43] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reward-mixing MDPs with few latent contexts are learnable. International Conference on Machine Learning, pages 18057-18082. PMLR, 2023.
[44] Kwon, Y. Efroni, Mannor, and C. Caramanis. Prospective side information for latent MDPs. arXiv preprint a3Xiv: 10.07596, 2023.
[45] Q. Liu, A. Chung, C. Szepesvári, and C. Jin. When is partially observable reinforcement learning not scary? In Conference on Learning Theory, pages 5175-5220. PMLR, 2022.
[46] Q. Liu, P. Netrapalli, C. Szepesvari, and C. Jin. Optimistic MLE: A generic model-based algorithm for partially observable sequential decision making. In Proceedings of the 55th Annual ACM Symposium Theory of Computing, pages 363-376, 2023.
[47] H. Namkoong, R. Keramati, S. Yadlowsky, and E. Brunskill. Off-policy policy evaluation for sequential decisions under unobserved confounding, Advances in Neural Information Processing Systems,
2020.
[48] Z. Rimon, A. Tamar, and G. Adler Meta reinforcement learning with finite training tasks-a density estimation approach, Advances in Neural Information Processing Systems, 35:13640- 13653, 2022.
12

---

[49] D. Russo and B. Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances Neural Information Processing Systems, 26, 2013.
[50] C. Shi, M. Uehara, J. Huang, and N. Jiang. A minimax learning approach off- evaluation in confounded partially observable markov decision processes. In International Conference on Machine Learning, pages 20057-20094. PMLR, 2022.
[51] S. Singh, M. R. James, and M. R. Rudary. Predictive state representations: new theory for modeling dynamical systems. In Proceedings of the 20th conference on Uncertainty in artificial intelligence, pages 512-519, 2004.
[52] R. D. Smallwood and E. Sondik. The optimal control of partially observable markov processes over finite horizon. Operations research, 1165::071-11888 1973.
[53] L. N. Steimle, D. L.. Kaufman, and B. T. Denton. Multi-model markov decision processes. Optimization Online URL http://www.. optimization-online org/DB HEEDDBFFIE//211111 F111/2208////14344444 pdf. 2018.
[54] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[55] K. Tan and Z. Xu. A natural extension to online algorithms for hybrid RL with limited coverage. arXiv preprint arXiv:2403.00 09701, 2024.
[56) G. Tennenholtz, U. Shalit, and S. Mannor. Off-policy evaluation in partially observable envi- ronments. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10276-10283, 2020.
[57] M. Uehara, A. Sekhari, J. D. Lee, N. Kallus, and W. Sun. Provably efficient reinforcement learning in partially observable dynamical systems. Advances in Neural Information Processing Systems, 35:578-592 2022.
[58] M. Uehara, C. Shi, and N. Kallus. A review of off-policy evaluation in reinforcement learning. arXiv preprint arXiv: 22t.::2212.6635.. 0235555 2022.
[59] T. Xie, D. J. Foster, Y. Bai, N. Jiang, and S. M. Kakade. The role of coverage in online reinforcement learning. The Eleventh International Conference on Learning Representations, 2022.
[60] T. Xie, Y. Ma, and Y.-X. Wang. Towards optimal off-policy evaluation for reinforcement learning with marginalized importance sampling. Advances in neural information processing systems, 32, 2019.
[61] Y. Xu, J. Zhu, C. Shi, S. Luo, and R. Song. An instrumental variable approach to confounded off-policy evaluation. In International Conference Machine Learning, pages 18844-38888) PMLR, 2023.
[62] G. Yauney and P. Shah. Reinforcement learning with action-derived rewards for chemotherapy and clinical trial dosing regimen selection. In Machine Learning for Healthcare Conference, pages 161-226. PMLR, 2018
[63] W. Zhan, M. Uchara, W. Sun, and J. D. Lee. PAC reinforcement learning for predictive state representations. arXiv preprini 2rXii:2207.05738 2022.
[64] T. Zhang. From E-entropy to KL-entropy: Analysis of minimum information complexity density estimation. Annals of Statistics, (4(5):2180-2211) 2006.
[65] Y. Zhang and N. Jiang. On the curses of future and history in future-dependee value functions for off- off-policy evaluation. arXiv preprint 240.2 2024.
[66] R. Zhou, R. Wang, and S. S. Du. Horizon-free and variance-dependen reinforcement learning for latent markov decision processes. In International Conference on Machine Learning, pages 42698 42723. PMLR, 2023.
[67] I.. Zintgraf, K. Shiarlis, M. Igl. S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. VariBAD: A very good method for adapiive deep RL via meta-learning. In International Conference on Learning Representations, 2019.
13

---

Appendix A Additional Preliminaries
The difference between the values of any policy क II measured on two models 01, 02 Q bounded by the total-variation (TV) distance between trajectory distributions, that is,
IV
Vi'l - H TV(PT, (T)),
since the maximum reward that can be obtained in an episode is bounded by H due Assumption 2.2 Hence, if we can show that
TV(P (T) e/H =: CTV: VT E II.
(8)
then an optimal policy #* of the empirical model Ô is guaranteed to be 2e-optimal in the true model 0" Henceforth, we focus on finding an empirical model 0 that satisfies
To bound the TV-distance between trajectory distributions for all history-dependen policies π E II between any two LMDP models, start by unfolding the expression of statistical distance
(T)))=
Σ
(F((aa P2 (( re)
(Sre.Te)LE(n)
H EII (aa TT:e)
M
H
M
Σ
II t=0
(e) - Σ m==
2
II t=0
Th (241 (relzr))
x
w
m==
When the context is clear, we compare trajectory distributions between any two given model parameters 0,, 02 O, and denote the probability measure from each model following PP. (-), (-)
A.1
Additional Notation
To reduce the notation overload, we use (m((ye Tm (#14 (m When the context clear, we often use shorthand n, = x (a, (1:)), and denote tiita as shorthand of the product of time-consecuttvv sequence from & to 12, i.e,, ttyd Iī_,, Te- When implicitly mean that the &f part of yr, which denote s'(y,, must match to the 8 part of It-1, which we denote as 8(1++1). Using the notation, we rewrite the unfolded TV-disstace equation the following:
over both 2; and ye-
ΣΣ TI.H EM
W I PL
Phi (uellr)) - EM with II Ph (unll))))
We use shorthand (X) for (X) PX EM m=1 W (lP((()) Σ E P2. (X))
and thus Σx (X) TV(PT PT (X) where the summation is over all possible realizations of random variable X. Finally, we denote d 2M
A.2 Preliminaries for Lemma 4.2
Here we define a few more quantities that will be crucial for the proofs for Section In bounding the total variation distance in terms of tested policies without exponential blow-up H, the key is marginalize events across time steps. Let us first fix the checkpoini time-steps (TI,..,, Ty) for E [d], and a sequence of executable policies ub = (No, v1,,.,, va). Each 1th segment policy will be executed within time interval (r,, 1++ for i > 0.
To proceed, for the initial policy to,, let ( be the smallest quantity, among the contexts, the ratio between the state visitation probabilities in consecutive steps when th is executed (recall that we denote 2, = (st, a1). 8/1 = (re, 8141)):
too (st) pa (s++1)
( ((
min ne(1.)
min me(M,
pm (Fe, (Ti, (4t1))))))))
This quantity can be understood as the minimum (over latent contexts) of the sneudd-possterin probabilities of 1-step event given the future state if vo is memoryless, since
O (Z,, rei St+1) vo(a = min min pio pro (2x, re r1)) if vho E 1mmls- nE(1)) mE[M]
14

---

Henceforth we use a shorthand 0 == (1t We recursively define sequence of the above quantity.
Next, we fix the event (x+, Vr) ) at the event-log time-steps. For all i 0 and T; < Ti+1, we define ( Draa) and pm (Frae e Vr recursively as the following:
Pm
(R.M(W
(se|STc41 Pm
(r ( Set1l
([ Brea)
min nE(1.2) |m (xr, Srr )=0
min
(10)
Ph (st+1l57:+1)
and
P (Z Vrri)) x (riss( (Stu4) 18 +1 Pm (Ut4 (T(4))
Pm I
m
pn.oldie (8' (8'(Tr (17r)1)8844H M8rc+1) + (Friss Wras))
(11)
Here, we define [[o] = 1, where we recall that ((T;; t) means the memory reset of policy ı at time step t + 1. Zri. Vro.o 6, and 18 = 1 and ph wn,, for 1, 2. The key point here is that in this recursive construction, as increase, we have at least one such that either pm gl.+1 0 or P i.e., at least one context is removed from consideration at each checkpoint.
2,
In the subsequent steps our proof, we often omit the dependence on (Iir Vrne): . well as Tro in z and p* when the context is clear. Finally, we define
( V+) = Σ E LI
2/TT
2m=1 Pm
Now we are ready to state our key intermediate lemma:
Lemma A.1 For any target policy π 11 and sequence of segment policies ap (No, 1/1, my va). the following holds:
ΣΣ JI.
Em W/ IT P (Vblxe) - Em= w II P (3 (x:))
Z1:H BrEM
pa max, Mr (IT,
Σ Σ A(xr,87)) x
[1:-1-1 m((r.ve) (Srepa (80+1) Pm(err)) (rrep (Frez))
+EBubbeq H,d) ir
(12)
where m(x 9r) is the smallest m E [M] such that pm 1.171 > O.
A.3 Preliminaries for Lemma
We present additional tools that are useful for proving Lemma 4.4 We first define the notion of (latent)
segment coverage coefficient of the ser of test policies Utest with respect as the following:
Definition A.2 (LMDP Segment Coverage Coefficient) The coverage of π E I1 with respect set of test policies V teet S 11 is defined as:
maxp:|M| P (st, s/s/ s
P(V testi TT) = max max max f <3g 8,8' m
(13)
maxpew.... (2r(v:t)) (8t2 = 8/3/ 8))
The key lemma is bound the coverage coefficient of segmented policies consisting of mixture policies defined as the following:
Lemma A.3 Let Dh be mixture of a subset of behavioral policies the following set with fixed to E [H]:
arg max pe(vsto) (Seg+t = =/ss/, = =s), Vm, s, s' #E West test
teet-
(14)
Let (el. We define the mixture policy WE ;= Enve b. i.e. given the time interval i until the next checkpoint time, the first uniformly randomly picks one policy from Ve and executes the picked policy afterwards. Let ubs == (the, vhe,-.,, 1he) (of length d + 1). Then the following holds:
Clube T) < (nA. pl testi T))a
15

---

A.4 Auxiliary Concentration Lemmas
The following lemmas are the standard concentration of log log-likelihood values of the models within
the confidence set. The proofs are standard and can also be in e.g. [5] [4] and [20] 64]. We let D* be the dataset at the beginning of the hthh iteration Algorithm
We denote 5 == log(K
Lemma A.4 (Uniform Bound on the Likelihood Ratios) With probability 1 for any l for all k E [K] and for any 0 E O
Σ
log(P)(T)) - 3
Σ
log (T)).
(15)
(T.D)
(T,R)ED*
Lemma A.5 (Concentration of Maximum Likelihood Estimators) With probability 1 1 for all k E [K], & e [M]] and O we have
log Pg.(T) (T(
Σ
TP (P (T)
Σ (T,T)ED
+ +3..
(T.D)ED
Appendix B
Additional Related Work
The literature on reinforcement learning theory is fast growing. While learning in fully observable systems has been extensively studied in the past decades, relatively little has been understood for online exploration in partially observable systems until recently. We review recent theoretical advances in reinforcement learning with partial observations that are closely related to us.
Learning near-optimal policy in POMDPs is a notoriously hard problem [52] due to its full generality. In particular, the statistical complexity of online exploration in general POMDP fundamentally suffers from the curse of horizon [39]. A earlier breakthrough involved introducing structural assumptions system dynamics, which enable the recovery of underlying
Exploration in POMDPs
POMDPs model under the uniform ergodicity assumption 25
23
Recent theoretical breakthrough concerns the exploration problem in POMDPs without the ergodicity assumption 32 45), initiating remarkable progress in understanding the statistical complexity of reinforcement learning in POMDPs under proper structural assumptions. To this date, well-studied tractable POMDP classes (that overcome the curse of horizon) can be considered largely as a system with a "short-window" window assumption is the prior knowledge that the a short consecutive execution of purely random actions is enough obtain sufficient statistics of histories, i.e., the full-rankness of latent state-future observation emission matrices. A short sequence uniformly random actions become the set of core tests of the system [5]] In such systems, learning the optimal policy only incurs poly(S,A) A complexity [46) (with window-size L), breaking the curse of horizon.
for efficient exploration [16] 18 17
13) The crux of the short-
Unfortunately, the same story does not apply in LMDPs, as there is no such "short-window" assumption that allows us to learn the sufficient statistics of histories. This calls for a set of new ideas and concepts, which could be of independent interest.
Along with the fast progress in online reinforcement learning with partial observations, there a growing interest in off-policy evaluation with partial observations 47 (6) 61]. While the sample complexity of OPE has been extensively studied in MDPs under various model class assumptions [31 60, 58], most existing OPE results POMDPs are asymptotic sampling.
Off-Policy Evaluation in POMDPs
in nature, and often suffers the curse of horizon due to their use of importance-wett
Recently, several recent works have proposed an alternative measure of coverage in the latent space, breaking the curse of horizon [50 10, 65]. However, their results rely on the weakly-revealing assumptions that is often made in tractable POMDP classes [45], and can only evaluate within the
class of memoryless policies. Our results are developed for the off-policy evaluation in LMDPs with several new concepts, which can also be of independent interest to off off-policy evaluation problems in partially observed systems.
16

---

B.1 Additional Details on the Full-Rankness Assumption
We give more detailed explanation of the full-rankness assumption that has become popular in the POMDP literature [45]. As mentioned, the statistical sufficiency of core-tests, which is represented as the minimum singular value of the "latent state-future trajectory" emission matrix L. Such an assumption has been exploited in the earlier work of LMDPs in [11] where the matrix T. is defined in the following form:
[m[m, (1/eest: Te:H)] Pu Whettt (Tehls),
where Presst test is test policy, and Te: is the future trajectory after time step Therefore, direct application of POMDP approaches such as OMLE require the prior knowledge of Leest and minn (Lx) > 0 for 8 S. The rationale behind such assumptions to ensure that a distribution of future trajectories can be converted to belief of latent contexts, hence distribution of future observations can serve as an alternative a belief state. However, we are not given the set core-tests V teat, or even the existence of Veet that ensures min (T..) > O for general Latent MDPss
recent work by Chen et al. 14, polynomial upper bound in M has been established under a notion of atrong-separatin between contexts with sufficiently long time horizon H. In essence, their assumptions guarantee that min (L,) > 0 holds for most of the states with priori given test policy, along with additional analysis for the tail of episodes. is of great importance to identify such practical assumptions that lead to fully polynomial upper bounds, especially for instances with some proper notion of separations even when no prior knowledge of test policies is provided.
With Separation.
Appendix C Proofs for Section
C.1 Proof of Lemma 3.1
This base case corresponds to Lemma 4.2 with Mi M2 For convenience, we let 0, 0* and 0g 0, and thus, 1 = Pp- and IP: = Po. We can show the inequality recursively:
H
ΣΣ
II (1)lax))
I (2(()))) t=
FI VIe
r==
H-1
H-1
Σ
Σ TTH
II t=0
Pi (uelx))))
II P2(u))))))
Σ P2 P2(yn 31
V1:H-1
==
H-1
Σ TI-H ==0
II
Pi(ue Σ Pi(un P2(un(an))
Note that Egg P2 P2(yn[eH) and
Enn Emn THH 11. Pi = (T
since we implicitly sum over SH = (1H as we described in Appendix A.1 Thus,
ΣΣ THH
II P2(2)
ΣΣ
TI.H
II
P (1
II P2(
JELIM UI
t=
t=0
SI 91
t=0
Σ
PT [PP (P (P((nnoo))
s 21
Then show that
PT (T( (P P2(un(nn)
Σ (T (P((wH(H)) - Pa(un|zn)] = Σ
P. (œ
C())))
Σ (eH. PT(
Σ (T)()) (P(
17

---

x 20 ())()) (2H, (ZH.UH). VH).
Applying the same step inductively from t = ./ to = 1, we get the lemma.
C.2 Proof of Lemma 3.2
Let We == (> (a.t), Vx, t (2, 1) == arg ma.. 1 (2r 2)
Then, let dee E II be
policy that can adapt to the predetermined checkpoint 1, such that the [Wel Exere vi Le. mixture of policies in V... Note that |Vel s HSA. Lemma
1
3.1 tells us that
S Σ
C(Ve: TV(PP P (T)
EE
O(UEiT) VE
2 Σ
TV(PI (T)
E M) NEWS
< 2H
DV PE (T).
(16)
V
Then we apply Lemma A.4 and Lemma A.5 to deduce that
163
Σ TV (PT
P (T)
Rsest
On the other hand, note that
px* (or x)
VE ps (Ie x))
C(y:: mh) max max
max
tE (E|#] SE p (xq = 2)
#ESxA maxj<k p (Irt x))
Now using the while loop condition, we have
C(ve:T") Vl V Vreeet3
6TV < TV(PT" P. P (T) S 8H max (E[H]
HSAB
Px* (xt = x)
S 8H Peast
max max HE(H] #EX max Pri (I, =))
Rearranging the inequality, implies that there exists at [H] and
X such that
max T n (z+ = x² < Jck
8H
HSA3
p=* (z+ = =x).
CTV
Nleet
C.3
Proof of Theorem 3.3
We first show that Algorithm terminates after K = H. log(1/7) iterations where Gast\ H?
consider perturbed model 0" (w, T, R) where
f(1ss, a) (1 7)T (-18, + 71,
By simulation lemma
38
for any E II, note that TV(PI, (vlo)) s 2yS for all and thus
H
TV(IPT PT-)(
ΣΣ Σ
II
II P2(un)2))
t=
H-1 Pi P1(2)))))
H-1
Σ TI:-1
II 1=
Ru:H-1 Mi:H-1
t=0
Σ
P7 (Pi(v)(a)) P=(vn(=E))
18

---

H-1
H-1
Σ
Σ
TI:H-
II
Pi ())
II
P (2())())) + +2,5
#1.8 - Mi:H-
- 2
For an arbitrary & iteration, we check whether the coverage doubling argument (Lemma D.I still holds. To see this, first note that we can define Ee, the and Ö(pg; T) as in Lemma 3.2 in terms of 0":
pz (Ia x)) }
n (2.1) Vx, & (x, = arg max JE0.1....k-1
and de is a checkpoint-dependemt policy where the is a mixture of Gg.. and
PT. (Ee = x)
C(p)))
max max tE[A] ME
P (3t x3)
Now invoke Lemma A.4 and Lemma A.5 to show that
D2 TV
(T) - TV
IVve P
(7) + TV (P Dyee) (T)
0+
2SH +
Σ TV P P² (T),
TV
[ell
weV;
for all 9 C using the triangle inequality for TV distance and (a b)2 < 2(a + b²). Thus,
derive equation [66 terms of 0*
163
TV(PT G- (P) < 2H
(2SH)²²²
V
Vlteet
P (Ix: x))
16.9
S 22HHHAAL
+ (2SH))
max max
Neest
tE[M] sEX maxjck Ph (Ix. = x))
On the other hand,
PT))) max (TV(PTL g. PT. (T), TV(PTT. g.: P7t) (T))(T)
2ETv 2SH, > 1.5erv:
by setting = Etest/ (4SAH)* Let Tlees
= 643(H*SA) SA)/e2a 2est. and we have
P (It = x))
2εTv Eteet mmax max- (E(H] ZEX maxy< P (xt, x)
Hence the same doubling argument holds, and MDP-OMLE (Algorithm) will terminate after at most
K O(HSA log( (HSA A/te)))))))
iterations. We note that all inequalities hold for all K iterations with probability at least 1,. Finally, by setting étast = énv and êry t/H, we can conclude that the total number of trajectories that was generated by MDP-OMLE is bounded by
0(1) (og
Appendix D Proofs for Section
D.1
Proof Sketch
In this section, we provide overview of the proof for Theorem 4.5
Compared to Algorithm | the
main differences LMDPss from the MDP cases are two-fold:
(a) Our goal is to find the optimal historr- (non-Markovian) policy.
19

---

(b) We cannot observe any context-segment pair that previous behavioral policies could not cover.
For the first point, (a), we already have reduced the problem from matching all history-dependent policies to a set of behavioral policies generated by the concatenation of memoryless policies in Lemma 4.4 For the second point, (b), even though we cannot observe the latent context under which each segment is covered, we can improve the coverage each context-segment pair given the test every iteration:
test set V/
Lemma D.I Lel Reest 2 38MP(SHF) (M Then with probability at least 1 at every tth iteration in Algorithm 2 there must exist at least one (m, t1, 12. 8, s'), such that
pm (stz = s|ss =s') max P =|ss = s').
v text
That is, we can ensure that the coverage of at least one context-segment pair is being exponentially improved despite the unobservability of latent contexts m.
For a moment, to simplify the discussion, we first assume that the uniformly random policy Unif(A) has non-zero y > 0 probability for covering all segments in all contexts:
plnit(A) () (St2 = S2 St. = $1) 2 % Pm
vt < f2, m, $1,
(17)
We note that this assumption will be eventually removed in our final result. Thus, if we start from the initial coverage 7 over all context-segment pairs, then since every probability the range of [0, 1]. this doubling-up event can happen at most log(1/7) times for every context-segment pair. Therefore, Algorithm| must terminate after at most = M H log(1/7) iterations.
Separately from the coverage improvement in Lemma D.I the standard concentration of the confidence set on the generated trajectory data is given by the maximum-likelihood estimation:
Lemma D.2 With probability at least - for all kth iterations in Algorithm 2 let Veg (vh...., w S test be any subset of candidate test policies, and let uhe = (We, 2/e (ni)(AA) where We = Eiein) th is mixture of policies in Me.. Then for any model in the confidence set 0 E C, the following holds:
Σ
Σ
TV (P pp(bessit)) (Iz: V+) S nd. Tess
43
TESubSeq(H.d) =C (0
Equipped with Lemma D.2 With probability at least 7 we terminate the while-loop after at most K H log(1/7) iterations, and each while-loop generates Kd-1 .Ttatt new trajectories, leading to a total
Kdd Nrest = 2 log(1/>)) 3//
sample complexity. The netr-optimality guarantce for the final empirical model is given by Lemma MSA)
4.4 where we set Étest = M (2H+2 ETy to obtain an (Herv)-optimal policy.
Without Initial Coverage. Now remove the initial 7 coverage assumption (17). To do so, consider virtual LMDP model (fw Ti, R }) 3) M== with perturbed transition models, i.e. (-Is,, a) = (111))T=(44 (-Is, a) +~1 for all (s,a) a) E S x A. For it is easy to see that for all π E II, we have TV(P7 (T) 2HS Thus, now with sufficiently small and it is straightforward that the perturbed model has the y >> 0 segment coverage for any policy, which concludes Theorem 4.5
shift arguments to this perturbed model
D.2 Proof of Lemma
We start from equation 12 in Lemma A.1
We refer the reader to Appendix A.2 for the used notation
here, with additional notation we define here: for a subset of indices / C [| we write T = (Ti)ser to refer to subsequence of at positions 1, and T/I subsequence positions outside of I.
20

---

We can first bound (x Mr) as the following:
(o)
Σ II
[[[|=]]
VET
Em wit (Ilier Pil Phe (07441/77 .ti)) (Thеπиi7 Pil Phe (8r,107, 1-1) Phh (Ur, (ET())
- Em w (IlEz
Pm 10 (Ieelt 1/1
(8 Pm 1 1 P (V+, (VrT
The ( (Frm-11-ranll 2
Σ |[[[=]]
(1E/A)/((1AA)
where =(;:T) satisfies z(I;T); = 0 j and otherwise. Then we can observe that
Hherr 1 (Z (1
(1/A4
pu(viiit)) m((r. wer
(Sr44s (8r+1) P Hrr (Vrees (#
IIirr
pul (8t+1 (t1)t1) IIEE polos--iit--) (I (IT118T111 wt) (Vrz m()
m((e
(18)
- (s IT/ V7,)) (()
using inequality that can be derived by definition of l3:
p (8 I+1)
m(
p-1
Pm(aa U+ (xr.)
(M(V--17:-)) (17+4197,_11)))
m(Zr.37)
We are now ready to use this inequality bound the TV distance between trajectory distribution via equation (12). To do so, we rearrange the summation orders as follows:
Σ
Mi Dm=n wi I Pil (vdl2)) - EM= wht IT Pa (y mmn)
TIH
pJ (I 3r
Σ Σ
Σ
A(TT,
pu( (Sreta 18r. Pa (W (ETEE))
M [] me Mi] #t th mm (rrl)) mm
Pit (Zir Mr) Su( = ((1:T)) (S'r, teyss Merys
Σ
Σ
Σ
Im p (s',, It B7,1))
T H)] m mE|Mi] Fy Mr @rr Mr Jmm Icll
Pm (s Le, VT, 8 ( (1:T)) ( Teey Vry1))
Σ E Σ
Σ
( = (s #T, Ur,1)
T H) me
Mil 15 o
$ M -C(b)))) Σ
Σ
Σ TV [P((ATT((T))))
(1( (s'r) Bays Vrr))
(EdrC(H).I [g
S M. A.MC((pn
Σ
Σ
Σ
(IE(USTT p(( 17),
aZd H |T [C
proving Lemma
4.2
D.3 Proof of Lemma 4.4
Before proceeding to the proof, we refer the reader to Appendix A.3 for additional preliminaries. Recall the definition of Ue in Lemma A.3
V = arg max pu(Vito) = =s/o,, 5'), s, C v teet-
de V m
21

---

Note that the above definition is invariant to to.
Now consider Vtest miss- Then, VE is the of policies that maximize the probability po( (Stg+r 81%6, = s'), i.e., memoryless policy that aims to reach after & time steps (this policy is invariant to the starting state s'). Therefore, we can first induce that | S MHS. Then, the Markovian policy maximizing the probability reach certain state under fixed context m is also best within the history-dependens policy class II, and therefore
maxTrag Pm (Se4to = 81660 = = Tisto) S 1.
p(llmlsi R) == max max
m,,,8' maXpev, p(((vto) (sL- s/15 5²)
We can now invoke Lemma and Lemma A.3 and noting that
TV(PT. P2
(6 TTyp VT,)) S TV(PP, (T)
Let be be as defined in Lemma A.3 and for any π E II, we can conclude that
TV(PI,P5)() PT) s M ()
Σ
Σ
TV (P (P
P (
(T
Eubboog H.d) rell
S M (MHSA)
Σ
Σ
TV (PI pP( (T).
TESubeeq( H.d) =E(0,1)
Finally, only remains to bound the total variation distance with n(bbii T, 2). To see this, for given q S d, Subseq (H, and E {0,13/T easy check that
TV (pr(I)eiti))) 1 (T)
max TV (PI(WIT pP(W)))))) (T) S Cteat. (
Noting that Eqsa (#) < min (H 2#), assuming we are in the regime H >> d, we conclude that
TV(PT,P))( s M(MSA) (2H2)4 Etest:
concluding the proof.
D.4 Proof of Theorem 4.5
We continue from the conclusion of Lemma D.I and the remaining step to ensure that LMDP-OMLE terminates after K S log(1/7) iterations where 7 / (112d without the initial coverage
assumption 17 We consider perturbed model 0 ((wm, T R J. where
fm (-18, a) = (1->)T% (-18, a) + +71,
for all m. By the performance difference lemma [35], for any π E II,
TV(Pgg. PT))(T)= S Em w TV(PT ge Pz. (T S 2SyH.
For every tthh iteration, we check whether the coverage doubling argument (Lemma D.I still holds. To see this, first note that we can define ő( v/ tati т) and V.. vie and ((bbe T)) as in Lemma A.3 in terms of 0r. Then Lemma D.2 can be modified to guarantee that
Σ
Σ
TV² (pr(abeiT .p p(lleettt))) (T) nd Neest
16,3
+ (2HH))
+ HESubSeq(H.d) 2E(0.1)l=)
for all E C* using the triangle inequality for TV distance and (a + 6)2 < 2(a² + B2). Thus, we can derive (19 in terms of Ô- as
(2H)43
TV(PF g., PT)(T) SM(nA)(()))))
nd Mtest + (2H))(48H)22.
22

---

History (with Possible Contexts)
02
dg
2 =
Q1 = 1,
(m = 1)
41 = 2, = (m = 2)
E r2
"1 2, # (m 1 or 3)
r2
d1 = 1,71 # -1 (m = or 3)
r 0
Table 1: The first step generates of four possible belief's of history. Measured elements the table denote the expected rewards actions at following behavioral policy. We can see that all m E [M a it holds that IP m (82, a2) > for all m E [M d2 E A: yet, we cannot estimate E[ra] given some histories.
On the other hand,
TV(PPN (0 (V)) Poz) (T))(T) > 26TV -28H > 1.56rv:
max
6.
Now we set y Geor (8HnA)-2d(1)) (AMSH)) with Rrast 64M (4 /Lest and n MH52 and we have
26ty < 4 d. ((
teat T Creat
Hence the same doubling argument holds, and Algorithm 2 will terminate after at most
K OMM5 log( M HMM/StAaaa A//taa))))
iterations. We note that all inequalities hold for all K iterations with probability least 1-17. Finally, we invoke 4. éteet = é eroyp we can conclude that the total number trajectories we generated is bounded by
( H, M, S, crv 6/44SAHY and d = 2M - 1,
poly(M H, H, S, log (ogss
D.5 Counter Example for Remark
One may wonder why it is not sufficient to consider a single latent-state coverability analogous to Lemma 3.1 analogous to equation 2 defined as the following:
P PT = x)
max max max
[E #EA mE[M]
PR (x, = =x)
Here we present a counter example where the multi-step events must be considered even the latent
state space: the LMDP consists of M 3 MDPs with S {1, A = {1, R 0, 1)
and Hl = 2. All MDP starts from = and with a transition kernel Tw. (82 = 2[s1,0) =
for all m [M] and a A. Reward models are given by R (r -1ss 1,a = 1)
Ra(r = 1 = 1, a 2) = 1, and Rm (r|s = 1,a=1) = 0.5 for r E (0, 11. m E (2,3), and
(rls 1,a = 2) = 0.5 for E (-, m E (1,3).
Now we target to measure the expected rewards of actions executed by a behavioral policy at $2 = 2 as in Table [ In this example, all actions covered under all contexts following the behavior policy, i.e., (m 02) > 0. Yet, we cannot estimate the expected reward of action d2 = under context m 1. Therefore, the speculated single latent-state coverage coefficient finite for this problem, however, off-policy evaluation guarantee cannot be established only with the single latent-state coverability.
Appendix E Deferred Proofs
E.1
Proof of Lemma D.1
For any fixed checkpoint los recall the definition Vgg as defined equation
14
V
arg max ven
pelerto) (S1+ = =/10/ s'), Vm, s, s,
Note that [Vll < 2 and invariant to since V C Mmls- Now for any memoryless policy E 1mla. we recall Lemma 4.2 with the = (1h ... Wee Unif (A)), where the e 1I is mixture policy of L.g. We have
TV(PT.,P5)(T & MC
Σ
Σ
TV (pr())TT))) " p (T)
TESubeeq( H.d) xE (0,11))
23

---

s MC (211)..
(2
E
Σ
TV (poleera)) Po p
(T)
TEBubaeq (1))) =e (0,1)lm)
(2H)*3
S 4M(nAp(Vtost test: m))a
(19)
nd Dreat
where the last inequality follows by Lemma have
D.2
By the while-loop condition, for mm we
max (TV(P g- P2 (T), TV(P (T)) > 2 by the triangle inequality for TV distance. Therefore, we can conclude that leati x()) 16(2n
4e (Ap((
fleast
Plugging Reeett = 64M ( /eesst with = MHS2 we have 4dt < p( Vitat π ) = pl ((Viee -1 teeti )
Thus, p( (Wart m
> 2, which in turn implies Lemma D.1
tasti
E.2 Proof of Lemma D.2
By the construction of confidence set in equation ① and Lemma A.5 for all & [K] and @ E Ck. have
Σ (7) s 33,
(Tix)eDs
where 3 = l00g(K))l/)) Let 2is be a mixture of a subset candidate policies Ve With the = (N's, ve, te, Unif(A)) and for every 7 E
(th1, Vi2,.., Dn S Vd--1 test.
Subseq H, d), E 1), we have
TV (pHLRCTT pp p)((EtTti)) (T) g+
Σ
TV 1 or. un. (( Pp W (nn (a) (T)
nd Ee
Σ
TV2 Ph ( Und (T).
nd
Therefore,
Σ
Σ
TV (P((VeiTia(sT)) ( p (8 1//1://71.
r TESbbag H.d) [[[l(]
Σ
Σ
Σ
TV (1
pe p (01 (T)
nd
TEBubsoq(H.d) De (0.1))
33
nd -Teeat
where the last inequality is due to the construction of our dataset Di and the concentration guarantee for the confidence set Ck.
E.3 Proof of Lemma A.1
First note that rewrite atomic probability P. (ve(aa)) Pm (re, 8141/2)) as
(re, s141 pn. ($141)
pn.vo (st)
(rz, s1,1 - 1 ((r, + & ((0111s411)
pn. (8.)
pn. (s
In turn, moving from conditioning on the event (Ir, y.). we view the LMDP model after induction as if there are at most M1 1 contexts in the first LMDP model (or M2 - in the second LMDP model if (x,, ry; 8,) is attained with n = 2). We often use a shorthand 70) (*+, yt) = ( Te 8141) to reduce the burden on the notation.
24

---

Proof. The basic strategy is to apply iteratively the triangle inequality. To illustrate, we first expand the equation
EE
DM with Il
Pi (uel - Em w II
P (Mel (2))
THH
Σ
E
THS EM= wit (il
Pal (
P.l (111111)) 1 (21, v1) + (11. v1) I P. (vel
11.81 23.8 H MaIN
Σ
m=
Ma
with P2 (82)
P (o1) Pm
pu (s2))
(m2)
p (1111133 (1), + (x11
I Pa (ve(22))
ΣΣ
THe (01 vn)
Em wit P.l I Pa (vel EM= w m (s2)
IT== P (v 21)
21.81 22
N2:H
ΣΣ
E TIl Em w/ (s) (111 (x1, 91) P1) (.. I P.l. (vel
21.01 #2.8
-EMe wh (s)) (n)lar) - (2 y1 P (22))()2) I P (unlare)
(20)
We can continue applying triangle inequalities all possible first event-logging time step, and we can start with the following incquality:
ΣΣ
M
wit Il
Pal (uelor)) Em win II Pa (vel
TIl
Σ
Tin 1on---n111
TLE(H) marri Mhrr
Σ Te 11:8
#ntlil
Em
pii ( Un IH
Pl
Ph. (ull2))
Em Em=1 22
(Zr Vr, (2 P
I
Ph (v)l
Recall that least one phy 1211 is O, that is, one of the contexts in either of the systems is eliminated.
Now for (i), we can pick the next event-logging time step T2 > 71, and apply the triangle incquality similarly, and repeat such event-logging until all contexts are exhausted. Since there are at most 2.M contexts, cannot repeat this process more than d 2M times. Hence, we arrive to the following inequality:
ΣΣ
TI:H
ΣM
w.
I
Pil (Uellz) - EM w
I P2 (v.l
Σ
Σ (xT.Vr)x Σ
Σ
II
rESubaeq( H
20:r1-1
+ 1=0
Woirg
Mrje|
1::8
where we set 7 H 1. To proceed from here, we first observe that for any E [M] and i 2 0 such that p1 0,
11=++1 Pm (velae))
1 1 + & Aetto
pulliits))
Pm
(8 (8,+1)
11111 1- T(a Pm
S (a (TTr))
(1(()))))) (ST (7141)
pa (I + 1or Ur. 1:74 -1 (1
(21)
25

---

Thus, we can summarize that
EE TIH
Em w II Pu (vcc22)) - E My
win I P (ye (z))
P mixx fr (Tr Br, + -1
Σ E A(a+8++)) x
E Σ
II
pu(vitt) (g 1 の) m(z+
rESubsagl H.d) 2T.le
We note that each term in the product sequence is equivalent
P m((r, .8) (b //4
m(x+.87) (Sr, 187 11 1) mi(r ur) (Mrre (xT11))
and thus
|T|
p m(xr.ur (Ire 1ir Vos +e4 1
pulpsite) (Rr (87;+1) m((
1=0
mla (11:H Vi:H)
IT
pu( m(Z+.()))
(877 1x,+1) mm(x- (Nr. (mr))))
Here we set Tq+1 H + 1 and Ps (5H+11)) for any and conditional event. Since the denominator does not depend on events outside event-logging time-steps т, we can marginalize the probabilities in the inner summation and conclude the lemma.
E.4 Proof of Lemma A.3
Let us slightly extend the lemma such that we consider different sets of behavioral policies for different checkpoint time-steps.
Proof.
Note that for all m, s', s,,1,, t2. by the construction of 2/g. it follows that
pelvedt (st, = 5/3/, = =') 2 max
Pi
(st2 s|s s')
PEV;
We can observe that for any m,
pi ( ( Zr Vr,))
(87,41
p p (5,,15'_,, (Br., 1
A
iE!
On the other hand, for any x E T S [H], I [|T|], Zr. Vr.
PM (5 Tr (r S II max pa
(87,4165 Tis .)
II
max
P (1 Tie Pm (xT:)
Tim
sE/g//I
Applying the inequality with the definition of (( we have
p (5/ S "T,TT I. DT,VT y)
C(bE)))= = max max IS
max 3'E8 !I (xx)))( 111-11
max ME[M]
p (1:T)) (s s FT,1 x. VT,, = y)
maxT
Pa (874+110r. TIST.
< max
max
max 8' 12! ( 1n
max ME[M] iET
P P (ST44188/ ,)
max
po (s 5 1n TI TI:T-a)
A.
/ pp
(s,, |s/ (
KE(0)T
S (nA ((V
concluding Lemma A.3
26

---

E,5
Proof of Lemma A.4
This is by now standard MLE technique for constructing confidence sets RL
Proof.
The proof follows a Chernoff bound type of technique:
(T(T)
( (T) Po (7)
Po-
Σ
log
2 Eg.
log
B
P Pg-(T) (T)
(TT
(Mia)ED
(T) (T)
S Po.
exp
Σ
log
2 exp (3)
P (T)
(TT))DD
(T()) cxp(-3).
I Eg-
cxp
Σ
log
P.. (T)
(T.R)EDE
The last inequality is by Markov's inequality. Note that random variables are (T, T) in the trajectory dataset D, and
P (T) Py (T)
)
Eg-
Σ
log
(DPPP((")))))) s 0.
(AT))DDA
Therefore,
)
(T) Pz.
PT(T)
Eg-
exp
Σ
log
Eeg-
Σ Po(DD) =1.
(T)(T)
(F.T)EDE
DR
Combining the above, taking union bound over & E [K] rounds and E O, letting log(K((1/m) with probability 1 17, the inequality in Lemma A.4 holds.
E.6 Proof of Lemma A.5
Proof. By the TV-distance and Hellinger distance relation, for any i, T, x and E [M].
TV (P5(T), < 2H
(T))
(T( PT-
P7 1P-
- E ErrPE
[V
2-2 log
Erney.
To bound the summation over samples, we start from
(T( PT- (T)
Σ
H2 (P (
Σ (r.*)ED
log Erwpy.
(T.))EDs
On the other hand, by the Chernoff bound,
(T(T) P (7)
(T(T)
Paa-
Σ
log
Σ (T,
log Eropy. g"
B
PIT (7)
exp E(TT))EDE log
PI(T) PT-(T)
< Eg.
exp(-B)
exp [(TT))DD LogEEaPE
PH(T) F. (r)
PJ(T) (T,)) V
(T-(7)
= Eg-
cxp(-8)
PE(T) (T.(T)
I1 (l E.apy
27

---

Po(
IT (F,T))
Eg-
Ey-
(r)
D
exp(-33
Pole.)
11 (r,a)) EENPE
Pr- (1,7)
PG(T)
(
(T)
Eg-
exp(-3) exp(-8))
((,x)))D E.r....
Polc.) of (1.T)
where in the last line, we used the tower property of expectation. Thus, again by setting 3 = log(K' @l/n),, with probability at least 17, we have
Σ (P) (T.))ED
1
Σ
log
P (T) P (T)
3
(T.T)ED
(T)
(T)
Σ
log
Σ
log
B,
2
P
JP (T)
(T,T)ED
(T.
for all 1s [K] and E O. Now apply Lemma A.4 and finally have
Σ PT
1
Σ
log
(ī(T)
3 2.8
P (T)
(F.
(r.m)EDe
Since TV 2H², we get the lemma.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract present the problem consider in this work designing efficient learning algorithm for the RL setting and elaborate, in highlevel, on our new approach for solving this open problem.
Justification:
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: In the conclusion part we highlight the limitations of our work the main ones are: (i) our algorithm is optimal only up to polynomial factors and closing this gap is left future work, (ii) designing computationally and sample efficient algorithm, under some oracle assumptions, is an open problem which is left for future work.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
The complete proof is given in the Appendix. Further, we made substantial effort to provide intuition for the proof in the main paper: by providing analysis for the simpler MDP problem (and complementary analysis in the Appendix), as well as by connecting the analysis of this simpler setting the analysis of the LMDP setting.
Justification:
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [NA]
Justification: There are no experimental results in this paper, but resolution of algorithmic open problem.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: There are no experimental results in this paper, but resolution an algorithmic open problem.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: There are no experimental results in this paper, but resolution of algorithmic open problem.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification:
There are no experimental results in this paper, but resolution of algorithmic open problem.
29

---

8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [NA]
Justification: There are no experimental results in this paper, but resolution of algorithmic open problem.
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurIPS Code of Ethics https:
Answer: [Yes]
Justification: There are no potential harms caused by the research process or required mitigation measures that should have been taken for this work. Further, this work focuses on mathematical framework with no immediate societal impact or potential harmful consequences, to our opinion.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work focuses on establishing the learnability of the LMDP setting, was not established by previous works. The LMDP setting has been investigated in the past, we mentioned in the introduction section, by both empirical and theoretical communities. We do not see immediate socictal impacts of our work between the new results we derived, and the promise for improving algorithms for the LMDP, and POMDP settings future works.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer:[NA]
Justification: There are no experimental results in this paper, but resolution of an algorithmic open problem.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: There are experimental results that make use of data models in this paper, but resolution of an algorithmic open problem.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: There are no experimental results in this paper, but resolution of algorithmic open problem.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
30

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
31


Reviews:
[{"review_id": "vAUIxBdkoU", "reviewer": "Reviewer_jHm3", "summary": "The paper studies latent MDPs, an MDP framework with a set of MDPs and the environment samples a random MDP at the beginning of each episode. To avoid a $A^H$ sample complexity, previous works either assume separation or similarity of transitions. This work removes these conditions, and provide an algorithm with $(SA)^{O(M)}$ upper bound, which matches the lower bound with $M = O(1)$. Speficially, the algorithm adapts the information theoretic Optimistic MLE algorithm for the LMDP setting, by collecting data with all possible segments of all previous policies (+random actions). The key analysis tool is the LMDP coverage coefficient.", "questions": "see above", "limitations": "n/a", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper proposes the first algorithm that matches the lower bound of LMDPs (to poly order) without the assumption of separation and similar transition dynamics.\n\n2. The proposed LMDP coverage coefficient, although exponential in $d$ ($M$), seems to capture the right complexity of the LMDPs. \n\n3. The detailed and intutive introduce of Optimistic MLE algorithm makes the proposed algorithm easy to understand.", "weaknesses": "1. The techincal innovation seems rather limited -- the main techiques are quite similar as OMLE. \n\n2. I do not understand the relation to OPE: which part of the algorithm relates to OPE? It seems like the only change to OMLE is the way of data collection so there is no additional OPE subroutine?\n\n3. The \"without any additional structural assumption\" claim in the abstract seems inaccurate: even OMLE is limited with the assumption of finite SAIL (bli-linear rank), and the proposed analysis only applies to tabular? Also, does the analysis generate to general function approximation if we ues the correponding coverage coefficient?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for a thoughtful review and constructive feedback on our paper. We would like to start by emphasizing our technical novelty.\n\n\n**About Technical Novelties**\n\nOur LMDP-OMLE algorithm builds on the general framework of OMLE. The key novelty of the algorithm is the design of data-collection (exploration) policies (we also refer the reviewer to our general responses on the novelties). In designing the data-collection policy, we removed the need for core-tests and well-conditionedness. Further, in our analysis, we introduced the novel concept of LMDP coverage coefficient. These innovations result from efforts to formalize *moment-exploration*, as suggested in previous work [1]. Due to these, we were able to remove the restrictive assumptions of shared transition dynamics and separation.\n\n\nHence, our key results deviate substantially from existing approaches, and, hopefully, establish a new perspective on exploration in partially observed environments. Further, we believe our results contain a sufficient amount of novel concepts and analysis techniques.  \n\n\n\n\n**Off-Policy Evaluation for Online Guarantees**\n\nOur off-policy evaluation guarantee guides the data-collection policy of the LMDP-OMLE algorithm. In particular, the key novelty lies in the proposal of LMDP coverage coefficient (Definition 4.1): this quantity measures the data-coverage in terms of the visitation probability to *tuples* of different states and actions. Algorithmically, this suggests that our exploration algorithm should aim to visit \"uncovered\" tuples of state-action pairs, reaching one target state (and action) and moving to the next target state (and action) in the target tuple. Once we can design the exploration policy in such a way, we can show the sample-complexity upper bound via the coverage-doubling argument (Theorem 4.5) as we have illustrated in the MDP case.\n\n\nHere we note that establishing online guarantees via off-policy evaluation is just one possible approach for exploration in LMDPs. It would be very interesting to see if other approaches such as going through the problem-specific complexity measures such as Eluder-dimension [2], Bellman-rank [3] or Decision-Estimation Coefficient (DEC) [4] can be applied using the results provided in this work.\n\n\n\n\n**Why Our Approach Doesn't Make Any Assumptions?**\n\nWe clarify that this work studies LMDP \"without any assumptions\" in the sense that we consider the class of *all possible LMDP models*, unlike in previous work that assumes certain separations [5,6] or shared transitions dynamics [1].\n\n\n-------------\n\nWe hope our responses have clarified any doubts or questions. If there are any remaining issues, please let us know. Otherwise, we would highly appreciate it if the reviewer could consider raising the score based on our responses.\n\n\n\n\n[1] Kwon et al., \"Reward-mixing MDPs with few latent contexts are\nlearnable\", ICML 2023\n\n[2] Russo and Van Roy, \"Eluder dimension and the sample complexity of optimistic exploration\", NeurIPS 2013\n\n[3] Jin et al., \"Bellman eluder dimension: New rich classes of RL problems, and sample-efficient algorithms\", NeurIPS 2021\n\n[4] Foster et al., \"The statistical complexity of interactive decision\nmaking\", arXiv 2021\n\n[5] Hallak et al., \"Contextual markov decision processes\", arXiv 2015\n\n[6] Chen et al., \"Near-optimal learning and planning in separated\nlatent MDPs\", COLT 2024", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "4Ucx7laqFu", "reviewer": "Reviewer_8u1Y", "summary": "This paper introduces a new version of the coverage coefficient for analyzing latent Markov Decision Processes (MDPs). It demonstrates how to link the proposed coverage coefficient with sample complexity using MDPs. Additionally, the paper presents an algorithm and provides a bound on the sample complexity of this algorithm.", "questions": "See \"Weaknesses\" part for the questions.", "limitations": "See \"Weaknesses\" part for the limitations.", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "The use of MDPs to illustrate the concept is effective, making it easier for the audience to follow.", "weaknesses": "1. The comparison with related work is insufficient. Including a table that compares the results in this paper with those in [1] would be helpful.\n\n2. The connection and comparison between this work and existing studies on the coverage coefficient [2] should be more concrete.\n\n3. The meaning and design of the segment policy are unclear, making it difficult for the audience to understand the intuition behind why this coverage coefficient is helpful for analysis. Although a counterexample is provided in the appendix, it is still not intuitive enough.\n\n4. In my understanding, the core idea behind the analysis is that complexity depends on the longest policy sequence $\\pi_1,\\ldots,\\pi_k$, where $C(\\pi_i,\\pi_j)=\\infty$. If this understanding is correct, there are two follow-up questions:\n(a) Can we define a complexity measure similar to the Eluder dimension?\n(b) Can the complexity measure generalize to cases beyond latent MDPs?\n(c) Is this complexity measure equivalent to or weaker than the Eluder dimension?\nAnswers to these questions might help improve this paper.\n\n5. Although the results in this paper do not require further structural assumptions, the significance of an exponential upper bound is questionable.\n\n6. Lack of discussion with works that have complexity measure without Markovian assumption [3].\n\n[1]. J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reinforcement learning in reward-mixing MDPs. Advances in Neural Information Processing Systems, 34, 2021.\n\n[2]. P. Amortila, D. J. Foster, and A. Krishnamurthy. Scalable online exploration via coverability. arXiv preprint arXiv:2403.06571, 2024.\n\n[3]. Zhong, Han, et al. \"Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your insights and suggestions, which will help us improve our work. Below, we address the weaknesses you mentioned.\n\n\n**Comparison to Previous Work on LMDPs**\n\nWe would like to highlight that our work is the first to propose a general exploration algorithm applicable to the entire class of LMDP models. In contrast, previous works [1] focus solely on the subclass of LMDPs known as \"Reward-Mixing\" MDPs. These models assume no context ambiguity in transition dynamics, only in reward models. Under this assumption, data collection becomes less challenging. One can first learn the transition model separately through reward-free exploration, and then explore specific state-action tuples in order to obtain samples of \"moments\". However, when transition dynamics also vary across contexts, data-collection becomes significantly more complex, and is the main challenge in our work.\n\nWe will further clarify this distinction and the associated challenges in Section 1.1 - Challenge 2 in our revised manuscript.\n\n\n\n**Coverage Coefficient in MDPs vs LMDPs**\n\nExisting work that studies \"the role of coverage in online exploration\" is focused on fully observed settings (Low-Rank MDP, Block MDP are also fully-observed, although latent dynamics is much lower dimensional). In contrast, LMDP, and more broadly POMDP, involve dynamics and policies that are history-dependent -- an optimal policy of an LMDP may depend on the entire history. This means that a single latent-state coverability does not adequately capture the complexity of offline learning in these settings (as we illustrated in the counterexample in Appendix). \n\nConsequently, off-policy evaluation in LMDPs requires a new notion of coverage that measures coverability over *sequences of state-action pairs*. Our main contribution is the proposal of this LMDP-coverage concept, which we connect to off-policy evaluation and online exploration. We believe this not only represents a significant advancement in learning LMDPs but also offers a conceptual contribution to the learning of general POMDPs. We will make this point clearer and more accessible to readers in our revision. \n\n\n\n\n\n\n**Eluder Dimension vs Our Approach**\n\nThe Eluder dimension essentially measures the longest possible policy sequence that can incur a large prediction error. However, in partially observed settings, this definition may fail to capture the true complexity of online exploration. For instance, in multi-step weakly-revealing POMDPs (or PSRs), without executing core-tests, we may still suffer from the curse of horizon [2].\n\nOne might modify the definition of the Eluder dimension as proposed in the GEC paper you shared (details will follow below). However, similarly to the limitation of known tractable POMDP classes, the GEC definition also relies on the prior knowledge of how to construct the data-collection policies (using core-tests). Our technical contribution lies in the design and understanding of data-collection policies in LMDPs. We believe there is significant potential to develop a more general, fine-grained notion of complexity for a broader class, potentially including both revealing POMDPs and LMDPs. We see this as an exciting direction for future research.\n\n\n\n\n\n\n**Upper Bounds Exponential in M?**\n\nWe accept this criticism, though due to the lower bound established in [3] we cannot avoid this dependence in M without any assumptions. However, we would like to highlight that our results represent a substantial advancement from previous work. Specifically, our algorithm is the first to achieve sample efficiency in a partially observed setup without relying on standard assumptions such as weakly-revealing or low-rankness, which are well-studied in the literature. Even in cases where $M=2$ or $3$, the solutions to LMDPs were previously unknown, making our findings significant. While we acknowledge that our current results are achieved with $M=O(1)$, we believe it is important to emphasize the positive aspect of our contributions in pushing the boundaries of what is known in this field (we also refer the reviewer to our general response on the importance of studying LMDPs).\n\n\n\n\n\n**Comparison/Connection to Other Complexity Measures (GEC)**\n\nWe thank the reviewer for drawing our attention to this work, and we will add this reference in our revision. Yes, the idea behind GEC resonates with us at a very high-level -- it aims to capture the largest discrepancy in prediction errors from small training errors. However, as with other complexity measures such as (generalized) Eluder-dimension or Decision estimation coefficient, the main bottleneck is to show the boundedness of GEC in LMDPs, as there has been no upper bound established up-to-date. Furthermore, our construction of exploration (data-collection) policy is very different from known methods. \n\nWhile all the above is in part described in our Section 1.1, we will be more explicit about comparison to existing complexity measures in our revision. \n\n\n-------------\n\nWe hope our responses clarify any misunderstandings and resolve the issues identified. Please let us know if there are any remaining questions or concerns. Otherwise, we kindly request a reevaluation of our work in light of the provided clarifications. Thank you for your consideration and effort.\n\n\n\n[1] Kwon et al., \"Reinforcement Learning in Reward-Mixing MDPs\", NeurIPS 2021\n\n[2] Chen et al., \"Lower bounds for learning in revealing POMDPs\", ICML 2023\n\n[3] Kwon et al., \"RL for Latent MDPs: Regret Guarantees and a Lower Bound\", NeurIPS 2021", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "jOBrenug4N", "reviewer": "Reviewer_AN8Q", "summary": "This paper studies latent MDP, where the underlying dynamics and rewards are controlled by some latent states (not revealed to the learner), and the learner attempts to do learning and planning based on the trajectory data. \n\nThe algorithm builds from the optimistic MLE algorithm, which iteratively checked whether whether there exists uncovered policies. To make the OMLE works for LMDP, this paper constructed segmented policies, and iteratively check whether such policies are all covered. \n\nAs a result, they showed the first polynomial sample complexity result for LMDP when the number of latent states are constant.", "questions": "Can these results somehow adapt to the POMDP setting?", "limitations": "Yes. The authors addressed all the limitations listed in the guidelines.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper is well written. The algorithm, theorems and proofs are clear.\n\nThe results of polynomial sample complexity for LMDP with constant latent states seem very interesting. \n\nThe idea of constructing segmented policies is novel.", "weaknesses": "I don't see any significant weaknesses of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the encouraging comments and positive assessment of our paper. Below, we address an interesting question you raised:\n\n**Can these results somehow adapt to the POMDP setting?**\n\nYes, we believe our results can be adaptable to the POMDP setting in two senses:\n\n- Algorithmically speaking, our algorithm overall falls in a general framework of OMLE; the key difference is in the data-collection policy part which may differ with different environments and assumptions. But at a higher-level, they can be viewed in a unified principle.\n \n- Technically speaking, our analysis based on the notion of coverage can be useful to give a different interpretation of solving the POMDP settings (that are known to be tractable). The key to understand in such a way is to define a suitable notion of coverage in general POMDPs -- which is in part done in some recent work on off-policy learning for POMDP [1]. This is an interesting future question. \n\n\n\n\n--------------------\n\nThank you once again for your thoughtful review and positive assessment. If you have any further questions or suggestions, please let us know.\n\n[1] Zhang and Jiang, \"On the curses of future and history in future-dependent value functions for off-policy evaluation\", arXiv 2024", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "FAZ9eNVjs6", "reviewer": "Reviewer_Ch5D", "summary": "This paper studies latent Markov decision processes (LMDP) with M = O(1) number of MDPs. In other words, there exists a set of MDPs (unknown to the agent) and the environment randomly selects one MDP at the beginning of each episode. The selected MDP is not revealed to the agent and therefore the agent must infer which MDP is selected from the feedback. This paper proposes an algorithm with sample complexity poly(S,A)^M that matches the \\Omega((SA)^M) lower bound.\n\nOn the technical side, this paper designs a model-based exploration algorithm that actively collects new data if the collected data can shrink the confidence interval of the model parameters. On standard tabular RL, this paper shows that (a) the algorithm terminates after poly(SAH) rounds of exploration, and (b) any remaining model parameter can be used to construct a near-optimal policy after the exploration stage. Because of the generality of the algorithm, this paper further extends the algorithm to the LMDP setting and proves a poly(S,A)^M sample complexity.", "questions": "-\tIt seems to me that the OMLE algorithm is very general. Is it straightforward to prove some general sample complexity results with problem-dependent bounds? Intuitively, the bound would include terms like the statistical complexity of the model class and some complexity measure of the MDP family. \n-\tThe data collection step in Algorithm 2 (Line 5-7) seems to be quite brute-force since it has to enumerate all the segmented policy. Could the authors elaborate on its reason?", "limitations": "The authors adequately addressed the limitations and potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-\tThe exposition of the main ideas is very clear because of the instantiation on tabular MDPs.\n-\tThe OMLE algorithm is neat, and it provides a general framework for model-based exploration. \n-\tThe sample complexity result matches the \\Omega(S,A)^M lower bound and requires no additional assumption on the structure of the LMDP family.", "weaknesses": "-\tThe algorithm is not computationally efficient even for standard tabular MDPs because it has to enumerate both the parameters and the policy in the confidence interval, though I understand that computational efficiency is not the focus of this paper.\n-\tSection 4 is hard to follow without background knowledge of several prior works. E.g., the choice d = 2M-1 is not well-justified. I encourage the authors to revise Section 4 to make it more self-contained.\n-\tIt seems to me that the title of this paper somewhat overclaims the result since the sample complexity is still exponential in M, and the assumption that M is finite is not intrinsic to the LMDP setup (although it is necessary due to the sample complexity lower bound).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review and constructive feedback on our paper. We address the mentioned weaknesses below. \n\n\n**Practicality of the Proposed Algorithm**\n\nAs the reviewer correctly pointed out, our focus was indeed on the purely statistical learning aspect of the problem. Nevertheless, we acknowledge and appreciate the reviewer's perspective on the computational aspects of the algorithm. Our long-term goal is to design an oracle-efficient algorithm for LMDPs that can be implemented with general function classes and ERM-style oracles. While we believe this is a feasible direction, achieving it will require significant effort from both theoretical and practical standpoints.\n\n\n**Intuition on Policy-Segmentation Regarding Prior Work**\n\nThank you for your suggestion. Indeed, we agree that it would be useful to provide the intuition on how we connect moment-exploration ideas presented in [1] to our off-policy evaluation results. We will add the following paragraph in our revision at the beginning of Section 4:\n\n*Before we dive into our key results, let us provide our intuition on how we construct the OPE lemma for LMDPs. Our construction is inspired by the moment-exploration algorithm proposed in [1]: when state-transition dynamics are identical across latent contexts, {\\it i.e.,} $T_1 = T_2 = ... = T_M$, we can first learn the transition dynamics with any reward-free type exploration scheme for MDPs [2], and then set the exploration policy that sufficiently visits certain tuples of state-actions of length at most $d=2M-1$. For the exploration policy, the work in [1] set a memoryless exploration policy $\\psi \\in \\Pi_{\\texttt{mls}}$ which sets the visitation probability to certain tuples sufficiently large. We note that the same moment-exploration strategy cannot be applied to general LMDPs with different state-transition dynamics since learning the transition dynamics itself involves latent contexts. Nevertheless, the intuition from [1] suggests that our key statistics are this visitation probabilities to all tuples of state-actions within a trajectory.*\n\n\n\n**Significance of Our Results on $M$**\n\nWe accept this criticism and appreciate the feedback. That being said, we would like to highlight that our results represent a substantial advancement from previous work on LMDPs. Specifically, our algorithm is the first to achieve sample efficiency in a partially observed setup without relying on standard assumptions such as weakly-revealing or low-rankness, which are well-studied in the literature. Even in cases where $M=2$ or $3$, the solutions to LMDPs were previously unknown, making our findings significant. While we acknowledge that our current results are achieved with $M=O(1)$, we believe it is important to emphasize the positive aspect of our contributions in pushing the boundaries of what is known in this field.\n\n\n\n\n\n\n**Possible Extensions to General Function Classes**\n\nThis is a great question. We assume that by \"problem-dependent bound\", the reviewer is referring to sample-complexity bounds with general function approximation. We believe that extending our approach to general function approximation is feasible, particularly given recent advances in RL; however, this is not entirely straightforward and requires careful consideration. For Latent MDPs specifically, our starting point will be to define what is the good notion of \"coverage\" in LMDPs with large state/action spaces -- that will direct us to the question of what are the \"moments\" with general function approximation. If we can define this notion well, then we believe that we should be able to come up with a notion similar to the LMDP-coverage with function approximation. Another possibility is to come up with a general notion of statistical complexity measure for Latent MDPs/POMDPs such as DEC [3] (see also our discussion on GEC with Reviewer 8u1Y). Such an extension is beyond the scope of this paper, and we think it is an exciting future direction. \n\n\n\n\n\n\n**More Efficient Ways to Construct Data-Collection Policies**\n\nThis is another great point. The current brute-force construction is due to the nature of worst-case analysis, where joint events at all length-$d$ tuples of state-actions need to be observed under all contexts. Since we do not have the observability over which policy drives us to the next state-action pair under latent contexts of interest, our solution brute-force all possibilities. \n\nThat being said, we are hopeful that under some practical assumptions, we may expect improvement in two ways: \n\n- We can rule out some redundant combinations, as they might already be covered by other combination (e.g., this idea can be related to G-optimal design as in the Latent Multi-Armed Bandit setting [4]) \n\n- Other possibility is to identify that for some combinations, at some point, it becomes obvious that the segmentation of significantly smaller than $d$ segments would be sufficient. \n \nThese are the next big questions and important future work. \n\n\n\n------------------------\n\n\nWe hope that our answers to the raised questions are satisfactory. Please let us know if you have any other concerns or questions. We would greatly appreciate it if you could consider reevaluating our work, taking into account the strengths and improvements we've outlined.\n\n\n[1] Kwon et al., \"Reward-mixing MDPs with few latent contexts are learnable\", ICML 2023\n\n[2] Jin et al., \"Reward-free exploration for reinforcement learning\", ICML 2020\n\n[3] Foster et al., \"The statistical complexity of interactive decision making\", arXiv 2021\n\n[4] Kwon et al., \"Tractable Optimality in Episodic Latent MABs\", NeurIPS 2022", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their effort, time, and their valuable feedback. While we respond to each reviewer on the specific concerns raised, we would like to emphasize the importance of studying the Latent MDP setting as well as the technical novelties in our work. 



**Why Solving LMDPs is Important:** 


It is often useful to model the population of environments as a mixture of simpler distributions. Latent MDP is an interactive version of mixture modeling, and has a potential for many real world problems, *e.g.,* in dialogue, recommender or in healthcare systems, when complete information on a user or patient is not given. Nevertheless, no algorithm has been known beyond the setting when clustering is relatively straight-forward. We believe that studying general cases when clustering is not so straight-forward has many potentials to advancing the field both in theory and practice, as we detail below. 




- *LMDP requires rethinking exploration strategies in partially observed environments:*  Several known approaches for general POMDPs overcome the challenge by assuming (1) the well-conditionedness of the system (latent state-observation emission matrix must be in full-rank), and (2) the prior knowledge of core-tests (those policies that guarantee the full rankness of the state-observation matrix). Not only does this require strong domain knowledge, but also there is no known approach to learn the optimal policy beyond such cases. 



- *Potential to enlarge the scope beyond known tractable POMDP classes and general approaches:* Our work is the first to propose the sample-efficient algorithm for an important and broad class of POMDPs that do not provide the common assumptions of well-conditionedness and known core-tests. Furthermore, the proposed LMDP-OMLE algorithm is designed with as much flexibility and generality as possible without specific domain knowledge and assertions, unlike previous solutions that rely on certain clustering-enabling or separation assumptions. We believe our result paves the way to push existing theories further to classes of other potentially tractable POMDPs (in particular, to general function approximation with large state/action spaces in POMDPs). 








**Technical Novelties in Our Work** 

While we follow the principles of Optimistic Maximum Likelihood Estimation (OMLE) [1], we emphasize that our key novelty lies in the algorithmic design and analysis of data-collection policies given the confidence set, without the given knowledge of core-tests. Our approach is the first that goes beyond the known assumptions for POMDPs, and integrates the method-of-moments in the sequential decision-making setup. In particular, we formalize the concept of moment-matching in our off-policy evaluation lemma (Lemma 4.2) by introducing the visitation probability to tuples of different state-action pairs and learning their joint probabilities (this captures the correlations and moments of the system). Utilizing this off-policy evaluation lemma, we derive a sample-efficient algorithm through our coverage-doubling argument, detailed in Lemma D.1 of the Appendix. Although the mathematics involved is not particularly fancy and is purely algebraic, we believe that the conceptual and algorithmic innovations are substantial.


[1] Liu et al., "Optimistic MLE: A Generic Model-Based Algorithm for Partially Observable Sequential Decision Making", STOC 2023

Meta-review: The paper presents a novel approach to latent Markov decision processes (LMDPs). The work's primary contributions include the introduction of the LMDP coverage coefficient and adapting the OMLE algorithm to this setting, offering a detailed analysis of sample complexity in tabular settings. Reviewers appreciated the novelty of the segmentation policies. However, some raised concerns about the limited technical innovation compared to existing methods like OMLE and a lack of sufficient comparison with related work. Despite these concerns, most reviewers lean towards acceptance.
Decision: Accept (poster)