Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: I become more positive after reading the rebuttal and other reviews. In particular,

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

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

The following is the original review.
***************************************************************************

This 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:

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

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

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

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

    - if the upper bound developed in Theorem 3.2 has a uniform dependence on \epsilon
    - 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)
limitations_and_societal_impact: I have discussed the limitations of this paper above.

This is a theoretical paper, and it is hard to speculate its negative societal impacts.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
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:

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

* The paper explicitly considers the (realistic) case of inexact, incomplete oracles.

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

* The use of LPs is clean and tidy.  Love it.

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

In terms of NeurIPS criteria:

* Originality: the paper seems original.
* Quality: the paper is of very high quality, with a solid logical flow, good references, and clear writing.
* Clarity: the paper is heavily theoretical, and therefore is heavy on notation.
* 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.

As 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?


limitations_and_societal_impact: This is not applicable, given the nature of the work.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: 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. 
main_review: Originality

This paper studies to a novel theoretical setting to the best of my knowledge, albeit an arguably narrow setting. 

Quality

Overall, 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. 

Clarity

Overall, the presentation of the paper is very good. I found no issues in terms of clarity.

Significance

The 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. 
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_title: null

final_score: rating: 7: Good paper, accept
review_content: 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. 

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

The paper describes the algorithm and its theoretical analysis in terms of the number of episodes required to learn a near-optimal policy.
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.

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

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

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

3. 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.
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_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: 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.

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.

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

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

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

limitations_and_societal_impact: ok
review_title: null

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)