Title: RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluation
Venue: NeurIPS 2024
Accepted: yes
Keywords: Latent Markov Decision Process

Online Exploration

Off-Policy Evaluation

Partially Observed Systems

Sample Complexity

Information Theory
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:

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. In the last decade, there has been significant progress in solving LMDPs under different structural assumptions. However, for general LMDPs, even in the tabular case, no algorithm is known to provably match the existing lower bound [41]. We introduce the first sample-efficient algorithm for LMDPs without any additional distributional assumptions . Our result builds off a new perspective on the role of offpolicy evaluation guarantees and coverage coefficients in LMDPs, a perspective, that 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 LMDPs, and especially, for partially observed environments.

1 Introduction

In Reinforcement Learning (RL) [54], an agent aims to maximize the long-term cumulative rewards through interactions within an 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 on the Markovian assumption that requires the state to be fully observable. However, many real-world decision problems involve critical partially observed or latent information, such as sensitive or 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 a 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, 11, 21, 41, 37] within the framework that we here collectively refer to as Latent Markov Decision Processes (LMDP) following [41]. In LMDPs, nature selects an MDP from a 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 an episode (the horizon). However, the identity of the chosen MDP is not given to the agent. We call this unknown identity the latent context .

Most prior work on LMDPs has relied on strict separation assumptions ( e.g., [11, 24, 37]). The applicability of these approaches is limited to scenarios where the horizon is sufficiently large and identification of the latent model can be guaranteed, i.e., H ≫ Ω( 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 Ω( A H ) - 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].

Recently, a series of works [40, 42, 43] proposed sample efficient algorithms without separation assumptions when M = O (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 LMDP setting with different transition dynamics (see Section 1.1). Consequently, to date, the following question has remained open:

Can we break the curse of horizon in LMDPs if M = O (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 > 2 M , and focus on whether we can improve the trivial upper bound that incurs complexity Ω( A H ) . Since a Ω( SA ) M lower bound for LMDPs has been established [41], our goal is to achieve an upper bound of poly( S, 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 a confidence set - a set of candidate models (hypothesis) that can explain the observed data with high probability - and execute a policy that will shrink the volume of the confidence set is produced and executed [6, 38, 27, 7, 33, 36]. The entirety of the statistical problem is to analyze the decaying rate of confidence sets under proper model class assumptions [49, 34, 19].

Challenge 1: Limitation of Existing POMDP Algorithms. Existing approaches for online exploration 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, 8, 16, 17, 45, 57, 13, 22, 46, 26]. Without specifying the proper core-tests, the volume of confidence sets may not decay in a desired rate, leading to the curse of horizon Ω( A H ) [39, 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 (ii) it assumes the full-rankness of the state-observation emission matrix when the core-tests are being executed.

While LMDPs are a special class of POMDPs, neither the existence of a set of core-tests is known a priori , nor it is possible to recover the belief over latent contexts from distribution of trajectories (see Section B for details). This creates a fundamental challenge for existing approaches when applied to LMDPs. Further, little is understood on learning a near-optimal policy among 'doubly-exponential' number of candidate history-dependent 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-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 := 2 M -1 state-action pairs x = ( ( s [1] , a [1] ) , ..., ( s [ d ] , a [ d ] ) ) . This in turn suggests that the exploration algorithm must learn how to visit these lengthd state-action tuples simultaneously, i.e., find a policy that ensures that x appears as a subsequence of the entire trajectory with high enough probability.

When the transition dynamics of different latent contexts is similar, reaching optimally to d state-action pairs is a minor challenge; e.g., we can first learn the shared transition kernel with any reward-free

Wtest

←Add a new test policy Tt:

Genarate

e0(e-2)

sampleswithnew

30 E C, TV(P, P) >∈

segmented policies v E Ilexp(Ytest)

UpdatetheconfidencesetC

Add newsample

(loss constraints on dataset)

trajectoriestodataset

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 in 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 are not equal, is impossible, since even if we obtain the samples of correlations, different policies may result in 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 settings [30, 34, 19]. 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 to 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, 2, 29, 9, 55, 3, 4]. 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 a sample-efficient algorithm for the LMDP setting. This further showcases the usefulness of OPE for online exploration in POMDPs.

Arguably, the fundamental question in OPE is the following: how much does a behavioral policy ψ tell about a target policy π ? The simplest form of the OPE guarantee in MDPs relies on the notion of coverage coefficient given by:

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 ( ψ ; π ) = ∞ implies there exists a state-action pair, at some time-steps, that cannot be reached under ψ , but can be reached with π .

The algorithmic framework we 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, we are required to develop a new notion of coverage coefficient and new OPE tools. We propose a coverage coefficient that can be informally described as follows:

where m is the unobserved latent context, T = ( s t , a t , r t ) t ∈ [ H ] is a sampled trajectory, E 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 a at the t th time step regardless of ψ (for the formal definition, see Definition 4.1). Note that the coverage coefficient cannot be 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 4.2):

where TV ( P 1 , P 2 )( · ) is the total-variation (TV) distance between two probability measures P 1 , P 2 .

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 1 for highlevel illustration), repeats the following: (i) find a 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 a set of (obtained) test policies and interventions, an exploration strategy for LMDPs that we introduce.

2 Preliminaries

We consider an 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 , θ, H ) with a state space S ; action space A ; reward space R , and a finite-time horizon H . θ is a model parameter consisting of multiple MDPs in the model θ := ( { w m , T m , R m } ) M m =1 . In each m th MDP, T m : S × A × S → [0 , 1] maps a state-action pair and a next state to a probability; R m : S × A × R → [0 , 1] is a probability of rewards; { w m } M m =1 are the mixing weights such that at the beginning of every episode the m th model is chosen with probability w m .

Without loss of generality, we assume that there exists a null state that represents the starting and terminal state s 0 = s H +1 = ∅ , and a null action at the beginning of an episode a 0 = ∅ , even though actual policies do not take any action at the beginning. T m ( ·| s 0 , a 0 ) is the initial state distribution of the m 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 H > 2 M . Further, we assume the reward values are finite and bounded:

Assumption 2.2 (Finite and Bounded Reward) The reward distribution has finite support with (arbitrarily large) cardinality, and each reward is bounded: | r | ≤ 1 for all r ∈ R .

We also note that this concept can be easily generalized to instantaneous observations that include rewards, and thus, we do not lose much generality due to Assumption 2.2. We consider a policy class Π which contains all history-dependent policies π : Ξ × ( S , A , R ) ∗ × ( S × [ H ]) → ∆( A ) , where Ξ is the space of independent variables decided at the beginning of execution. As a special case, we consider the class of memoryless policies: Π mls : ( S × [ H ]) → ∆( A ) We are interested in finding an optimal history dependent policy π ∈ Π that maximizes the expected reward: V ∗ θ ∗ := max π ∈ Π E π θ ∗ [ ∑ H t =1 r t ] , where θ ∗ ∈ Θ is the true model parameter and E π θ ∗ [ · ] is expectation taken over the true LMDP model M ∗ when policy π is executed.

Notation We use [ n ] := { 1 , . . . , n } and [ n ] + := { 0 } ∪ [ n ] . We define d := 2 M -1 and assume H > 2 M . Let SubSeq ( H,d ) be the set of subsequences of (1 , 2 , ..., H ) with length less than or equal to d , i.e., SubSeq ( H,d ) := { ( τ 1 , τ 2 , ..., τ q ) | q ∈ [ d ] , 1 ≤ τ 1 < ... < τ q ≤ H } . We often denote a state-action pair ( s, a ) as one symbol x = ( s, a ) ∈ X = ( S × A ) , and an reward-next state pair ( r, s ′ ) as one symbol y = ( r, s ′ ) ∈ Y = ( R×S ) . We often express the next state at time step t as either s t +1 or s ′ t , and the pair of instantaneous observation and next state as y t = ( r t , s t +1 ) = ( r t , s ′ t ) . For any segment of a sequence ( z 1 , z 2 , ..., z H ) from t 1 to t 2 , we often simplify the notation as z t 1 : t 2 . We denote the entire trajectory as T := ( s, a, r ) 1: H , and T 1: t = (( s, a, r ) 1: t -1 , s t ) for a history of length t . For any set S , we define S ⊗ k as a short-hand for the k -times Cartesian power of S .

Algorithm 1 MDP-OMLE

1: Input: n test ∈ N , β, ϵ TV , η > 0 , C 0 = Θ

2: Initialize k = 0

3: while there exists π k ∈ Π mls , and θ 1 , θ 2 ∈ C k such that TV ( P π k θ 1 , P π k θ 2 ) ( T ) > 4 ϵ TV do

4: Generate data {T k j } n test j =1 by executing π k , update D k ←D k -1 ∪ { ( T k j , π k ) } n test j =1

5: Refine the confidence set with the dataset:

6: end while

7: Pick any θ ∈ C k and return the optimal policy of M := ( S , A , O , θ ) .

We define SubTraj ( T , τ ) ⊆ ( X × Y ) ⊗ | τ | as a valid subsequence of trajectories at time-steps τ ∈ SubSeq ( H,d ) , i.e., if ( x τ , y τ ) ∈ SubTraj ( T , τ ) , for any i such that τ i = τ i +1 , y τ i = ( r τ i , s ′ τ i ) and x τ i +1 = ( s τ i +1 , a τ i +1 ) must have s ′ τ i = s τ i +1 .

For a tuple of state-action pairs (or states) of length q , we denote x = ( x [1] , ..., x [ q ] ) (or s = ( s [1] , ..., s [ q ] ) with bracketed indices for each element to distinguish from time steps. We use | x | for the length of sequence x . We denote the cardinality of the state and action space as S := |S| and A := |A| . For any two models θ 1 , θ 2 , we often denote P 1 ( · ) := P θ 1 ( · ) and P 2 ( · ) := P θ 2 ( · ) whenever the context is clear. We denote P m ( · ) for a probability measured conditioned on the context m ∈ [ M ] over the ground-truth model ( θ 1 when we compare θ 1 and θ 2 ). We denote Unif ( A ) as the uniform distribution over a set A . Let TV ( P 1 , P 2 )( X ) be the total-variation distance between two probability measures P 1 ( · ) , P 2 ( · ) over a random variable X .

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-OMLE depicted in Algorithm 1. MDP-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 a policy for which the TV distance between trajectory distributions of models in 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 its application, we present the result.

Lemma 3.1 (TV Bound via OPE for MDPs) For any behavioral and target policies ψ, π ∈ Π , let the coverage coefficient be defined by:

For any two models θ, θ ∗ ∈ Θ , the TV distance between trajectory distributions following a target policy π ∈ Π is bounded as follows:

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 at each iteration by executing a policy π k on the

k th iteration, i.e., n test = ∞ . Further, assume the algorithm is at the beginning of its k +1 iteration. In the infinite sample regime all models in the confidence set must have matching event distribution relatively to the underlying model measured when policy π k is tested. Specifically, for all θ ∈ C k and t ∈ [ H ] it holds that TV ( P π k θ ∗ , P π k θ )( x t , y t ) = 0 . Then Lemma 3.1 implies the following: for all policies π for which C ( π k ; π ) < ∞ it also holds that TV ( P π θ ∗ , P π θ )( T ) = 0 . Conversely, assume the condition of the while loop at the beginning of the k +1 iteration holds true, namely, there exists a policy ¯ π for which TV ( P ¯ π θ ∗ , P ¯ π θ )( T ) > 0 . Then Lemma 3.1 also implies that C ( π k ; ¯ π ) = ∞ , namely, there exists an x ∈ X and t ∈ [ H ] such that P ¯ π θ ∗ ( x t = x ) > 0 whereas P π k θ ∗ ( x t = x ) = 0 . Next, recall that MDP-OMLE sets the data collection policy at the k +1 iteration to be π k +1 = ¯ π . Hence, the data collection policy at the k +1 iteration will visit some state-action pair at some time step π k did not visit. Hence, in the infinite sample regime, MDP-OMLE halts after at most HSA iterations, as there are at most HSA different state-action pairs in different time steps.

The intuition presented above is robust to sampling error, i.e, when n test < ∞ . To simplify the discussion, let us temporarily assume that P π θ ∗ ( x t = x ) > γ for all π ∈ Π and x ∈ S × 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 k > 0 in Algorithm 1, there exists at least one t ∈ [ H ] and x ∈ X such that

with some absolute constant c > 0 .

Therefore, by setting the number of samples to be n test ≥ (4 H 2 SAβ ) / ( cϵ TV ) 2 , we ensure that in every iteration MDP-OMLE doubles the coverage of at least one state-action pair at a certain time step. Therefore, the algorithm terminates within at most K = O ( HSA · log(1 /γ )) iterations with high probability . After termination, we are guaranteed that any two models in the confidence set are ϵ TV -close in TV-distance for any policy, hence we can obtain ϵ = ( Hϵ TV ) -optimal policy. To summarize, we state the following theorem:

Theorem 3.3 Let K = O ( HSA ) log( HSA/ϵ ) and β = log( K | Θ | /η ) . Then, with probability at least, 1 -δ , MDP-OMLE terminates after K iterations with at most N episodes being generated, where

and outputs an ϵ -optimal policy with probability at least 1 -η .

In a typical tabular MDP setting, we take O (log | Θ | ) = ˜ O ( SA ) , by discretizing the class of MDPs. Hence the sample complexity of MDP-OMLE is N = ˜ O ( H 6 S 3 A 3 /ϵ 2 ) . 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 1.1, Challenge 1).

4 Efficient Exploration in LMDPs

In previous section we presented a new approach to analyze the OMLE algorithm for MDPs. Next, we develop an analogous technique for the LMDP setting and design the LMDP-OMLE algorithm. Central to its design and analysis is an OPE lemma and a new coverage coefficient which we now present.

Intuition from moment-exploration algorithm in [43]. 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 [43]: when state-transition dynamics are identical across latent contexts, i.e., T 1 = T 2 = ... = T M , we can first learn the transition dynamics with any reward-free type exploration scheme for MDPs [33], and then set the exploration policy that sufficiently visits some tuples of state-actions x of length at most d . Specifically, they set a memorlyess exploration policy ψ ∈ Π mls which sets P ψ ( x τ = x ) sufficiently large for some τ ∈ SubSeq ( H,d )

Algorithm 2 LMDP-OMLE 1: Input: K,d,n test ∈ N , ϵ test , η > 0 , β = log( K | Θ | /η ) , Ψ 0 test = { Unif ( A ) } , D 0 = { ( T j , Unif ( A )) } n test j =1 where each T j is generated by executing Unif ( A ) , C 1 as defined in (1) 2: Initialize k = 1 3: while there exists π k ∈ Π mls , and θ 1 , θ 2 ∈ C k such that TV ( P π k θ 1 , P π k θ 2 ) ( T ) > 4 ϵ test do 4: Ψ k test ← Ψ k -1 test ∪ { π k } , initialize D k = D k -1 5: for all ψ = ( ψ 0 , ψ 1 , ..., ψ d ) ∈ Ψ k test × ... × Ψ k test ×{ Unif ( A ) } with at least one i ∈ [ d -1] + being ψ i = π k , and τ ∈ SubSeq ( H,d ) , z ∈ { 0 , 1 } | τ | do 6: Generate data {T j } n test j =1 by executing ν ( ψ ; τ , z ) 7: Update D k ←D k ∪ { ( T j , ν ( ψ ; τ , z ) } n test j =1 8: end for 9: Update the confidence set C k +1 according to equation (1) 10: k ← k +1 11: end while 12: Pick any θ ∈ C k and return the optimal policy of M := ( S , A , R , θ ) .

and x ∈ X ⊗ | τ | . 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 [43] suggests that our key statistics are this visitation probabilities to all tuples of state-actions within a trajectory.

4.1 Off-Policy Evaluation in LMDPs

The OPE lemma we derive in this section makes use of a behavior policy of a special form which we refer to as a 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 ν ( ψ ; τ , z ) , takes as an input a sequence of history-dependent policies, ψ = ( ψ 0 , ..., ψ d ) , a sequence of time steps, we call checkpoints, τ = ( τ 1 , ..., τ | τ | ) ∈ SubSeq ( H,d ) , and a sequence of binary numbers z = ( z 1 , ..., z | τ | ) ∈ { 0 , 1 } | τ | where | τ | ≤ d , and returns a history-dependent policy.

The segmented policy ν ( ψ ; τ , z ) switches sequentially between different policies in ψ . The time steps in which the switch occurs are determined by τ : starting from time step τ i +1 policy ψ i will be executed. Finally, the sequence z determines whether an intervention with a random action will occur at the τ i time-step. If z i = 1 the executed action at time step τ i is the uniform action, Unif ( A ) , and, otherwise, the policy ψ i -1 is executed. The segmented policy is also denoted by

where ' π a ◦ ( t,z ) π b ' means switch to policy π b at starting from time step t +1 , and at time step t take random action if z = 1 and otherwise execute π a .

We are now ready to define a coverage coefficient for the LMDP class of models. This new coverage coefficient is central to the analysis and design of LMDP-OMLE .

Definition 4.1 (LMDP Coverage Coefficient) The LMDP coverage coefficient of a sequence of policies ψ ∈ Π ⊗ ( d +1) with respect to a target policy π ∈ Π in is given by:

The LMDP coverage coefficient C ( ψ ; π ) between a sequence of policies, ψ , and a target, historydependent, policy π , depends on the worst-case way to generate a segmented policy, ν ( ψ ; τ , z ) from ψ . Further, it is a worst-case ratio of the probability of a sequence of observations within | τ | = d different time steps, namely, x τ , y τ . This is different than the standard coverage coefficient (see equation (2)), that depends on observation from a single time. Fortunately, C ( ψ ; π ) requires only a partial set of observations, instead of using full trajectories. This is crucial towards developing

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 ( ψ ; π ) 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 d = 2 M -1 . For any two models θ, θ ∗ ∈ Θ , and for any π ∈ Π and ψ ∈ Π ⊗ ( d +1) , let C ( ψ ; π ) be defined as (4) over θ ∗ . Then the following holds:

This result is analgous to the OPE result for MDPs (see Lemma 3.1). It is a tool that allows us to bound the TV distance between trajectory distributions of a history-dependent policy π by a term that depends on a segmented policy ν ( ψ ; τ , z ) and an LMDP coverage coefficient. Importantly, the term on the RHS that depends on the segmented policy, ν ( ψ ; τ , z ) , is a sum of distributions of partial trajectories of size | τ | ≤ 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.1, namely an analogous to (2) defined as:

In Appendix D.5 we provide a 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 to 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 to the above issue: we reduce the search space from history-dependent policies to memoryless policies. This allows us 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:

where ν ( ψ i ; τ i ) denotes a segmented policy executing ψ i after the τ th i time-step with memory reset, hence ignoring the history up to τ i (the conditioning event s ′ τ 0 = s ′ [0] at i = 0 can be ignored). Inspired by the form in denominator, we aim to double the following probability defined over a context-segment pair:

for at least one m ∈ [ M ] , s, s ′ ∈ S and t 1 < t 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 to apply the coverage doubling argument.

The restriction to the class of memoryless policies allows us to resolve these issues since

if π, ψ ∈ Π mls since P m represents the latent Markovian transition dynamics. Thus, we can aim to double up the quantity in equation (6). To apply this argument, we establish our second key lemma, a crucial building block for the coverage doubling argument:

Lemma 4.4 (Sufficiency of Memoryless Polices for LMDPs) Suppose the following holds:

Then for any history-dependent policies π ∈ Π , the following holds:

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 ϵ test > 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 a high-level, LMDP-OMLE follows similar recipe to MDP-OMLE , and is summarized in Algorithm 2. It can be described as follows:

Find a memoryless policy π k ∈ Π mls whose prediction on trajectory distributions does not match between two models in the confidence set C k . Add π k to the collection of test policies Ψ k test , that forms each segment of (segmented) exploration policies.

Collect new sample trajectories following the new set of segmented policies for exploration, generated by different combinations of collected test policies and switching operations.

Update the confidence set C k with Maximum Likelihood Estimation (MLE) on the updated dataset D k by equation (1).

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 4.4, we know that when the while loop terminates, any optimal policy of a model contained in the confidence set is a 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 ) Let d = 2 M -1 and assume H > 2 M . After at most K = O ( MS 2 H ) · log( MSAH/ϵ ) iterations, LMDP-OMLE (Algorithm 2) terminates with at most N episodes being generated where

and outputs an ϵ -optimal policy with probability at least 1 -η .

Note that in tabular LMDPs with finite support rewards, we have log( | Θ | ) = O ( S 2 A |R| log(1 /ϵ )) . The appeal of log( | Θ | ) dependence is an flexible extension of the same result to parameterized reward distributions. In Section D.1 and D.4 we provide a proof overview and a full proof of Theorem 4.5.

5 Conclusion and Future Work

In this work, we presented the first sample-efficient algorithm for LMDPs, resolving an 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 are a few remaining questions for the LMDP setting.

Tightness of the Result. The upper bound in Theorem 4.5 scales with ˜ O ( MSAH · log(1 /ϵ )) O ( M ) , while the existing lower bound is Ω( SA ) M . Closing this polynomial gap in the exponent, and having a matching upper and lower bounds can be valuable for a 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 is a segmented policy with intervention at different checkpoints. While this result allows us to provide guarantees on LMDP-OMLE and prove it learns a near-optimal policy, this result is restrictive, in that it does not provide general guarantees for OPE nor makes it possible to derive regret guarantees. In particular, can we evaluate π ∈ Π without policy-switching or intervention when the behavioral policy is a generic history-dependent policy ψ ∈ Π ? Further, is there an algorithm with provable poly( S, A ) M · √ T 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, 44] or additional structural assumptions [41, 63, 14]. Deriving new conditions can be of great importance for real-world applications, e.g., there could be more practical notion of separation, or the set of instances that allows the notion of coverage-coefficient 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.

Acknowledgment

This research was partially funded by AFOSR/AFRL grant no. FA9550-18-1-0166, NSF Grants 2019844 and 2112471, and Israel Science Foundation Grant No. 2199/20.

References

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.

A. Al-Marjani, A. Tirinzoni, and E. Kaufmann. Active coverage for PAC reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory , pages 5044-5109. PMLR, 2023.

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.

P. Amortila, D. J. Foster, and A. Krishnamurthy. Scalable online exploration via coverability. arXiv preprint arXiv:2403.06571 , 2024.

A. Anandkumar, R. Ge, D. Hsu, S. M. Kakade, and M. Telgarsky. Tensor decompositions for learning latent variable models. The Journal of Machine Learning Research , 15(1):2773-2832, 2014.

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.

M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning , pages 263-272. PMLR, 2017.

K. Azizzadenesheli, A. Lazaric, and A. Anandkumar. Reinforcement learning of POMDPs using spectral methods. In Conference on Learning Theory , pages 193-256, 2016.

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.

A. Bennett and N. Kallus. Proximal reinforcement learning: Efficient off-policy evaluation in partially observed markov decision processes. Operations Research , 2023.

E. Brunskill and L. Li. Sample complexity of multi-task reinforcement learning. In Uncertainty in Artificial Intelligence , page 122. Citeseer, 2013.

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 on Artificial Intelligence (AAAI-12) , 2012.

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 Representations , 2022.

F. Chen, C. Daskalakis, N. Golowich, and A. Rakhlin. Near-optimal learning and planning in separated latent MDPs. arXiv preprint arXiv:2406.07920 , 2024.

F. 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.

C. Dann, N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. On oracleefficient PAC RL with rich observations. In Advances in neural information processing systems , pages 1422-1432, 2018.

Y. Efroni, C. Jin, A. Krishnamurthy, and S. Miryoosefi. Provable reinforcement learning with a short-term memory. In International Conference on Machine Learning , pages 5832-5850. PMLR, 2022.

Y. Efroni, D. Misra, A. Krishnamurthy, A. Agarwal, and J. Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Conference on Learning Representations , 2021.

D. J. Foster, S. M. Kakade, J. Qian, and A. Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487 , 2021.

S. A. Geer. Empirical Processes in M-estimation , volume 6. Cambridge university press, 2000.

C. Gentile, S. Li, P. Kar, A. Karatzoglou, G. Zappella, and E. Etrue. On context-dependent clustering of bandits. In International Conference on Machine Learning , pages 1253-1262. PMLR, 2017.

N. Golowich, A. Moitra, and D. Rohatgi. Learning in observable POMDPs, without computationally intractable oracles. Advances in neural information processing systems , 35:1458-1473, 2022.

Z. D. Guo, S. Doroudi, and E. Brunskill. A PAC RL algorithm for episodic POMDPs. In Artificial Intelligence and Statistics , pages 510-518, 2016.

A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.02259 , 2015.

D. Hsu, S. M. Kakade, and T. Zhang. A spectral algorithm for learning hidden markov models. Journal of Computer and System Sciences , 78(5):1460-1480, 2012.

R. Huang, Y. Liang, and J. Yang. Provably efficient UCB-type algorithms for learning predictive state representations. In The Twelfth International Conference on Learning Representations , 2023.

T. Jaksch, R. Ortner, and P. Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research , 11:1563-1600, 2010.

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.

Z. Jia, G. Li, A. Rakhlin, A. Sekhari, and N. Srebro. When is agnostic reinforcement learning statistically tractable? Advances in Neural Information Processing Systems , 36, 2023.

N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. Contextual decision processes with low bellman rank are PAC-learnable. In International Conference on Machine Learning , pages 1704-1713. PMLR, 2017.

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.

C. Jin, S. M. Kakade, A. Krishnamurthy, and Q. Liu. Sample-efficient reinforcement learning of undercomplete POMDPs. arXiv preprint arXiv:2006.12484 , 2020.

C. Jin, A. Krishnamurthy, M. Simchowitz, and T. Yu. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning , pages 4870-4879. PMLR, 2020.

C. Jin, Q. Liu, and S. Miryoosefi. Bellman eluder dimension: New rich classes of RL problems, and sample-efficient algorithms. Advances in neural information processing systems , 34:1340613418, 2021.

S. Kakade and J. Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning , pages 267-274, 2002.

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.

C. Kausik, K. Tan, and A. Tewari. Learning mixtures of markov chains and MDPs. In International Conference on Machine Learning , pages 15970-16017. PMLR, 2023.

M. Kearns and S. Singh. Near-optimal reinforcement learning in polynomial time. Machine learning , 49:209-232, 2002.

A. Krishnamurthy, A. Agarwal, and J. Langford. PAC reinforcement learning with rich observations. In Advances in Neural Information Processing Systems , pages 1840-1848, 2016.

J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reinforcement learning in reward-mixing MDPs. Advances in Neural Information Processing Systems , 34, 2021.

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.

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.

J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reward-mixing MDPs with few latent contexts are learnable. In International Conference on Machine Learning , pages 18057-18082. PMLR, 2023.

J. Kwon, Y. Efroni, S. Mannor, and C. Caramanis. Prospective side information for latent MDPs. arXiv preprint arXiv:2310.07596 , 2023.

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.

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 on Theory of Computing , pages 363-376, 2023.

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 , 33:18819-18831, 2020.

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:1364013653, 2022.

D. Russo and B. Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances in Neural Information Processing Systems , 26, 2013.

C. Shi, M. Uehara, J. Huang, and N. Jiang. A minimax learning approach to off-policy evaluation in confounded partially observable markov decision processes. In International Conference on Machine Learning , pages 20057-20094. PMLR, 2022.

S. Singh, M. R. James, and M. R. Rudary. Predictive state representations: a new theory for modeling dynamical systems. In Proceedings of the 20th conference on Uncertainty in artificial intelligence , pages 512-519, 2004.

R. D. Smallwood and E. J. Sondik. The optimal control of partially observable markov processes over a finite horizon. Operations research , 21(5):1071-1088, 1973.

L. N. Steimle, D. L. Kaufman, and B. T. Denton. Multi-model markov decision processes. Optimization Online URL http://www. optimization-online. org/DB_FILE/2018/01/6434. pdf , 2018.

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018.

K. Tan and Z. Xu. A natural extension to online algorithms for hybrid RL with limited coverage. arXiv preprint arXiv:2403.09701 , 2024.

G. Tennenholtz, U. Shalit, and S. Mannor. Off-policy evaluation in partially observable environments. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 10276-10283, 2020.

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.

M. Uehara, C. Shi, and N. Kallus. A review of off-policy evaluation in reinforcement learning. arXiv preprint arXiv:2212.06355 , 2022.

T. Xie, D. J. Foster, Y. Bai, N. Jiang, and S. M. Kakade. The role of coverage in online reinforcement learning. In The Eleventh International Conference on Learning Representations , 2022.

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.

Y. Xu, J. Zhu, C. Shi, S. Luo, and R. Song. An instrumental variable approach to confounded off-policy evaluation. In International Conference on Machine Learning , pages 38848-38880. PMLR, 2023.

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.

W. Zhan, M. Uehara, W. Sun, and J. D. Lee. PAC reinforcement learning for predictive state representations. arXiv preprint arXiv:2207.05738 , 2022.

T. Zhang. From ε -entropy to KL-entropy: Analysis of minimum information complexity density estimation. Annals of Statistics , 34(5):2180-2210, 2006.

Y. Zhang and N. Jiang. On the curses of future and history in future-dependent value functions for off-policy evaluation. arXiv preprint arXiv:2402.14703 , 2024.

R. Zhou, R. Wang, and S. S. Du. Horizon-free and variance-dependent reinforcement learning for latent markov decision processes. In International Conference on Machine Learning , pages 42698-42723. PMLR, 2023.

L. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. VariBAD: A very good method for bayes-adaptive deep RL via meta-learning. In International Conference on Learning Representations , 2019.

Appendix A Additional Preliminaries

The difference between the values of any policy π ∈ Π measured on two models θ 1 , θ 2 ∈ Θ are bounded by the total-variation (TV) distance between trajectory distributions, that is,

since the maximum reward that can be obtained in an episode is bounded by H due to Assumption 2.2. Hence, if we can show that

then an optimal policy ˆ π ∗ of the empirical model ˆ θ is guaranteed to be 2 ϵ -optimal in the true model θ ∗ . Henceforth, we focus on finding an empirical model ˆ θ that satisfies (8).

To bound the TV-distance between trajectory distributions for all history-dependent policies π ∈ Π between any two LMDP models, we start by unfolding the expression of statistical distance

Whenthe context is clear, we compare trajectory distributions between any two given model parameters θ 1 , θ 2 ∈ Θ , and denote the probability measure from each model following π as P π 1 ( · ) , P π 2 ( · ) .

A.1 Additional Notation

To reduce the notation overload, we use P m ( y t | x t ) = T m ( s t +1 | x t ) R m ( r t | x t ) . When the context is clear, we often use a shorthand π t = π ( a t |T 1: t ) , and denote π t 1 : t 2 as a shorthand of the product of a time-consecutive sequence from t 1 to t 2 , i.e., π t 1 : t 2 = ∏ t 2 t = t 1 π t . When we sum over both x t and y t , we implicitly mean that the s ′ t part of y t , which we denote as s ′ ( y t ) , must match to the s t +1 part of x t +1 , which we denote as s ( x t +1 ) . Using the notation, we rewrite the unfolded TV-distance equation as the following:

We use a shorthand δ π ( X ) for | P π 1 ( X ) -P π 2 ( X ) | = | ∑ M m =1 w 1 m P 1 ,π m ( X ) -∑ M m =1 w 2 m P 2 ,π m ( X ) | , and thus ∑ X δ π ( X ) = TV ( P π 1 , P π 2 )( X ) where the summation is over all possible realizations of a random variable X . Finally, we denote d = 2 M -1 .

A.2 Preliminaries for Lemma 4.2

Here we define a few more quantities that will be crucial for the proofs for Section 4. In bounding the total variation distance in terms of tested policies without exponential blow-up in H , the key is to marginalize events across time steps. Let us first fix the checkpoint time-steps τ = ( τ 1 , ..., τ q ) for q ∈ [ d ] , and a sequence of executable policies ψ = ( ψ 0 , ψ 1 , ..., ψ d ) . Each i th segment policy will be executed within time interval ( τ i , τ i +1 ] for i ≥ 0 .

To proceed, for the initial policy ψ 0 , let l 0 be the smallest quantity, among the contexts, of the ratio between the state visitation probabilities in consecutive steps when ψ 0 is executed (recall that we denote x t = ( s t , a t ) , y t = ( r t , s t +1 ) ):

This quantity can be understood as the minimum (over latent contexts) of the 'pseudo-posterior' probabilities of 1-step event given the future state if ψ 0 is memoryless, since

Henceforth we use a shorthand l 0 t := l 0 ( x t , r t ; s t +1 ) . We recursively define a sequence of the above quantity.

Next, we fix the event ( x τ , y τ ) at the event-log time-steps. For all i ≥ 0 and τ i < t ≤ τ i +1 , we define l i t ( x τ 1: i , y τ 1: i ) and p n,i m ( x τ 1: i , y τ 1: i ) recursively as the following:

and

Here, we define t [0] ≡ -1 , where we recall that ν ( π ; t ) means the memory reset of a policy π at time step t +1 . x τ 1:0 , y τ 1:0 ≡ ϕ , and l 0 0 ≡ 1 and p n, 0 m ≡ w n m for n = 1 , 2 . The key point here is that in this recursive construction, as i increase, we have at least one i such that either p 1 ,i +1 m = 0 or p 2 ,i +1 m = 0 , i.e., at least one context is removed from consideration at each checkpoint.

In the subsequent steps in our proof, we often omit the dependence on ( x τ 1: i , y τ 1: i ) , as well as π τ 0: i in l i t and p n,i +1 m when the context is clear. Finally, we define

Now we are ready to state our key intermediate lemma:

Lemma A.1 For any target policy π ∈ Π and a sequence of segment policies ψ = ( ψ 0 , ψ 1 , ..., ψ d ) , the following holds:

where m ( x τ , y τ ) is the smallest m ∈ [ M ] such that p 1 , | τ | m > 0 .

A.3 Preliminaries for Lemma 4.4

Wepresent additional tools that are useful for proving Lemma 4.4. We first define the notion of (latent) segment coverage coefficient of the set of test policies Ψ test with respect to π as the following:

Definition A.2 (LMDP Segment Coverage Coefficient) The coverage of π ∈ Π with respect to a set of test policies Ψ test ⊆ Π is defined as:

The key lemma is to bound the coverage coefficient of segmented policies consisting of mixture policies defined as the following:

Lemma A.3 Let ψ ξ ∈ Π be a mixture of a subset of behavioral policies for the following set with a fixed t 0 ∈ [ H ] :

Let n = | Ψ ξ | . We define the mixture policy as ψ ξ := 1 n ∑ ψ ∈ Ψ ξ ψ , i.e., given the time interval l until the next checkpoint time, ψ ξ first uniformly randomly picks one policy from Ψ ξ and executes the picked policy afterwards. Let ψ ξ := ( ψ ξ , ψ ξ , ..., ψ ξ ) (of length d +1 ). Then the following holds:

A.4 Auxiliary Concentration Lemmas

The following lemmas are the standard concentration of log-likelihood values of the models within the confidence set. The proofs are standard and can also be in e.g., [1, 45, 44] and [20, 64]. We let D k be the dataset at the beginning of the k th iteration in Algorithm 2. We denote β := log( K | Θ | /η ) .

Lemma A.4 (Uniform Bound on the Likelihood Ratios) With probability 1 -η for any η > 0 , for all k ∈ [ K ] and for any θ ∈ Θ ,

Lemma A.5 (Concentration of Maximum Likelihood Estimators) With probability 1 -η , for all k ∈ [ K ] , t ∈ [ H ] and θ ∈ Θ , we have

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.

Exploration in POMDPs Learning a 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 a general POMDP fundamentally suffers from the curse of horizon [39]. A earlier breakthrough involved introducing structural assumptions on system dynamics, which enable the recovery of underlying POMDPs model under the uniform ergodicity assumption [25, 5, 8, 23].

Recent theoretical breakthrough concerns the exploration problem in POMDPs without the ergodicity assumption [32, 45], initiating a 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' for efficient exploration [16, 1, 18, 17, 22, 45, 13]. The crux of the shortwindow assumption is the prior knowledge that the a short consecutive execution of purely random actions is enough to obtain sufficient statistics of histories, i.e., the full-rankness of latent state-future observation emission matrices. A short sequence of uniformly random actions become the set of core tests of the system [51]. In such systems, learning the optimal policy only incurs poly( S, A ) · A L complexity [46] (with window-size L ), breaking the curse of horizon.

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.

Off-Policy Evaluation in POMDPs Along with the fast progress in online reinforcement learning with partial observations, there is a growing interest in off-policy evaluation with partial observations [47, 56, 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 in POMDPs are asymptotic in nature, and often suffers the curse of horizon due to their use of importance-weight sampling.

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-policy evaluation problems in partially observed systems.

B.1 Additional Details on the Full-Rankness Assumption

We give a 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 [41], where the matrix L is defined in the following form:

where ψ test ∈ Ψ test is a test policy, and T t : H is the future trajectory after time step t . Therefore, direct application of POMDP approaches such as OMLE require the prior knowledge of Ψ test and σ min ( L s ) > 0 for all s ∈ S . The rationale behind such assumptions is to ensure that a distribution of future trajectories can be converted to a belief of latent contexts, hence a distribution of future observations can serve as an alternative to a belief state. However, we are not given the set of core-tests Ψ test , or even the existence of Ψ test that ensures σ min ( L s ) > 0 for general Latent MDPs.

With Separation. In a recent work by Chen et al. [14], a polynomial upper bound in M has been established under a notion of strong-separation between contexts with a sufficiently long time horizon H . In essence, their assumptions guarantee that σ min ( L s ) > 0 holds for most of the states with a priori given test policy, along with additional analysis for the tail of episodes. It 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.

Appendix C Proofs for Section 3

C.1 Proof of Lemma 3.1

This base case corresponds to Lemma 4.2 with M 1 = M 2 = 1 . For convenience, we let θ 1 = θ ∗ and θ 2 = θ , and thus, P 1 = P θ ∗ and P 2 = P θ . We can show the inequality recursively:

since we implicitly sum over s H = s ′ ( y H -1 ) as we described in Appendix A.1. Thus,

Then we can show that

Applying the same step inductively from t = H to t = 1 , we get the lemma.

C.2 Proof of Lemma 3.2

Let Ψ ξ := { π j ∗ ( x,t ) , ∀ x, t | j ∗ ( x, t ) := arg max j ∈ 0 , 1 ,...,k -1 P π j ( x t = x ) } . Then, let ψ ξ ∈ Π be a policy that can adapt to the predetermined checkpoint l , such that ψ ξ = 1 | Ψ ξ | ∑ ψ ∈ Ψ ξ ψ , i.e., a mixture of policies in Ψ ξ . Note that | Ψ ξ | ≤ HSA . Lemma 3.1 tells us that

Then we apply Lemma A.4 and Lemma A.5 to deduce that

On the other hand, note that

Now using the while loop condition, we have

Rearranging the inequality, implies that there exists a t ∈ [ H ] and an x ∈ X such that

C.3 Proof of Theorem 3.3

We first show that Algorithm 1 terminates after K = HSA log(1 /γ ) iterations where γ = ϵ 2 test /H 2 . We consider a perturbed model ˆ θ ∗ = ( w, ˆ T, R ) where

By simulation lemma [38], for any π ∈ Π , note that TV ( P π 1 , P π 2 )( y | x ) ≤ 2 γS for all x, y , and thus

For an arbitrary k iteration, we check whether the coverage doubling argument (Lemma D.1) still holds. To see this, first note that we can define ˆ Ψ ξ , ˆ ψ ξ and ˆ C ( ˆ ψ ξ ; π ) as in Lemma 3.2 in terms of ˆ θ ∗ :

and ˆ ψ ξ is a checkpoint-dependent policy where ˆ ψ ξ is a mixture of ˆ Ψ ξ , and

Now we invoke Lemma A.4 and Lemma A.5 to show that

for all θ ∈ C k using the triangle inequality for TV distance and ( a + b ) 2 ≤ 2( a 2 + b 2 ) . Thus, we can derive equation (16) in terms of ˆ θ ∗ as

On the other hand,

by setting γ = ϵ test / (4 SAH ) 4 . Let n test = 64 β ( H 3 SA ) /ϵ 2 test , and we have

Hence the same doubling argument holds, and MDP-OMLE (Algorithm 1) will terminate after at most

iterations. We note that all inequalities hold for all K iterations with probability at least 1 -η . Finally, by setting ϵ test = ϵ TV and ϵ TV = ϵ/H , we can conclude that the total number of trajectories that was generated by MDP-OMLE is bounded by

Appendix D Proofs for Section 4

D.1 Proof Sketch

In this section, we provide an overview of the proof for Theorem 4.5. Compared to Algorithm 1, the main differences in LMDPs from the MDP cases are two-fold:

(a) Our goal is to find the optimal history-dependent (non-Markovian) policy.

(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 m under which each segment is covered, we can improve the coverage of each context-segment pair given the test set Ψ k test every iteration:

Lemma D.1 Let n test ≥ 3 βM 2 (8 H 2 ) d ( MS 2 A 2 ) d /ϵ 2 test . Then with probability at least 1 -η , at every k th iteration in Algorithm 2, there must exist at least one ( m,t 1 , t 2 , s, s ′ ) , such that

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 a non-zero γ > 0 probability for covering all segments in all contexts:

We note that this assumption will be eventually removed in our final result. Thus, if we start from the initial coverage γ > 0 over all context-segment pairs, then since every probability is in the range of [0 , 1] , this doubling-up event can happen at most log(1 /γ ) times for every context-segment pair. Therefore, Algorithm 2 must terminate after at most K = MS 2 H log(1 /γ ) iterations.

Separately from the coverage improvement in Lemma D.1, 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 1 -η , for all k th iterations in Algorithm 2, let Ψ ξ = { ψ 1 , ..., ψ n } ⊆ Ψ k -1 test be any subset of candidate test policies, and let ψ ξ = ( ψ ξ , ..., ψ ξ , Unif ( A )) where ψ ξ := 1 n ∑ i ∈ [ n ] ψ i is a mixture of policies in Ψ ξ . Then for any model in the confidence set θ ∈ C , the following holds:

Equipped with Lemma D.2 With probability at least 1 -η , we terminate the while-loop after at most K = HS 2 M log(1 /γ ) iterations, and each while-loop generates K d -1 · n test new trajectories, leading to a total

sample complexity. The near-optimality guarantee for the final empirical model is given by Lemma 4.4 where we set ϵ test = M -1 (2 H 2 MSA ) -d · ϵ TV to obtain an ( Hϵ TV ) -optimal policy.

Without Initial Coverage. Now we remove the initial γ > 0 coverage assumption (17). To do so, consider a virtual LMDP model ˆ θ = ( { w ∗ m , ˆ T m , R ∗ m } ) M m =1 with perturbed transition models, i.e., ˆ T m ( ·| s, a ) = (1 -γ ) T ∗ m ( ·| s, a ) + γ 1 for all ( s, a ) ∈ S ×A . For ˆ θ , it is easy to see that for all π ∈ Π , we have TV ( P π ˆ θ , P π θ ∗ )( τ ) ≤ 2 HSγ . Thus, now we can shift our arguments to this perturbed model with sufficiently small γ , and it is straightforward that the perturbed model has the γ > 0 segment coverage for any policy, which concludes Theorem 4.5.

D.2 Proof of Lemma 4.2

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 I ⊆ [ | τ | ] , we write τ I := ( τ i ) i ∈ I to refer to a subsequence of τ at positions I , and τ /I a subsequence at positions outside of I .

We can first bound ∆( x τ , y τ ) as the following:

where z ( I ; τ ) satisfies z ( I ; τ ) j = 0 if j ∈ I and 1 otherwise. Then we can observe that

using inequality that can be derived by definition of l i t :

We are now ready to use this inequality to bound the TV distance between trajectory distribution via equation (12). To do so, we rearrange the summation orders as follows:

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 Ψ ξ in Lemma A.3:

Note that the above definition is invariant to t 0 .

Now consider Ψ test = Π mls . Then, Ψ ξ is the set of policies that maximize the probability P ν ( ψ ; t 0 ) m ( s t 0 + t = s | s ′ t 0 = s ′ ) , i.e., a memoryless policy that aims to reach s after t time steps (this policy is invariant to the starting state s ′ ). Therefore, we can first induce that | Ψ ξ | ≤ MHS . Then, the Markovian policy maximizing the probability to reach a certain state under a fixed context m is also best within the history-dependent policy class Π , and therefore

We can now invoke Lemma 4.2 and Lemma A.3, and noting that

Let ψ ξ be as defined in Lemma A.3, and for any π ∈ Π , we can conclude that

Finally, it only remains to bound the total variation distance with ν ( ψ ξ ; τ , z ) . To see this, for a given q ≤ d , τ ∈ Subseq ( H,d ) , and z ∈ { 0 , 1 } | τ | , it is easy to check that

Noting that ∑ q ≤ d ( H q ) ≤ min( H d , 2 H ) , assuming we are in the regime H ≫ d , we conclude that

concluding the proof.

D.4 Proof of Theorem 4.5

We continue from the conclusion of Lemma D.1, and the remaining step is to ensure that LMDP-OMLE terminates after K = MS 2 H log(1 /γ ) iterations where γ = ϵ 2 test / ( H 2 d ) without the initial coverage assumption (17). We consider a perturbed model ˆ θ ∗ = ( { w ∗ m , ˆ T m , R ∗ m } M m =1 ) where

for all m . By the performance difference lemma [35], for any π ∈ Π ,

For every k th iteration, we check whether the coverage doubling argument (Lemma D.1) still holds. To see this, first note that we can define ˆ ρ (Ψ k test ; π ) and ˆ Ψ ξ , ˆ ψ ξ and ˆ C ( ψ ξ ; π ) as in Lemma A.3 in terms of ˆ θ ∗ . Then Lemma D.2 can be modified to guarantee that

for all θ ∈ C k using the triangle inequality for TV distance and ( a + b ) 2 ≤ 2( a 2 + b 2 ) . Thus, we can derive (19) in terms of ˆ θ ∗ as

̸

On the other hand,

Now we set γ = ϵ 2 test (8 HnA ) -2 d (4 MSH ) -2 with n test = 64 M 2 β ( HnA 2 ) d /ϵ 2 test and n = MHS 2 , and we have

Hence the same doubling argument holds, and Algorithm 2 will terminate after at most

iterations. We note that all inequalities hold for all K iterations with probability at least 1 -η . Finally, we invoke Lemma 4.4 with ϵ test = ϵ TV · poly( H,M,S,A ) -d , ϵ TV = ϵ/ (4 MSAH ) d and d = 2 M -1 , we can conclude that the total number of trajectories we generated is bounded by

D.5 A Counter Example for Remark 4.3

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:

Here we present a counter example where the multi-step events must be considered even in the latent state space: the LMDP consists of M = 3 MDPs with S = { 1 , 2 } , A = { 1 , 2 } , R = {-1 , 0 , 1 } , and H = 2 . All MDP starts from s 1 = 1 and with a transition kernel T m ( s 2 = 2 | s 1 , a ) = 1 for all m ∈ [ M ] and a ∈ A . Reward models are given by R 1 ( r = -1 | s = 1 , a = 1) = 1 , R 2 ( r = 1 | s = 1 , a = 2) = 1 , and R m ( r | s = 1 , a = 1) = 0 . 5 for r ∈ { 0 , 1 } , m ∈ { 2 , 3 } , and R m ( r | s = 1 , a = 2) = 0 . 5 for r ∈ {-1 , 0 } , m ∈ { 1 , 3 } .

Now we target to measure the expected rewards of actions executed by a behavioral policy at s 2 = 2 as in Table 1. In this example, all actions are covered under all contexts following the behavior policy, i.e., P m ( s 2 , a 2 ) > 0 . Yet, we cannot estimate the expected reward of action a 2 = 1 under context m = 1 . Therefore, the speculated single latent-state coverage coefficient is 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 t 0 , recall the definition Ψ ξ as defined in equation (14):

Note that | Ψ ξ | ≤ MS 2 H and invariant to t 0 since Ψ k -1 test ⊂ Π mls . Now for any memoryless policy π ∈ Π mls , we recall Lemma 4.2 with ψ ξ = ( ψ ξ , ..., ψ ξ , Unif ( A )) , where ψ ξ ∈ Π is a mixture policy of Ψ ξ . We have

̸

Plugging

n

test

where the last inequality follows by Lemma D.2. By the while-loop condition, for π k ∈ Π mls , we have

by the triangle inequality for TV distance. Therefore, we can conclude that

(8

HnA

)

/ϵ

2

test with

n

=

MHS

, we have

Thus, ρ (Ψ k -1 test ; π k ) > 2 , which in turn implies Lemma D.1.

E.2 Proof of Lemma D.2

By the construction of confidence set in equation (1) and Lemma A.5, for all k ∈ [ K ] and θ ∈ C k , we have

where β = log( K | Θ | /η ) . Let ψ ξ be a mixture of a subset of candidate policies Ψ ξ = { ψ 1 , ψ 2 , ..., ψ n } ⊆ Ψ k -1 test . With ψ ξ = ( ψ ξ , ψ ξ , ..., ψ ξ , Unif ( A )) and for every τ ∈ Subseq ( H,d ) , z ∈ { 0 , 1 } | τ | , we have

Therefore,

where the last inequality is due to the construction of our dataset D k and the concentration guarantee for the confidence set C k .

E.3 Proof of Lemma A.1

First note that we can rewrite an atomic probability P n m ( y t | x t ) = P n m ( r t , s t +1 | x t ) as

In turn, moving from conditioning on the event ( x t , y t ) , we view the LMDP model after induction as if there are at most M 1 -1 contexts in the first LMDP model (or M 2 -1 in the second LMDP model if l 0 ( x t , r t ; s t ) is attained with n = 2 ). We often use a shorthand l 0 ( x t , y t ) = l 0 ( x t , r t ; s t +1 ) to reduce the burden on the notation.

2

d

2

= 64

M

2

β

Proof. The basic strategy is to apply iteratively the triangle inequality. To illustrate, we first expand the equation at t = 1 :

We can continue applying triangle inequalities to all possible first event-logging time step, and we can start with the following inequality:

Recall that at least one p 1 , 1 m or p 2 , 1 m is 0, that is, one of the contexts in either of the two systems is eliminated.

Now for ( i ) , we can pick the next event-logging time step τ 2 > τ 1 , and apply the triangle inequality similarly, and repeat such event-logging until all contexts are exhausted. Since there are at most 2 M contexts, we cannot repeat this process more than d = 2 M -1 times. Hence, we arrive to the following inequality:

where we set τ | τ | +1 ≡ H +1 . To proceed from here, we first observe that for any m ∈ [ M ] and i ≥ 0 such that p 1 ,i m > 0 ,

Thus, we can summarize that

We note that each term in the product sequence is equivalent to

and thus

Here we set τ q +1 = H + 1 and P π m ( s H +1 |· ) = 1 for any m , π and conditional event. Since the denominator does not depend on events outside of 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, t 1 , t 2 , by the construction of ψ ξ , it follows that

We can observe that for any m ,

On the other hand, for any π ∈ Π , τ ⊆ [ H ] , I ∈ [ | τ | ] , x τ , y τ ,

Applying the inequality with the definition of ρ (Ψ ξ ; π ) , we have

E.5 Proof of Lemma A.4

This is by now a standard MLE technique for constructing confidence sets in RL [1].

Proof. The proof follows a Chernoff bound type of technique:

The last inequality is by Markov's inequality. Note that random variables are ( τ, π ) in the trajectory dataset D , and

Therefore,

Combining the above, taking a union bound over k ∈ [ K ] rounds and θ ∈ Θ , letting β = log( K | Θ | /η ) , with probability 1 -η , 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 ι, τ , π and t ∈ [ H ] ,

To bound the summation over samples, we start from

On the other hand, by the Chernoff bound,

where in the last line, we used the tower property of expectation. Thus, again by setting β = log( K | Θ | /η ) , with probability at least 1 -η , we have

for all k ∈ [ K ] and θ ∈ Θ . Now we can apply Lemma A.4, and finally have

Since TV 2 ≤ 2 H 2 , we get the lemma.

□

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]

Justification: The abstract present the problem we consider in this work - designing an efficient learning algorithm for the RL setting - and elaborate, in a highlevel, on our new approach for solving this open problem.

2. Limitations

Question: Does the paper discuss the limitations of 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 as 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 a complete (and correct) proof?

Answer: [Yes]

Justification: 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 to the analysis of the LMDP setting.

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 of whether the code and data are provided or not)?

Answer: [NA]

Justification: There are no experimental results in this paper, but a resolution of an 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 a resolution of an algorithmic open problem.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a resolution of an algorithmic open problem.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: There are no experimental results in this paper, but a resolution of an algorithmic open problem.

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 a resolution of an algorithmic open problem.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a 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, as we mentioned in the introduction section, by both empirical and theoretical communities. We do not see immediate societal impacts of our work between the new results we derived, and the promise for improving algorithms for the LMDP, and POMDP settings in 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 (e.g., pretrained language models, image generators, or scraped datasets)?

Answer:[NA]

Justification: There are no experimental results in this paper, but a resolution of an algorithmic open problem.

12. Licenses for existing assets

Question: 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?

Answer: [NA]

Justification: There are no experimental results that make use of data or models in this paper, but a 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 a resolution of an 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.

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n-\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.\n\nWEAKNESSES:\n-\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).\n\nQUESTIONS:\n-\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?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: 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. 

The 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. 

As a result, they showed the first polynomial sample complexity result for LMDP when the number of latent states are constant.

STRENGTHS:
The paper is well written. The algorithm, theorems and proofs are clear.

The results of polynomial sample complexity for LMDP with constant latent states seem very interesting. 

The idea of constructing segmented policies is novel.

WEAKNESSES:
I don't see any significant weaknesses of this paper.

QUESTIONS:
Can these results somehow adapt to the POMDP setting?

LIMITATIONS:
Yes. The authors addressed all the limitations listed in the guidelines.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe use of MDPs to illustrate the concept is effective, making it easier for the audience to follow.\n\nWEAKNESSES:\n1. 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.\"\n\nQUESTIONS:\nSee \"Weaknesses\" part for the questions.\n\nLIMITATIONS:\nSee \"Weaknesses\" part for the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}

comments: 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.

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.

2. The proposed LMDP coverage coefficient, although exponential in $d$ ($M$), seems to capture the right complexity of the LMDPs. 

3. 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. 

2. 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?

3. 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?

QUESTIONS:
see above

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-26