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

arXiv: https://arxiv.org/abs/2303.03284 · HTML: https://mlsanity.com/paper/2303.03284/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 76 references that are in the map.
- First posted on arXiv 6 Mar 2023.

## Builds on (most cited first, 28 of 28)

- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Conditional Image Generation with PixelCNN Decoders](https://mlsanity.com/paper/1606.05328/) (2016)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Deep Variational Reinforcement Learning for POMDPs](https://mlsanity.com/paper/1806.02426/) (2018)
- [Partially Observable Markov Decision Processes in Robotics: A Survey](https://mlsanity.com/paper/2209.10342/) (2022)
- [Machine Learning in Python: Main developments and technology trends in data science, machine learning, and artificial intelligence](https://mlsanity.com/paper/2002.04803/) (2020)
- [Shaping Belief States with Generative Environment Models for RL](https://mlsanity.com/paper/1906.09237/) (2019)
- [POPGym: Benchmarking Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2303.01859/) (2023)
- [Temporal Difference Variational Auto-Encoder](https://mlsanity.com/paper/1806.03107/) (2018)
- [Particle Filter Recurrent Neural Networks](https://mlsanity.com/paper/1905.12885/) (2019)
- [Steady State Analysis of Episodic Reinforcement Learning](https://mlsanity.com/paper/2011.06631/) (2020)
- [Flow-based Recurrent Belief State Learning for POMDPs](https://mlsanity.com/paper/2205.11051/) (2022)
- [Optimal Cost Almost-sure Reachability in POMDPs](https://mlsanity.com/paper/1411.3880/) (2014)
- [Recurrent networks, hidden states and beliefs in partially observable environments](https://mlsanity.com/paper/2208.03520/) (2022)
- [Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes (Technical Report)](https://mlsanity.com/paper/2112.09655/) (2021)

## Built on by (most cited first, 9 of 9)

- [Informed POMDP: Leveraging Additional Information in Model-Based RL](https://mlsanity.com/paper/2306.11488/) (2023)
- [PIGDreamer: Privileged Information Guided World Models for Safe Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2508.02159/) (2025)
- [A Theoretical Justification for Asymmetric Actor-Critic Algorithms](https://mlsanity.com/paper/2501.19116/) (2025)
- [An Empirical Study on the Power of Future Prediction in Partially Observable Environments](https://mlsanity.com/paper/2402.07102/) (2024)
- [CoEx -- Co-evolving World-model and Exploration](https://mlsanity.com/paper/2507.22281/) (2025)
- [Deep SPI: Safe Policy Improvement via World Models](https://mlsanity.com/paper/2510.12312/) (2025)
- [Informed Asymmetric Actor-Critic: Leveraging Privileged Signals Beyond Full-State Access](https://mlsanity.com/paper/2509.26000/) (2025)
- [Descriptive History Representations: Learning Representations by Answering Questions](https://mlsanity.com/paper/2506.02125/) (2025)
- [Reinformed Dreamer: An Asymmetric World Model Efficiently Trained through Latent Guidance](https://mlsanity.com/paper/2607.26040/) (2026)

## Abstract (arXiv)

Partially Observable Markov Decision Processes (POMDPs) are used to model environments where the full state cannot be perceived by an agent. As such the agent needs to reason taking into account the past observations and actions. However, simply remembering the full history is generally intractable due to the exponential growth in the history space. Maintaining a probability distribution that models the belief over what the true state is can be used as a sufficient statistic of the history, but its computation requires access to the model of the environment and is often intractable. While SOTA algorithms use Recurrent Neural Networks to compress the observation-action history aiming to learn a sufficient statistic, they lack guarantees of success and can lead to sub-optimal policies. To overcome this, we propose the Wasserstein Belief Updater, an RL algorithm that learns a latent model of the POMDP and an approximation of the belief update. Our approach comes with theoretical guarantees on the quality of our approximation ensuring that our outputted beliefs allow for learning the optimal value function.