The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models

We thank the reviewer for their insightful comments, and answer their questions in the following.

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

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.

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.

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.

References:

• Florent Delgrange, Ann Nowe, and Guillermo Perez. Wasserstein auto-encoded MDPs: Formal verification of efficiently distilled RL policies with many-sided guarantees. (ICLR 2023)

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.

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

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.

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.

References :

• Bojun Huang. Steady state analysis of episodic reinforcement learning. (NeurIPS 2020)
• D. Revuz. Markov Chains. North-Holland mathematical library. Elsevier Science Publishers B.V., second (revised) edition, 1984.
• Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming

We would like to thank the reviewer for their comments, having attentively checked the Appendix, and their very positive appreciation of our work.

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.

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.

References:

• Bojun Huang. Steady state analysis of episodic reinforcement learning. (NeurIPS 2020)
• Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym (2016)
• Fabio Pardo, Arash Tavakoli, Vitaly Levdik, and Petar Kormushev. Time limits in reinforcement learning. (ICML 2018)

We appreciate the reviewer's insightful comments, and in the subsequent discussion, we aim to answer their engaging technical questions and address their concerns.

Responses to specific questions:

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.

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

Sampling states from the true belief. In the section related to local losses, sampling a state from the true belief function w.r.t. a history is achieved in practice via a replay buffer, where states perceived during the interaction are stored. Sampling a state from the replay buffer is therefore equivalent to sampling (i) a history, and then (ii) a state from the resulting belief computed through the true belief update function $\tau^*$.

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

Representation quality relies on optimality. The reviewer rightly noticed that the applicability of the Theorem is limited to optimal policies. At the moment, we conjecture that the Theorem may not apply to arbitrary policies. Technically, this is due to the fact that we rely on almost Lipschitzness of the optimal value function of the POMDP (see Lemma E.9 in Appendix for additional details), which does not necessarily hold in general, for arbitrary policies. Nevertheless, in the context of RL we are particularly interested in optimal policies, so it is natural to focus on them.

Negative belief loss. The negative values in Figure 3 result from the Monte Carlo approximation of the belief loss. This actually may occur when minimizing KL for two distributions $P, Q$ with densities $p, q$ via Monte Carlo sampling, since the log-ratio $\log \frac{p(x)}{q(x)}$ of the two distributions is averaged from samples produced by $x \sim P$ and not integrated over the whole space. We made clearer in the caption of Figure 3 that the belief loss is estimated.

Nature of our guarantees

We argue that our algorithm does enable an “end-to-end” learning in the sense that, at any time of the optimization process, the current measures of the losses of our algorithm serve as indicators of the proximity between the latent and original models. In essence, this is exactly what we present in our main text; and in addition to end-to-end learning it allows a posteriori model checking, if desired. Our work is not oriented towards sample complexity or regrets, but rather on the representation guarantees that are induced by our algorithm based on the optimization of tractable losses. This tractability is related to their measurability, and the feasibility of their optimization and estimation at any time of the interaction (a.k.a. “on-policy” learning).

Precisely, we focus on representation quality guarantees when function approximators are used to generate a suitable representation of the space over which decisions are based. To do so, we provide tractable losses, whose minimizations are compliant with stochastic gradient descent algorithms used by neural networks. These losses are guaranteed to lead to (i) a latent space model where the agent behaves closely to the way it behaves in the original environment, and (ii) a latent representation of the history space (over which decisions are based) that is guaranteed to be almost-Lipschitz continuous in the value function of the optimal policy. While (i) allows the agent to maintain a belief on the latent state space and take its decisions accordingly, (ii) guarantees that the representation induced by our approach is suitable to optimize the expected return, while ensuring that collapsing issues does not occur under optimality (any histories having close representation never lead to divergent returns).

We stress that we designed the losses to be locally tractable, i.e., measurable and estimable under latent policies, which is in stark contrast to losses whose optimality relies on the whole state, action, and observation spaces, which is infeasible to optimize in practice. This means that at any time of the training process, the current measures of the losses are informative on the closeness of the latent and original models (Theorem 3.2).

We agree with Reviewer Rz6s that the representation quality (Theorem 3.3) holding specifically for the optimal latent policy can be seen as a limitation since it could be interesting to assess the representation quality offered by our representation for any policy, but (i) we specifically focus on RL, where we seek for optimality, and (ii) we currently conjecture that such a bound does not hold for arbitrary policies, under the current assumptions. We defer as future work the refinement of the quality of the representation under relaxed assumptions.

References

• Lee, J., Agarwal, A., Dann, C. and Zhang, T., 2023, July. Learning in pomdps is sample-efficient with hindsight observability. ICML 2023
• Liu, Q. et al. (2022). When is partially observable reinforcement learning not scary?. In Conference on Learning Theory. PMLR.
• Huang, R. et al. Provably efficient ucb-type algorithms for learning predictive state representations. (2023).
• SeyedSalehi E. el al., Approximate information state based convergence analysis of recurrent Q-learning, EWRL 2023
• Larsen and Skou, Bisimulation through probabilistic testing. Proceedings of the 16th ACM SIGPLAN-SIGACT symposium on Principles of programming languages. 1989.
• Givan, R. et al. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence (2003)
• De Alfaro, L., et al.. Discounting the future in systems theory. International Colloquium on Automata, Languages, and Programming. 2003.
• Desharnais, J., et al. Metrics for labelled Markov processes. Theoretical computer science (2004)
• Chatterjee, et al. Algorithms for game metrics. Logical Methods in Computer Science (2010).
• Ferns, N.,et al.. Metrics for finite Markov decision processes. (2012).
• Gelada, et al. Deepmdp: Learning continuous latent space models for representation learning. ICML. PMLR, 2019.
• Van der Pol E, et al.. Plannable Approximations to MDP Homomorphisms: Equivariance under Actions, AAMAS 2020.
• Castro, PS, et al. MICo: Improved representations via sampling-based state similarity for Markov decision processes. Advances in Neural Information Processing Systems 34 (2021):
• Zhang, A., et al.(2021). Learning invariant representations for reinforcement learning without reconstruction.
• Zang, H., et al.. SimSR: Simple distance-based state representations for deep reinforcement learning. AAAI 2022.
• F. Delgrange, et al.. Wasserstein auto-encoded MDPs: Formal verification of efficiently distilled RL policies with many-sided guarantees. (ICLR 2023)
• Castro, PS, et al. Equivalence Relations in Fully and Partially Observable Markov Decision Processes. IJCAI. 2009.