RL for Latent MDPs: Regret Guarantees and a Lower Bound

1 Introduction . 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 can not 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 can not 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 . 2 Preliminaries . 2.1 Problem Setup : Latent MDPs . 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 . 2.2 Predictive State Representation ( PSR ) . 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 ) . 2.2.1 Spectral Learning of PSRs in LMDPs . 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 .

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