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

We thank the reviewer for the thoughtful review and constructive feedback on our paper. We address the mentioned weaknesses below.

Practicality of the Proposed Algorithm

As the reviewer correctly pointed out, our focus was indeed on the purely statistical learning aspect of the problem. Nevertheless, we acknowledge and appreciate the reviewer's perspective on the computational aspects of the algorithm. Our long-term goal is to design an oracle-efficient algorithm for LMDPs that can be implemented with general function classes and ERM-style oracles. While we believe this is a feasible direction, achieving it will require significant effort from both theoretical and practical standpoints.

Intuition on Policy-Segmentation Regarding Prior Work

Thank you for your suggestion. Indeed, we agree that it would be useful to provide the intuition on how we connect moment-exploration ideas presented in [1] to our off-policy evaluation results. We will add the following paragraph in our revision at the beginning of Section 4:

Before we dive into our key results, let us provide our intuition on how we construct the OPE lemma for LMDPs. Our construction is inspired by the moment-exploration algorithm proposed in [1]: when state-transition dynamics are identical across latent contexts, {\it i.e.,} $T_1 = T_2 = ... = T_M$, we can first learn the transition dynamics with any reward-free type exploration scheme for MDPs [2], and then set the exploration policy that sufficiently visits certain tuples of state-actions of length at most $d=2M-1$. For the exploration policy, the work in [1] set a memoryless exploration policy $\psi \in \Pi_{`mls`}$ which sets the visitation probability to certain tuples sufficiently large. We note that the same moment-exploration strategy cannot be applied to general LMDPs with different state-transition dynamics since learning the transition dynamics itself involves latent contexts. Nevertheless, the intuition from [1] suggests that our key statistics are this visitation probabilities to all tuples of state-actions within a trajectory.

Significance of Our Results on $M$

We accept this criticism and appreciate the feedback. That being said, we would like to highlight that our results represent a substantial advancement from previous work on LMDPs. Specifically, our algorithm is the first to achieve sample efficiency in a partially observed setup without relying on standard assumptions such as weakly-revealing or low-rankness, which are well-studied in the literature. Even in cases where $M=2$ or $3$, the solutions to LMDPs were previously unknown, making our findings significant. While we acknowledge that our current results are achieved with $M=O(1)$, we believe it is important to emphasize the positive aspect of our contributions in pushing the boundaries of what is known in this field.

Possible Extensions to General Function Classes

This is a great question. We assume that by "problem-dependent bound", the reviewer is referring to sample-complexity bounds with general function approximation. We believe that extending our approach to general function approximation is feasible, particularly given recent advances in RL; however, this is not entirely straightforward and requires careful consideration. For Latent MDPs specifically, our starting point will be to define what is the good notion of "coverage" in LMDPs with large state/action spaces -- that will direct us to the question of what are the "moments" with general function approximation. If we can define this notion well, then we believe that we should be able to come up with a notion similar to the LMDP-coverage with function approximation. Another possibility is to come up with a general notion of statistical complexity measure for Latent MDPs/POMDPs such as DEC [3] (see also our discussion on GEC with Reviewer 8u1Y). Such an extension is beyond the scope of this paper, and we think it is an exciting future direction.

More Efficient Ways to Construct Data-Collection Policies

This is another great point. The current brute-force construction is due to the nature of worst-case analysis, where joint events at all length-$d$ tuples of state-actions need to be observed under all contexts. Since we do not have the observability over which policy drives us to the next state-action pair under latent contexts of interest, our solution brute-force all possibilities.

That being said, we are hopeful that under some practical assumptions, we may expect improvement in two ways:

• We can rule out some redundant combinations, as they might already be covered by other combination (e.g., this idea can be related to G-optimal design as in the Latent Multi-Armed Bandit setting [4])

• Other possibility is to identify that for some combinations, at some point, it becomes obvious that the segmentation of significantly smaller than $d$ segments would be sufficient.

These are the next big questions and important future work.

---

We hope that our answers to the raised questions are satisfactory. Please let us know if you have any other concerns or questions. We would greatly appreciate it if you could consider reevaluating our work, taking into account the strengths and improvements we've outlined.

[1] Kwon et al., "Reward-mixing MDPs with few latent contexts are learnable", ICML 2023

[2] Jin et al., "Reward-free exploration for reinforcement learning", ICML 2020

[3] Foster et al., "The statistical complexity of interactive decision making", arXiv 2021

[4] Kwon et al., "Tractable Optimality in Episodic Latent MABs", NeurIPS 2022

We sincerely appreciate the encouraging comments and positive assessment of our paper. Below, we address an interesting question you raised:

Can these results somehow adapt to the POMDP setting?

Yes, we believe our results can be adaptable to the POMDP setting in two senses:

• Algorithmically speaking, our algorithm overall falls in a general framework of OMLE; the key difference is in the data-collection policy part which may differ with different environments and assumptions. But at a higher-level, they can be viewed in a unified principle.

• Technically speaking, our analysis based on the notion of coverage can be useful to give a different interpretation of solving the POMDP settings (that are known to be tractable). The key to understand in such a way is to define a suitable notion of coverage in general POMDPs -- which is in part done in some recent work on off-policy learning for POMDP [1]. This is an interesting future question.

---

Thank you once again for your thoughtful review and positive assessment. If you have any further questions or suggestions, please let us know.

[1] Zhang and Jiang, "On the curses of future and history in future-dependent value functions for off-policy evaluation", arXiv 2024

We are grateful for your insights and suggestions, which will help us improve our work. Below, we address the weaknesses you mentioned.

Comparison to Previous Work on LMDPs

We would like to highlight that our work is the first to propose a general exploration algorithm applicable to the entire class of LMDP models. In contrast, previous works [1] focus solely on the subclass of LMDPs known as "Reward-Mixing" MDPs. These models assume no context ambiguity in transition dynamics, only in reward models. Under this assumption, data collection becomes less challenging. One can first learn the transition model separately through reward-free exploration, and then explore specific state-action tuples in order to obtain samples of "moments". However, when transition dynamics also vary across contexts, data-collection becomes significantly more complex, and is the main challenge in our work.

We will further clarify this distinction and the associated challenges in Section 1.1 - Challenge 2 in our revised manuscript.

Coverage Coefficient in MDPs vs LMDPs

Existing work that studies "the role of coverage in online exploration" is focused on fully observed settings (Low-Rank MDP, Block MDP are also fully-observed, although latent dynamics is much lower dimensional). In contrast, LMDP, and more broadly POMDP, involve dynamics and policies that are history-dependent -- an optimal policy of an LMDP may depend on the entire history. This means that a single latent-state coverability does not adequately capture the complexity of offline learning in these settings (as we illustrated in the counterexample in Appendix).

Consequently, off-policy evaluation in LMDPs requires a new notion of coverage that measures coverability over sequences of state-action pairs. Our main contribution is the proposal of this LMDP-coverage concept, which we connect to off-policy evaluation and online exploration. We believe this not only represents a significant advancement in learning LMDPs but also offers a conceptual contribution to the learning of general POMDPs. We will make this point clearer and more accessible to readers in our revision.

Eluder Dimension vs Our Approach

The Eluder dimension essentially measures the longest possible policy sequence that can incur a large prediction error. However, in partially observed settings, this definition may fail to capture the true complexity of online exploration. For instance, in multi-step weakly-revealing POMDPs (or PSRs), without executing core-tests, we may still suffer from the curse of horizon [2].

One might modify the definition of the Eluder dimension as proposed in the GEC paper you shared (details will follow below). However, similarly to the limitation of known tractable POMDP classes, the GEC definition also relies on the prior knowledge of how to construct the data-collection policies (using core-tests). Our technical contribution lies in the design and understanding of data-collection policies in LMDPs. We believe there is significant potential to develop a more general, fine-grained notion of complexity for a broader class, potentially including both revealing POMDPs and LMDPs. We see this as an exciting direction for future research.

Upper Bounds Exponential in M?

We accept this criticism, though due to the lower bound established in [3] we cannot avoid this dependence in M without any assumptions. However, we would like to highlight that our results represent a substantial advancement from previous work. Specifically, our algorithm is the first to achieve sample efficiency in a partially observed setup without relying on standard assumptions such as weakly-revealing or low-rankness, which are well-studied in the literature. Even in cases where $M=2$ or $3$, the solutions to LMDPs were previously unknown, making our findings significant. While we acknowledge that our current results are achieved with $M=O(1)$, we believe it is important to emphasize the positive aspect of our contributions in pushing the boundaries of what is known in this field (we also refer the reviewer to our general response on the importance of studying LMDPs).

Comparison/Connection to Other Complexity Measures (GEC)

We thank the reviewer for drawing our attention to this work, and we will add this reference in our revision. Yes, the idea behind GEC resonates with us at a very high-level -- it aims to capture the largest discrepancy in prediction errors from small training errors. However, as with other complexity measures such as (generalized) Eluder-dimension or Decision estimation coefficient, the main bottleneck is to show the boundedness of GEC in LMDPs, as there has been no upper bound established up-to-date. Furthermore, our construction of exploration (data-collection) policy is very different from known methods.

While all the above is in part described in our Section 1.1, we will be more explicit about comparison to existing complexity measures in our revision.

---

We hope our responses clarify any misunderstandings and resolve the issues identified. Please let us know if there are any remaining questions or concerns. Otherwise, we kindly request a reevaluation of our work in light of the provided clarifications. Thank you for your consideration and effort.

[1] Kwon et al., "Reinforcement Learning in Reward-Mixing MDPs", NeurIPS 2021

[2] Chen et al., "Lower bounds for learning in revealing POMDPs", ICML 2023

[3] Kwon et al., "RL for Latent MDPs: Regret Guarantees and a Lower Bound", NeurIPS 2021

We thank the reviewer for a thoughtful review and constructive feedback on our paper. We would like to start by emphasizing our technical novelty.

About Technical Novelties

Our LMDP-OMLE algorithm builds on the general framework of OMLE. The key novelty of the algorithm is the design of data-collection (exploration) policies (we also refer the reviewer to our general responses on the novelties). In designing the data-collection policy, we removed the need for core-tests and well-conditionedness. Further, in our analysis, we introduced the novel concept of LMDP coverage coefficient. These innovations result from efforts to formalize moment-exploration, as suggested in previous work [1]. Due to these, we were able to remove the restrictive assumptions of shared transition dynamics and separation.

Hence, our key results deviate substantially from existing approaches, and, hopefully, establish a new perspective on exploration in partially observed environments. Further, we believe our results contain a sufficient amount of novel concepts and analysis techniques.

Off-Policy Evaluation for Online Guarantees

Our off-policy evaluation guarantee guides the data-collection policy of the LMDP-OMLE algorithm. In particular, the key novelty lies in the proposal of LMDP coverage coefficient (Definition 4.1): this quantity measures the data-coverage in terms of the visitation probability to tuples of different states and actions. Algorithmically, this suggests that our exploration algorithm should aim to visit "uncovered" tuples of state-action pairs, reaching one target state (and action) and moving to the next target state (and action) in the target tuple. Once we can design the exploration policy in such a way, we can show the sample-complexity upper bound via the coverage-doubling argument (Theorem 4.5) as we have illustrated in the MDP case.

Here we note that establishing online guarantees via off-policy evaluation is just one possible approach for exploration in LMDPs. It would be very interesting to see if other approaches such as going through the problem-specific complexity measures such as Eluder-dimension [2], Bellman-rank [3] or Decision-Estimation Coefficient (DEC) [4] can be applied using the results provided in this work.

Why Our Approach Doesn't Make Any Assumptions?

We clarify that this work studies LMDP "without any assumptions" in the sense that we consider the class of all possible LMDP models, unlike in previous work that assumes certain separations [5,6] or shared transitions dynamics [1].

---

We hope our responses have clarified any doubts or questions. If there are any remaining issues, please let us know. Otherwise, we would highly appreciate it if the reviewer could consider raising the score based on our responses.

[1] Kwon et al., "Reward-mixing MDPs with few latent contexts are learnable", ICML 2023

[2] Russo and Van Roy, "Eluder dimension and the sample complexity of optimistic exploration", NeurIPS 2013

[3] Jin et al., "Bellman eluder dimension: New rich classes of RL problems, and sample-efficient algorithms", NeurIPS 2021

[4] Foster et al., "The statistical complexity of interactive decision making", arXiv 2021

[5] Hallak et al., "Contextual markov decision processes", arXiv 2015

[6] Chen et al., "Near-optimal learning and planning in separated latent MDPs", COLT 2024