Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes

Abstract

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

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

“…The optimization process relies on a temperature 𝜆 ∈ [0, 1) that controls the continuity of the latent space learned, the zero-temperature corresponding to a discrete latent state space. This procedure guarantees M 𝜃 to be probably approximately bisimilarly close [12,16,27] to M as 𝜆 → 0: in a nutshell, bisimulation metrics imply the closeness of the two models in terms of probability measures and expected return [13,14].…”

“…Wasserstein Auto-encoded MDPs (WAE-MDPs) [11] are latent space models that are trained based on the optimal transport from trajectory distributions, resulting from the execution of the RL agent policy in the real environment M, to that reconstructed from the latent model M 𝜃 . The optimization process relies on a temperature 𝜆 ∈ [0, 1) that controls the continuity of the latent space learned, the zero-temperature corresponding to a discrete latent state space.…”

“…Precisely, to accommodate with the WAE-MDP framework and leverage the guarantees presented in Sect. 3.2, we use the MAF presented in [11] that learns multivariate, latent, relaxed distributions that become discrete (binary) in the zero-temperature limit of the WAE-MDP. We define the sub-belief as the vector representation of the belief distribution, the MAF allowing the transformation of sub-beliefs into beliefs.…”

“…We learn the latent model of the POMDP with Wasserstein-MDP (WAE-MDP) [12]. We then learn the belief-update network (BUN) by minimizing the Wasserstein distance with the exact belief-update rule in the latent POMDP.…”

MSVIPER: Improved Policy Distillation for Reinforcement-Learning-Based Robot Navigation

“…Policy distillation involves transformation of a policy from one format to another while keeping the essential inputoutput as similar as possible. This could be transforming a neural net or Markov Decision Process (MDP) into a smaller network or MDP [20], or into a different form entirely, such as a saliency map [21] or tree [22], [13]. For policy distillation, our approach is based on VIPER [13] because it is flexible in terms of being able to use any method to learn the expert policy.…”