RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluation

Summary

This paper studies latent Markov decision processes (LMDP) with M = O(1) number of MDPs. In other words, there exists a set of MDPs (unknown to the agent) and the environment randomly selects one MDP at the beginning of each episode. The selected MDP is not revealed to the agent and therefore the agent must infer which MDP is selected from the feedback. This paper proposes an algorithm with sample complexity poly(S,A)^M that matches the \Omega((SA)^M) lower bound. On the technical side, this paper designs a model-based exploration algorithm that actively collects new data if the collected data can shrink the confidence interval of the model parameters. On standard tabular RL, this paper shows that (a) the algorithm terminates after poly(SAH) rounds of exploration, and (b) any remaining model parameter can be used to construct a near-optimal policy after the exploration stage. Because of the generality of the algorithm, this paper further extends the algorithm to the LMDP setting and proves a poly(S,A)^M sample complexity.

Strengths

- The exposition of the main ideas is very clear because of the instantiation on tabular MDPs. - The OMLE algorithm is neat, and it provides a general framework for model-based exploration. - The sample complexity result matches the \Omega(S,A)^M lower bound and requires no additional assumption on the structure of the LMDP family.

Weaknesses

- The algorithm is not computationally efficient even for standard tabular MDPs because it has to enumerate both the parameters and the policy in the confidence interval, though I understand that computational efficiency is not the focus of this paper. - Section 4 is hard to follow without background knowledge of several prior works. E.g., the choice d = 2M-1 is not well-justified. I encourage the authors to revise Section 4 to make it more self-contained. - It seems to me that the title of this paper somewhat overclaims the result since the sample complexity is still exponential in M, and the assumption that M is finite is not intrinsic to the LMDP setup (although it is necessary due to the sample complexity lower bound).

Questions

- It seems to me that the OMLE algorithm is very general. Is it straightforward to prove some general sample complexity results with problem-dependent bounds? Intuitively, the bound would include terms like the statistical complexity of the model class and some complexity measure of the MDP family. - The data collection step in Algorithm 2 (Line 5-7) seems to be quite brute-force since it has to enumerate all the segmented policy. Could the authors elaborate on its reason?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations and potential negative societal impact.

Summary

This paper studies latent MDP, where the underlying dynamics and rewards are controlled by some latent states (not revealed to the learner), and the learner attempts to do learning and planning based on the trajectory data. The algorithm builds from the optimistic MLE algorithm, which iteratively checked whether whether there exists uncovered policies. To make the OMLE works for LMDP, this paper constructed segmented policies, and iteratively check whether such policies are all covered. As a result, they showed the first polynomial sample complexity result for LMDP when the number of latent states are constant.

Strengths

The paper is well written. The algorithm, theorems and proofs are clear. The results of polynomial sample complexity for LMDP with constant latent states seem very interesting. The idea of constructing segmented policies is novel.

Weaknesses

I don't see any significant weaknesses of this paper.

Questions

Can these results somehow adapt to the POMDP setting?

Rating

8

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

Yes. The authors addressed all the limitations listed in the guidelines.

Summary

This paper introduces a new version of the coverage coefficient for analyzing latent Markov Decision Processes (MDPs). It demonstrates how to link the proposed coverage coefficient with sample complexity using MDPs. Additionally, the paper presents an algorithm and provides a bound on the sample complexity of this algorithm.

Strengths

The use of MDPs to illustrate the concept is effective, making it easier for the audience to follow.

Weaknesses

1. The comparison with related work is insufficient. Including a table that compares the results in this paper with those in [1] would be helpful. 2. The connection and comparison between this work and existing studies on the coverage coefficient [2] should be more concrete. 3. The meaning and design of the segment policy are unclear, making it difficult for the audience to understand the intuition behind why this coverage coefficient is helpful for analysis. Although a counterexample is provided in the appendix, it is still not intuitive enough. 4. In my understanding, the core idea behind the analysis is that complexity depends on the longest policy sequence $\pi_1,\ldots,\pi_k$, where $C(\pi_i,\pi_j)=\infty$. If this understanding is correct, there are two follow-up questions: (a) Can we define a complexity measure similar to the Eluder dimension? (b) Can the complexity measure generalize to cases beyond latent MDPs? (c) Is this complexity measure equivalent to or weaker than the Eluder dimension? Answers to these questions might help improve this paper. 5. Although the results in this paper do not require further structural assumptions, the significance of an exponential upper bound is questionable. 6. Lack of discussion with works that have complexity measure without Markovian assumption [3]. [1]. J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Reinforcement learning in reward-mixing MDPs. Advances in Neural Information Processing Systems, 34, 2021. [2]. P. Amortila, D. J. Foster, and A. Krishnamurthy. Scalable online exploration via coverability. arXiv preprint arXiv:2403.06571, 2024. [3]. Zhong, Han, et al. "Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond."

Questions

See "Weaknesses" part for the questions.

Rating

3

Confidence

3

Soundness

3

Presentation

2

Contribution

1

Limitations

See "Weaknesses" part for the limitations.

Summary

The paper studies latent MDPs, an MDP framework with a set of MDPs and the environment samples a random MDP at the beginning of each episode. To avoid a $A^H$ sample complexity, previous works either assume separation or similarity of transitions. This work removes these conditions, and provide an algorithm with $(SA)^{O(M)}$ upper bound, which matches the lower bound with $M = O(1)$. Speficially, the algorithm adapts the information theoretic Optimistic MLE algorithm for the LMDP setting, by collecting data with all possible segments of all previous policies (+random actions). The key analysis tool is the LMDP coverage coefficient.

Strengths

1. The paper proposes the first algorithm that matches the lower bound of LMDPs (to poly order) without the assumption of separation and similar transition dynamics. 2. The proposed LMDP coverage coefficient, although exponential in $d$ ($M$), seems to capture the right complexity of the LMDPs. 3. The detailed and intutive introduce of Optimistic MLE algorithm makes the proposed algorithm easy to understand.

Weaknesses

1. The techincal innovation seems rather limited -- the main techiques are quite similar as OMLE. 2. I do not understand the relation to OPE: which part of the algorithm relates to OPE? It seems like the only change to OMLE is the way of data collection so there is no additional OPE subroutine? 3. The "without any additional structural assumption" claim in the abstract seems inaccurate: even OMLE is limited with the assumption of finite SAIL (bli-linear rank), and the proposed analysis only applies to tabular? Also, does the analysis generate to general function approximation if we ues the correponding coverage coefficient?

Questions

see above

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

n/a

Thank you for the response

Thanks for the response and positive feedback!