Provable Representation with Efficient Planning for Partially Observable Reinforcement Learning

We provided some necessary introduction and common knowledge in RL and addressed the questions below.

• RL vs. Planning

  In RL, the model learning refers to estimating the transition and reward, and planning is referring to finding the optimal policy with the given or learned transition and reward (as explained in the Introduction section). We rigorously follow these definitions, where in our paper the algorithm learns a proper representation during the learning phase that can be used to perform planning (finding the optimal policy) easily.

• Real-world problems

  We tested our algorithm on MuJoCo environments with partial observations, and image-based RL problems, which are considered as real-world benchmarks in the RL community [1, 2]. We will rephrase this claim. Again, we would emphasize the proposed µLV-Rep is the first practical representation learning algorithm for POMDPs with tractable planning and exploration, and justified by rigorously theoretical guarantees and strongly empirical performances, as far as we know. It is unrealistic to ask one single conference paper to complete every aspect from theory, to simulation, to real robot deployment.

• Offline POMDPs

  The pessimism mechanism in offline setting is similar to the optimism in the online setting, and the major difference between the optimism in online setting and pessimism in offline setting is that, in the offline setting we minus the same bonus from the reward while in the online setting we add the bonus into the reward. We already added the corresponding discussion in Appendix D in our updated version.

• Novelty and Significance comparing to Ren et al. (2023a)

  We strongly disagree with this claim about the novelty and significance of our work. Based on our personal communication with the authors of Ren et al (2023a), how to extend LV-Rep to POMDP is “extremely difficult, which cost them one year already”, as we also discussed in Section 4.
  The extension is highly nontrivial and significant as we discussed in Section 3, especially considering the research on planning in POMDPs has stuck for decades. We emphasize the simplicity of the algorithm should be considered as an advantage, which enables the LV-Rep to more practical settings to handle partial observations, rather than a drawback.

• L-step decodable POMDPs vs L-order MDPs

  We are not sure what is the concrete definition of L-order MDPs. We assume the reviewer is asking about the megastate MDPs defined in [3]. In fact, the connection between L-step decodable POMDPs and L-step megastate MDPs have been discussed extensively in [3]. One can convert L-step decodable POMDPs to a L-step megastate MDPs. This reduction is useful in justifying the difficulty of the problem by characterizing the lower bound dependency on the problem size. However, such a simple reduction will induce extra dependency in the megastate, as the successive megastates have overlap of the state trajectories. This extra dependency does not reduce the difficulty in planning, and it does not provide justification for low-rank linear representation in MDPs. We understand the reviewer would like to use this reduction to justify the triviality of applying the existing representation learning for MDPs to POMDPs. However, as our personal communication and discussion with the authors of [1, 3, 4, 5], this argument does not hold.

• Latent Variable z vs. Belief of state

  No. Such understanding is incorrect. Recall that by definition that belief is the posterior of the true state, while our latent variable z does not have to be the state but an arbitrary variable, as long as it satisfies the factorization in Eq. 16. Meanwhile, belief is a distribution, while z is a variable.
  These properties of latent variable representation unleash us from identifiability restriction, which is one of our benefits to overcome the difficulty we discussed in Section 3 about recovering belief. We also discussed this in our main text in page 5 the remark about identifiability.

• Low-rank POMDPs

  The low-rank structure is indicating the latent variable factorization can be finite, or even if it is infinite, the eigendecay speed is fast, as we discussed in Appendix. The L-step decodability property ensures the rank is smaller than the exponential of the horizon, as discussed in [3].

• Non-Markovian in Visual Control

  In the image-based visual control benchmarks, although we can have the images, the control states, e.g., velocity, can not be inferred from one single image. There is no doubt that the Markovian methods can be applied for these tasks, but the performances will be significantly degenerated, due to the missing of important information. To further demonstrate the non-markovian property, we tested the algorithms with markovian assumption (L=1) vs non-markovian (L>1) (Figure 6 of Appendix H.1). As we can see, the non-Markvoian algorithm performs significantly better.

Firstly, I would like to address the tone of your response, which appears somewhat condescending. The phrase "We provided some necessary introduction and common knowledge in RL" suggests an assumption about the reviewer's level of expertise, akin to that of a student recently introduced to RL. Additionally, I noticed the absence of a customary 'thank you' in your response to my review, a courtesy extended to other reviewers. This inconsistency in approach is somewhat dismissive.

However, as a professional reviewer, I still respond to your disconcerting rebuttal:

RL vs. Planning

I definitely know what they are. So what is the planning part in your algorithm, like Dreamer's planning? The paper is too dense, as the other reviewer pointed out, and it's hard to grasp the core component of the algorithm.

The problem on this point is writing. The title and abstract are too broad. The introduction focuses on the motivation and prior work too much.

By “efficient" we mean the statistical and computational complexity avoids an exponential dependence on history length, while the computational components of learning, planning and exploration are computationally feasible; while by “practical" we mean that every component of an algorithm can be easily implemented and applied to a real-world problem. In this paper, we provide a affirmative answer to these questions.

Don't you feel these arguments are too strong? If so, RL is solved. I don't know anyone will call dm-control a real-world problem. To clarify, I don't ask you to work on real robots, but ask you to soften your overall writing. Unfortunately, you did not make any changes.

Offline POMDPs

Ok, during rebuttal you have added it and the proof required the readers for some external references.

The claim “state-of-the-art”

Unfortunately again, you did not respond to this.

Significance

I'm not familiar with Ren et al.'s work. I also think the anecdotal evidence like personal communication is not a professional way to show some work is generally hard. Furthermore, the fact that the work is hard does not imply the work is significant.

But at least, section 4.1 is very straightforward once you have L-step decodability assumption. Let the other reviewers and AC decide the main algorithm's significance.

L-order MDPs

This concept stems from L-order HMM. You can view it as the recent L observations (and actions) in a POMDP compose of the state in the corresponding MDP.

Latent Variable z vs. Belief of state

You did not directly answer my question. Can you compare Eq 17 with Dreamer's model learning objective? A belief state is a posterior, but it can be also viewed as a vector, expressed with the parameters of the posterior. In Dreamer or other variational belief approaches, they make it stochastic, i.e., we get a belief state z sampled from a distribution.

Low-rank POMDPs

So L-step decodability belongs to low-rank POMDPs?

Non-Markovian in Visual Control

Sounds good!

Thank you for your response. However, I must express that my primary concerns remain unaddressed. I find the tone of your response to be still condescending, as highlighted in the following points:

“Sec 4.1 is straightforward”

The phrase “Sec 4.1 is straightforward” was not meant to imply triviality. Simplicity in scientific approaches is often commendable and can be non-trivial. My intention was not to undermine the significance of your work, which I am uncertain about.

Your criticism on me that "lacks the basic respect not only for us, but also for the previous researchers, who are working on this problem." is unfounded and quite troubling. Such an accusation is not conducive to a constructive scientific discourse.

We never claim we solved RL problem, as we emphasize in the main text, the paper is trying to answer the question.

When I referred to your strong claims, it was my interpretation of the tone and implications in your paper, not an accusation that you claimed to have "solved the RL problem".

cannot “be easily implemented and applied to a real-world problem"

To be specific, you can check the assumptions made in your approach. For example, the definitions 1,2 are on the structures of POMDPs. How does your algorithm scale to very large L, which is common in real-world problem?

L-order MDPs

Thanks and I got it now. However, your previous claim hat

We understand the reviewer would like to use this reduction to justify the triviality of applying the existing representation learning for MDPs to POMDPs. However, as our personal communication and discussion with the authors of [1, 3, 4, 5], this argument does not hold.

Your assumption that my critique is intended to trivialize is again groundless. My goal was to seek clarity.

Belief of state

To clarify, I was asking about approximate belief state, like what these variational approaches learn.

"State-of-the-art”

I did not see that you have updated your PDF.

We sincerely believe the paper should be reviewed by evaluating the contribution and significance, instead of criticizing the usage of terminology.

It is essential to support claims of being SOTA with solid empirical evidence. The usage of such terminology is not just a semantic issue; it reflects the scope and impact of your claims.

Finally, my objective is not to criticize but to ensure the clarity, accuracy, and significance of your contributions to the field.

Thanks for being professional. The following are our replies.

• Significance The claim that “Sec 4.1 is straightforward” after seeing the results lacks the basic respect not only for us, but also for the previous researchers, who are working on this problem. The L-step decodability assumption has been proposed by Efroni et al, in 2020, and a recent followup paper trying to figure out computation tractable planning (Guo et al, 2023) is just published, while still leaving the planning problem partially open. If the planning problem is indeed “straightforward” as the reviewer claimed, the problem should be solved far earlier than our paper.
  Again, we emphasize that simplicity does not mean straightforward. The simplicity of an algorithm should be considered as an advantage, rather than a drawback.

• RL vs. Planning

  • Re: Planning.

    With the proposed representation, the Q-function can be represented linearly. And the planning is executing Bellman backup in the linear space as we described. While in Dreamer, it is not clear in which space the Bellman backup should be conducted to obtain a theoretical rigorous algorithm.

  • Re: RL.

    We never claim we solved RL problem, as we emphasize in the main text, the paper is trying to answer the question: “How can we design an efficient and practical RL algorithm for structured partial observations”, instead of solving RL, which we never claimed.

    This question is right before the paragraph you cited. The paragraph you cited is written to explain the problem in detail. In fact, the proposed algorithm indeed avoids exponential dependence on history length, while with tractable planning and exploration. Meanwhile, if you read the paragraph you cited carefully, it claims that our algorithm can be applied to real-world problems. If you think there is any step of the proposed algorithm that cannot “be easily implemented and applied to a real-world problem“, please point it out and we are open to discuss.

• L-order MDPs

  In the RL community, this is known as L-step mega-state MDPs. We already explained the connection between L-decodable and L-step mega-state MDP in our last response, and more importantly, the difficulties in planning still exist even with the form of so-called “L-order MDPs”, due to the observation overlaps in mega-state.

• Latent Variable z vs. Belief of state

  We would like to clarify that belief refers to the posterior of true state, which requires extra identifiability assumption, while as we discussed in the remark in page 5.

• Terminology confusion

  • “Real-world problems”

    We are sorry that the term “real-world” makes the reviewer uncomfortable. In fact, we already modified the claim in our experiment setting. While there are only three places in maintext, where the words show up:

    • “In real-world reinforcement learning, state information is often only partially observable,” in abstract;
    • “Such algorithms have been applied to many real-world applications with image- or text-based observations (Berner et al., 2019; Jiang et al., 2021), sometimes even surpassing human-level performance (Mnih et al., 2013; Kaufmann et al., 2023).” in page 1 introduction;
    • “while by “practical" we mean that every component of an algorithm can be easily implemented and applied to a real-world problem.” in page 2.

    If read carefully, you may see the first one is describing the problem setting, the second is describing dreamer and chatgpt, the third one is describing the target for which we are pursuing.

  • "State-of-the-art”

    We have softened our claim and removed “state-of-the-art”, which is used for our competitors, not our method.

We sincerely believe the paper should be reviewed by evaluating the contribution and significance, instead of criticizing the usage of terminology.

Thanks for your replies. As the discussion is close to end, I do not expect further responses to my following comment and summary. I believe with additional efforts, your work has the potential for publication in a top-tier conference.

We believe “straightforward” and no “technical significance” might not be appropriate to be used here.

I'd like to clarify my previous statement: I mentioned the work has "not much significance", not "no significance". I trust this misquotation was accidental.

Regarding the use of "straightforward", I maintain its appropriateness. The key observations, as the paper stated, are straightforward in replacing belief state with history, and history with $L$-step memory, based on your assumptions.

the claim of real-world problem, large $L$

Concerning your assumption on $L$-step and its real-world applicability, testing with $L$ values between 2 to 5 is insufficient to claim easy implementation in real-world problems without empirical evidence. The validity of the linear structure assumption in real-world problems also remains unclear.

I recommend softening the claim about real-world applicability to "has the potential to be applied", unless you provide empirical support.

“Belief of state”

Let me be straightforward. Here I rewrite Eq. 17 with $l=1$:

$$\max_q E_{q(z_h \mid x_h,a_h,o_{h+1})}[\log P(o_{h+1} \mid z_h)] - D_{KL}(q(z_h \mid x_h,a_h,o_{h+1}) \mid\mid p(z_h\mid x_h))$$

The first term is to reconstruct observations, and the second term is KL divergence between the posterior of $z$ and prior of $z$. The DreamerV2 model learning objective is very similar to this, except that they have a reward prediction term (which I am unsure about your method) and the prior also conditions on $a_h$ (which I believe should be added).

state-of-the-art

I found you changed it in the main PDF, but not supplementary material. Please also update the supp accordingly.

Finally, I want to conclude my review. I maintain my rating due to the seemly limited significance of the work, dense writing style, unclear contributions and connections with prior work, some overreaching claims, and the tone and unprofessionalism of responses to my feedback (e.g., several anecdotal evidences and assumed negative intent on the reviewer) which I flagged for future ethics review.

We thank the reviewer for the suggestion on the presentation and other feedback. We address the concerns in the following:

• Presentation

  Despite the complicatedness of the related techniques and the subtleness in planning for POMDPs, which makes the paper technically dense, we try our best to make the paper self-contained, and design a simple algorithm. We have revised our manuscript accordingly and try our best to make the paper more convenient to be understood. We would like to clarify some key typos in the previous manuscript

• • Reward $r(s_h, a_h)$

  In fact the agent can only observe a reward $r(o_h, a_h)$, which can be defined as $r(o_h, a_h) = \int r(s_h, a_h) P(o_h|s_h) ds_h$ and can be viewed as a sample from $r(s_h, a_h)$, which is commonly assumed in the existing POMDP literature [1, 2]

• • $P(s_1|o_1)$ and $\tau$

    $P(s_1|o_1)$ can be defined as the posterior distribution induced by the initial distribution and emission. $\tau$ is the history defined in the second paragraph of the preliminaries.

• • $\theta$ in Equation 5

    $\theta$ is just a vector that can induce the reward.

• Parameter l in Eq. 16

  For any $l \in N_+$, $p({o_{h+1:h+l}}|x_h, a_h) \int_{\mathcal{Z}} p(z_h|x_h, a_h) p^\pi ({o_{h+1:h+l}}|z_h) d z_h$ Is factorizable and share the same $p(z_h|x_h, a_h)$, but with different $p^\pi ({o_{h+1:h+l}}|z_h)$. As we discussed that $p(z_h|x_h, a_h)$ forms the function space for Q-function, and thus, is the latent variable representation we are seeking, we can perform MLE on arbitrary l-step future observations conditional distribution.

• Definition for the moment matching policy $\nu_\pi$

  The formal definition of νπ is complicated and we provide a formal definition at the beginning of the proof for Lemma 8. We also add more concrete discussion in Appendix C about the moment matching policy.

• Ablation Study

  We have added more ablation studies (Appendix H.1) to demonstrate the effects of the major components, including representation dimension and window size.

• Experimental setting

  We set up the experiment setting exactly following the benchmark [1], which has been widely adopted in [2,3,4] for fairness, especially for the comparison to LV-Rep with full observation in Table 1 (denoted as Best-FO), to demonstrate the performance gap in MDP and POMDP.

[1] Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforcement learning. arXiv preprint arXiv:1907.02057, 2019.

[2] Tianjun Zhang, Tongzheng Ren, Mengjiao Yang, Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Making linear mdps practical via contrastive representation learning. In ICML, 2022

[3] Tongzheng Ren, Chenjun Xiao, Tianjun Zhang, Na Li, Zhaoran Wang, Sujay Sanghavi, Dale Schuurmans, and Bo Dai. Latent variable representation for reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023b.

[4] Tongzheng Ren, Tianjun Zhang, Lisa Lee, Joseph E Gonzalez, Dale Schuurmans, and Bo Dai. Spectral decomposition representation for reinforcement learning. ICLR, 2023c.

We thank the reviewer for the feedback. Please see our response below:

• Assumptions

We first emphasize that the assumptions are merely for the proofs, which are commonly used in the literature (e.g. [1, 2, 3, 4]) and are thought to be mild. While the proposed µLV-Rep is still widely applicable even without the assumptions in practice. In fact, in our empirical experiments, the assumptions are not guaranteed to be held. But the proposed algorithm still outperforms or is comparable to the existing strong competitors, including DreamerV2, which demonstrates the robustness of the proposed algorithm in practice.

• Novelty and Significance comparing to Ren et al. (2023a)

As we mentioned, our core contribution is not theoretical analysis. Instead, our contribution lies on extending LV-Rep for POMDP that can perform provable representation learning and planning with function approximations (as illustrated in Section 4), which is highly nontrivial and significant as we discussed in Section 3, especially considering the research on planning in POMDPs has stuck for decades. We emphasize the simplicity of the algorithm should be considered as an advantage, which enables the LV-Rep to more practical settings to handle partial observations, rather than a drawback.

• Pessimism in Offline

The pessimism mechanism in offline setting is similar to the optimism in the online setting, and the major difference between the optimism in online setting and pessimism in offline setting is that, in the offline setting we minus the same bonus from the reward while in the online setting we add the bonus into the reward. We already added the corresponding discussion in Appendix D our updated version.

• Tightness of Bounds

First, we emphasize that our sample complexity bound is for POMDPs, which matches the lower bound dependency on the horizon of actions [1]. Regarding the tightness of the dependency on the other quantities in the sample complexity, it is still an open problem even for the existing representation learning methods including [2, 3, 4]. The proposed µLV-Rep is the first practical representation learning algorithm for POMDPs with tractable planning and exploration, with strong theoretical guarantees, as far as we know. How to provide better representation learning algorithms with better sample complexities is still an open problem and beyond the scope of this manuscript.

[1] Efroni, Yonathan, et al. Provable reinforcement learning with a short-term memory. In International Conference on Machine Learning, 2022.

[2] Agarwal, Alekh, et al. ”Flambe: Structural complexity and representation learning of low rank mdps.” Advances in neural information processing systems 33 (2020): 20095-20107.

[3] Uehara, Masatoshi, et al. ”Representation Learning for Online and Offline RL in Low-rank MDPs.” International Conference on Learning Representations. 2022.

[4] Ren, Tongzheng, et al. ”Latent Variable Representation for Reinforcement Learning.” The Eleventh International Conference on Learning Representations. 2023