{'fraction_non_alphanumeric': 0.049043471263905375, 'fraction_numerical': 0.01659592444328217, 'mean_word_length': 3.783679753656659, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes (Technical Report)

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.


Introduction
While reinforcement learning (RL) has been applied to a wide range of challenging domains, from game playing (Mnih et al. 2015) to real-world applications such as effective canal control (Ren et al. 2021), more widespread deployment in the real world is hampered by the lack of guarantees provided with the learned policies. Although there are RL algorithms which have limit-convergence guarantees in the discrete setting (Tsitsiklis 1994) -and even in some continuous settings with function approximation, e.g., Nowe (1994) -these are lost when applying more advanced techniques which make use of general nonlinear function approximators (Tsitsiklis and Roy 1997) to deal with continuous Markov decision processes (MDPs) such as deep-RL (e.g., Mnih et al. 2015). In this paper, we apply such advanced RL algorithms to unknown continuous MDPs with (i) reachability, (ii) safety-constrained reachability, or (iii) discounted-reward objectives. To recover the formal guarantees, we use the obtained policy to train a variational autoencoder (VAE) which gives us a discrete latent model that approximates the unknown environment. We build upon the DeepMDP framework (Gelada et al. 2019) to provide guarantees on the quality of the abstraction induced by this model. DeepMDPs are provided with such guarantees when their loss functions are minimized. These can be defined on the entire state space (global) or on states visited under a given policy (local). The guarantees concern a state embedding function, linking the latent and original MDPs and are defined as bounds on the difference of their value function and bisimulation distance. The latter was only developed for global losses. While these are interesting in theory, they are often infeasible to measure in practice. In contrast, we introduce such bounds in the local setting and further consider an action embedding function to handle continuous actions. Importantly, we focus on general MDPs and do not restrict our attention to deterministic ones as was done by Gelada et al. to enable the approximation and minimization of their losses via neural networks. We also give PAC approximation schemes to compute both the losses and said bounds.
Our VAE is trained by maximizing a lower bound on the likelihood of traces generated by executing the RL policy in the environment. We derive a loss function, incorporating variational versions of the local losses, that enables learning (i) discrete state and action spaces, (ii) an MDP defined over these spaces, (iii) state and action embedding functions, linking the original and discrete MDPs, and (iv) a distilled version of the RL policy set over the discrete spaces which can be executed in both models via the embedding functions. An important challenge for our approach is the posterior collapse problem which often occurs when optimizing a variational model (e.g., Alemi et al. 2018). We present a novel approach based on prioritized experience replay (Schaul et al. 2016) to resolve this when learning a discrete latent model.
All of the above result in an efficient way of training a VAE to obtain a discrete latent model that is provably approximately bisimilar to the unknown MDP, further yielding a distilled version of the RL policy. These enable the application of formal methods and tools that have been developed for discrete MDPs: for instance, PRISM (Kwiatkowska, Norman, and Parker 2011), MODEST (Hartmanns and Hermanns 2014), and STORM (Hensel et al. 2021). Other related work. Frameworks providing formal guarantees during the RL process include the work of Junges et al. (2016), Shielded-RL (Alshiekh et al. 2018;, and AlwaysSafe (Simão, Jansen, and Spaan 2021). These all require an abstract model of the safety aspect of the environment. Our approach is complementary in that we assume no prior knowledge and learn an abstraction. Notably, our goal is not the same: they aim at verifying whether the exploration is safe while our goal is to verify policies learned via any RL technique. Other approaches share ours in the particular case of verifying deep-RL policies, but rely on a known (abstraction of the) environment model. Bacci and Parker (2020) require the neural network (NN) specifying the policy, the environment to be deterministic, and a formal description of the probability of action failures. Carr, Jansen, and Topcu (2020) require the environment to be discrete and focus on policies represented as recurrent NNs by discretizing their hidden states via quantized autoencoders, in the same spirit as our policy distillation. Finally, Alamdari et al. (2020) focus on tree-based policies distilled from deep-RL, without considering abstraction quality guarantees.
VAEs have been used in the context of (model-based) RL to learn latent representations of the unknown environment and train simpler policies from the features extracted (e.g., Corneil, Gerstner, and Brea 2018;Freeman, Ha, and Metz 2019;Lee et al. 2020;Burden, Siahroudi, and Kudenko 2021). In particular, Corneil, Gerstner, and Brea (2018) focused on learning discrete latent MDPs from continuousstate environments with discrete actions (without guarantees nor distilled policies) to plan via prioritized sweeping.

Background
We write rT s " tn P N | n ď T u. For A Ď X, we denote by 1 A : X Ñ r1s the indicator function: 1 A paq " 1 iff a P A. Let X be a complete and separable space and ΣpX q denote the set of all Borel subsets of X . We write PpX q for the set of measures P defined on X and SupppP q " tx P X | P pxq ą 0u to denote their support. Discrepancy measures. Let P, Q P PpX q with density functions p and q. Their discrepancy can be measured via • Kullback-Leibler (KL) divergence: D KL pP Qq " Ex"P rlogp ppxq {qpxqqs.
• Wasserstein: W d pP, Qq " inf λPΛpP,Qq Ex,y"λ dpx, yq, where d : XˆX Ñ r0, 8r is a distance metric over X and ΛpP, Qq is the set of all couplings of P and Q. • Total Variation (TV): d T V pP, Qq " sup APΣpX q |P pAqQ pAq|. If X is equipped with the discrete metric 1 ‰ , TV coincides with the Wasserstein measure.
Markov decision processes. A Markov decision process (MDP) is a tuple M " xS, A, P, R, , AP, s I y where S is a set of states; A, a set of actions; P : SˆA Ñ PpSq, a probability transition function; R : SˆA Ñ R, a reward function; : S Ñ 2 AP , a labeling function over a set of atomic propositions AP; and s I P S, the initial state. The set of enabled actions of s P S is Actpsq Ď A. We assume Actpsq ‰ H for all s P S. If |Actpsq| " 1 for all s P S, M is a fully stochastic process called a Markov chain (MC).
Let T Ď AP, we write T " t s | psq X T ‰ H u Ď S and T " Sz T . We assume AP and labels being respectively one-hot and binary encoded. We write M s for the MDP obtained when we replace the initial state of M by s P S, M ' R 1 when we replace the reward function by R 1 , and M oe B when we make absorbing states from B Ď S, i.e., by changing Pp¨| s, aq to P 1 p¨| s, aq such that P 1 pB | s, aq " 1 for all s P B, a P Actpsq. We refer to MDPs with continuous states or actions spaces as continuous MDPs. In that case, we assume S and A are complete separable metric spaces equipped with a Borel σ-algebra and ´1 pTq P ΣpSq for any T Ď AP.

Trajectories.
A trajectory τ of M is a sequence of states and actions τ " xs 0:T , a 0:T´1 y where s 0 " s I , s t`1 " Pp¨| s t , a t q and a t P Actps t q for t P rT´1s. The set of infinite trajectories of M is Traj M . An execution traceτ of M is a trajectory that additionally records labels and rewards encountered. The set of execution traces of M is Traces M .
Policies. A (memoryless) policy π : S Ñ PpAq of M is a stochastic mapping from states to actions such that Supppπp¨| sqq Ď Actpsq. The set of memoryless policies of M is Π ml M . An MDP M and π P Π ml M induce an MC M π along with a unique probability measure P M π on the Borel σ-algebra over measurable subsets E Ď Traj M (Puterman 1994). We drop the superscript when the context is clear. For π P Π ml M , we denote by P π : S Ñ PpSq the probability transition distribution of M π , and by R π : S Ñ R its reward function. We writeτ " xs 0:T , a 0:T´1 , r 0:T´1 , l 0:T y " M π for xs 0:T , a 0:T´1 y " P M π withτ P Traces Mπ .
Stationary distributions. Let π P Π ml M , ξ t π : S Ñ PpSq with ξ t π ps 1 | sq " P Ms π`t s 0:8 , a 0:8 | s t " s 1 u˘be the distribution giving the probability for the agent of being in each state of M s after t steps, and B Ď S. B is a strongly connected component (SCC) of M π if for any pair of states s, s 1 P B, ξ t π ps 1 | sq ą 0 for some t P N. It is a bottom SCC (BSCC) if (i) B is a maximal SCC, and (ii) for each s P B, P π pB | sq " 1. The unique stationary distribution of B is ξ π P PpBq. We write s, a " ξ π as shorthand for first sampling s from ξ π and then a from π. An MDP M is ergodic if for all π P Π ml M , the state space of M π consists of a unique aperiodic BSCC with ξ π " lim tÑ8 ξ t π p¨| sq for all s P S.
Events and functions. Let C, T Ď AP, we define the constrained reachability (resp. reachability) event as C U T " t s 0:8 , a 0:8 | Di P N, @j ă i, s j P C ^s i P T u P ΣpTraj M q (resp. ♦T " H U T). Safety w.r.t. a set of failure states T can be expressed as a safe-constrained reachability event to a safe destination C (resp. safety event) through T U C (resp.˝ T " Traj M z♦T). Let γ P r0, 1r, ϕ P t , C U T, ♦T u where is the empty symbol, and R T " p1´γq1 T ˆA , the value obtained by running π P Π ml When ϕ P t C U T, ♦T u, observe that V ϕ π ptq = 1 for t P T and lim γÑ1 V ϕ π psq " P Ms π pϕq for s P S. The action-value function is Q ϕ π ps, aq " R 1 ps, aq`E s 1 "Pp¨|s,aq rγV ϕ π ps 1 qs, with R 1 " R if ϕ " and R 1 " R T otherwise.
Given the original (continuous, possibly unknown) environment modeled as an MDP, a latent space model is another (simpler, smaller, and explicit) MDP with state-action space linked to the original one via embedding functions. The latter can be learned to optimize an equivalence criterion between the two models. Formally, fix MDPs M " xS, A, P, R, , AP, s I y and M " @ S, A, P, R, , AP, s I D such that S is equipped with metric d S . Let φ : S Ñ S and ψ : SˆA Ñ A be respectively state and action embedding functions. We refer to @ M, φ, ψ D as a latent space model of M and M as its latent MDP. We write Π ml " Π ml M , and Q π for the action-value function of a policy π P Π ml in M. We also consider π as a policy in M: states passed to π are embedded with φ, then actions executed are embedded with ψ. Let π P Π ml and s P S, we write a " πp¨| sq for a " πp¨| φpsqq and Q π ps, aq as shorthand for Q π ps, ψps, aqq.
A particular point of interest is to focus on discrete latent models, where d S " 1 ‰ . In the following, we adopt the latent space model formalism of Gelada et al. (2019).
Notations. Let π P Π ml , we write |R ‹ π | for sup sPSˇR π psqˇˇ. We say that M is xK π R , K π P y-Lipschitz if for all s 1 , s 2 P S,ˇR π ps 1 q´R π ps 2 qˇˇď K π R d S ps 1 , s 2 q, W d S`P π p¨| s 1 q, P π p¨| s 2 q˘ď K π P d S ps 1 , s 2 q.
Assuming M is discrete, all W d S terms can be replaced by d T V since Wasserstein coincides with TV when using the discrete metric. In that case, (optimal) constants K π R and K π P can be computed in polynomial time in M for any π P Π ml . Henceforth, we make the following assumptions. Assumption 3.1. MDP M is ergodic. Assumption 3.2. Rewards of M are scaled in the interval r´1 2 , 1 2 s, i.e., R : SˆA Ñ r´1 2 , 1 2 s. Assumption 3.3. The embedding function preserves the labels, i.e., φpsq " s ùñ psq " psq for s P S, s P S.
Seemingly restrictive at first glance, Assumption 3.1 is compliant with RL environments and a wide range of continuous learning tasks (every episodic RL process is ergodic, see Huang 2020). Discarding it restricts the upcoming guarantees to BSCCs. Assumption 3.2 basically requires rewards to be bounded and re-scalable in this interval. This is a reasonable assumption in practice and re-scaling is straightforward if bounds are known (otherwise, dynamical rescaling is still feasible but complicates the implementation).
In Sect. 4.2, we show that Assumption 3.3 can be made trivial. Note that our approach requires Assumption 3.2 and 3.3.
Proofs of the claims made in the following two subsections are provided in supplementary material (Appendix B).

Bisimulation and Value Difference Bounds
We aim now at formally checking whether the latent space model offers a good abstraction of the original MDP M. To do so, we present bounds that link the two MDPs. We extend the bounds from Gelada et al. (2019) to discrete spaces while additionally taking into account state labels and discounted reachability events. Moreover, we present new bisimulation bounds in the local setting. Bisimulation. A (probabilistic) bisimulation is a behavioral equivalence between states. Formally, a bisimulation on M is an equivalence relation B Φ such that for all s 1 , s 2 P S and Φ Ď t R, u, s 1 B Φ s 2 iff PpT | s 1 , aq " PpT | s 2 , aq, ps 1 q " ps 2 q if P Φ, and Rps 1 , aq " Rps 2 , aq if R P Φ, for each action a P A and (Borel measurable) equivalence class T P S{B Φ . Properties of bisimulation include trace, trajectory, and value equivalence (Larsen and Skou 1989;Givan, Dean, and Greig 2003). The relation can be extended to compare two MDPs (in our case M and M) by considering the disjoint union of their state space. We denote the largest bisimulation relation by " Φ . Pseudometrics. Desharnais et al. (2004) introduced bisimulation pseudometrics for continuous Markov processes, generalizing the notion of bisimilariy by assigning a bisimilarity distance between states. A pseudometricd satisfies symmetry and the triangle inequality.
Probabilistic bisimilarity can be characterized by a logical family of functional expressions derived from a logic L. More specifically, given a policy π P Π ml M , we consider a family F L γ pπq of real-valued functions f , parameterized by the discount factor γ and defining the semantics of L in M π . Such a logic allows formalizing discounted properties, including reachability (e.g., Chatterjee et al. 2008). The related pseudometricd π is defined as:d π ps 1 , s 2 q " sup f PF L γ pπq |f ps 1 q´f ps 2 q|, for all s 1 , s 2 P S. We distinguish between pseudometricsd R π , characterized by functional expressions including rewards (e.g., Ferns, Precup, and Knight 2014), andd π , whose functional expressions are based on state labels (e.g., Chen, van Breugel, and Worrell 2012). LetP be the space of pseudometrics on S and ϕ P t R, u. Define ∆ :P ÑP so that ∆pd ϕ π qps 1 , s 2 q is p1´γq |R π ps 1 q´R π ps 2 q|¨1 t R u pϕq`M where: M " max " γW d ϕ π pP π p¨| s 1 q, P π p¨| s 2 qq , 1 ‰ p ps 1 q, ps 2 qq¨1 t u pϕq, thend ϕ π is its unique fixed point whose kernel is " tϕu , i.e., d ϕ π ps 1 , s 2 q " 0 iff s 1 " tϕu s 2 (van Breugel and Worrell 2001; Ferns, Precup, and Knight 2014), and |V π ps 1 q´V π ps 2 q| ď d R π ps1,s2q {1´γ (Ferns, Panangaden, and Precup 2005). Bisimulation bounds. While computing this distance is intractable in practice, especially when continuous spaces are involved, we emphasize that local losses can be used to evaluate the original and latent model bisimilarity. Fix π P Π ml and assume M is discrete. Given the induced stationary distribution ξ π in M, where M is xK π R , K π P y-Lipschitz. The result provides us a general way to assess the quality of the abstraction: the bisimulation distance between states and their embedding is guaranteed to be small in average whenever local losses are small. Analogously, local losses can further be used to check whether two states with the same representation are indeed bisimilarly close: for s 1 , s 2 P S with φps 1 q " φps 2 q, π ps 1 q`ξ´1 π ps 2 q˘, d π ps 1 , s 2 q ď γL ξπ P 1´γ`ξ´1 π ps 1 q`ξ´1 π ps 2 q˘. (2) Remark 1. Our goal is to quantify the gap between the behaviors induced by the latent policy executed in the original environment compared to those induced when executed in the latent MDP. This allows for example to lift the guarantees obtained by model-checking M π to M π . For this reason, notice that the expectations of Eq. 1 are set over states produced by π in the original environment: this way, states that are likely to be seen under π in M are ensured to have this behavioral gap bounded by a factor of L ξπ R and/or L ξπ P . Value difference bounds. Considering discounted returns or a specific event, the quality of the latent abstraction can be in particular formalized by means of value difference bounds. These bounds can be intuitively derived by taking the value function as a real-valued function from P˘, then, E s,a"ξπˇQ π ps, aq´Q π pφpsq, aqˇˇď Moreover, for any states s 1 , s 2 P S with φps 1 q " φps 2 q, |V π ps 1 q´V π ps 2 q| ď 1´γ`ξ´1 π ps 1 q`ξ´1 π ps 2 q˘, |V ϕ π ps 1 q´V ϕ π ps 2 q| ď γL ξπ P 1´γ`ξ´1 π ps 1 q`ξ´1 π ps 2 q˘.
Intuitively, when local losses are sufficiently small, then (i) the value difference of states and their embedding that are likely to be seen under a latent policy is also small, and (ii) states with the same embedding have close values.

Checking the Quality of the Abstraction
While bounding the difference between values offered by π P Π ml in M and M is theoretically possible using L ξπ R and L ξπ P , we need to accurately approximate these losses from samples to further offer practical guarantees (recall that M is unknown). Although the agent is able to produce execution traces by interacting with M, estimating the expectation over the Wasserstein is intractable. Intuitively, even if approximating Wasserstein from samples is possible (e.g., Genevay et al. 2019), this would require access to a generative model for Pp¨| s, aq (e.g., Kearns, Mansour, and Ng 2002) from which we would have to draw a sufficient number of samples for each s, a drawn from ξ π to then be able to estimate the expectation. Gelada et al. (2019) overcome this issue by assuming a deterministic MDP, which allows optimizing an approximation of L ξπ P through gradient descent. To deal with general MDPs, we study an upper bound on L ξπ P that can be efficiently approximated from samples: where ξ π ps, a, s 1 q " ξ π ps, aq¨Pps 1 | s, aq and φps | sq " 1 " pφpsq, sq. We now provide probably approximately correct (PAC) guarantees for estimating local losses for discrete latent models, derived from the Hoeffding's inequalities.
This yields the following Theorem, finally allowing to check the abstraction quality as a bounded value difference.
Theorem 3.5. Let π P Π ml and assume M is discrete and P˘. Let ε, δ P s0, 1r and ξ π be the stationary distribution of M π . Then, after T ě Q´l We now provide a framework based on variational autoencoders (Kingma and Welling 2014) that allows us to learn a discrete latent space model of M through the interaction of the agent executing a pre-learned RL policy π P Π ml M with the environment. Concretely, we seek a discrete latent space model We propose to learn the parameters xι, θy of an encoder Q ι and a behavioral model P θ from which we can retrieve (i) the embedding functions φ ι and ψ ι,θ , (ii) the latent MDP components P θ , R θ , and θ , and (iii) a latent policy π θ P Π ml , via min θ DpM π , P θ q, where D is a discrepancy measure. Intuitively, the end goal is to learn (i) a discrete representation of S and A (ii) to mimic the behaviors of the original MDP over the induced latent spaces, thus yielding a latent MDP with a bisimulation distance close to M, and (iii) to distill π into π θ . In the following, we distinguish the case where we only learn S, with A " A (in that case, A is assumed to be discrete), and the one where we additionally need to discretize the set of actions and learn A.

Evidence Lower Bound
In this work, we focus on the case where D KL is used as discrepancy measure: the goal is optimizing min θ D KL pM π P θ q or equivalently maximizing the marginal log-likelihood of traces of M, i.e., Eτ"M π rlog P θ pτ qs, where τ " xs 0:T , a 0:T´1 , r 0:T´1 , l 0:T y, z P Z with Z " S if A " A and Z " SˆA otherwise, P π θ ps 0:T q " ś T´1 t"0 P π θ ps t`1 | s t q, and P π θ ps 0:T , a 0:T´1 q " ś T´1 t"0 π θ pa t | s t q¨P θ ps t`1 | s t , a t q. The dependency ofτ on Z in Eq. 5 is made explicit by the law of total probability.
Optimizing Eτ"M π rlog P θ pτ qs through Eq. 5 is typically intractable (Kingma and Welling 2014). To overcome this, we use an encoder Q ι pz 0:T |τ q to set up a lower bound on the log-likelihood of produced traces, often referred to as evidence lower bound (ELBO, Hoffman et al. 2013): The purpose of optimizing the ELBO is twofold. First, this allows us learning φ ι via (i) Q ι ps 0:T |τ q " ι being an action encoder. We assume here that encoding states and actions to latent spaces is independent of rewards. We additionally make them independent of l 0:T by assuming that is known. This allows φ ι to encode states and their labels directly into the latent space (cf. Sect. 4.2).
Second, we assume the existence of latent reward and label models, i.e., P R θ and P θ , allowing to recover respectively R θ and θ , as well as a generative model P G θ , enabling the reconstruction of states and actions. This allows decomposing the behavioral model P θ into: P θ ps 0:T , a 0:T´1 , r 0:T´1 , l 0: Model ψ θ allows learning the action embedding function via ψ ι,θ pa | s, aq " Es"φ ι p¨|sq ψ θ pa | s, aq for all s P S, a P A, a P A. We also argue that a perfect reconstruction of labels is possible, i.e., P θ pl 0:T | s 0:T q " 1, due to the labels being encoded into the latent space. From now on, we thus omit the label term.
Deterministic embedding functions φ ι , ψ ι,θ and R θ can finally be obtained by taking the mode of their distribution. Back to the local setting. Taking Assumption 3.1 into account, drawing multiple finite tracesτ " M π can be seen as a continuous interaction with M along an infinite trace (Huang 2020). This observation allows us to formulate the ELBO in the local setting 1 and connect to local losses: where D and R denote respectively the distortion and rate of the variational model (Alemi et al. 2018), given by We omit the subscripts when the context is clear. The optimization of ELBO`M θ , φ ι , ψ ι,θ˘a llows for an indirect optimization of the local losses through their variational versions: (i) L ξπ R via the log-likelihood of rewards produced, and (ii) 9 L ξπ P where we change the Wasserstein term to the KL divergence. Note that this last change means we do not necessarily obtain the theoretical guarantees on the quality of the abstraction via its optimization. Nevertheless, our experiments indicate KL divergence is a good proxy of the Wasserstein term in practice. In particular, in the discrete setting, Wasserstein matches TV and one can relate the proxy with the original metric using the Pinkster's inequality.

VAE Distributions
Discrete distributions. We aim at learning discrete latent spaces S and A, the distributions φ ι , Q A ι , P θ , and π θ are thus supposed to be discrete. Two main challenges arise: (i) gradient descent is not applicable to learn ι and θ due to the discontinuity of S and A, and (ii) sampling from these distributions must be a derivable operation. We overcome these by using continuous relaxation of Bernoulli distributions to learn a binary representation of the latent states, and the Gumbel softmax trick for the latent action space (Jang, Gu, and Poole 2017; Maddison, Mnih, and Teh 2017). Labels. To enable log P θ pl 0:T | s 0:T q " 0, we linearly encode ps t q " l t into each s t via φ ι . Recall that labels are binary encoded, so we allocate them |AP| bits in S. Then, φ ι ps | sq ą 0 implies psq " θ psq, for all s P S, s P S, satisfying Assumption 3.3 if φ ι is deterministic. Decoders. For P G θ , ψ θ , and P R θ , we learn the parameters of multivariate normal distributions. This further allows linking all s P S to the parameters of P G θ p¨| sq for explainability.

Posterior Collapse
A common issue encountered while optimizing variational models via the ELBO is posterior collapse. Intuitively, this results in a degenerate local optimum where the model learns to ignore the latent space. With a discrete encoder, this translates into a deterministic mapping to a single latent state, regardless of the input. From an information-theoretic point of view, optimizing the ELBO gives way to a tradeoff between the minimization of R and D, where the feasible region is a convex set (Alemi et al. 2018). Posterior collapse occurs when R « 0 (auto-decoding limit). On the other hand, one can achieve D « 0 (auto-encoding limit) at the price of a higher rate. Regularization terms. Various solutions have been proposed in the literature to prevent posterior collapse. They include entropy regularization (via α P r0, 8r, e.g., Burke, Hristov, and Ramamoorthy 2019;Dong et al. 2020) and KLscaling (via β P r0, 1s, e.g., Alemi et al. 2018), consisting in changing ELBO to xα, βy -ELBO "´pD`β¨Rqὰ¨H pQ ι q, where HpQ ι q denotes the entropy of an encoding distribution. We choose to measure the entropy of the marginal encoder, given by Q ι psq " Es"ξ π φ ι ps | sq. Intuitively, this encourages the encoder to learn to make plenty use of the latent space. The parameter β allows to interpolate between auto-encoding and auto-decoding behavior, which is not possible with the standard ELBO objective.
A drawback of these methods is that we no longer optimize a lower bound on the log-likelihood of the input while optimizing xα, βy -ELBO. In practice, setting up annealing schemes for α and β allows to eventually recover ELBO and avoid posterior collapse (α " 0 and β " 1 matches ELBO). Prioritized replay buffers. To enable meaningful use of latent space to represent input state-actions pairs, Q ι should learn to (i) exploit the entire latent space, and (ii) encode wisely states and actions of transitions yielding poor ELBO, being generally sensitive to bad representation embedding. This motivates us to use a prioritized replay buffer (Schaul et al. 2016) to store transitions and sample them when optimizing ELBO. Draw xs, a, r, s 1 y " ξ π and let p xs,a,r,s 1 y be its priority, we introduce the following priority functions.
• Bucket-based priority: we partition the buffer in |S| buckets. Let N P N and b : S Ñ N be respectively step and latent state counters. At each step, let s " φ ι p¨| sq, we assign p xs,a,r,s 1 y to N {bpsq, then increment N and bpsq of one. This allows φ ι to process states being infrequently visited under π and learn to fairly distribute S. • Loss-based priority: we set p xs,a,r,s 1 y to its individual transition loss, which enables to learn improving the representation of states and actions that yield poor ELBO.

Experiments
The goal of our experiments 2 is to evaluate the quality of the latent space model learned and the policy distilled via our VAE-MDP framework. This evaluation consists of: an analysis of the training of the latent space model and the benefits of our method to avoid posterior collapse, assessing the quality of the abstraction learned via PAC local losses bounds, and testing the performance of the distilled policy. This allows to assess if the latent model learned yields a sound compression of the state-action space that retains the necessary information to optimize the return. We evaluate our method on classic OpenAI environments (Brockman et al. 2016) with (i) continuous states and discrete actions (Cart-Pole, MountainCar, and Acrobot), and (ii) where both, states and actions, are continuous (Pendulum and LunarLander). We distill RL policies π learned via DQN (Mnih et al. 2015) for case (i), and SAC (Haarnoja et al. 2018) for case (ii). The Latent spaces. Since latent spaces are trained to enable formal verification, we choose log 2 |S| and |A| ranging from 9 bits and 2 actions (coarser) to 16 bits and 5 actions (finer) to make them tractable for model checkers -see Budde et al. (2020) for a performance comparison of modern tools for a range of instances with large model size. Depending on the property to verify, the latent space may have to be finer since (i) we need to reserve |AP| bits in the representation of S for labels, and (ii) this allows the agent to take more precise decisions over a finer partition of the latent space. We then select the model with the best trade-off between abstraction quality (measured viaL ξπ θ R andL ξπ θ P ) and performance (i.e., the return approximated by running π in M).
ELBO optimization. In order to allow ELBO to be trained efficiently, posterior collapse has to be tackled from the very first stages of training. In fact, we found that the KL-scaling and entropy regularization annealing schemes (cf. Sect. 4.3) were necessary to avoid the latent space to collapse into a single state-action pair after only a few training steps. We found the most efficient to start with an autoencoder behavior (β 0 " 0) and a large entropy regularizer (e.g., α 0 " 10) during the 10 4 first steps, and then anneal them via α t " α 0¨p 1´τ q t , β t " 1´p1´τ q t (e.g., τ " 10´5) to fully recover the original ELBO in a second training phase. We also found that prioritized experience replays further prevent posterior collapse during training (cf. Fig. 1). Fig. 2a shows that training ELBO this way results in a stable learning procedure that successfully minimizes D while preventing an auto-decoding behavior by keeping R away from 0.
Local losses. We compute the PAC local losses bounds presented in Sect. 3 along training steps (Fig. 2b). When the policy is eventually distilled (Fig. 2c) and π θ achieves performance similar to those of π, the learning curves stabilize, while maximizing ELBO successfully allows minimizing the local losses. In every environment where we tested our method, a low reward loss is maintained while the transition loss reaches values in the interval r 1 {5, 3 {5s for most of the instances. These values are thus guaranteed to upper bound the expected bisimulation distance between M and M as well as their value difference. We do not expect our approach to reach zero reward and transition loss though, since we pass from continuous to discrete spaces: the abstraction induced by our approach is always coarser than the original spaces, which translates in general to precision loss. This precision loss is often encoded through the discrete probability transitions. For instance, a state s which deterministically transitions to a close state s 1 in M such that φ ι psq " s " φ ι ps 1 q induces P π θ ps | sq ą 0. Observe that the PAC computation ofL ξπ θ P is sensitive to the entropy of P θ , which thus may induce a residual transition loss. L ξπ θ P is on the contrary not sensitive to this. In practice, its value will thus often be lower than its approximated upper boundL ξπ θ P . Policy distillation. The guarantees derived in Sect. 3 are only valid for π θ , the policy under which formal properties can be verified. We further need evaluate if it achieves sound performance in the RL environment: checking the expected value difference via Thm. 3.5 is most significant when V π θ psq is worth in any state s likely to be reached in M π θ . We compare the episode return achieved by π θ against π in the environment M (Fig. 2c). Our framework allows learning to improve π θ along training steps and eventually achieving the return of the original policy π. This illustrates that training our latent model enables distilling π into a latent policy π θ that achieves similar performance in M.

Conclusion
In this work, we presented VAE-MDPs, a framework for learning discrete latent models of unknown, continuousspaces environments with bisimulation guarantees. We detailed how such latent models can be learned by executing an RL policy in the environment and showed that the procedure yields a distilled version of the latter as a side effect. To provide the guarantees, we introduced new local losses bounds aimed at discrete latent spaces with their PAC-efficient approximation algorithm derived from the execution of the distilled policy. All this enables the verification of RL policies for unknown continuous MDPs.
Experimental results demonstrate the feasibility of our approach through the PAC bounds and the performance of the distilled policy achieved for various environments. Our tool can also be used to highlight the lack of robustness of input policies when the distillation fails.
Complementary to safe-RL approaches addressed via formal methods, we emphasize the ability of our tool to be coupled with such algorithms when no model of the environment is known a priori. The applicability of our method enables its use in future work for real-world case studies and more complex settings, such as multi-agent systems.

A Details on Discrepancy Measures
Let P, Q P PpX q be probability distributions over X with density functions p and q.
Wasserstein. The Kantorovich-Rubinstein duality allows formulating the Wasserstein distance between P and Q as where F d is the set of 1-Lipschitz functions, i.e., F d " t f : |f pxq´f pyq| ď dpx, yq u.
Total Variation. One can reformulate the total variation distance between P and Q as When X is continuous, TV might thus give an overly strong measure of the numerical differences across probability distributions to all measurable sets since the distance between two point masses dx, dy is 1 unless x " y (Ferns, Panangaden, and Precup 2005).

B Latent Space Models
From now on, fix MDP M " xS, A, P, R, , AP, s I y with latent space model @ M, φ, ψ D , M " @ S, A, P, R, , AP, s I D . We write V π for the value function of a policy π P Π ml in M.

Smooth-valuations.
A policy π P Π ml is said to be K V -(d T V -)smooth-valued if sup sPSˇV π psqˇˇď K V and if for all a P A, sup sPSˇQπ ps, aqˇˇď K V . Observe that all policies π P Π ml are p |R ‹ π | {1´γq-smooth-valued, where |R ‹ π | " sup sPSˇR π psqˇˇ.
Lipschitzness. A policy π P Π ml is K V -Lipschitz-valued if for all s 1 , s 2 P S,ˇˇV π ps 1 q´V π ps 2 qˇˇď K V d S ps 1 , s 2 q and for all a P A,ˇˇQ π ps 1 , aq´Q π ps 2 , aqˇˇď K V d S ps 1 , s 2 q. The MDP M is xK π R , K π P y-Lipschitz if for all s 1 , s 2 P S,ˇR π ps 1 q´R π ps 2 qˇˇď K π R d S ps 1 , s 2 q, W d S`P π p¨| s 1 q, P π p¨| s 2 q˘ď K π P d S ps 1 , s 2 q.
Moreover, all latent policies π P Π ml with K π P ď 1 γ are p K π R{1´γK π P q-Lipschitz-valued (Gelada et al. 2019, Lem. 1). In the discrete setting, we replace the Wasserstein term by TV and we omit the d S terms.
This finally yields

B.4 Value Difference Bound
Lemma B.3. Let π P Π ml and assume M is discrete and xK π R , K π P y-Lipschitz. Let C, T Ď AP, ϕ P t C U T, ♦T u and P˘. Then, given the induced stationary distribution ξ π in M, E s,a"ξπˇQ π ps, aq´Q π pφpsq, aqˇˇď Proof. The first result for discounted return follows directly from (Gelada et al. 2019). The choice of K V is derived from the fact that K π P ď 1 γ since M is discrete: d T V is bounded by one as well as the diameter of S, equipped with the discrete metric. For the reachability cases, first observe that Q ϕ π ps,¨q " 1 when s P T Ď S and Q ϕ π ps,¨q " 1 when s P T Ď S. Moreover, if ϕ " C U T, Q ϕ π ps,¨q " 0 when s P C Ď S and Q ϕ π ps,¨q " 0 when s P C Ď S. By Assumption 3.3, s P C Y T implies that the value difference is zero since the labels of each state and those of its latent abstraction are the same. Therefore, let I be 1 C X T if ϕ " C U T and 1 T if ϕ " ♦T, E s,a"ξπˇQ ϕ π ps, aq´Q ϕ π pφpsq, aqˇ" E s,a"ξπ « IpsqˇˇˇˇE

B.6 Transition-reward function
In general RL environments, it is quite common for rewards to be generated from a reward function defined over transitions rather than over state-action pairs., i.e., R : SˆAˆS. In that case, we have Rps, aq " Es 1 "Pp¨|s,aq Rps, a, s 1 q. Taking transitionrewards into account, we set up an upper bound on L ξπ R that can be efficiently estimated from samples (using Lemma 3.4). Lemma B.5. Let ξ π be a stationary distribution of M under π P Π ml , then L ξπ R ď E s,a,s 1 "ξπˇR`s , a, s 1˘´R pφpsq, aqˇˇ" 9 L ξπ R .
Proof. The upper bound 9 L ξπ R on the local reward loss is given as follows:

B.7 Local Transition Loss Upper Bound
Recall that to deal with general MDPs, we study an upper bound on L ξπ P (namely, 9 L ξπ P ) that can be efficiently approximated from samples. We derive it in the following Lemma. Lemma B.6. Let ξ π be a stationary distribution of M under π P Π ml , then L ξπ P ď E s,a,s 1 "ξπ W d S`φ`¨| s 1˘, Pp¨| φpsq, aq˘" 9 L ξπ P where we write ξ π ps, a, s 1 q " ξ π ps, aq¨Pps 1 | s, aq and φps | sq " 1 " pφpsq, sq.
Proof. The upper bound 9 L ξπ P on the local transition loss is obtained as follows. Proof. Since M is ergodic, the agent acts in the BSCC induced by π and each tuple xs t , a t , r t , s t`1 y can be considered to be drawn independently from the stationary distribution ξ π of the BSCC in the limit. The result follows then from the Hoeffding's bound (Hoeffding 1963): • consider the sequence of independent random variables X R 0:T´1 such that X R t " |r t´R pφps t q, a t q| P r0, 1s, thenL ξπ R " 1 T ř T´1 t"0 X R t and Pr´ˇˇL ‹ R´L ξπ Rˇě ε¯ď δ R , with δ R " 2e´2 T ε 2 and L ‹ R " L ξπ R if R is defined over state-action pairs and L ‹ R " 9 L ξπ R if it is defined over transitions. • consider the sequence of random variables X P 0:T´1 such that X P t " 1´Ppφps t`1 q | φps t q, a t q P r0, 1s, The latent MDP M being discrete, we have 9 L ξπ P " E s,a,s 1 "ξπ 1´P`φ`s 1˘| φpsq, a˘‰ thenL ξπ P " 1 T ř T´1 t"0 X P t and Pr´ˇˇ9 L ξπ P´L ξπ Pˇě ε¯ď δ P , with δ P " 2e´2 T ε 2 . Therefore, We thus consider δ ě 4e´2 T ε 2 , i.e., log`δ 4˘ě´2 T ε 2 and T ě´l , and E s,a"ξπˇQ ϕ π ps, aq´Q ϕ π pφpsq, aqˇˇď with probability at least 1´δ. To offer an error of at most ε, we need which yields the result.

C Details on Variational MDPs
ELBO. In this work, we focus on the case where D KL is used as divergence measure. The goal is thus to optimize min θ D KL pM π P θ q, i.e., maximizing the marginal log-likelihood of traces of the original MDP Eτ"M π rlog P θ pτ qs. Recall that whereτ " xs 0:T , a 0:T´1 , r 0:T´1 , l 0:T y, z P Z such that Z " S if A " A and Z " SˆA otherwise, P π θ ps 0:T q " ś T´1 t"0 P π θ ps t`1 | s t q, and P π θ ps 0:T , a 0:T´1 q " ś T´1 t"0 π θ pa t | s t q¨P θ ps t`1 | s t , a t q. The dependency ofτ on Z in Eq. 6 is made explicit by the law of total probability.
Proof. First, observe that f ps t , a t , r t , s t`1 q.

C.1 Latent Distributions
Discrete latent state space. W.l.o.g., we assume latent states to have a binary representation, we further see each bit as a Bernoulli random variable. Let α P R and λ P s0, 1s, s λ P r0, 1s has a relaxed Bernoulli distribution s λ " RelaxedBernoullipα, λq with logit α and temperature parameter λ iff (a) if l is a logistic sample with location parameter αλ´1 and scale parameter λ´1, i.e., l " Logisticpαλ´1, λ´1q, then s λ " σplq, σ being the sigmoid function, (b) lim λÑ0 RelaxedBernoullipα, λq " Bernoullipαq, meaning Pplim λÑ0 s λ " 1q " σpαq, and (c) let p α,λ be the density of RelaxedBernoullipα, λq, then p α,λ ps λ q is log-convex in s λ . Since the logistic distribution belongs to the location-scale family, it is fully reparameterizable, i.e., l " Logisticpαλ´1, λ´1q " x " Logisticp0, 1q and l " α`x λ . In practice, we train φ ι and P θ to infer and generate log 2 |S| logits, and we anneal λ to 0 during training while using Bernoulli distributions for evaluation. As suggested by Maddison, Mnih, and Teh (2017), we use two different temperature parameters and we fix their initial values to λ φι " 2 {3 for the encoder distribution and λ P θ " 1 {2 for the latent transition function.
Discrete latent action space. We learn discrete latent distributions for Q A ι and π θ by learning the logits of continuous relaxation of discrete distributions, through the Gumbel softmax trick. Let n " |A|, and A " t a 1 , . . . , a n u. Drawing a i " Categoricalpαq from a discrete (categorical) distribution in P`A˘with logits parameter α P R n is equivalent to applying an arg max operator to the sum of Gumble samples and these logits, i.e., i " arg max k rα k´l ogp´log k qs, where P r0, 1s n is a uniform noise. The sampling operator being fully reparameterizable here, arg max is however not derivable, preventing the operation to be optimized through gradient descent. The proposed solution is to replace the arg max operator by a softmax function with temperature parameter λ P s0, pn´1q´1s.
Concretely, x P r0, 1s n has a relaxed discrete distribution x " RelaxedCategoricalpα, λq with logits α, temperature parameter λ iff (a) let P r0, 1s n be a uniform noise and G k "´logp´log k q, then x k " exp ppαk`Gkqλ´1q ř n i"1 exp ppαi`Giqλ´1q for each k P rns, (b) lim λÑ0 RelaxedCategoricalpα, λq " Categoricalpαq, meaning Pplim λÑ0 x k " 1q " exp α k ř n i"1 exp αi , and (c) let p α,λ be the density of RelaxedCategoricalpα, λq, then p α,λ pxq is log-convex in x. In practice, we train Q A ι and π θ to infer and generate logits parameter α, and we anneal λ from pn´1q´1 to 0 during training while using discrete distributions for evaluation. Again, as prescribed by Maddison, Mnih, and Teh (2017), we use two different temperature parameters, one for Q A ι and another one for π θ .

C.2 Prioritized Experience Replay
Modern RL techniques allow to learn the parameters θ of a policy, model, or value function, by processing encountered experiences η via a loss function Lpη, θq. If these are processed sequentially and discarded after all parameter updates, the latter are strongly correlated and rare events are quickly forgotten. Prioritized experience replay buffers (Schaul et al. 2016) are finite data structures D that allow to overcome both issues by (i) storing experiences η with some assigned priority p η , and (ii) producing η " D according to P pηq " p ς η { ř η 1 PD p ς η 1 , ς P r0, 1s, which allows minimizing Eη"D rw η¨L pη, θqs for importance sampling weight w η ą 0. The replay buffer D is uniform if ς " 0. Otherwise, when ς ą 0, a bias is introduced because priorities change the real distribution of experiences and consequently the estimates of the expectation of L. This bias is alleviated via w η " p|D|¨P pηqq´ω for each η P D, ω P r0, 1s, where ω " 1 fully corrects it.
We use prioritized replay buffers to store transitions and sample them when optimizing our loss function, i.e., ELBO. This allows to alleviate the posterior collapse problem, as shown in Figure 3. We introduced two priority functions, yielding a bucket-based and a loss-based prioritized replay buffer. In the following, we elaborate how precisely we assign priorities via the loss-based priority function. Loss-based priorities. Draw xs, a, r, s 1 y " ξ π , we aim at assigning p xs,a,r,s 1 y to its individual transition loss, i.e., L ELBO : for s, s 1 " φ ι p¨| s, s 1 q, a " Q A ι p¨| s, aq if A ‰ A and a " a else, L ELBO`@ s, a, r, s 1 D , @ s, a, s 1 D"´l og P G θ`s 1 | s 1˘´l og P R θ pr | s, aq´log π θ pa | sq`log φ ι`s 1 | s 11 "`A , A˘¨log P π θ`s 1 | s˘´1 ‰`A , A˘¨`log ψ θ pa | s, aq`log P θ`s 1 | s, a˘´log Q A ι pa | s, aq˘.
Notice that we cannot assign directly p xs,a,r,s 1 y to L ELBO since priorities require to be strictly positive to compute the replay buffer distribution. A solution is to pass L ELBO pxs, a, r, s 1 y , xs, a, s 1 yq to the logistic function, i.e., p xs,a,r,s 1 y " x ‹¨σ`k¨`L ELBO`@ s, a, r, s 1 D , @ s, a, s 1 D˘´x 0˘w here x ‹ ą 0 is the scale, k ą 0 is the growth rate, and x 0 is the location of the logistic. In practice, we set x ‹ to the desired maximum priority and we maintain upper (L max ) and lower bounds (L min ) on the loss during learning. Along training steps, we tune these values according to the current loss and we set x 0 " L max´Lmin 2 and k " x ‹ L max´Lmin .

D Details on RL Policy Distillation
Concretely, the goal of π θ is to mimic the behavior of π through the latent spaces. However, executing π θ in M via φ ι and ψ ι,θ as described in Sect. 3 to evaluate our latent space model (e.g., Lem. 3.4) may result in performance loss compared to that offered under π. This is due to the abstraction learned via the ELBO: we do not expect a zero-distortion by encoding the state-action space through the VAE (and M π θ to behave exactly like M π ), but to minimize it. Due to this remaining distortion, decisions taken according to π θ may result in unreachable states under π, causing π θ to behave poorly in that particular states since the learning process only allows to learn from transitions produced via π. We propose two approaches that can be used alone or together to alleviate this problem. Globally-robust RL policy. In general, policies π learned through deep-RL provide local performance, meaning they provide good performance in states s P S that are likely to be reached from s I by executing π. It is, however, often suitable to distil and verify globally-robust RL policies, i.e., policies trained to provide good global performance, in a wider range of states, in particular those that are likely to be reached through the distillation. Training globally robust RL policies can be achieved by allowing the agent to learn (or pursuing its training) in a modified version of M, where s I is picked at random from S. Training π this way will usually take longer since this may require more exploration. The degree of randomness can be decided according to the exploration/exploitation trade-off and the additional training time allowed. Moreover, we argue that this additional time is acceptable in our context since enabling and applying model checking are time-consuming tasks by nature. ε-mimic. Assuming π is sufficiently robust for the distillation, we allow the encoder to process states reachable under both π θ and π by learning from a mixtureπ ε θ of the two, defined asπ ε θ pa | sq " πpa | sq¨p1´εq`π θ pa | sq¨ε for all s P S, a P A, given ε P r0, 1s. Given the robustness assumption, states entered according to a decision a " π θ p¨| sq (produced with probability ε) should be efficiently processed by π. In practice, we start with a high value of ε (e.g., ε " 1 {2) to encourage exploration of states possibly reachable under π θ and we anneal it to 0 during training. Let θ1, θ2 P r0, 2πs be the angles of the two rotational joints, • s1 " cospθ1q • s2 " sinpθ1q • s3 " cospθ2q • s4 " sinpθ2q • s5: angular velocity 1 • s6: angular velocity 2 x1ąp´s1´s3¨s1`s4¨s2, 1q, 1ěps1, 0q, 1ěps2, 0q 1ěps3, 0q 1ěps4, 0q 1ěps5, 0q 1ěps6, 0qy • p1: RL agent target • p2: θ1 P r0, π {2s Y r 3π {2, 2πs • p3: θ1 P r0, πs • p4: θ2 P r0, π {2s Y r 3π {2, 2πs • p5: θ2 P r0, πs • p6: positive angular velocity (1) • p7: positive angular velocity (2) Pendulum R 3 Let θ P r0, 2πs be the joint angle • s1 " cospθq • s2 " sinpθq • s3: angular velocity x1ěps1, cosp π {3qq, 1ěps1, 0q, 1ěps2, 0q 1ěps3, 0qy  Table 2: Labeling functions used in our experiments. We provide a short description of the state space and the meaning of the atomic propositions forming the labels. Recall that labels and AP are respectively binary and one-hot encoded. Let n " |AP| with AP " t p 1 , . . . , p n u, we write p i for the one-hot vector of size n with the i th entry set to 1, with i P rnsz t 0 u. The LunarLander labeling function is based on the heuristic function of the OpenAI's implementation (Brockman et al. 2016).
Gerstner, and Brea (2018) and we reconstruct the latent dynamics via frequency estimation (e.g., Bazille et al. 2020) to check the abstraction quality in a PAC manner (cf. Fig. 2): for each environment, we evaluated the latent space model quality based on frequency-estimated latent dynamics from transitions collected by (repeatedly) simulating the environment for 2¨10 5 steps.

E.4 Labeling functions
The labeling functions we used for each environment are detailed in Table 2.

F Reproducibility of the Results
Seeding. The code we provide allows for setting random seeds to run the experiments, making the overall shape of the Figures we present in this paper reproducible. However, the exact values obtained in each particular instance are not: our prioritized experience replays (and those of TF-AGENTS) rely on REVERB replay buffers, handling randomness via the ABSEIL C++ library 3 which does not allow for manual seeding. All other libraries using randomness are seeded properly.