RL for Latent MDPs: Regret Guarantees and a Lower Bound

Abstract:
In this work, we consider the regret minimization problem for reinforcement learning in latent Markov Decision Processes (LMDP). In an LMDP, an MDP is randomly drawn from a set ofM possible MDPs at the beginning of the interaction, but the identity of the chosen MDP is not revealed to the agent. We first show that a general instance of LMDPs requires at least Ω((SA) ) episodes to even approximate the optimal policy. Then, we consider sufficient assumptions under which learning good policies requires polynomial number of episodes. We show that the key link is a notion of separation between the MDP system dynamics. With sufficient separation, we provide an efficient algorithm with local guarantee, i.e., providing a sublinear regret guarantee when we are given a good initialization. Finally, if we are given standard statistical sufficiency assumptions common in the Predictive State Representation (PSR) literature (e.g., [6]) and a reachability assumption, we show that the need for initialization can be removed.

heading: 1 Introduction
text: Partially observable Markov decision processes (POMDPs) [42] give a general framework to describe partially observable sequential decision problems. In POMDPs, the underlying dynamics satisfy the Markovian property, but the observations give only partial information on the identity of the underlying states. With the generality of this framework comes a high computational and statistical price to pay: POMDPs are hard, primarily because optimal policies depend on the entire history of the process. But for many important problems, this full generality can be overkill, and in particular, does not have a way to leverage special structure. We are interested in settings where the hidden or latent (unobserved) variables have slow dynamics or are even static in each episode. This model is important for diverse applications, from serving a user in a dynamic web application [18], to medical decision making [45], to transfer learning in different RL tasks [8]. Yet, as we explain below, even this area remains little understood, and challenges abound. Thus, in this work, we consider reinforcement learning (RL) for a special type of POMDP which we call a latent Markov decision process (LMDP). LMDPs consist of some (perhaps large) numberM of MDPs with joint state space S and actions A. In episodic LMDPs with finite time-horizon H , the static latent (hidden) variable that selects one ofM MDPs is randomly chosen at the beginning of each episode, yet is not revealed to the agent. The agent then interacts with the chosen MDP throughout the episode (see Definition 1 for the formal description).
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Related Work. The LMDP framework has previously been introduced under many different names, e.g., hidden-model MDP [11], Multitask RL [8], Contextual MDP [18], Multi-modal Markov decision process [45] and Concurrent MDP [9]. Learning in LMDPs is a challenging problem due to the unobservability of latent contexts. For instance, the exact planning problem is P-SPACE hard [45], inheriting the hardness of planning from the general POMDP framework. Nevertheless, the lack of dynamics of the latent variables, offers some hope. As an example, if the number of contextsM is bounded, then the planning problem can be at least approximately solved (e.g., by point-based value iteration (PBVI) [39], or mixed integer programming (MIP) [45]). The most closely related work studying LMDPs is in the context of multitask RL [47, 8, 34, 18]. In this line of work, a common approach is to cluster trajectories according to different contexts, an approach that guided us in designing the algorithms in Section 3.4. However, previous work requires very long time-horizon H SA in order to guarantee that every state-action pair can be visited multiple times in a single episode. In contrast, we consider a significantly shorter time-horizon that scales poly-logarithmic with the number of states, i.e.,H = poly log(MSA). This short time-horizon results in a significant difference in learning strategy even when we get a feedback on the true context at the end of episode. We refer the readers to Appendix A for additional discussion on related work. Main Results. To the best of our knowledge, none of the previous literature has obtained sample complexity guarantees or studied regret bounds in the LMDP setting. This paper addresses precisely this problem. We ask the following:
Is there a sample efficient RL algorithm for LMDPs, with sublinear regret?
The answer turns out to be not so simple. Our results comprise a first impossibility result, followed by positive algorithmic results under additional assumptions. Specifically:
• First, we find that for a general LMDP, polynomial sample complexity cannot be attained without further assumptions. That is, to find an approximately optimal policy we need at least Ω ( (SA)M ) samples, i.e., at least exponential in the number of contextsM (Section 3.1). This lower bound even applies to instances with deterministic MDPs.
• We find that there are several natural assumptions under which optimal policies can be learned with polynomial sample complexity. Similarly to mixture problems without dynamics, the key link is a notion of separation between the MDPs. With sufficient separation, we show that there is a planning-oracle efficient RL algorithm with polynomial sample complexity. A critical development is adapting the principle of optimism as in UCB, but to the partially observed setting where value-iteration cannot be directly applied, and thus neither can the UCRL algorithm for MDPs.
• Finally, under additional statistical sufficiency assumptions that are common in the Predictive State Representation (PSR) literature (e.g., [6]) and a reachability assumption, we show that the need for initialization can be entirely removed.
• Finally, we perform an empirical evaluation of the suggested algorithms on toy problems (Section 4), while focusing on the importance of the made assumptions.

heading: 2 Preliminaries

heading: 2.1 Problem Setup: Latent MDPs
text: We start with the definition of episodic reinforcement learning in latent Markov decision process:
Definition 1 (Latent Markov Decision Process (LMDP)) Consider a set of MDPsM with joint state spaceS and joint action spaceA in a finite time horizonH . LetM = |M|, S = |S| andA = |A|. Each MDPMm ∈M is a tuple (S,A, Tm, Rm, νm) where Tm : S ×A× S → [0, 1] a transition probability maps a state-action pair and a next state to a probability, Rm : S ×A× {0, 1} → [0, 1] a probability measure for rewards that maps a state-action pair and a binary reward to a probability, and νm is an initial state distribution. Let w1, ..., wM be the mixing weights of LMDPs such that at the start of every episode, one MDPMm ∈M is randomly chosen with probability wm.
We assume the mixing weights are uniform and known a priori, i.e., w1 = ... = wM = 1/M . This is only for the ease of presentation, and does not affect any algorithmic idea or main results in this paper. The goal of the problem is to find a (possibly non-Markovian) policy π in a policy class Π that maximizes the expected return: V ∗M := maxπ∈Π ∑M m=1 wmEπm [∑H t=1 rt ] ,where Eπm[·] is
expectation taken over the mth MDP with a policy π. If not specified otherwise, we find the best policy in a set of history-dependent policies π : (S,A, {0, 1})∗ × S → ∆(A) that map an entire history to a probability distribution over actions. We define the notion of regret relative to a (possibly approximate) planning oracle. Thus, suppose we have a planning-oracle with the following approximation guarantee: V πM ≥ ρ1V ∗M − ρ2, where π is a returned policy whenM is given to the planning-oracle, and ρ1, ρ2 are multiplicative and additive approximation constants respectively. We then define the regret as the comparison to the best approximation guarantee that the planning-oracle can achieve:
Regret(K) = K∑ k=1 (ρ1V ∗ M − ρ2)− V πk M , (1)
where πk is a policy executed in the kth episode. For example, we can use point-based value-iteration (PBVI) [39] as a planning-oracle:
Example 1 PBVI [39] with discretization level d > 0 in the belief space (over RM ) returns an dH2 additive approximate policy. That is, equation (1) holds with ρ1 = 1 and ρ2 = dH2.

heading: 2.2 Predictive State Representation (PSR)
text: A partially observable dynamical system can be viewed as a model that generates a sequence of observations from observation space O with (controlled) actions from action space A. A predictive state representation (PSR) is a compact description of a dynamical system with a set of observable experiments, or tests [41]. Specifically, a test of length t is a sequence of action-observation pairs given as τ = aτ1oτ1oτ2 ...aτt oτt . A history h = ah1oh1ah2oh2 ...aht oht is a sequence of action-observation pairs that has been generated prior to a given time. A predictionP(τ |h) = P(oτ1:t|h||do aτ1:t) denotes the probability of seeing the test sequence from a given history, given that we intervene to take actions aτ1a τ 2 ...a τ t . In latent MDPs, the observation space can be considered as a pair of next-states and rewards, i.e., O = S × {0, 1} and ot = (st+1, rt). As we work with a special class of POMDPs, we customize the formulation for LMDPs. The set of histories consists of a subset of histories that end with different states, i.e.,H = ⋃ sHs, where each element h ∈ Hs is a short sequence of state-action-rewards of length l ending with state s:
h = sh1a h 1r h 1 s h 2 ...s h l−1a h l−1r h l−1s = (s, a, r) h 1:l−1s.
We define Pπm(Hs) a vector of probabilities where each coordinate is a probability of sampling each history in Hs in the mth MDP with a policy π. Likewise, each element in tests τ ∈ T is a short sequence of action-reward-next states of length at most l:
τ = aτ1r τ 1s τ 2 ...a τ l r τ l s τ l+1 = (a, r, s ′)τ1:l.
We denote Pm(T |s) as a vector of probability where each coordinate is a success probability of each test in themth MDP starting from a state s. That is,
Pm(T |s)i = Pm(rτi1 s τi 2 ...r τi l s τi l+1|s||do a τi 1 ...a τi l ).

heading: 2.2.1 Spectral Learning of PSRs in LMDPs
text: In spectral learning, we build a set of observable matrices that contains the (joint) probabilities of histories and tests, and then we can extract parameters from these matrices by performing singular value decomposition (SVD) and regressions [7]. In order to apply spectral learning techniques, we require the following technical conditions on statistical sufficiency of tests:
Condition 1 (Sufficient Tests) For all s ∈ S, for the test set T , let Ls = [P1(T |s)|P2(T |s)|...|PM (T |s)]. Then σM (Ls) ≥ στ for all s ∈ S with some στ > 0.
Here, σM (·) is the minimum (M th) singular value of a matrix. Another technical condition for spectral learning method is a rank non-degeneracy condition for sufficient histories:
Condition 2 (Sufficient Histories) For all s ∈ S, for the history set Hs ending with a state s, let Hs = [Pπ1 (Hs)|Pπ2 (Hs)|...|PπM (Hs)]> with a sampling policy π. Then σM (LsHs) ≥ Pπ(end state = s) · σh with some σh > 0.
Here Pπ(end state = s) is a probability of sampling a history ending with s. Along with the rank condition for tests, pairs of histories and tests can be thought as many short snap-shots of long trajectories obtained by external experts or some exploration policy (e.g., random policy in uniformly ergodic MDPs). Following the notations in [7], let PT ,Hs = LsHs. Conditions 1 and 2 ensure σM (PT ,Hs) > 0. Under these conditions, the goal of spectral learning algorithm is to output PSR parameters which are used to compute P̂(τ |h), the estimated probability of any future observations (or tests τ ) given any sampled histories h. We refer to Appendix E.1 for a detailed procedure.

heading: 2.3 Notations
text: We denote the underlying LMDPwith true parameters asM∗. With slight abuse of notation, we denote the l1 distance between two probability distributions D1 and D2 on a random variable X conditioned on event E as ‖(PX∼D1 − PX∼D2)(X|E)‖1 = ∑ X∈X |PX∼D1(X|E)− PX∼D2(X|E)|, where X is a support of X . When we do not condition on any event, we omit the conditioning on E. When we measure a transition or reward probability at a state-action pair (s, a), we use T or R instead of P. We use Pm to refer to the probability of any event measured in the mth context (or in mth MDP). In particular, Pm(s′, r|s, a) = Tm(s′|s, a)Rm(r|s, a). If we use P without any subscript, it is a probability of an event measured outside of the context, i.e., P(·) = ∑M m=1 wmPm(·). If the probability of an event depends on a policy π, we add superscript π to P. Similarly, Em[·] is expectation taken over themth context and π is added as superscript if the expectation depends on π. We use ·̂ to denote any estimated quantities. a . b implies a is less than b up to some constant and logarithmic factors. poly(·) means the order of polynomial complexity (up to logarithmic factors) in referenced parameters. We interchangeably use o, an observation, to replace a pair of next-state and immediate reward (s′, r) to simplify the notation. We occasionally express a length t > 0 history (s1, a1, r1, ..., st, at, rt) compactly as (s, a, r)1:t.

heading: 3 Main Results
text: In this section, we first obtain a hardness result for the general case. We then consider sample- and computationally efficient algorithms under additional assumptions.

heading: 3.1 Fundamental Limits of Learning General LMDPs
text: We first study the fundamental limits of the problem. In particular, we are interested in whether we can learn the optimal policy after interacting with the LMDP for a number of episodes polynomial in the problem parameters. We prove a worst-case lower bound, exhibiting an instance of LMDP that requires at least Ω ( (SA)M ) episodes:
Theorem 3.1 (Lower Bound) There exists an LMDP such that for finding an -optimal policy π for which V π M ≥ V ∗M − , we need at least Ω ( (SA/M)M/ 2 ) episodes.
The hard instance consists of fully deterministic MDPs with possibly stochastic rewards, indicating an exponential lower bound in the number of contexts even for the easiest types of LMDPs. The example is constructed such that, in the absence of knowing true contexts, all wrong action sequences of length M cannot provide any information with zero reward, whereas the only correct action sequence gets a total reward of 1 under one specific context. The construction is given in Appendix B. Theorem 3.1 prevents a design of efficient algorithms with growing number of contexts. We note here that Theorem 3.1 holds even for restricted classes of policies, e.g., memoryless policies. Furthermore, our construction of hard instances does not allow to find any approximate policy with ρ1 = ω((SA) −M ) within a polynomial number of episodes either. To the best of our knowledge, this
Algorithm 1 Latent Upper Confidence Reinforcement Learning (L-UCRL) Initialize visit countsNm(s, a),N(m) and parameters (T̂m, R̂m, ν̂m) properly 1: for each kth episode do 2: Get a policy πk for M̃k in Lemma 3.2 3: Play policy πk and get the trajectory τ = (s, a, r)1:H 4: Get an estimated belief over contexts b̂ with either Algorithm 2 (when contexts are given), or Algorithm 3 (when we infer contexts) 5: for m = 1, ...,M and t = 1, ...,H do 6: Nm(st+1|at, st)← Nm(st+1|at, st) + b̂(m) 7: Nm(rt|st, at)← Nm(rt|st, at) + b̂(m) 8: Nm(s1)← Nm(s1) + b̂(m) 9: Update empirical parameters T̂m, R̂m, ν̂m 10: end for 11: end for
is the first lower bound of its kind for LMDPs. Next, we investigate natural assumptions which help us to develop an efficient algorithm when only polynomial number of episodes are available.

heading: 3.2 The Critical First Step: Contexts in Hindsight
text: Suppose the true context of the underlying MDP is revealed to the agent at the end of each episode. We do not require any assumptions on the environments in this scenario. Note that this scenario is different from fully observable settings (i.e., knowing the true context at the beginning of an episode). In the latter scenario, we would simply haveM -decoupled RL problems in standard MDPs. While this can be considered as a “warm-up” for the sequel, it is motivated by real-world examples. Moreover, the key technical insight here will prove important for the sequel as well. Knowing contexts in hindsight allows us to construct a confidence set for parameters:
C = {M |‖(νm − ν̂m)(s)‖1 ≤ √ cν/N(m), ‖(Tm − T̂m)(s′|s, a)‖1 ≤ √ cT /Nm(s, a),
‖(Rm − R̂m)(r|s, a)‖1 ≤ √ cR/Nm(s, a), ∀m, s, a}, (2)
where Nm(s, a) is the number of times each state-action pair (s, a) in mth MDP is visited, and N(m) is the number of episodes we interact with the mth MDP. With properly set parameters cT = O(S log(K/η)), cR = O(log(K/η)) and cν = O(S log(K/η)) for the confidence intervals, M∗ ∈ C with high probability for allK episodes. With the construction of confidence sets, it is then natural to try to design an optimistic RL algorithm, as in UCRL [22]. An obvious optimistic value in light of (2) is maxπ,M∈C V πM. However, solving this optimization problem is more general than solving an LMDP. In fully observable settings, we could replace the complex optimization problem by adding a proper exploration bonus to obtain an optimistic value function [4]. In partially observable environments, value iteration is only defined in terms of belief-states and not the observed states. For this reason, existing techniques solely based on the value-iteration cannot be directly applied for LMDPs. Yet, we find that proper analysis of the Bellman update rule over the belief state reveals that an empirical LMDP with properly adjusted hidden rewards is optimistic:
Proposition 3.2 We construct an optimistic LMDP M̃ whose parameters are given such that:
T̃m(s ′|s, a) = T̂m(s′|s, a), R̃obsm (r|s, a) = R̂m(r|s, a), ν̃m(s) = ν̂m(s), R̃hidinit(m) = min ( 1, √ cν/N(m) ) R̃hidm (s, a) = H min ( 1, √ 5(cR + cT )/Nm(s, a) ) ,
where R̃hidinit(m) is an initial hidden reward given when starting an episode with a contextm, and R̃obsm (·|s, a) is a probability measure of an observable immediate reward r whereas R̃hidm (s, a) is a hidden immediate reward (that is not visible to the agent) for a state-action pair (s, a) in a contextm. Then for any policy π, the expected long-term reward is optimistic, i.e., V π
M̃ ≥ V πM∗ .
Algorithm 2 Access to True Contexts Input: Receive true context m∗ in hindsight Output : Belief over hidden contexts:
b̂(m) = {
1, form = m∗ 0, form 6= m∗
Algorithm 3 Inference of Contexts Input: Trajectory τ = (s1, a1, r1, ..., sH , aH , rH) Output: Return an estimate of belief over contexts b̂: p̂m(τ) = Π H t=1(α+ (1− 2αS)P̂m(st+1, rt|st, at)),
b̂(m) = p̂m(τ)∑M m=1 p̂m(τ) .
Here, hidden reward is a deterministic reward that happens for every state-action pair, but does not appear in any observation history during the episode. We note that most existing planning algorithms can incorporate the hidden-reward structure without changes to maximize the long-term observed + hidden rewards. For instance, the PBVI algorithm [39] can be executed as it is in the planning step. Hence in each episode, we can build one optimistic model from Proposition 3.2, and call the planning-oracle to get a policy to execute for the episode. Then we simply run the policy and update model parameters in a straight-forward manner. The algorithm can be efficiently implemented as long as some efficient (approximate) planning algorithms are available. To establish Proposition 3.2, we make use of the ‘alpha vector’ representation [42] of the value function of general POMDPs. Detailed analysis is deferred to Appendix C.1. With the optimistic model constructed in Proposition 3.2, planning-oracle efficient implementation is straightforward. The resulting latent upper confidence reinforcement learning (L-UCRL) algorithm is summarized in Algorithm 1. Based on the established optimism in Proposition 3.2 and by carefully bounding the on-policy errors we arrive to the following regret guarantee of L-UCRL.
Theorem 3.3 Let N = HK. The regret of the Algorithm 1 is bounded by: Regret(K) ≤ K∑ k=1 (V πk M̃k − V πkM∗) . HS √ MAN.
Proof of Theorem 3.3 is given in Appendix C.4. The central result of this section, Theorem 3.3 leads to the following observation: a polynomial sample complexity is possible for the LMDP model assuming the context of the underlying MDP is supplied at the end of each episode. Next, we explore ways to relax this assumption, while keep supplying with a polynomial sample complexity guarantee.
3.3 When we can Infer Contexts?
Without explicit access to the true context at the end of an episode, it is natural to estimate the context from the sampled trajectory. One sufficient condition to infer the context is the following:
Assumption 1 (δ-Strongly Separated MDPs) For all m, m1,m2 ∈ [M ] such that m1 6= m2, for all (s, a) ∈ S ×A, l1 distance between probability of observations o = (s′, r) of two different MDPs in LMDP is at least δ > 0, i.e., ‖(Pm1 − Pm2)(o|s, a)‖1 ≥ δ for some constant δ > 0.
In order to reliably infer the true contexts the seperatedness between MDPs alone is not sufficient, since we need to estimate the contexts from the current empirical estimates of LMDPs. In order to reliably estimate the context from empirical estimate of LMDPs, we need a well-initialized empirical transition model of the LMDP: ‖(T̂m − Tm)(s′|s, a)‖1, ‖(ν̂m − νm)(s)‖1, ‖(R̂m −Rm)(r|s, a)‖1 ≤ init, ∀(s, a), (3)
for some initialization error init > 0. Note that while the initialization error is relatively small, it can be still not good enough to obtain a near-optimal policy (i.e., it will result in a linear regret). We can consider as if the state-action pairs are already visit at least N0 = cT / 2init times in each context. Once the initialization is given along with separation between MDPs, we can modify Algorithm 1 to update the empirical estimate of LMDP using the estimated belief over contexts computed in Algorithm 3. Note that when we update the model parameters, we increase the visit count of stateaction pair (s, a) atmth MDP by b̂(m). With Assumption 1, it approximately adds a count for the
Algorithm 4 (Informal) Recovery of LMDP parameters Learn PSR parameters up to precision o(δ) Get clusters {T̂m(·|s, a), R̂m(·|s, a)}(s,a)∈S×A,m∈[M ] with learned PSR parameters Build each MDP model by correctly assigning contexts to estimated model parameters ReturnWell-initialized model {T̂m, R̂m}m∈[M ]
correctly estimated context, but even without Assumption 1, the update steps can still be applied. In fact, this is equivalent to an implementation of the so-called (online) expectation-maximization (EM) algorithm [10] for latent MDPs. Thus Algorithm 1 with Algorithm 3 essentially results in combining L-UCRL and the EM algorithm. In terms of performance guarantees, using Algorithm 3 as a sub-routine for L-UCRL gives the same order of regret as in Theorem 3.3 as long as the true context can be almost reliably inferred:
Theorem 3.4 Suppose Assumption 1 holds withH > C · δ−4 log2(1/α) log(N/η) for some absolute constants C, δ > 0, and a parameter α > 0 such that α ln(1/α) ≤ δ2/(200S). If the initialization parameters satisfy equation (3) with some initialization error init ≤ δ2/(200 ln(1/α)), then with probability at least 1− η, the regret of Algorithm 1 is bounded by:
Regret(K) . HS √ MAN.
The proof of Theorem 3.4 is given in Appendix D.3. The provable guarantees are only given for wellseparated LMDPs. Nevertheless, we empirically evaluate Algorithm 1 as a function of separations and initialization (see Figure 1). An interesting consequence of Assumption 1 is that the length of episode can be logarithmic in the number of state and actions. With much longer time-horizons H ≥ Ω(S2A/δ2), [8, 18] assumed similar δ-separation only for some (s, a) pairs. While Assumption 1 requires a stronger assumption of δ-separation for all state-actions, the requirement on the time-horizon can be significantly weaker with large state and action spaces. For a more discussion on the separation condition, we refer the readers to Appendix D.1.

heading: 3.4 Learning LMDPs without Initialization
text: Finally, we discuss efficient initialization with some additional assumptions. Clustering trajectories is the cornerstone of all our technical results, as this allows us to estimate the parameters of each hidden MDP and then apply the techniques of Section 3.2. The challenge is how to cluster when we have short trajectories, and no good initialization. The key is again in Assumption 1. In Section 3.3, we use a good initialization to obtain accurate estimates of the belief states. These can then be clustered, thanks to Assumption 1, allowing us to obtain the true label in hindsight. Without initialization, we cannot accurately compute the belief state, so this avenue is blocked. Instead, our key idea is to leverage a predictive state representation (PSR) of the POMDP dynamics, and then show that Assumption 1 allows us to cluster in this space. Algorithm 4 gives our approach. We first explain the high-level idea, and subsequently detail some of the more subtle points. Suppose we have PSR parameters allowing us to estimate P(o|h‖do a), (the probabilities of any future observations o = (s′, r) given a history h and intervening action a) to within accuracy o(δ). We then show that we can again apply Assumption 1, to (almost) perfectly cluster the MDPs by true context at the end of the episode. After we collect transition probabilities at all states near the end of episode, we can construct a full transition model for each MDP. Learning the PSR to sufficient accuracy requires an additional assumption. We show that the following standard assumption on statistical sufficiency of histories and tests, is sufficient for our purposes (see also Section 2.2.1 and Appendix E.1):
Assumption 2 (Sufficient Tests/Histories) Let T andH be the set of all possible tests and histories of length l = O(1) respectively, with a given sampling policy π (e.g., uniformly random policy) for historiesH. T andH satisfy Condition 1 and 2 respectively.
While the worst-case instance may require l ≥M to satisfy the full-rank conditions, we assume that the length of sufficient tests/histories is l = O(1). In fact, l = 1 has been (implicitly) the common assumption in the literature on learning POMDPs [20, 5, 17, 25]. Empirically, we observe that the more MDPs differ, the more easily they satisfy Assumption 2. See Figure 2. At this point, we are not aware whether sample-efficient learning is possible with only Assumption 1. Though the main idea and key assumption are above, a few important details and technical assumptions remain to complete this story. The primary guarantee still required is that we have access to an exploration policy with sufficient mixing, to guarantee we can collect all required information to perform the PSR-based clustering. The following assumption ensures that additional Õ(M/α2) sample trajectories obtained with the exploration policy π can provide M clusters of estimated one-step predictions Pm(o|s, a) for every state s and intervening action a.
Assumption 3 (Reachability of States) There exists a priori known exploration policy π such that, for allm ∈ [M ] and s ∈ S, we have Pπm(sH−1 = s) ≥ α2 for some α2 > 0.
A subtle point here is that we still have an ambiguity issue in the ordering of contexts (or labels) assigned in different states, which prevents us from recovering the full model for each context. We resolve this issues ambiguity assuming the MDP is connected, and give the full description of Algorithm 4 in Appendix E.2. We conclude this section with an (informal) end-to-end guarantee:
Theorem 3.5 (Informal) Let Assumption 2 hold for an LMDP instance with a sampling policy π. Furthermore, assume the LMDP satisfies Assumptions 1 and 3. Then there exists an algorithm such that with probability at least 2/3, it returns a good initialization of LMDP parameters that satisfies (3) in time poly(Al, S,H,M, σ−1h , σ−1τ , α −1 2 , δ, init).
Theorem 3.5 completes the pipeline for learning in latent MDPs: we initialize the parameters by the estimated PSR and clustering (see Appendix E) up to some accuracy, and then we run L-UCRL to refine the model and policy up to arbitrary accuracy (Algorithm 1). Full version of Theorem 3.5 can be found in Theorem E.3.

heading: 4 Experiments
text: In this section, we evaluate the proposed algorithm on synthetic data. Our first two experiments illustrate the performance of L-UCRL (Algorithm 1) for various levels of separation and quality of initialization. Then, we empirically study the performance of the PSR-Clustering algorithm for randomly generated LMDPs for different levels of separation and time-horizon.

heading: 4.1 The Value of True Contexts in Hindsight
text: We first study the importance of receiving the true contexts in hindsight for the approach analyzed in this work, by comparing the performance of Algorithm 1 when instantiating it with Algorithm 2 or 3 as a sub-routine. We generate random instances of LMDPs of sizeM = 7, S = 15, A = 3 and set the time-horizon H = 30. The reward distribution is set to be 0 for most state-action pairs. We compare when we give a true context to the algorithm (Algorithm 2) and when we infer a context with random initialization or good initialization (Algorithm 3). In the latter, it is equivalent to
running the EM algorithm for the model estimation. For the planning algorithm, we use Q-MDP heuristic [32] which shows good empirical performance. We measure the model estimation error as error := minσ∈PermM ∑ (m,s,a) ‖(Pm − P̂σ(m))(s′, r|s, a)‖1, where PermM denotes all lengthM permutation sequences. The measured errors are averaged over 10 independent experiments. The experimental results are given in Figure 1(a). When the true context is given at the end of episode (with Algorithm 2), L-UCRL converges to the optimal policy as our theory suggests. On the other hand, if the true context is not given (with Algorithm 3), the quality of initialization becomes crucial; when the model is poorly initialized, the estimated model converges to a local optimum which leads to a sub-optimal policy. When the model is well-initialized, L-UCRL performs as well as when true contexts are given in hindsight.

heading: 4.2 Performance of L-UCRL with Good Initialization
text: In our second experiment, we focus on the performance of L-UCRL (Algorithm 1) along with Algorithm 3 under different levels of separation (δ in Assumption 1) when approximately good model parameters are given. For various levels of δ, we generate the parameters for transition probabilities randomly while keeping the distance between different MDPs to satisfy δ ≤ ‖(Tm1 − Tm2)(s
′|s, a)‖1 ≤ 2δ form1 6= m2. We show the error in the estimated model and average long-term rewards in Figure 1(b). When the separation is sufficient (larger δ or H), the estimated model converges fast to the true parameters. When the separation gets small (smaller δ or H), the convergence speed gets slower. This type of transition in the convergence speed of EM (the update of model parameters with Algorithm 3) is observed both in theory and practice when the overlap between mixture components gets larger (e.g., [29]). On the other hand, the policy steadily improves regardless of the level of separation. We conjecture that this is because the optimal policy would only need the model to be accurate in the total-variation distance, not in the actual estimated parameters.

heading: 4.3 Initialization with PSR and Clustering
text: In the third experiment, we evaluate the initialization algorithm (Algorithm 4) for randomly generated LMDP instances. Since PSR learning requires a (relatively) large number of short sample trajectories, we evaluate this step on smaller instances with S = 7, A = 2,M = 3. The LMDP instances are generated similarly as in the second experiment with different levels of δ and H . The reward and initial distributions are set the same across all MDPs. To learn the parameters of PSR, we run 106 episodes with H = 4. We assume histories and tests of length 1 are statistically sufficient with the uniformly random policy. In the clustering step, we run an additional 5 · 103 episodes to obtain longer trajectories of length H = 20, 40 and 80. We report the experimental results in Figure 2. We first observe how the level of separation δ between MDPs impacts trajectory separation, i.e., belief state vs true label (left). Recall that this separation property is the key for clustering trajectories. We then examine the performance of Algorithm 4 (see full Algorithm 5) for various levels of separation. Empirically, it succeeds to get a good initialization of an LMDP model when we have sufficient separation. As the separation level decreases, the algorithm starts to fail (Right). There are two possible sources of the failure: (1) the belief state is far from the true context, and (2) the similarity between MDPs drops theM th singular value of PT ,H (Middle). We can compensate for (1) if we
have a longer time-horizon to infer true contexts, as in the leftmost graph. For (2), if theM th singular value drops, we require more samples for the estimation of PSR parameters. In our experiments, as we decreased δ we found that failure in the spectral learning step was the more significant of the two.

heading: 5 Future Work
text: There are several interesting research avenues in continuation of this work. An interesting direction is to study RL algorithms for LMDPs with no underlying assumptions. Although our lower bound suggests such an algorithm necessarily suffers an exponential dependence in the number of contexts, if this number is small, such dependence might be acceptable. Specifically, we conjecture that general LMDPs can be learned with sample complexity of poly ( (HSA)M , −1 ) . For a special case when MDPs are deterministic, we show that the exponential dependence inM is sufficient In Appendix G. The case for general LMDPs is an interesting open question. Furthermore, a needed empirical advancement is to design efficient ways to learn the set of sufficient histories/tests for learning predictive state representation of LMDPs. This can dramatically improve the performance of our algorithms when a sufficiently good initial model needs to be learned.

heading: Acknowledgement
text: The research was funded by NSF grant 2019844, and by the Army Research Office and was accomplished under Cooperative Agreement Number W911NF-19-2-0333. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.

Review summaries:
This paper studies episodic RL in the regret minimization setup in Latent MDPs (LMDPs), which belong to the family of POMDPs. The paper presents a sample complexity lower bound for general LMDPs (i.e., in the absence of any restrictive assumption on the LMDP) that grows exponential with the number $M$ of contexts (or base MDPs) --- that is, learning a near-optimal policy in the worst-case may require an exponential (in $M$) number of episodes. It then presents L-UCRL, a UCRL2-style algorithm for LMDPs, which is proven to achieve $O(\sqrt{T})$ regret with high probability under some assumptions, where the regret bound depends polynomially on the number of state-action pairs and on the number of base MDPs.  

The authors in this work consider RL under the latent MDP model, where an unknown MDP is sampled from a known set of $M$ MDPs each episode. In this setting, the authors first show a general statistical lower bound which indicates that (approximately) learning optimal policy in the L-MDP setting requires exponential number of episodes (exponential in $M$). Given the hardness result, the authors first resolve the problem in the standard *online* framework where each episode is revealed to the learner after the episode itself: in this scenario, they provide an optimistic parameter estimation algorithm that eventually yields an efficient algorithm with sublinear regret. Finally, the authors show that the learner can actually infer these episodes after they are played assuming separation between the different MDPs and benign initialization (which is also investigated and relaxed into other assumptions such as test sufficiency and reachability of MDPs).

In this paper, the authors study efficient algorithms and hardness results for Latent MDPs. In a latent MDP, in each iteration, an MDP is randomly drawn from a set of M possible MDPs, and the agent interacts with that MDP without knowing its identity. The goal is to find a policy so that the expected cumulative reward is maximized, where the expectation is taken over the choice of the MDP and its stochasticity. The authors first show that without further assumptions, LMDPs are hard to learn. The authors then give a set of sufficient conditions so that LDMPs are learnable with polynomial sample complexity. One possible condition that makes LMDPs learnable with polynomial number of samples is the case when the identity of the randomly chosen MDP is revealed to the agent at the end of each trajectory. The authors further show that such identity can be learned if certain seperatedness condition holds. Preliminary experimental results are attached.