Research map / Reinforcement learning / MDPs and planning
The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Wasserstein Belief Updater (WBU), a model-based RL method that learns a WAE-MDP latent POMDP model with privileged state access during training and a feed-forward, normalizing-flow belief encoder trained to match the belief update, with value-difference guarantees.
Problem
Learning reliable belief representations for POMDPs without relying on RNNs with BPTT, which lack guarantees that their history encodings suffice for optimal return.
Approach
Learn a latent model (WAE-MDP) of the POMDP using state access during training, then train a feed-forward sub-belief encoder with a normalizing flow (MAF) to approximate the exact belief update via a Wasserstein/KL proxy, avoiding BPTT, with theoretical bounds on value differences and representation quality.
Key results
- WBU is the only evaluated method that shows mid- to long-term memorization on RepeatPrevious, where R-A2C and DVRL fail.
- On SpaceInvaders with hidden alien direction and confounded fire, WBU earns the highest rewards among WBU, R-A2C and DVRL.
- On StatelessCartPole, R-A2C learns faster early, while WBU eventually reaches R-A2C's final performance; under noise WBU converges faster than R-A2C, which has higher variance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use WBU instead of R-A2C or DVRL | RepeatPrevious · RepeatPrevious | undiscounted cumulative return | Helped |
| Use WBU instead of R-A2C | StatelessCartPole · StatelessCartPole | undiscounted cumulative return | Mixed |
| Use WBU instead of DVRL | StatelessCartPole · StatelessCartPole | undiscounted cumulative return | Helped |
| Use WBU instead of R-A2C and DVRL | SpaceInvaders · SpaceInvaders | undiscounted cumulative return | Helped |
| Use WBU instead of R-A2C | StatelessCartPole with Gaussian observation noise · StatelessCartPole with Gaussian noise | undiscounted cumulative return | Mixed |
| Use WBU instead of R-A2C | SpaceInvaders with binary observation noise · SpaceInvaders with binary noise | undiscounted cumulative return | Helped |
Surprising findings
- Feed-forward belief encoder trained without backpropagation through time performs competitively with RNN-based R-A2C.
Limitations
- Assumes access to the true state during training (privileged information).
- Scaling to high-dimensional observations such as images may be computationally intensive due to observation filtering.
- Wasserstein optimization is challenging; KL is used as a practical proxy with weaker guarantees.
Experimental setup
Datasets: POPGym RepeatPrevious POPGym StatelessCartPole partially observable MinAtar SpaceInvaders
Benchmarks: POPGym RepeatPrevious POPGym StatelessCartPole MinAtar SpaceInvaders (partially observable)
Metrics: undiscounted cumulative return belief loss
Baselines: R-A2C DVRL
Builds on: WAE-MDP R2D2-style recurrent policies A2C Masked Autoregressive Flow DVRL
Closest papers by meaning
- Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy (NeurIPS 2024 · similarity 0.88)
- Learning POMDP World Models from Observations with Language-Model Priors (NeurIPS 2026 · similarity 0.88)
- Provable Representation with Efficient Planning for Partially Observable Reinforcement Learning (ICML 2024 · similarity 0.87)
- Provable Partially Observable Reinforcement Learning with Privileged Information (NeurIPS 2024 · similarity 0.87)
- Partially Observable Reinforcement Learning with Memory Traces (ICML 2025 · similarity 0.86)
- ESCORT: Efficient Stein-variational and Sliced Consistency-Optimized Temporal Belief Representation for POMDPs (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.