# Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 61 references that are in the map.
- First posted on arXiv 22 Mar 2023.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Safe Reinforcement Learning via Shielding](https://mlsanity.com/paper/1708.08611/) (2017)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [MADE: Masked Autoencoder for Distribution Estimation](https://mlsanity.com/paper/1502.03509/) (2015)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Lagging Inference Networks and Posterior Collapse in Variational Autoencoders](https://mlsanity.com/paper/1901.05534/) (2019)
- [Environment-Independent Task Specifications via GLTL](https://mlsanity.com/paper/1704.04341/) (2017)
- [The Probabilistic Model Checker Storm](https://mlsanity.com/paper/2002.07080/) (2020)
- [MICo: Improved representations via sampling-based state similarity for Markov decision processes](https://mlsanity.com/paper/2106.08229/) (2021)
- [From optimal transport to generative modeling: the VEGAN cookbook](https://mlsanity.com/paper/1705.07642/) (2017)
- [Safety-Constrained Reinforcement Learning for MDPs](https://mlsanity.com/paper/1510.05880/) (2015)
- [Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems](https://mlsanity.com/paper/1910.09588/) (2019)
- [Steady State Analysis of Episodic Reinforcement Learning](https://mlsanity.com/paper/2011.06631/) (2020)
- [Efficient Model-Based Deep Reinforcement Learning with Variational State Tabulation](https://mlsanity.com/paper/1802.04325/) (2018)
- [Deep reinforcement learning for large-scale epidemic control](https://mlsanity.com/paper/2003.13676/) (2020)
- [Probabilistic Guarantees for Safe Deep Reinforcement Learning](https://mlsanity.com/paper/2005.07073/) (2020)
- [SimSR: Simple Distance-based State Representation for Deep Reinforcement Learning](https://mlsanity.com/paper/2112.15303/) (2021)
- [Formal Methods with a Touch of Magic](https://mlsanity.com/paper/2005.12175/) (2020)
- [Model-predictive control and reinforcement learning in multi-energy system case studies](https://mlsanity.com/paper/2104.09785/) (2021)
- [Probabilistic Model Checking and Autonomy](https://mlsanity.com/paper/2111.10630/) (2021)
- [Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes (Technical Report)](https://mlsanity.com/paper/2112.09655/) (2021)
- [LTLf Synthesis on Probabilistic Systems](https://mlsanity.com/paper/2009.10883/) (2020)
- [Wasserstein-Wasserstein Auto-Encoders](https://mlsanity.com/paper/1902.09323/) (2019)

## Abstract (arXiv)

Although deep reinforcement learning (DRL) has many success stories, the large-scale deployment of policies learned through these advanced techniques in safety-critical scenarios is hindered by their lack of formal guarantees. Variational Markov Decision Processes (VAE-MDPs) are discrete latent space models that provide a reliable framework for distilling formally verifiable controllers from any RL policy. While the related guarantees address relevant practical aspects such as the satisfaction of performance and safety properties, the VAE approach suffers from several learning flaws (posterior collapse, slow learning speed, poor dynamics estimates), primarily due to the absence of abstraction and representation guarantees to support latent optimization. We introduce the Wasserstein auto-encoded MDP (WAE-MDP), a latent space model that fixes those issues by minimizing a penalized form of the optimal transport between the behaviors of the agent executing the original policy and the distilled policy, for which the formal guarantees apply. Our approach yields bisimulation guarantees while learning the distilled policy, allowing concrete optimization of the abstraction and representation model quality. Our experiments show that, besides distilling policies up to 10 times faster, the latent model quality is indeed better in general. Moreover, we present experiments from a simple time-to-failure verification algorithm on the latent space. The fact that our approach enables such simple verification techniques highlights its applicability.