Paper: The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models (ICLR-2024)

Reviewer comment (Theory):
Even though the method is called Wasserstein Belief Updater and the theoretical analysis is for the Wasserstein belief loss, the actual algorithm uses a KL loss for the belief updater part. Hence, the guarantees do not necessarily hold for the actual algorithm.

Author response:
**Theoretical Justification for KL as a Proxy.** The KL divergence, while more tractable to optimize, also offers theoretical backing. As per the Pinsker's inequality, in the zero-temperature limit, the KL divergence bounds the Wasserstein distance. This relationship allows us to leverage the easier optimization landscape of the KL divergence while still maintaining a connection to the Wasserstein framework. Essentially, optimizing the KL divergence serves as a practical approach to approximate the Wasserstein distance, especially given the challenges associated with the latter in our specific learning context. We refer to the answer to Reviewer Gzjw (paragraph *Wasserstein proxy*) for the practical challenges of Wasserstein optimization.

Reviewer comment (Reproducibility):
The algorithm doesn't have any step for updating the MAF. Is there something missing or is the MAF pre-trained by some other way?

Author response:
**Masked autoregressive flow.** The MAF is trained with the sub-belief encoder to minimize Equation 8. This is mentioned in the training paragraph of Section 4 and in the algorithm pseudo code in Appendix F. Details concerning the sampling and learning procedure of the particular MAF we use are given in Delgrange et. al, 2023, Section 3.2 (last paragraph).

Reviewer comment (Theory):
In Section 4, the authors claim that the training of the sub-belief encoder does not require the challenging BPTT as in RNN. But given the recursive structure of the sub-belief encoder, it is not clear why the claim is true. Even when the sub-belief encoder is updated only by the belief loss, the gradient flow might still need to propagate back if the previous sub-belief is also generated by the same sub-belief encoder. Back-propagation might stop if the previous sub-belief is actually generated by other networks, like using (3), but this part is a bit unclear.

Author response:
**Backpropagation through time (BPTT)** is a powerful tool for learning value functions due to their backward recurrence nature; an update in the value of a history $h_t$ ​ under policy $\pi$ can propagate changes back to previous histories. This approach is efficient as it reduces the number of necessary updates by allowing changes to flow through past timesteps. However, in the context of belief learning, the situation is different. Beliefs progress through forward recurrence, implying that a modification in belief $b_t$ affects subsequent beliefs but not preceding ones. Consequently, BPTT does not offer the same utility for belief learning as it does for value functions. To align with this forward recurrence nature and avoid the unnecessary complications of BPTT, we apply a stop gradient in our model, effectively preventing the backward flow of updates. This approach is illustrated in Figure 2 and detailed in the algorithm presented in Appendix F, where the use of the stop gradient __sg__ is explicitly implemented and explained.

Reviewer comment (Reproducibility):
Can you provide more details on MAF training?

Author response:
**Masked autoregressive flow.** The MAF is trained with the sub-belief encoder to minimize Equation 8. This is mentioned in the training paragraph of Section 4 and in the algorithm pseudo code in Appendix F. Details concerning the sampling and learning procedure of the particular MAF we use are given in Delgrange et. al, 2023, Section 3.2 (last paragraph).

Reviewer comment (Reproducibility):
Can you provide more details on how the sub-belief encoder is trained and why BPTT is not needed with some gradient analysis?

Author response:
We thank the reviewer for their reactivity and allowing us to expand on why BPTT is not needed in the context of Belief Learning.

The reviewer noted that in the RNN-A2C the hidden state $z_t$ is also computed recursively, which is in the same spirit as how the belief $b_t$ is computed. The main difference between beliefs and hidden states lies in what they represent. 

On the one hand, beliefs are state distributions which can be inferred from the previous belief and be themselves seen as Markovian states (cf. the belief MDP construction). Therefore, one can see the belief update function as the transition function of an MDP where an observation is further provided. Disabling BPTT is therefore consistent with the literature of model based RL for learning Markovian transition functions (e.g., Gelada et al. 2019, Vincent François Lavet et al. 2019, Delgrange et al. 2022). Notably, our goal is not targeted towards learning directly a history representation (from sequences) *but rather on learning to reproduce the belief update rule*, based on the latent stochastic transition function $\bar{P}$, and observation function $\bar{O}$ of the latent model, which are learned separately and considered fixed in this phase. Therefore, stopping the gradients does not produce approximate gradients, as the error linked to this recursive update can be fully determined through the error of our latent model (Theorem 3.2, Equation 6), and using BPTT would increase the risk of accumulating errors across time.

On the other hand, hidden states are latent representations of the history and it is up to the RL algorithm to learn a useful representation of the history for the control task. This is in stark contrast to our approach: we separately learn a representation which is guaranteed to be well suited for the RL agent, and we use it as input of the policy. In RNN-A2C the main objective is to find a policy that optimizes the return conditioned on the full history, making it sound to do BPTT in that case. In our case the RL agent uses the beliefs as a given state-signal to condition the value function and the policy, and therefore the RL objective is not backpropagated through the belief encoder.

Finally, it is worth mentioning that as beliefs are computed with a forward recurrence, errors may accumulate over time. This means that if the belief $b_{t+k}$ diverges significantly from its correct value, then $b_{t+k+1}$ will diverge even more. Then, allowing gradients from the learning of $b_{t+k+1}$ to be computed based on an incorrect $b_{t+k}$ and further flow to $b_t$ may hinders the learning procedure. 

**References**
- C. Gelada et al. "DeepMDP: Learning Continuous Latent Space Models for Representation Learning". ICML 2019.
- V. François-Lavet et al. "Combined reinforcement learning via abstract representations". AAAI 2019.
- F. Delgrange et al. "Distillation of RL policies with formal guarantees via variational abstraction of Markov decision processes". AAAI 2022.

Reviewer comment (Writing):
Some possible typos and missing definitions: 
  - Is $\varphi^*$ defined somewhere?
  - In equation (5), should the Wassertein distance be given by $D$ or?

Author response:
**Typos and missing definitions.** We did not explicitly define $\varphi^*$ which corresponds to the recursive application of $\varphi$ ( analogously to $\tau^*$). We modified the paper to make it explicit. The distance $D$, defined in the box page 5 as the distance between the real belief updater and our approximation, is set to the Wasserstein distance in the paragraph Belief Losses above Equation 5.

Reviewer comment (Theory):
Wasserstein Belief Updater finally chooses the objective function to be KL divergence as a Wasserstein proxy. I agree with the difficulty of optimizing Wasserstein (which is quite singular). But probably the authors can think about the entropy-regularized version of Wasserstein distance? Hope the work after the following paper could help.  
[ Marco Cuturi, Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances ]  
[ Marco Cuturi, et. al., Fast Computation of Wasserstein Barycenters ]

Author response:
**Wasserstein proxy.** We thank the reviewer for the pointers to entropy-regularized versions of the Wasserstein distance. As pinpointed in the main text, Wasserstein optimization is known to be challenging, and one can indeed avoid the optimal coupling (primal) or Lipschitz (dual) constraints by considering entropy-regularized or relaxed distances. Still, the major challenge in learning to minimize a Wasserstein-like discrepancies between our latent belief encoder and the true latent belief update function is that *one need be able to sample from both distributions*, even via an entropy-regularized version of the distance. As mentioned in the main text, sampling from our belief encoder $\varphi$ is done via a masked autoregressive flow, but it is non-trivial to sample batches of states from the true latent belief yielded by $\overline{\tau}$: indeed, in that case, we need to infer the distribution from the closed form of Equation 3. To that aim, one could consider Monte Carlo Markov Chain sampling techniques such as Metropolis-Hastings, but this complicates the learning process as such procedures are computationally expensive, and further brings technical challenges. In contrast, Kullback Leilbler (KL) is easier to optimize since we only need to sample from our belief encoder to minimize the divergence (see Equation 8), while still offering guarantees (KL bounds Wasserstein by the Pinsker’s inequality in the zero-temperature limit, as mentioned in the main text). We defer as future work the study of alternative Wasserstein-like discrepancy to our KL loss.

Reviewer comment (Theory):
I did not see what are the requirements of the sets we discuss? I guess they are at least discrete to make certain Markov chain theoretical results hold. But as the paper talked about Cartpole, I guess the theory would meant to be possibly applied to continuous state spaces (real manifold at least).

Author response:
**General state space POMDPs.** We thank the reviewer for pointing out that in general state spaces, ergodicity notions are different. We are aware of the issue and, while there is a precedent of allowing this small disconnection between theory and practice, we have (for completeness) confirmed the guarantees do hold in general state spaces under reasonable assumptions. Below we elaborate on these two points.  

As the reviewer mentioned, the theory and results developed by Huang Bojun (2020) hold for finite or countable infinite state spaces. Nevertheless, for their experiments, they evaluated their approach on uncountable infinite state space MDPs, on RobotSchoolBullet environments. Arguably, their evaluation agrees with the theory developed since everything is discrete in practice: the computer simulates the environment whose episodes are derived from a finite branching of histories unrolled through a machine-induced discretization of the observations.  

Assuming that uncountable infinite state spaces are endowed with a usual Borel sigma-algebra, we can still show that all our results hold in any episodic RL setting. We have updated the Appendix accordingly. Intuitively, the policies considered in our work all induce a *recurrent Harris chain* (Definition C.6 in Appendix), where the interaction is guaranteed to eventually result in a stationary distribution. Therefore, the approach of Huang 2020 (which facilitates rapid convergence to the stationary distribution) is still valid in our setting. We added a complete, formal discussion about this setting in Appendix C.3 (part (b) of the proof).

Reviewer comment (Theory):
In the theoretical part, the existence of $\mathcal{H}_{\overline{\pi}}$ is relying on the ergodicity of an episodic MDP (Huang Bojun, 2020), where the state space $\mathcal{S}$ is assumed to be discrete (finite or countably infinite). I want to know the topology (and geometry) of the history space.

Author response:
**General state space POMDPs.** We thank the reviewer for pointing out that in general state spaces, ergodicity notions are different. We are aware of the issue and, while there is a precedent of allowing this small disconnection between theory and practice, we have (for completeness) confirmed the guarantees do hold in general state spaces under reasonable assumptions. Below we elaborate on these two points.  

As the reviewer mentioned, the theory and results developed by Huang Bojun (2020) hold for finite or countable infinite state spaces. Nevertheless, for their experiments, they evaluated their approach on uncountable infinite state space MDPs, on RobotSchoolBullet environments. Arguably, their evaluation agrees with the theory developed since everything is discrete in practice: the computer simulates the environment whose episodes are derived from a finite branching of histories unrolled through a machine-induced discretization of the observations.  

Assuming that uncountable infinite state spaces are endowed with a usual Borel sigma-algebra, we can still show that all our results hold in any episodic RL setting. We have updated the Appendix accordingly. Intuitively, the policies considered in our work all induce a *recurrent Harris chain* (Definition C.6 in Appendix), where the interaction is guaranteed to eventually result in a stationary distribution. Therefore, the approach of Huang 2020 (which facilitates rapid convergence to the stationary distribution) is still valid in our setting. We added a complete, formal discussion about this setting in Appendix C.3 (part (b) of the proof).

Reviewer comment (Theory):
Is it the collection of all $(a,o)$ sequences such that $(\mathcal{A}\cdot \Omega)^*=Map(\mathbb{N}, \mathcal{A}\times \Omega)$, or it is the collection of all finite-length tuples that $(\mathcal{A}\cdot \Omega)^*=\bigsqcup_{n\in\mathbb{N}}(\mathcal{A}\times \Omega)^n$? Categorically speaking, is it a direct product (product) or direct sum (coproduct)? See https://en.wikipedia.org/wiki/Product_(category_theory).

Author response:
**History space.** We admit that we are not experts in category theory, but relying on the sigma-algebra defined by Revuz (1984) and Putterman (1994), the history space is $(\mathcal{A}\cdot \Omega)^*=\bigsqcup_{n\in\mathbb{N}}(\mathcal{A}\times \Omega)^n$. This should clarify what topology we are working in.

Reviewer comment (Theory):
Is there any theoretical result on the continuous spaces such as CartPole? The definition of ergodicity is slightly different there.

Author response:
**General state space POMDPs.** We thank the reviewer for pointing out that in general state spaces, ergodicity notions are different. We are aware of the issue and, while there is a precedent of allowing this small disconnection between theory and practice, we have (for completeness) confirmed the guarantees do hold in general state spaces under reasonable assumptions. Below we elaborate on these two points.  

As the reviewer mentioned, the theory and results developed by Huang Bojun (2020) hold for finite or countable infinite state spaces. Nevertheless, for their experiments, they evaluated their approach on uncountable infinite state space MDPs, on RobotSchoolBullet environments. Arguably, their evaluation agrees with the theory developed since everything is discrete in practice: the computer simulates the environment whose episodes are derived from a finite branching of histories unrolled through a machine-induced discretization of the observations.  

Assuming that uncountable infinite state spaces are endowed with a usual Borel sigma-algebra, we can still show that all our results hold in any episodic RL setting. We have updated the Appendix accordingly. Intuitively, the policies considered in our work all induce a *recurrent Harris chain* (Definition C.6 in Appendix), where the interaction is guaranteed to eventually result in a stationary distribution. Therefore, the approach of Huang 2020 (which facilitates rapid convergence to the stationary distribution) is still valid in our setting. We added a complete, formal discussion about this setting in Appendix C.3 (part (b) of the proof).

Reviewer comment (Theory):
The biggest drawback to this work is the latent observability condition. The methodology being presented is fairly dependent on this condition, since it attempts to learn the latent transition model, observation model, and belief states. This simplifies the problem setting quite a bit by making the POMDPs more or less a complicated involved MDP. It is not clear how much applicability this method has beyond the settings where this assumption holds.

Author response:
**Access to the latent model vs Access to states during training.** 
1. **Access to the latent space model:** In Section 3, we present how to learn the latent space model, which is then used later on, in Section 4. Specifically, this includes using the learned latent transition and latent observation functions to train our belief encoder to replicate latent belief update function. *In practice, the latent model and belief encoder are learned in a round robin fashion* (see Figure 1). 
2. **Access to the states during training:** For a detailed discussion on the applicability of Assumption 1, we kindly refer the reviewer to our Introduction and Remark 1, in the main text. Notice that despite the observability of the state during learning, the problem remains formalized as a POMDP, since the resulting policy is defined over past observations and actions, and not on the MDP states. We note that in the field of planning in POMDPs, the access to the model of the environment is typically assumed, which is a much stronger assumption, while keeping a POMDP formulation. Therefore, we believe our approach is not a simplification of the problem and does not reduce the problem into a complicated involved MDP.

Reviewer comment (Novelty):
In terms of related work, there is a fair amount of literature on provable guarantees for exploration in POMDPs under latent observability (e.g. Lee et al. ICML 2023 and related/followup work) or under structural conditions for POMDPs (e.g. Liu et al. 2022 and related/followup work) which is not mentioned here.

Author response:
**Related work.** We thank the reviewer for the pointers to relevant works. We have included references to Lee et al. in our revised related  work section. We would like to highlight the fact that *Lee et al.'s assumption of hindsight observability aligns closely with our Assumption 1.*

Reviewer comment (Theory):
The current theoretical guarantees are similar to the other papers on learning with bisimulation-like metrics, in that the extent of the theoretical results is a bound on the approximation error incurred via the population loss functions. I would argue that this is not an "end-to-end" theoretical guarantee that is useful, since there are many questions about how one can minimize such a loss when only given access to samples from the distribution (i.e. how close are the empirical losses and the population losses, are there any double-sampling issues, is one sampling from distributions one does not have access to, etc.) and if we can minimize it on policies that actually matter for downstream performance rather than only on the on-policy distributions (currently it is simply an "on-policy" bound which holds only for the distribution under which the data was collected). As an example of the weakness of such a bound, the usual bisimulation metric in tabular MDPs satisfies a similar bound (Theorem 5.1, Ferns et al. 2004), but there is still an exponential lower bound for learning any bisimulation-like abstraction (Appendix B of Modi et al., 2020). Thus, the argument that this paper presents theory that can be used to give guarantees is either incomplete or inaccurate, depending on one's perspective. The authors could perhaps support their theory by experimentally verifying that the (population) losses are indeed minimized and verifying that the quantities on the LHS of their bounds are controlled by their procedure.

Author response:
**Theoretical guarantees.** Our focus is on developing deep RL theories for POMDPs, which is not necessarily oriented towards sample complexity results. We indeed intentionally targeted on-policy metrics, which explains the expectation operation over histories perceived during the interaction in our bounds. While WAE-MDPs are linked to bisimulation (pseudo-)metrics with guarantees, our belief encoder learning does not directly relate to bisimulation metrics (and definitely not to strong bisimulation concepts). However, we acknowledge the potential to strengthen our guarantees by incorporating bisimulation metrics in future work, as mentioned in our Conclusion, and we agree that this represents a considerable challenge.

Reviewer comment (Writing):
In terms of writing and presentation, I feel that more effort should be put into streamlining the notations and explanations. For example, even after several re-reads, it is still difficult to keep track of the notation for the different models MDPs/POMDPs being passed between, to type-check things (e.g. the state embedding functions depend on different things at different points), and to grok which ideas are essential/novel and which are simply details.

Author response:
**Notations and formalities.** To assist the reader with the notations, we added an Index of notations in Appendix. Note that for the sake of presentation, we needed sometimes to find the right balance between formalities and intuitions in the paper. However, all the statements and rigorous proofs are clearly detailed and formalized in the Appendix. Furthermore, we also kept into consideration the remarks of the other reviewers to improve on this balance and the overall presentation of the results.

Reviewer comment (Theory):
There is a fair amount of statements which are "fake rigorous" i.e. do not check out formally. Here is a list:  
1. "Unlike in MDPs, stationary policies that are based solely on the current observation of P do not induce any probability space on trajectories of M". This is probably not what the authors meant to say, since any policy (stationary or not) induces a distribution over trajectories.  
2. "This procedure guarantees \bar{M} to be probably approximately bisimlarly close to M as \lambda \rightarrow 0". Is this a formal claim? There is no PAC bound given here nor in the cited prior work on WAE-MDPs.  
3. Lemma 3.1: "There is a well defined probability distribution [...] over histories likely to be perceived at the limit by the agent when it executes \bar{\pi} in P". What does this mean? I can try to guess what this mean, but as a formal statement this does not make sense.  
4. Theorem 3.3: there is a "for any \epsilon > 0" in the statement and in the bound but no other quantities in the bound depend on $\varepsilon$.

Author response:
**Distribution over histories.** Lemma 3.1 ensures the existence of a (stationary) distribution over histories occurring under the latent policy, when the latter is executed in the original POMDP. We have modified the main text of the paper to improve the readability of the Lemma. Please, further notice that Appendix C is dedicated to the proof of the statement, where an extended and detailed version of this Lemma is formalized (Lemma C.1).

**Almost Lipschitzness.** In Theorem 3.3, the $\epsilon$ represents an error factor. Intuitively, when the losses go to zero, given $\epsilon > 0$, there always exists a constant $K < \infty$ so that the value function is almost surely $K$-Lipschitz w.r.t. the latent belief space, up to the error factor $\epsilon$.  The error $\epsilon$ is tied to $K$ as the proof relies on the _almost-Lipschitzness_ of the value function (see Definition E.8 and Lemma E.9 in Appendix). The epsilon-K notation is standard when referring to almost Lipschitzness (cf . Vanderbei 1991).

Reviewer comment (Experiments):
The experimental results show mild improvement over the baselines in some environments, and no improvement over the baselines in other environments. What is the source of this varying degree of improvement? Can more experiments be performed to understand when this specific approach yields benefits?

Author response:
**Experimental performance.** The source of the varying degree of improvement compared to the baseline is detailed in the experimental section. In short, whenever only short memorization is required, RNNs tend to perform well and learn fast, although they do so without any guarantees. Their learning capability and stability decrease as the size of required memory increases or when observations are noisy.

Reviewer comment (Presentation):
Assumption(2) requiring an episodic RL process appears unexpectedly in Section 3.2. It would help to anticipate the need for this assumption earlier in the introduction.

Author response:
**Episodic assumption.** Assumption 2 holds in a wide range of RL scenarios: intuitively this means that the agent is trained in an episodic manner, i.e., the environment is reset when the agent succeeds, fails, or after a finite number of time steps.This actually holds in the majority of RL environments relying on a simulator (see, e.g., Brockman et al., 2016; Pardo et al., 2018). Continual RL is one of the few settings where this Assumption does not hold. However, in our setting, the episodic assumption is used to prove the existence of a stationary distribution over the histories perceived under the learned policy (Lemma 3.1). We note that we can relax the episodic assumption if we can assume the existence of a stationary distribution, which is common in continual RL (see, e.g., Huang 2020 for a discussion). We added a remark (Remark 2) in the paper related to this episodic assumption, aiming to explore the constraints of this assumption and identify scenarios in which it can be relaxed.

Reviewer comment (Theory):
It would help to discuss how restrictive Assumption(2) is, as was done for Assumption(1), and whether it is mainly a technicality for some of the theoretical derivations.

Author response:
**Episodic assumption.** Assumption 2 holds in a wide range of RL scenarios: intuitively this means that the agent is trained in an episodic manner, i.e., the environment is reset when the agent succeeds, fails, or after a finite number of time steps.This actually holds in the majority of RL environments relying on a simulator (see, e.g., Brockman et al., 2016; Pardo et al., 2018). Continual RL is one of the few settings where this Assumption does not hold. However, in our setting, the episodic assumption is used to prove the existence of a stationary distribution over the histories perceived under the learned policy (Lemma 3.1). We note that we can relax the episodic assumption if we can assume the existence of a stationary distribution, which is common in continual RL (see, e.g., Huang 2020 for a discussion). We added a remark (Remark 2) in the paper related to this episodic assumption, aiming to explore the constraints of this assumption and identify scenarios in which it can be relaxed.

Reviewer comment (Theory):
Perhaps related to Assumption(2), the future work section mentioned adaptations under relaxed assumptions, citing Lambrechts et al. 2022. It would help to elaborate on some of those potential relaxations.

Author response:
**Episodic assumption.** Assumption 2 holds in a wide range of RL scenarios: intuitively this means that the agent is trained in an episodic manner, i.e., the environment is reset when the agent succeeds, fails, or after a finite number of time steps.This actually holds in the majority of RL environments relying on a simulator (see, e.g., Brockman et al., 2016; Pardo et al., 2018). Continual RL is one of the few settings where this Assumption does not hold. However, in our setting, the episodic assumption is used to prove the existence of a stationary distribution over the histories perceived under the learned policy (Lemma 3.1). We note that we can relax the episodic assumption if we can assume the existence of a stationary distribution, which is common in continual RL (see, e.g., Huang 2020 for a discussion). We added a remark (Remark 2) in the paper related to this episodic assumption, aiming to explore the constraints of this assumption and identify scenarios in which it can be relaxed.

Reviewer comment (Theory):
The variance assumption in Thm.3.3 appears to be a technicality, and may hinder the actual message. Looking at the proof, it seems only needed for the value function to be almost Lipschitz. Smoothing the observation — introduced in S3.1 to enable training — would make the variance vanish by design as learning converges. It is suggested to amend the statement of Thm.3.3 to clarify this, for example by stating “as learning converges to the deterministic mapping ensured by ...”.

Author response:
**Comments on the presentation.** We have updated the paper to implement the presentation suggestions. As the reviewer rightfully noticed, the discussion about the observation variance is technical and only required to learn the belief updater, in Section 4. Therefore, we removed this discussion from Section 3, we modified Theorem 3.3, and refined Remark 3 in Section 4. Furthermore, to help the reader to keep track of the notations, we added an index of notations in Appendix.