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