Reinforcement Learning in Reward-Mixing MDPs
Authors: Jeongyeol Kwon
Yonathan Efroni
Constantine Caramanis
Shie Mannor
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
Partially Observable
Sample-Complexity
Method-of-Moments
Latent Confounders
Mixture Models

Abstract: Learning a near optimal policy in a partially observable system remains an elusive challenge in contemporary reinforcement learning. In this work, we consider episodic reinforcement learning in a reward-mixing Markov decision process (MDP). There, a reward function is drawn from one of $M$ possible reward models at the beginning of every episode, but the identity of the chosen reward model is not revealed to the agent. Hence, the latent state space, for which the dynamics are Markovian, is not given to the agent. We study the problem of learning a near optimal policy for two reward-mixing MDPs. Unlike existing approaches that rely on strong assumptions on the dynamics, we make no assumptions and study the problem in full generality. Indeed, with no further assumptions, even for two switching reward-models, the problem requires several new ideas beyond existing algorithmic and analysis techniques for efficient exploration. We provide the first polynomial-time algorithm that finds an $\epsilon$-optimal policy after exploring $\tilde{O}(poly(H,\epsilon^{-1}) \cdot S^2 A^2)$ episodes, where $H$ is time-horizon and $S, A$ are the number of states and actions respectively. This is the first efficient algorithm that does not require any assumptions in partially observed environments where the observation space is smaller than the latent state space. 

Reinforcement Learning in Reward-Mixing MDPs
Jeongyeol Kwon The University Texas at Austin maanchhungiuutexass....
Yonathan Efroni
Microsoft Research, NYC jonathan efroni wweethan.effroiippp. @gmail. com
Constantine Caramanis The University of Texas at Austin constantine@utexas. edu
Shie Mannor Technion, NVIDIA
shie@ee technion. ac i1, smannor@nvidh www.rriiiiiiii.ii. com
Abstract
Learning a near optimal policy in a partially observable system remains an clusive challenge in contemporary reinforcement learning. In this work, consider episodic reinforcement learning in a reward-mixing Markov decision process (MDP). There, a reward function is drawn from one of multiple possible reward models at the beginning of every episode, but the identity of the chosen reward model is not revealed to the agent. Hence, the latent state space, for which the dynamics are Markovian, is not given to the agent. We study the problem of learning a near optimal policy for two reward-mixing MDPs. Unlike existing approaches that rely on strong assumptions on the dynamics, we make no assumptions and study the problem in full generality. Indeed, with further assumptions, for two switching reward-modess the problem requires several new ideas beyond existing algorithmic and analysis techniques for efficient exploration. We provide the first polynomial-tin -omome algorithm that finds an e-optimal policy after exploring Opoly(( H, + 5 A²) episodes, where is time-horizon and S, A are the number of states and actions respectively. This the first efficient algorithm that does not require any assumptions in partially observed environments where the observation space is smaller than the latent state space.
1 Introduction
In reinforcement learning (RL), an agent solves a sequential docision-making problem in an unknown dynamic environment maximize the long-term reward [46]. The agent interacts with the environment by receiving feodback on its actions in the form state-dependent reward and observation.
One of the key challenges in RL is exploration: in the absence of auto-exploratory properties such as ergodicity the system, the agent has devise a clever scheme collect data from under-explored parts of the system. In a long line of work, efficient exploration in RL has been extensively studied for fully observable environments, i.e., under the framework Markov decision process (MDP) with a number of polynomial-tine algorithms proposed [34, 27, 43, 40, 18]. For instance, in tabular MDPs, ie., environments with a finite number of states and actions, we can achieve sample-optimality or minimax regret without any assumptions system dynamics [3, 50].
In contrast fully observable environments, very little is known about the exploration partially observable MDPs (POMDPs). In general, RL in POMDPs may require an exponential number of samples (without simplifying structural assumptions) [35, 31]. Therefore, is important consider natural sub-classes of POMDPs which admit tractable solutions. Previous studies focus on a special class of POMDPs where observation spaces are large enough such that a single observation provides
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

---

sufficient information on the underlying state [22, 35, 4. 12. 16, 30]. However, there are many applications in robotics, topic modeling and beyond, where the observations space is smaller than the latent state space. This "small observation space" setting is critical, as most prior work does not apply. Our work seeks to provide one of the first results to tackle a class of problems in this space.
In particular, we tackle a sub-class of POMDPs inspired by latent variable or mixture problems. Specifically, we consider reward-mixing MDPs (RM-MDPs). In RM-MDPs, the transition kernel and initial state distribution are defined as in standard MDPs, while the reward function is randomly chosen from one of M -reward models at the beginning of every episode (see the formal definition 1). RM- RM-MDPs are important in their own right. Consider, for instance, user-interaction model in a dynamical web system, where reward models vary across users with different characteristics [38, 23). A similar setting has been considered in a number of recent papers [8. 6, 23, 45, 7]. and most recently, and most directly related to our setting, [36].
Even for M = 2, the RM-MDP setting shares some of the key challenges of general POMDPs. This is because the best policy for the RM-MDP problem should account for not just the current state, but also an entire sequence of previous observations since some previous events might be correlated with a reward of the current action. And in the small observation space setting, prior work as in [22, 35, 4, 12, 16, 30] cannot be applied. The work in [36] develops an algorithmic framework for solving the RM-MDP problem in this small observation setting, but requires strong assumptions without which the algorithm falls apart. We discuss this in more detail below. Indeed, no work to date, has been able to address the RM-MDP problem, even for M 2. without significant further assumptions. This is precisely the problem we tackle in this paper.
Main Results and Contributions We focus on the problem of learning near-optimal policies in two reward-mixing MDPs, i.e., RM-MDPs with two reward models M = 2. To the best of our knowledge. no prior work has studied sample complexity of RL in RM-MDPs even for M 2 without any assumptions on system dynamics. Specifically, we provide the first polynomial-time algorithm which learns an c-optimal policy after exploring Oppoly H, c
A²) episodes. We also show that D25 A /2) number of episodes is necessary, and thus our result tight in S, A. This is the first efficient algorithm in a sub-class of partially observable domains with small observation spaces and without any assumptions on system dynamics.
On the technical side, we must overcome several new challenges that are not present in fully observable settings. One key hurdle relates to identifiability. In standard MDPs, an average of single observations from a single state-action pair is sufficient to get unbiased estimators of transition and reward models for the state-action. However, in RM-MDPs, the same quantity would only give an averaged reward model, where the average is taken over the reward model distribution. However, an average reward model alone not sufficient to obtain the optimal policy which depends on a sequence of previous rewards.
Our technique appeals to the idea of uncertainty in higher-order moments, bringing inspiration from algorithms for learning a mixture of structured distributions [19, 9, 26, 15). In such problems, the key for efficient learning is to leverage information from higher-order moments. Following this spirit, we consider estimating correlations of rewards at multiple different state-actions. A central challenge we must overcome is that in finite-horizon MDPs, it may be not possible to estimate all higher-order correlations with uniformly good accuracy. In fact, it may not be possible to estimate some correlations all when some pairs of state-actions are not reachable in the same episode. Therefore, we cannot simply rely on existing techniques that require good estimations of all elements in higher-order moment matrices. The main technical challenge is, therefore, to show that a near-optimal policy can still be obtained from uncertain and partial corvelations of state-actions. Our technical contributions are thus two-fold: (1) design of an efficient exploration algorithm to estimate each correlation up to some required accuracy. and (2) new technical tools to analyze errors from different levels of uncertainties in estimated higher-order correlations.
Related Work Efficient exploration in fully observed environments has been extensively studied in the framework of MDPs In tabular MDPs, a long line of work has studied efficient exploration algorithms and their sample complexity based on the principle of optimism [27, 43, 3, 18, 50, 47, 44] or posterior-sample [1, 43, 40, 41]. Beyond tabular settings, recent work seeks for efficient exploration in more challenging settings such as with large continuous state spaces [28, 32, 17].
2

---

All aforementioned work considers only fully observable environments where the complexity of exploration is significantly lower as there is no need to consider observation histories.
Most previous study on exploration in POMDPs focuses on large observation spaces such that a set of single observations is statistically sufficient to infer internal states, i.e,, the observation matrix has non-zero minimum singular value. Under such environments, previous work either considers a strong assumption on the uniform ergodicity of the latent state transitions [25, 22, 4] or the uniqueness of underlying states for every observation [35, 12, 16]. The exception is a recent work by Jin et al. [30] which studied efficient exploration without any assumptions on latent system dynamics.
One of the most closely related to problem to our work is a reinforcement learning problem with latent contexts (also referred multitask RL), which is considered in [6, 23, 36]. Previous approaches relied on several strong assumptions such as very long time horizon [6, 23], revealed contexts in hindsight or ergodic properties of system dynamics [36]. Our work takes a first step without such assumptions by focusing on environments where only the reward model differs across different tasks.
RM-MDP might be also viewed as special case of adversarial MDPs (e.g., [49, 39, 29I) with an oblivious adversary who can only play within a finite set of reward models. However, our goal is to find optimal policy in broader class of history-dependent policies, whereas in adversarial MDP literature they compare only to the best Markovian policy in hindsight. RM-MDD can be also considered an instance of latent variable models (e.g.. g., [14, 20, 13, 25, 24)), where sample reward sequences are influenced by unobserved confounders. R.M-MDP poses several new challenges as both challenges from reinforcement learning and statistical inference are interlaced by latent contexts.
2 Problem Setup
We define the problem of episodic reinforcement learning problem with time-horizon H in reward- mixing Markov decision processes (RM-MDDss) defined as follows:
Definition 1 (Reward-Mixing Markov Decision Process (RM-MDP)) Let M be tuple m= (Rm, De m==1) be a RM-MDP on a state space S and action space A where T x A S > [0, 1] is common transition probubility measures that maps state-action pair and a next state to a probability, and v is common initial state distribution. Let = |S| and reward A = =|A]. Wi...., WM are the mixing weights such that at the beginning of every episode, model Rm randomly chosen with probability Wm. reward model Rm, S x Ax {0, 1} + [0, 1] is probability measure for rewards that maps state-action pair and binary reward to a probability for m E [M].
(8, A, T, v, [m LM
In this work, we focus on a special case of RM-MDPs when M = 2, i.e., 2RM-MDPs. For the ease of presentation, assume all random rewards take values from {0, 1}, i.e. rewards are binary random variables. In this work, we assume a uniform-prior over each latent context such that w1 = w and W2 1 with found in Section 5. We consider a policy class II which contains all history-dependen policies (S.A, (0, 1)) S + A. The goal of the problem is to find a near optimal policy II that has near optimal value w.r.t. the optimal one, VM= := max Σm
1/2. Detailed discussions on these simplifying assumptions can be
- Ft] where E, -! is expectation taken over the th MDP with policy æ. More formally, wish our algorithm to return an (c, m) wwaablyapproximatllyy..o optimal policy, defined as follows:
Wm Ey
(PAC) optimal policy, which we also refer as a near
Definition ((c, -PCC optimal policy) An algorithm is (c. PPCC if it returns a policy à s.t.
P(VM VĀ 2e)2-11
2.1 Notation
We often denote a state-action pair (s, a) as one symbol x == (s, E S A. For any sequence (y1, ..., vt) of length I, we often simplify the notation as (1)1. for any symbol y. We denote 1 sum of any function f over random variable X conditioned on an event as [[[(X)) = Exex [f), where X is support of X. 1 {E} is an indicator function for any event E and
3

---

Algorithm 1 Learning Two Reward-Mixture MDPs
1: Run pure-exploration (Algorithm 2) to estimate second-order correlations
2: Estimate 2RM-MDP parameters M from the collected data (Algorithm 3)
3: Return मt the (approximately) optimal policy of M
P(E) is probability of any event & measured without contexts. If we use P without any subscript, it a probability of event measured outside of the context, i.e., P(-) = Em= Wmm Wmm (-))) We refer m the probability of any event measured in the m context (or in th MDP). If probability of an event depends on policy we add superscript æ to P. We denote VM as an expected long-term reward for model M with policy æ. We use denote empirical counterparts. For any set and d E N. A is d-ary Cartesian product over A. We use V to refer max(a, b) and b mean min(a,) b) for a, b E I..
Specifically for .M- 2: We use shorthand Pm (z) := Rm (r 1|x) for 1,2. Let expected averaged reward (x) == (x) + P2())), differences in rewards p- (z) := (pi - P2 (x)) and A p (x)
3 Algorithms
Before developing a learning algorithm for 2RM-MDP, let us provide intuition behind our algorithm. At a high-level, our approach lies on the following observation:
The latent reward model of 2RM-MDP can be recovered from reward correlations and the averaged reward.
Consider a simpler interaction model, in which we have perfect and exact access to correlations of the reward function. That is, suppose we can query for any (g)) E (S x A) 1)@2 the reward correlation function
p( mT) = [r' ry ar, my = (1) (z:)pp (12) + p2
Assume we can also query the exact expected average reward p+ (x) for all x E S x A. Suppose we can construct a matrix B E SAXSA indexed by state-actions is given by:
A, such that at its (i.j) entry
Br, p( x5)) p+ (x;).
Simple algebra shows that B = qq" where q E SAA is a vector indexed by x E S x A such that q. p- (z;). Hence, we compute the top principal component of B, from which we can recover
p- (x) for all x E S x A. Through the access to p- (I)) (ignore sign ambiguity issue for now) and the latent contexts via
p+ (I) can recover the probabilities conditioned
R1(r= 1|x) (+) + (z) and R2(r 1 px) p--(x.
(2)
Once we have the latent reward model
any approximate planning algorithm (e.g., point- based value iteration) to find an c-optimal policy.
However, when such exact oracle is not supplied, and the interaction with the 2RM-MDP is based on trajectories and roll-in policies, getting a point-wise good estimate of B is not generally possible. For example, some state-action pairs cannot be visited in the same episode, in which case we cannot get any samples for the correlations of such pairs of state-actions: p(z,, ) is completely unknown. In such a case, recovery of the true model parameters is not possible in general even if we estimate all reachable pairs without errors. Hence the main challenge is as follows: can we estimate a model from set of empirical second-order correlations such that the estimated model is close to the true model in terms of the optimal policy, even if is not close in model parameters?
For usual MDPs, states that cannot be reached pose no problems: we do not need to learn transition or reward probabilities of these states, since no optimal policy will use that state. Indeed, with this observation at hand, reward free exploration techniques [31, 33] can be utilized to learn a good model w.r.t. all possible reward functions (given fixed initial distribution). Our main conceptual contribution is to show how these techniques can be leveraged for the RM-MDP problem.

---

Overview of our approach. Our approach develops the following ideas.
1. Section 3.1: We define what we call a Second Order MDP, an augmented MDP whose states are pairs of states of the original MDP. Borrowing technology from single MDPs, we show that reward free exploration on the augmented MDP can be used to approximate second order moments of the original MDP, along with confidence intervals, obtained by how often a partieular pair of states can be reached.
2. Section 3.2: We show how find the parameters of a model whose second order moments are in the confidence set obtained from the augmented MDP. To do this, we show that we linear programming select an element of the uncertainty set that corresponds to valid second order moments for a model.
3. Section 3.3: The main contribution of the analysis, is then to show that for any model whose true second order statistics lie in this set obtained. approximate planning oracle is guaranteed to return an e-optimal solution to the true model. Theorem 3.2 shows that sample complexity depends quadratically in (S-A). This is well-expected, given that our algorithm must compute statistics on the augmented graph. We next give a matching lower bound: Theorem 3.3 shows that this quadratic dependence is unavoidable.
3.1 Pure Exploration of Second-Order Correlations
Our first goal is to collect samples for pairs of state-actions through exploration. This will ultimately allow us to estimate the correlations of the reward function, as well as the average reward per-state p+ (z). The backbone of pure-exploration is adaptive reward-free exploration schemes studied in
standard MDP settings [33]. We first consider surrogate model that helps to formulate data collection process for correlations of pairs. Specifically, we first define the augmented MDP:
Definition 3 (Augmented Second-Order MDPs) An augmented second-order MDP M is defined on a state-space S and action-space À where
$ (i, (1,2,3), u: (v', (0) E ((S A) (null))
Ā = a E (0,
Then in M, an augmented state (ie, Ut, 81) evolves under an action (at, z4) as follows:
i 1, = (roull, rull), , 81 ~ v()), 8141 ~ (1)),,
ui is Ej (2)/(t,)
fs i 2 = 0 ir 3 oft1 as)
if = 0 or in = 3
441
(3)
i, +1
else
else
The second-order augmented MDP consists of an extended state space where the first two coordinates i, v store previously selected state-actions followed by the coordinate s for the current state. The action space is product of original actions and choice actions 2, where 2 = 1 means that select the current state-action to be explored as a part of a pair in the episode. Specifically, when 2 = 1, by design we increase the count If i reaches 3, it means we collected sample of a pair (u', u²) in the episode.
Initially the true model parameters are unknown and thus are not aware of how to collect samples for any pair of state-actions or how much samples we need for each pair. order to resolve this, we use ideas from reward free exploration for MDPs: we consider the upper confidence error bound function Q that follows from the Bellman-equation for M with (pure) exploration bonus:
(1,0,) = 1 (i= 2 (IA n(or)) br(s, = =(1A
'T
n(ss
Ve((i,
1 (b,(i, v,22 + +E (g ((+r(((, (l((),)))
(4)
where 1 (i = =0==1} is an indicator of whether to collect samples for correlations between state- actions v/ Here, v' andd v' are first and second coordinates the next state following the transition rule (3), 12 O(log(K/n)). LT O(S log K (/n)) are properly set confidence interval parameters,
5

---

Algorithm 2 Pure Exploration of Second-Order Correlations
1: Initialize (()() = Vo = 1, n(v) = for all v E (S 2
2: while Vo <pe do
3: Get an initial state 91 for the kth episode. Let 01 (nulll, null),
4: for t = 1, ..., H do
5:
Pick (ac, 22) = arg (a.se) (i st), (a, 2)). Play action dN, observe next state 8 and reward Tr-
6:
7:
i - and = = 1 then
8:
Te f- To-
9:
end if
10: Update (i++1, Me+1, S1+1) according to the choice of and following the rule in (3)
11:
end for
12:
if i = 3 then
13:
n(UH+1) + n(V=+1) +1, ((n++1) + (1 (1-1/n(vn+1)) ((H+1)) +re
14:
end if
15:
Update D, f. 2+ from the trajectory (s, a, r)HH, and then update Q and V using (4), (5)
end while
K is the total number of episodes to be explored and QH+ 1 (+) greedy action with respect to Q and collect sufficient data for pairs of state-actions by pure exploration on the augmented MDP. We repeat the pure-exploration process until Vo becomes less than the threshold épc where
O. At every episode, we choose a
Vi v, 8) max (1) v, s), (a), Vo Vv/K + Ee(s) (1,1,,1,)) (w',=))e
(5)
and fr = (//n)) a confidence interval parameter for initial states. The eure-exploration procedure is summarized in Algorithm 2. The main purpose of Algorithm 2 is to balance the amount of samples for correlations.
3.2 Recovery of Empirical Model with Uncertainty
Once we gather enough samples to estimate correlations of pairs of state-actions we are assured that Vo S fpe- Given this, we describe an efficient algorithm to recover good estimate of the true 2RM-MDP parameters. We formulate the model recovery problem linear program (LP). Recall that for any 24,2y 2y (S A) Di; Zy:
compute the multiplication of differences two positions
u(,,, 2,) == p- (z;)p- -
(6)
Ideally with exact oracles for getting values of 21)) can recover p- (xx through LP: let ((x) = log lp- (x)l for all x E S x A. Then we first find a solution for the following LP:
((x;) + (æ,)) = log [u( - (x.)) Vz,, E
After solving for 1(x), we can find consistent assignments of signs sign(z) such that p- (x) = sign(x) exp(((x)), which can be formulated as 2-Satisfiability problem [2].
However, after the pure-exploration phase, only have estimates of correlations ((z;, 2zy) and confidence intervals b(2,, .2y) := /ig/n(z,,,;)+tt where Eg e/ H2 is a small tolerance parameter (as well as estimates of the expected average reward (z) and its confidence interval b(x) == (12/n(z)). Note that (()) is the number of counts that pair (II, is counted during pure- exploration phase. Let ú(2,, 24)) == (x:)pp_aa (x,, be an empirical estimate of u(x,, We want to recover (z) sign(x) exp(((x)) such that the following is satisfied:
(f,(z, 23)) + 24 (n:) 10 (αz)) ú(x, 25)1 S b(,,, 2y), Vz,, J E x A.
We present in Appendix B.1 the LP formulation to recover ?(xx) and sign(x). The outline of procedure is stated in Algorithm 3 and more details are presented in Algorithm 4. In order for Algorithm 3 to succeed, we need the following lemma on the existence of feasible solutions in the L.P with high probability:

---

Algorithm 3 Reward Model Recovery from Second- Correlations
1: Solve an LP for ((x) and find sign(æ) for a E x A that satisfies (7).
2: Clip [(x) within (-oo, û(z)] where ((x) = log (min( (I), ,, 1 P+ (x)) for all z E 8 x
3: Set (=(()) = sign(x)) exp(((x)) for all x x
4: Let (z) = (x) + p_(()_()))))))))))() (2(2) )= = (+(x) (- for all a 8 A.
5: Return M = (S,A, T, D, Rm 1 m= an empirical 2RM-MDP model.
Lemma 3.1 Suppose that we set the confidence parameters as 11 = 12 (llog(SA/n) and by = LT = (S log(SA/m) Then with probability at least 1 - 17, Algorithm returns a model M that satisfies constraints for all Iu, 3. E S x A.
Given Lemma 3.1, Algorithm is able to find at least one feasible solution, i.e., model M, in polynomial-ting -ime.. Via such a solution, we obtain p- IRSA and can recover an estimate of the latent model Pı (s, a) R (r 1|s, a) and p2 a) R2(r 1|s, a) for all state-action pairs by a plug-in estimate (2). We refer to such model as the empirical model. The proof of Lemma 3.1 is given in Appendix B.2.
3.3 Main Theoretical Results
Upper Bound We first prove the correctness of our main algorithm (Algorithm 1) and compute an upper bound on its sample complexity. Our upper bound has a quadratic dependence on S A. This is well-expected because must estimate quantities related to the augmented second order MDP. Our lower bound shows that this dependence is in fact tight. Additionally, our bounds exhibit an interesting dependence the error threshold, c, that in fact depends the minimum separation in distinguishable state-actions:
8 = 1
min A(x).
(8)
F =E))):A())>0
where A(x) (x)) This is reminiscent of similar results for mixture models (e.g., [10,37]) We state the following theorem on the sample-complextty of 2RM-MDPs.
Theorem 3.2 There exists a universal constant C > O such that if we run Algorithm 1 with Epe
17 (DV ) then Algorithm 2 terminates after at most K episodes where.
K K-C-H² S (1 (H+
( v -lgg(HSA/e)) log (H/e).
with probability at least 1 - 11. Farthermore, under the same high probabilistic event, Algorithm I returns an e-optimal policy.
Theorem 3.2 shows that we can find a e-optimal policy for any 2RM-MDP instance with at most O (poly (H, e --- S2 A²) episodes of exploration. We leave it as future work to investigate whether dependency on e can be uniformly O(1/e²) regardless of 8, as well as whether the dependency on H can be tightened.
Lower Bound Theorem 3.2, the sample complexity guaranteed by Algorithm is A2) We show that the dependency on 82 A2 is also necessary for two reward-mixing MDPs:
Theorem 3.3 There exists class of 2RM-MDPs such that order to obtain e-optimal policy, we need at least (S5 A2 episodes.
The proof of the lower bound can be found in Appendix C.
7

---

Computational Complexity Main bottlenecks of Algorithm 1 are in two parts: (1) solving the LP in Algorithm 3 and (2) computing the optimal policy for M. For (I), with SA variables and O(5² A2) constraints, solving the LP in Algorithm 3 takes poly(S, A) times (see [11] and the references therein for some known computational complexity results for solving LPs). For (2), point-based value-iteration algorithm [42] runs in O ( time to obtain e-optimal policy for any specified 2RM-MDP model. Therefore, the overall running time of Algorithm is polynomial in all parameters.
4 Analysis Overview
For the ease presentation, we assume that the true transition and initial state probabilities are known in advance. In Appendix B, we complete the proof of Theorem I without assuming known transition and initial probabilities but instead using and T estimated during the pure-exploration phase.
We first note that given any policy π II, difference expected rewards can bounded by 41- statistical distance in trajectory distributions. More specifically, consider the set of all possible trajectories T = (S x A x {0, 1}) f that is, any aine--actioo--wwadd sequence of length H. Then,
VV V
H (PT ((x))))
Σ TET
P (T) PT((T))
where 1pT (.) is the probability of an event measured in empiriia model M with policy æ. Therefore, for any policy π E II, our goal is to show that | PP*(T) & ((e/H)))))))))) i... the true and empirical models are close in ( statiiistaal distance in trajectory distributions.
model
(T)
The main challenge in the analysis is to bound the 1 distance when estimated first and second-order moments within the range of confidence intervals. We show that when all pairs of state-actions in were visited sufficient amount of times then we can bound the / difference. Let us first divide set of trajectories into several groups depending on the number of times that every pair in each trajectory has been explored during the pure-exploration phase. We first define the following sets:
2 = {(x;, E (S x A)@ @2 n( x3)) Z my): 3, = (I T Vt,, E [H], t + to s.t. (Ze, F12) E X;),
(9)
where n4 C. 3,,2 for some absolute constant C > O, c0 = H2 and e 26 for 2 1. 5, = max E1) is a threshold for distinguishable state-actions such that A(x) 2 6, In nut-shell, & is set of sequences of state-actions in which every pair of state-actions has been explored least my times. Then we split a set of trajectories into disjoint sets & En and & & n N& and for I E [ 1], ie., a set of trajectories with all pairs visited more than and at least one pair visited less than 74-1, where let L, O(log( Hf))) be the largest integer such that ny C 12- Let NL+11 = 0 and EL+11 1.
Recall that our goal is to control the I stattiiiti distance between distributions of trajectories for any history-dependent policies that resulted from true and empirical models. With above machinery, we can, instead, bound the I statistical distance between all trajectories as the following:
L+1 Σ
L+1
[ - -P-((T)
Σ TEEEEEE
(P - PTT(T))
<
Σ I
sup PT (Z1H E E) -(Ha), (10)
2
xell
where (Z E E),, is the probability that the random trajectory т observed when following policy is contained within the set Ej. The following lemma is critical in bounding the above inner sum for each fixed l:
Lemma 4.1 (Eventwise Total Variance Discrepancy) For E (0, 1,.,, L+1] and any history-
dependent policy π E IT, we have:
Σ -"(rr-1P"(1) S sup Pa (* E EEL
(11)
nell
The second key connection the relation between sup (8) and the data collected in pure- exploration phase. The following lemma connects them:

---

Lemma 4.2 (Event Probability Bound by Pure Exploration) With probability at least 1 - Algo rithm 2 terminates after at most episodes where
K 2 C C AHH + A)c log (SSAA/(cppT)))
(12)
with some absolute constant C > Furthermore, for every E [L + we have
sup P" (11. E EE) S 50) nell
(13)
Given the above two lemmas, we set Epc = cdo/ H We then apply Lemma 4.2 and (10) to bound difference expected value of true and empirical models:
V Vil H - S (H S O(c).
Finally, plugging Epe edu/ L L into the required number of episodes (12) gives sample-complexity bound presented in Theorem 3.2. Full proof of Lemma 4.2 is given in Appendix B.3.
4.1 Key Ideas for Lemma 4.1
The core of analysis by which we establish Lemma 4. goes as follows. We split the state-action pairs into two sets distinguishable and indistinguishable state-action pairs. Fix an / {0. I...L 1] and consider then set E/. We call a state- state-action action pair x E S A 6-distinguishable if A(x) larger then some 8 0. Then, by the definition of (xx), implies that pi (x) and (2r)) are sufficiently different. For a trajectory r E E/, we say that a state-action æ within r is 6--distinguishable if A (x) 2 &t-
States that are distinguishable have the following property. If pairwise correlations between three distinguishable state-actions are well estimated, then we can recover individual reward probability for each of the three state-actions as in the following lemma:
Lemma 4.3 Suppose that (21,82), (22, rs). (03, 21) E 2 and A(xa) 2 for all i (1,2,3). Let i" arg max ((x;), and in arg max A(r;). Then we have
i ((x))) - A(x;)) S E/2,, and [(x;,) A(x))(()))))))))))))))
1/2
A
If all reward parameters in a trajectory are relatively well-estimated, then we can bound the error in the estimated and true probability of the trajectory. Thus, if trajectory E Eÿ contains more than two distinguishable state-actions, we can rely on the closeness in reward parameters to bound (T) - 1p (r)l- On the other hand, if there are no such three distinguishable Mate-actions in т, then we cannot guarantee the correctness of reward parameters in In such case, we take a different path to show the closeness trajectory distributions. Full proof given Appendix
5 Discussions and Future Work
We developed the first efficient algorithm for RL in two reward-mixing MDPs without any assumptions on system dynamics. We conclude this work with discussions on our assumptions and several future directions.
Mixing Weights. When the prior over contexts is not uniform, i.e., w of 1/2, parameter recovery procedure (Algorithm 3) may be more complicated due to the sign-ambiguity issue in p- (x). When w = 1/2, the estimated model is identical whether we use P- (x) ~ p- (x) or (z). However, the model no longer identical when we have correct signs p (=()) or not if w f 1/2. In Appendix E, we discuss the extension to non- non-uniform mixing weights in great detail. Once we can solve with any known E [o, 1). then one possible solution to handle unknown mixing weight is to sweep all possible mixing weight for from 0 to I with discretization level O(e/H), recover an empirical model obtained by solving an LP assuming known for all discretized weights and pick the best policy by testing each policy the environment.
Non-uniform/Unnn
p- (xx)

---

Structured Reward Distributions. Our approach can be extended to structured reward distribu- tions, e.g., distributions over finite supports. For instance, if reward can take different values in {gi. 92, q1), we redefine y( x2) := - Pm (x) Pm (x2) T as a & x I mooment-matrin with setting Pm. (x) = [R (r qı (1)z), R. (r ()).... Rm = 9r [z) Then, we can extend an LP in Algorithm 3 to find (m) to match ; within confidence intervals. This approach would also pay extra -olatt)-factorr both in sample and computational complexity. Then, once we can verify that the recovered difference in reward models is close, i.e., (1) p-(2)) O(+), then the remaining proof for Theorem 3.2 would follow similar to the Bernoulli reward case. We leave more thorough investigation in this direction as future work.
Extensions to More Than Two Reward Models M > 2. As mentioned earlier, RM-MDP with M reward models an MDP analogy learning mixture of M product distributions over reward distributions. If every higher-onder correlation of interest between multiple state-actions can be estimated, we may apply existing methods for learning a mixture of product distributions over binary field [20, 19, 26]. For instance, Jain and Oh [26] proposed a spectral method which requires all off- off-diagonal elements up to third order moments to learn parameters of . M-mixtures of ne-rodduct distributions. To apply their method, we may estimate correlations of all three different as similarly in Algorithm then reward model recovery of RM-MDP can be converted learning a mixture of S.A-prouuce reward distributions. However, as we have already seen in this paper, the real challenge is in recovering the model from correlations of unevenly reachable sets of state-actions, and then analyzing errors with different levels of uncertainty in correlations. We believe this is very interesting direction for future research.
In addition to the above mentioned issues, we believe there are a number of possible future directions. While we specifically focused on RM-MDP problems, the more general problem of Latent MDP [36] has not been resolved even for MM = 2 without any assumptions. We believe that it would be an interesting future direction to develop efficient algorithms for Latent MDPs with small M. We can also consider more technical questions in RM-MDPs, e.g., what would be the optimal sample-complxxng dependency on H, e, and whether we can develop a fully online algorithm to minimize regret. We leave them as future work.
Acknowledgement
The research was funded by NSF grant 2019844, and by the Army Research Office and was accom- plished under Cooperative Agreement Number W91 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.
The views and conclusions
References
[1] J. Asmuth, I.. Li, M. I.. Littman, A. Nouri, and D. Wingate. A bayesian sampling approach to exploration in reinforcement learning. In Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence, pages 19-26, 2009.
[2] B. Aspvall, M. F. Plass, and R. E. Tarjan. A linear-time algorithm for testing the truth of certain quantified boolean formulas. Information Processing Letters, 8(3):121-123, 1979.
[3] M. G. Azar, Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning, pages 263 263-272. PMLR, 2017.
[4] K. Azizzadenesheli, A. Lazaric, and A. Anandkumar. Reinforcement learning of POMDPs using spectral methods. In Conference Learning Theory, pages 193-256, 2016.
[5] B. Boots, S. M. Siddiqi, and G. J. Gordon. Closing the learning-planning loop with predictive state representations. The International Journal of Robotics Research, 30(7):954- 2011.
[6] E. Brunskill and I.. Li. Sample complexity of multi-task reinforcement learning. In Uncertainty in Artificial Intelligence, page 122. Citeseer, 2013.
[7] Buchholz. and Scheftelowitsch. Computation of weighted sums of rewards for concurrent MDPs. Mathematical Methods of Operations Research, 89(1):1-42, 2019.
10

---

[8] I. Chadès, J. Carwardine, T. Martin, S. Nicol, R. Sabbadin, and O. Buffet. Momdps: solution for modelling adaptive management problems. In Twenty-Sixth AAAI Conference on Artificial Intelligence (AAAI-12), 2012.
[9] S.-O. Chan, I. Diakonikolas, X. Sun, and R. A. Servedio. Learning mixtures of structured distributions over discrete domains. In Proceedings of the twenty-fourh annual ACM-SIAM symposium on Discrete algorithms, pages 1380-1394. SIAM, 2013.
[10] Y. Chen, X. Yi, and C. Caramanis. Convex and nonconvex formulations for mixed regression with two components: Minimax optimal rates. IEEE Transactions on Information Theory,
64(3):1738-1766 2017.
[11] M. B. Cohen, Y. T. Lee, and Z. Song. Solving linear programs in the current matrix multiplication time. Journal of the ACM (JACM), 68(1):1-39, 2021.
[12] C. Dann, N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. On oracle- efficient PAC RL with rich observations. In Advances in neural information processing systems, pages 1422-1432, 2018.
[13] S. Dasgupta. Learning mixtures of gaussians. In 40th Annual Symposium on Foundations of Computer Science (Cat. No. 99CB37039), pages 634-644, IEEE, 1999.
[14] R. D. De Veaux. Mixtures of linear regressions. Computational Statistics & Data Analysis, 8(3):227-245, 1989.
[15] I. Diakonikolas, G. Kamath, D. Kane, J. Li, A. Moitra, and A. Stewart. Robust estimators in high-dimensions without the computational intractability. SIAM Journal on Computing. 48(2):742-864, 2019.
[16] S. Du, A. Krishnamurthy, N. Jiang, A. Agarwal, M. Dudik, and J. Langford. Provably efficient with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665-16744 2019.
[17] S. S. Du, S. M. Kakade, R. Wang, and L. F. Yang. Is a good representation sufficient for sample efficient reinforcement learning? International Conference on Learning Representations, 2019.
[18] Y. Efroni, N. Merlis, M. Ghavamzadeh, and S. Mannor. Tight regret bounds for model-based reinforcement learning with greedy policies. Advances in Neural Information Processing Systems, pages 1222-11223 2019.
[19] J. Feldman, R. O'Donelll and R. A. Servedio. Learning mixtures of product distributions over discrete domains. SIAM Journal on Computing, 1566)115661 2008
[20] Y. Freund and Y. Mansour. Estimating a mixture of two product distributions. In Proceedings of the twelfth annual conference on Computational learning theory, pages 53-62, 1999.
[21] A. Garivier, P. Ménard, and G. Stoltz. Explore first, exploit next: The true shape of regret in bandit problems. Mathematics of Operations Research, 44(2):377-399, 2019.
[22] Z. D. Guo, S. Doroudi, and F. Brunskill. A PAC RL algorithm for episodic POMDPs. In Artificial Intelligence and Statistics, pages 510-518, 2016.
[23] A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv preprint 2015.
1502.
[24] J. Hoffmann, S. Basu, S. Goel, and C. Caramanis. Learning mixtures of graphs from epidemic cascades. In International Conference on Machine Learning, pages 4342-4352. PMLR, 2020.
[25] D. Hsu, S. M. Kakade, and T. Zhang. spectral algorithm for learning hidden markov models. Journal of Computer and System Sciences, 78(5):
2012.
[26] P. Jain and S. Oh. Learning mixtures of discrete product distributions using spectral decomposi- tions. In Conference on Learning Theory. pages 824 856. PMLR, 2014.
[27] T. Jaksch, R. Ortner, and P. Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Reseurch, 11 11:1563-1600, 2010.
[28] N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. Contextual decision processes with low bellman rank are pac-leaeeable. In International Conference on Machine Learning, pages 1704-1713. PMLR, 2017.

---

[29] C. Jin, T. Jin, H. Luo, S. Sra, and T. Yu. Learning adversarial mdps with bandit feedback and unknown transition. arXiv preprint arXiv:1912.01192, 2019.
[30] C. Jin, S. M. Kakade, A. Krishnamurthy, and Q- Liu. Sample-efficient reinforcement learning of undercomplete POMDPs. arXiv preprint arXiv:2006. /2484, 2020.
[31] C. Jin, A. Krishnamurthy, M. Simchowitz, and T. Yu. Reward-free exploration for reinforcement learning. arXiv preprint arXiv:2002.0222
2020.
[32] C. Jin, Z. Yang. Z. Wang, and M. I. Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137-2143. PMLR, 2020.
[33] E. Kaufmann, 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.
[34] M. Kearns and S. Singh. Near-optimal reinforcement learning in polynomial time. Machine learning, 49(2):209-223 232, 2002.
[35] A. Krishnamurthy, A. Agarwal, and J. Langford. PAC reinforcement learning with rich observa- tions. In Advances Neural Information Processing Systems, pages 1840-1848, 2016.
[36] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. RL for latent MDPs: Regret guarantees and a lower bound. arXiv preprint urXiv:2102.04939 2021.
[37] J. Kwon, N. Ho, and C. Caramanis. On the minimax optimality of the EM algorithm for learning two-component mixed linear regression. arXiv preprint arXiv:2006.00 02601, 2020.
[38] O.-A. Maillard and S. Mannor. Latent bandits. In International Conference on Machine Learning. pages 136-144, 2014.
[39] G. Neu, A. Gyorgy, and C. Szepesvári. The adversarial stochastic shortest path problem with unknown transition probabilities. In Artificial Intelligence and Statistics, pages 805-813, PMLR, 2012.
[40] I. Osband and B. Van Roy. Model-based reinforcement learning and the eluder dimension. In Advances in Neural Information Processing Systems, pages 1466-1474, 2014.
[41] I. Osband and B. Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In International Conference on Machine Learning, pages 2701-2710. PMLR, 2017.
[42] J. Pineau, G. Gordon, and S. Thrun. Anytime point-based approximations for large POMDPs. Journal of Artificial Intelligence Research, 27:335-380 380, 2006.
[43] D. Russo and B. Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances in Neural Information Processing Systems, 2013.
[44] M. Simchowitz and K. G. Jamieson. Non-asymptotic gap-dependent regret bounds for tabular mdps. In Advances in Neural Information Processing Systems, pages 1153-1162, 2019.
[45] L. N. Steimle, D. L. Kaufman, and B. T. Denton. Multi-model markov decision processes. Optimization Online URL http://www. optimization-online org/DB F1LE/2018001/6444 pdf. 2018.
[46] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[47] A. Tossou, D. Basu, and C. Dimitrakakis. Near-optimal optimistic reinforcement learning using empirical bernstein inequalities. arXiv preprini arXiv: 1205. 2019.
[48] M. J. Wainwright. High--imensional statistics: A non-asympiotic viewpoint, volume 48. Cam- bridge University Press, 2019.
[49] J. Y. Yu, S. Mannor, and N. Shimkin. Markov decision processes with arbitrary reward processes. Mathematics of Operations Research, 34(3):337-757 2009.
[50] A. Zanette and E. Brunskill. Tighter problem-dependeet regret bounds in reinforcement learning without domain knowledge using value function bounds. In International Conference on Machine Learning, pages 7304-7312, PMLR, 2019.
12


Reviews:
[{"review_id": "t9a0Xo57AtG", "reviewer": "Reviewer_TkXV", "summary": "This paper studies episodic reinforcement learning (RL) in a finite-horizon reward-mixing Markov decision process (RM-MDP). Specifically, the authors have set up the problem in Section 2; and motivated, developed, and analyzed an algorithm (Algorithm 1,2,3) in Section 3. The main theoretical results are summarized in Theorem 3.2 (upper bound on the sample complexity) and Theorem 3.3 (lower bound on the sample complexity). The authors have provided an overview of the analysis in Section 4, and discussed potential future work in Section 5.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I become more positive after reading the rebuttal and other reviews. In particular,\n\n1) I still think the M=2 limitation is a major issue of this paper. In particular, as the authors have discussed, I do not think the methodology developed in this paper can be easily extended to cases with M>2. However, as other reviewers have mentioned, M=2 is already a progress.\n\n2) For the simplifying assumptions (e.g., the rewards are Bernoulli, the transition model is fixed, and the mixing weights are known), if the authors can well explain how to relax them in a response and/or in a section in the appendices, then I am fine with these assumptions.\n\nThe following is the original review.\n***************************************************************************\n\nThis paper is interesting and in general well written. The idea of using higher-order moments to develop efficient RL algorithms for RM-MDPs is interesting and thought-provoking. However, I do not think the current version of this paper has met the bar of a top-tier machine learning conference, due to the reasons listed below:\n\n1) My main concern is that the RM-MDPs considered in this paper are too restrictive. Specifically, in this paper, the authors have assumed that (1) M=2, where M is the number of reward models, (2) the rewards are Bernoulli, (3) the transition model is fixed, and (4) the mixing weights are known. These assumptions, especially the first two, seem to be very strong. Moreover, as the authors have discussed in Section 5, these assumptions are essential for the proposed algorithm to work. For instance, relaxing Assumption (1) or (2) might require a complete redesign, rather than a straightforward modification, of the proposed algorithm.\n\nThe authors have already mentioned some ideas of extending this paper to more general settings. I think this paper will be much stronger if the authors can relax Assumption (1) to general M>=2 and Assumption (2) to structured reward distributions, and propose a more general algorithm under these relaxed assumptions.\n\n2) For Section 2 (Problem Setup), it would be great if the authors can further clarify which parts of the RM-MDP model are unknown to the agent before it interacts with the model. Also, for a history-dependent policy \\pi, it only depends on the history in the current episode, but not the history before the current episode, right? This part is not clear from the paper, please clarify.\n\n3) I think some numerical experiments can further strengthen this paper. For instance, the authors might provide a numerical experiment to show how to apply the developed algorithm in a simplified dynamic web problem mentioned in Section 1. The authors might also use numerical examples to check \n\n    - if the upper bound developed in Theorem 3.2 has a uniform dependence on \\epsilon\n    - the performance of the proposed algorithm when the mixing weights or M are misspecified (e.g. when M is actually 1 but the agent thinks M=2)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I have discussed the limitations of this paper above.\n\nThis is a theoretical paper, and it is hard to speculate its negative societal impacts.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "qvHFMSTBHLQ", "reviewer": "Reviewer_KvwV", "summary": "This paper studies exploration in a particular latent variable MDP setting where the reward function is random chosen from one of M-reward models at the beginning of each episode. The authors restrict themselves to the setting where M=2, and provide a polynomial time algorithm to learn a $\\epsilon$-optimal policy. The authors do this by developing an algorithm that approximately estimates correlations of rewards at different state-actions entries in order to build an approximate model under which they can perform planning. ", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality\n\nThis paper studies to a novel theoretical setting to the best of my knowledge, albeit an arguably narrow setting. \n\nQuality\n\nOverall, I found this paper to be above bar from an execution perspective. The majority of the analysis is provided in the appendix, but from a high level builds on the analysis in recent work by Kwon et al. 2021. It should be noted however that their polynomial bounds seem to be a consequence of their choice of considering the more restricted M=2 setting. \n\nClarity\n\nOverall, the presentation of the paper is very good. I found no issues in terms of clarity.\n\nSignificance\n\nThe two main concerns I have with this paper (and the main reason I am hesitant to recommend acceptance) are the following. First, this work considers a particular case (M=2) of a subtype of latent MDPS (where only the reward function changes at the beginning of each episode) which are themselves a special class of POMDPs. I found this particular setting to be somewhat lacking in motivation. Second, this work makes a subtle assumption that the mixing weights are provided a priori, which the authors suggest might be solved by a discretizing scheme. This would however be clearly challenging for a larger number of components (M>2) where as the authors point out there are additional conceptual difficulties. Despite this, there is a possibility that this work could be of interest to subsets of the NeurIPS community, so I would consider raising my score if the above concerns are addressed. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I felt this authors did an above average job when it came to being upfront about the limitations of their work. My overall impression of this paper is more borderline than my score in part due to this. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "cyEpotUXoPn", "reviewer": "Reviewer_qQka", "summary": "This paper proposes a novel learning algorithm in a specific class of POMDPs known as reward-mixing-MDPs.  The paper blends clever representations, theoretical bounds, and a clean policy learning algorithm to help identify the unknown reward function for the special case of M=2 possible reward functions.  The paper is purely theoretical, and presents no experiments.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I liked this paper.  While it is outside my area, and I did not understand it deeply, and while I have significant concerns about its significance, I think it makes some important contributions, and I recommend acceptance.  Specifically:\n\n* A strength of the paper is its generality.  The paper makes very few assumptions on the nature of the RM-MDP to be solved, and yet is able to give an algorithm with polynomial complexity.\n\n* The paper explicitly considers the (realistic) case of inexact, incomplete oracles.\n\n* The idea of the \"Second-Order MDP\" is intriguing, and a helpful auxiliary construct that leads to the algorithm; it may be useful in other contexts.\n\n* The use of LPs is clean and tidy.  Love it.\n\n* The paper is generally well-written and understandable.  I was grateful that the authors took the time to build intuition about the algorithm and the logical flow of the bounds.\n\nIn terms of NeurIPS criteria:\n\n* Originality: the paper seems original.\n* Quality: the paper is of very high quality, with a solid logical flow, good references, and clear writing.\n* Clarity: the paper is heavily theoretical, and therefore is heavy on notation.\n* Significance: it is not clear how significant this paper is, mostly because the RM-MDP formalism seems fairly contrived, and the results are further limited to the case of M=2.  The paper does not give even passing examples of what sort of real-world phenomena this might actually model, so the impact of the paper will largely be on a small community of like-minded scholars.  The best hope is that it will spark additional results on more realistic model classes, but this seems somewhat unlikely given the close connection between the solution methodology and the specific assumptions of the model class.\n\nAs a minor note, I wasn't sure how useful Thm. 3.2 is, mostly because I am interested in the \"other\" type of bound: assuming that the RM-MDP exhibits some sort of structure, can we do much better than S^2A^2 sample complexity?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This is not applicable, given the nature of the work.", "review_text": "", "extra_scores": {"ethical_concerns": "There are no ethical issues with this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "NFUcld5cv2r", "reviewer": "Reviewer_EDyq", "summary": "The paper proposes an algorithm to learn __reward-mixing MDPs__ in polynomial time. Reward mixing MDPs (RM-MDPs) are finite-horizon MDPs where, in each episode, the reward function is selected randomly from a set of possible reward functions. In this sense, RM-MDPs can be seen as a particular case of a POMDP where part of the state (the current reward function) is not observable. \n\nThe paper considers the case where the set of rewards contains only 2 binary reward functions, selected with uniform probability, and proposes a PAC algorithm to learn a near-optimal policy for RM-MDPs with large probability. To do so, the paper defines two quantities, $p_+$ and $p_-$ that, for each action pair $(s, a)$ correspond to the average reward ($p_+$) and the average difference of rewards ($p_-$). The individual rewards can be recovered from $p_+$ and $p_-$, which, in turn, can be computed by building an augmented MDP in which the states correspond to pairs of states in the original MDPs, and the rewards are the product of the corresponding rewards. The reward function for this augmented MDP can be efficiently estimated using, for example, reward-free exploration and then used to compute $p_+$ and $p_-$ through linear programming. Once $p_+$ and $p_-$ are computed, the original rewards can be recovered, and efficient planning can be used to derive a near-optimal policy. \n\nThe paper describes the algorithm and its theoretical analysis in terms of the number of episodes required to learn a near-optimal policy.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I start by saying that although I am mostly unfamiliar with the line of work pursued in this paper---aligned with PAC analysis of RL algorithms and efficient exploration. For this reason, I am unfamiliar with key pieces of relevant literature and the methods of analysis. For this reason, I have difficulty in properly assessing the novelty and technical soundness of the paper.\n\nThis said, overall, I liked the paper. The proposed motivation convinced me, and the overall methodology proposed in the paper strikes me as technically grounded and novel. I have several questions/observations that I hope may help to clarify the paper further.\n\n1. My first observation is that the paper is notationally heavy, and the notation selected is not always self-evident. For example, it took me some time to properly understand the purpose of the augmented MDP and how it meets the needs of the algorithm. Although I realize that the paper is technically involved and notation-heavy, I believe that a little effort could be made to provide some more intuition along with the math.\n\n2. If I understood correctly, each episode of the pure exploration is used to \"visit\" at most a single state of the augmented MDP and build the corresponding \"reward estimate\" $\\hat{\\mu}$. Wouldn't it make sense to reset $i$ once it reaches 3? This could allow (eventually) visiting more than one state in a single episode, or am I missing something?\n\n3. If I understood correctly, the RM-MDP framework discussed in this paper is a particular case of latent MDPs (LMDPs). RM-MDP corresponds to LMDPs the transition probabilities and initial distributions for the different MDPs are the same. In that case, they require some form of separation between the MDP dynamics. How does this relate to the separation in (7)? More generally, I would like to understand better how the assumptions made in the LMDP model are instantiated in the RM-MDP framework and how they relate to the implicit assumptions in the paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper's contributions are theoretical and relatively abstract in nature, so a discussion of societal impact is not very pertinent, in my opinion. ", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "8vezo7_hvD", "reviewer": "Reviewer_6RaD", "summary": "The authors consider a special case of latent MDPs where there are two unknown MDPs with identical state transition dynamics but different reward functions, and they refer to this setup as reward-mixing MDPs which they also define for M models, though only the case M=2 is studied. They show a lower bound for the sample complexity that is identical to the one proven by Kwon et al. when introducing latent MDPs (but with setting M=2), but for this special case. Further, they argue that they are the firsts to prove a polynomial (in relevant quantities) upper bound without extra assumptions beyond the MDP setting (besides the identical state transition dynamics that became part of the definition of RM-MDPs). Kwon et. al. had looked at a few other special cases like distinguishable latent MDPs.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors consider a special case of latent MDPs where there are two unknown MDPs with identical state transition dynamics but different reward functions, and they refer to this setup as reward-mixing MDPs which they also define for M models, though only the case M=2 is studied. They show a lower bound for the sample complexity that is identical to the one proven by Kwon et al. when introducing latent MDPs (but with setting M=2), but for this special case. Further, they argue that they are the firsts to prove a polynomial (in relevant quantities) upper bound without extra assumptions beyond the MDP setting (besides the identical state transition dynamics that became part of the definition of RM-MDPs). Kwon et. al. had looked at a few other special cases like distinguishable latent MDPs.\n\nThe authors claim that they had to introduce several new techniques for achieving such results and indeed, unlike the work by Kwon et. al., this article does not extend familiar techniques from the MDP setting, which also makes it much harder to read this paper. I would find it easier if they tried to reconnect by discussing what becomes of the method for M=1 and what kind of bounds would be achieved for MDPs with the resulting technique.\n\nI think what makes identical state transitions a bit easier to handle, is that they can rely on pure exploration ideas developed for handling a reward function that is chosen after such a pure exploration phase. What the authors have to handle besides such work, is that even after the pure exploration when it is time to act near-optimally, they will for each episode not know what reward function is chosen (identifying the context in the terms of Kwon et. al)  and they also have to estimate the two reward functions without ever knowing which one was in play, challenges shared with latent MDPs in general.\n\nI find that is would be easier to read this article if it were more conceptual structuring and connect as much as possible, not just in a related works section, to the existing work on MDPs and latent MDPs. It is interesting if this special case of latent MDPs comes with the possibility of desirable sample complexity bounds, though I have not succeeded in properly evaluating the proofs of this.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "ok", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3-4 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The main critique of the paper rests with the required assumptions for the theory.  Reviewer TkXV summarized them well.  The author response clarified that some of these assumptions were just out of simplicity in explication, and could relatively easily be relaxed.  The authors should absolutely make this more in clear in revisions of the paper, even if the details of the relaxations are better suited for supplemental material.  However, the sticking point is M=2.  This limitation is both substantial and with no easy path to relax it.  The authors argue that this is still a valuable step even if it's not settling everything.  The opposite perspective is that the theoretical machinery doesn't even lend itself toward the M>2 question, and so is a minor advance.  With the other restrictive assumptions addressed, I'm partial to theory moving forward in distinct (even if not large) steps, and this is indeed a distinct step.  
Decision: Accept (Poster)