# Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes (Technical Report)

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 52 references that are in the map.
- First posted on arXiv 17 Dec 2021.

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

- [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)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [Safe Reinforcement Learning via Shielding](https://mlsanity.com/paper/1708.08611/) (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)
- [TensorFlow Distributions](https://mlsanity.com/paper/1711.10604/) (2017)
- [The Probabilistic Model Checker Storm](https://mlsanity.com/paper/2002.07080/) (2020)
- [Safety-Constrained Reinforcement Learning for MDPs](https://mlsanity.com/paper/1510.05880/) (2015)
- [Reverb: A Framework For Experience Replay](https://mlsanity.com/paper/2102.04736/) (2021)
- [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)
- [Learning to Predict Without Looking Ahead: World Models Without Forward Prediction](https://mlsanity.com/paper/1910.13038/) (2019)
- [Probabilistic Guarantees for Safe Deep Reinforcement Learning](https://mlsanity.com/paper/2005.07073/) (2020)
- [Formal Methods with a Touch of Magic](https://mlsanity.com/paper/2005.12175/) (2020)
- [Hybrid system identification using switching density networks](https://mlsanity.com/paper/1907.04360/) (2019)

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

- [The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models](https://mlsanity.com/paper/2303.03284/) (2023)
- [Verified Probabilistic Policies for Deep Reinforcement Learning](https://mlsanity.com/paper/2201.03698/) (2022)
- [Interventionally Consistent Surrogates for Agent-based Simulators](https://mlsanity.com/paper/2312.11158/) (2023)
- [Model-based Offline Reinforcement Learning with Local Misspecification](https://mlsanity.com/paper/2301.11426/) (2023)
- [Deep SPI: Safe Policy Improvement via World Models](https://mlsanity.com/paper/2510.12312/) (2025)
- [MSVIPER: Improved Policy Distillation for Reinforcement-Learning-Based Robot Navigation](https://mlsanity.com/paper/2209.09079/) (2022)
- [Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees](https://mlsanity.com/paper/2303.12558/) (2023)
- [Meta Policy Switching for Secure UAV Deconfliction in Adversarial Airspace](https://mlsanity.com/paper/2506.21127/) (2025)
- [Adaptive state-action abstractions via rate-distortion](https://mlsanity.com/paper/2606.06123/) (2026)

## Abstract (arXiv)

We consider the challenge of policy simplification and verification in the context of policies learned through reinforcement learning (RL) in continuous environments. In well-behaved settings, RL algorithms have convergence guarantees in the limit. While these guarantees are valuable, they are insufficient for safety-critical applications. Furthermore, they are lost when applying advanced techniques such as deep-RL. To recover guarantees when applying advanced RL algorithms to more complex environments with (i) reachability, (ii) safety-constrained reachability, or (iii) discounted-reward objectives, we build upon the DeepMDP framework introduced by Gelada et al. to derive new bisimulation bounds between the unknown environment and a learned discrete latent model of it. Our bisimulation bounds enable the application of formal methods for Markov decision processes. Finally, we show how one can use a policy obtained via state-of-the-art RL to efficiently train a variational autoencoder that yields a discrete latent model with provably approximately correct bisimulation guarantees. Additionally, we obtain a distilled version of the policy for the latent model.