Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees
Authors: Florent Delgrange
Ann Nowe
Guillermo Perez
Venue: ICLR 2023 poster (2023)
Keywords: Reinforcement learning
Formal Verification
Representation Learning

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

Published as a conference paper at ICLR 2023
WASSERSTEIN AUTO-ENCODER MDPs
FORMAL VERIFICATION OF EFFICIENTLY DISTILLED RL POLICIES WITH MANY-SIDED GUARANTEES
Florent Delgrange
Ann Nowé
Guillermo A. Pérez
AI Lab, Vrije Universiteit Brussel (VUB) University of Antwerp
AI Lab, VUB
University of Antwerp
Flanders Make
florent.o dew.anndeeggagggaa...... vub. . ac be
ABSTRACT
Although deep reinforcement learning (DRL) has many success stories, the large- scale deployment policies learned through these advanced techniques safety- critical scenarios is hindered by their lack of formal guarantees. Variational Markov Decision Processes (VAE-MDPs) are discrete latent space models that provide 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 simple time-to-failure verification algorithm the latent space. The fact that our approach enables such simple verification techniques highlights its applicability.
INTRODUCTION
Reinforcement learning (RL) emerging as a solution of choice to address challenging real-word scenarios such as epidemic mitigation and prevention strategies (Libin et al.) 2020). multi-energy management (Ceusters et al.) [02T]. or effective canal control (Ren et al.) 2021). RL enables learning high performance controllers by introducing general nonlinear function approximators (such neural networks) to scale with high-dimensional and continuous state-action spaces. This introduction, termed deep-RL, causes the loss of the conventional convergence guarantees of RL (Tsitsiklis 1994) as well as those obtained in some continuous settings (Nowe] 1994), and hinders their wide roll-out in critical settings. This work enables the formal verification of any such policies, learned by agents interacting with unknown, continuous environments modeled as Markov decision processes (MDPs). Specifically, we learn a discrete representation of the state-action space of the MDP, which yield both a (smaller, explicit) latent space model and a distilled version of the RL policy, that are tractable for model checking (Baier & Katoen, 2008). The latter are supported by bisimulation guarantees: intuitively, the agent behaves similarly in the original and latent models. The strength of our approach not simply that we verify that the RL agent meets a predefined set of specifications, but rather provide an abstract model on which the user can reason and check any desired agent property.
Variational MDPs (VAE-MDPs, Delgrange et (022) offer a valuable framework for doing so. The distillation is provided with PAC-verifable bisimulation bounds guaranteeing that the agent behaves similarly (i) in the original and latent model (abstraction quality); (ii) from all original states embedded to the same discrete state (representation quality). Whilst the bounds offer a confidence metric that enables the verification of performance and safety properties, VAE-MDPs suffer from several learning flaws. First, training a VAE-MDP relies variational proxies to the bisimulation bounds, meaning there is no learning guarantee on the quality the latent model via its optimization. Second, variational autoencoders (VAEs) (Kingma & Welling 2014 Hoffman et al. 2013) are known

---

Published as a conference paper at ICLR 2023
to suffer from posterior collapse (e.g., Alemi al. (0118) resulting in a deterministic mapping to unique latent state in VAE-MDPs. Most of the training process focuses on handling this phenomenon and setting up the stage for the concrete distillation and abstraction, finally taking place in a second training phase. This requires extra regularizers, setting up annealing schemes and learning phases, and defining prioritized replay buffers to store transitions. Distillation through VAE-MDPs is thus meticulous task, requiring large step budget and tuning many hyperparameters.
Building upon Wasserstein autoencoders (Tolstikhin et al. 2018) instead of VAES, we introduce Wasserstein auto-encoded MDPs (WAE-MDPs), which overcome those limitations. Our WAE relies on the optimal transport (OT) from trace distributions resulting from the execution of the RL policy in the real environment to that reconstrueted from the latent model operating under the distilled policy. In contrast to VAFs which rely on variational proxies, we derive a novel objective that directly incorporate the bisimulation bounds. Furthermore, while VAEs learn stochastic mappings to the latent space which need be determinized or even entirely reconstructed from data at the deployment time to obtain the guarantees, our WAE has no such requirements, and learn all the necessary components to obtain the guarantees during learning, and does not require such post- post-procsssing operations.
Those theoretical claims are reflccedd in our experiments: policies are distilled up to 10 times faster through WAE- than VAE-MDPs and provide better abstraction quality and performance in general, without the need for setting up annealing schemes and training phases, nor prioritized buffer and extra regularizer. Our distilled policies are able to recover (and sometimes even outperform) the original policy performance, highlighting the representation quality offered by our new framework: the distillation is able to remove some non-robustness of the input RL policy. Finally, we formally verified time-to-failure properties (e.g.. Pnueli 1977) to emphasize the applicability of our approach.
Other Related Work. Complementary works approach safe RL via formal methods (Junges et al. 2016 Alshiekh et al. 2018 Jansen al., 2020 Simão et al. 2021), aimed at formally ensuring safety during RL, all of which require providing an abstract model of the safety aspects of the environment. They also include the work of Alamdari et al. (2020), applying synthesis and model checking on policies distilled from RL, without quality guarantees. Other frameworks share our goal of verifying Reep-RL policies (Bacci & Parker 2020 Carr al. 2020) but rely on a known environment model, among other assumptions (e.g.. deterministic or discrete environment). Finally, DeepSynth (Hasanbeig et al., (02T) allows learning formal model from execution traces, with the different purpose of guiding the agent towards sparse and non- non-Markovian rewards.
On the latent space training side, WWAEs (Zhag et al. 2019) reuse OT as latent regularizer discrepancy (in Gaussian closed form), whereas we derive two regularizers involving OT. These two are, in contrast, optimized via the dual formulation of Wasserstein, as Wassertein -Gass (Arjovsky] et al. 2017). Similarly to VQ-VAEs (van den Oord et al., (017) and Latent Bernoulli AEs (Fajtl et al., 2020). , our latent space model learns discrete spaces via deterministic encoders, but relies on a smooth approximation instead of using the straight-through gradient estimator.
Works on representation learning for RL (Gelada et al. 2019 Castro et al. 2021 Zhang et al. 2021 I
Zang et al. 2022) consider bisimulation metrics optimize the representation quality, and aim at learning (continuous) representations which capture bisimulation, so that two states close in the representation are guaranteed to provide close and relevant information to optimize the performance of the controller. In particular, as in our work, DeepMDPs (Gelada et al. 2019) are learned by optimizing local losses, by assuming deterministic MDP and without verifiable confidence measurement.
2 BACKGROUND
In the following, we write (XX) for the set of measures over (complete, separable metric space) X.
Markoy decision processes (MDPs) are tuples M = (S, A.P.A.P.R R. f, AP, where S is a set of states; A, set of actions; P: S x A - ASS), probability transition function that maps the current state and action to a distribution over the next states; R: S A + IR, a reward function; L: S - AAP a labeling function over set of atomic propositions AP; and 8, E S, the initial state. If [A| = M is a fully stochastic process called a Markov chain (MC). We write M, for the MDP obtained when replacing the initial state of M by S. An agent interacting in M produces trajectories, i.c., sequences of states and actions T = (Sper: aur-i) where so = Sy and 8411 P(: 8t, 0t) for T. The set of infinite trajectories of M is Traj. We assume AP and
2

---

Published as a conference paper at ICLR 2023
M.
Ea
M
M
M
Wer
5
π
A
a+
a
a
P
P
P
P
P
T2
[e R)
(K R)
g
(a) Execution of the latent policy in the original and latent MDPs, and local losses.
(b) Parallel execution of the original RL policy it in the original and latent MDPs, local losses, and steady-state regularizer.
Latent flows: arrows represent (stochastic) mappings, the original (resp. latent) state-action space is spread along the blue (resp. green) area, and distances are depicted red. Distilling æ into I via flow (b) by minimizing We. allows closing the gap between flows (a) and (b).
Figure 1:
labels being respectively one-hot and binary encoded, Given T AP, we write & }= T is labeled with T, i.e., {(s) o T * Ø. and }= -T for H T. We refer to MDPs with continuous state or action spaces as continuous MDPs. In that case, we assume S and A are complete separable metric spaces equipped with Borel a o-algebra, and 8-1 (T) is Morel-measurable for any T S AP.
Policies and stationary distributions. A (memoryless) policy т: S - ((A) prescribes which action to choose at each step of the interaction. The set of memoryless policies of M II. The MDP M and π e II induce an MC Ma with unique probability measure Wrt on the Borel o-algebra over measurable subsets & S Traj (Puterman 1994). We drop the superscript when the context is clear. Define (s') 8) P. ((Sore, Gpee 8t = 8') e'}) as the distribution giving the probability of being in each state of M,, after steps. B S is botlom strongly connected component (BSCC) of M₂ if (i) B is maximal subset satisfying (x') 8) > 0 for any 8, E B and some / > 0, and (ii) E-- P(B| s, s,a) = 1 for all E S. The unique stationary distribution of B is & A(B). We write s, ~ & for sampling from & then a from т. An MDP M is ergodic if for all x E 11, the
state space of M, consists of unique aperiodic BSCC with & lim 5/ 8) for all s
Value objectives. Given E II, the value of state S is the expected value of random variable obtained by running æ from s. For a discount factor e [o. 11. we consider the following objectives. (i) Discounted return: we write V=(s) E.M |Er= > (st, at) for the expected discounted rewards accumulated along trajectories. The typical goal of an RL agent is to learn a policy x² that maximizes Vz-(s) (51) through interactions with the (unknown) MDP; (ii) Reachability; let C,T = AP. the (constrained) reachability event is CUT (aany Gpen Bi N, Vj < i, } CAs T
Traj. We write (s) E,M. 1 (s0:00 for the discounted probability of satisfying CUT, where t" is the length of the shortest trajectory prefix that allows satisfying s. Intuitively, this denotes the discounted return of remaining a region of the MDP where states are labeled with C. until visiting for the firsi time goal state labeled with T, and the return is the binary reward signal capturing this event. Safety w.r.t. failure states can be expressed as the safety-constraing reachability to destination T through -CUT. Notice that (i( P.M. ((p) when ? =
Latent MDP. Given the original (continuous, possibly unknown) environment model M, latent space model is another (smaller, explicit) MDP M = AŠ,A, AP, 81) with state-action space linked the original one via state and action embedding functions: Φ: S - S and th: S x A A. We refer to (M, Φ, 1p) as a latent space model of M and M as its latent MDP. Our goal is to learn (M, &1 v/> by optimizing an equivalence criterion between the two models. We assume that dg is metric on S, and write II for the set of policies of M and V, for the values of running से in M. Remark 1 (Latent flow). The latent policy गे can be seen as a policy in M (cf. Fig. [a): states passed to मे are first embedded with @ the latent space, then the actions produced by म are executed via & in the original environment. Let E S, we write ā 7(- s) for क(. Φ(s)), then the reward and next state are respectively given by R(s, a) R(s, ()((()))))))))))))())) (()) a)) and s' P( s, a) P( s, (((()) a)).
3

---

Published as a conference paper at ICLR 2023
Local losses allow quantifying the distance between the original and latent reward/transition functions in the local setting, i.e., under a given state-action distribution Ç A(S x A):
Li
EE -,0==
[R(s, R(o(s), a)
L. = E D(OP( s, P( ((s), 9,00g
where p 8, is the distribution of drawing s' P(- 8, a) then embedding s' - o(s'), and D is a discrepancy measure. Fig Ia depicts the losses when states and actions are drawn from a stationary distribution &a resulting from running म E Il in M. In this work, we focus on the case where D is the Wasserstein distance Wag: given two distributions P, Q over a measurable set X equipped with a metric d, Wa is the solution of the optimal transport (OT) from P to Q, i.., the minimum cost of changing P into Q (Villani 2009): (P.QQQ) = xnfxeA(P.Q) Bz,y~> d(x; y), A(P, Q) being the set of all couplings of P and Q. The Kantorovich duality yields Wa (P, Q) = Se,, Ex-P (xx) Ez-Q f'y) where Fa is the set of I-Lipschizz functions. Local losses related to a well-established behavioral equivalence between transition systems, called bisimulation.
Bisimulation. A bisimulation B on M is a behavioral equivalence between states $1, E S so that, 81 82 iff (i) P(T) $1,0) = (T 82; a). (ii) {(s) {s22, and (iii) R(s1, a) R(82, for each action a E A and (Borel measurable) equivalence class T S/B. Properties of bisimulation include trajectory and value equivalence (Larsen & Skou 1989; Givan et al.) 2003). Requirements (ii) and (iii) can be respectively relaxed depending on whether we focus only on behaviors formalized through AP or rewards. The relation can be extended to compare two MDPs (e.g., M and M) by considering the disjoint union of their state space. We denote the largest bisimulation relation by-
Characterized by a logical family of functional expressions derived from a logic C. bisimulation pseudometrics (Desharnais et al. 2004) generalize the notion of bisimilariy. More specifically, given a policy π E II, we consider a family_ F real-vluued functions parameterized by discount factor y and defining the semantics of L in M,, Such functional expressions allow to formalize discounted properties such as reachability, safety, as well as general w-regular specifications (Chatterjee et al.) 2010) and may include rewards as well (Ferns et al. 2014 The pseudometric dz is defined as the largest behavioral difference (2) = sup [f()) - f(s2],, and its kernel is bisimilarity: (=()1,82) 0 iff 81 82- particular, value functions are lipschite-continuuu w.r.t. dzz [VG - (=2) & K dd 82), where K is 1/(1>)) if rewards are included in F and otherwise. To ensure the upcoming bisimulation guarantees, make the following assumptions:
Assumption 2.1. MDP M is ergodic, Im(R) is a bounded space scaled [-1/2, 1/2], and the embedding function preserves the labels, i.e.. Φ(s)
= E(s) E(5) for s S, S.
Note that the ergodicity assumption is compliant with episodic RL and a wide range of continuous learning tasks (see Huang 2020 Delgrange et al. 2022 for detailed discussions this setting).
Bisimulation bounds (Deigrange et al., 2022). MM being set over continuous spaces with possibly unknown dynamics, evaluating can turn out to be particularly arduous, if not intractable. solution is to evaluate the original and latent model bisimilarity via local losses: fix म E ΠI. assume M is discrete, then given the induced stationary distribution & in M, let 81, 82 E S with (81) = (xx2):
L + 1-7
+ +
L
E (s, S~EW
(3)) &
1- ) (E) (s)) (s2)). (I)
The two inequalities guarantee respectively the quality of the abstraction and representation: when local losses are small, (i) states and their embedding are bisimilarly close in average, and (ii) all states sharing the same discrete representation are bisimilarly close. The local losses and related bounds can be efficiently PAC-estimated. Our goal is to learn latent model where the behaviors of the agent executing मे can be formally verified, and the bounds offer a confidence metric allowing to lift the guarantees obtained this way back to the original model M, when the latter operates under न. We show in the following how to learn a latent space model by optimizing the aforementioned bounds, and distill policies π II obtained via any RL technique to a latent policy æ IT.
3 WASSERSTEIN AUTO-ENCODED MDPs
Fix M, - (S,A, Ry ł, A, si> and (Mg. d., vo) as a latent space model of M parameter- ized by and 0. Our method relies learning behavioral model Eo of M from which we can
4

---

Published as a conference paper at ICLR 2023
retrieve the latent space model and distill æ. This can be achieved via the minimization of a suitable discrepancy between &o and Mr. VAE-MDPs optimize a lower bound on the likelihood of the dynamics of M using the Kullback-Leible divergence, yielding (i) Mo. (ii) distillation No of and (iii) Φ and thg- Local losses are not directly minimized, but rather variational proxies that do not offer theoretical guarantees during the learning process. To control the local losses minimization and exploit their theoretical guarantees, we present a novel autoencoder that incorporates them in its objective, derived from the OT. Proofs of the claims made in this Section are provided in Appendix
3.1 THE OBJECTIVE FUNCTION
Assume that S, A. and Im(R) are respectively equipped with metrics ds, dA. and dR, we define the raw transition distance metric d as the component-wish wiss sum of distances between states, actions, and rewards occurring of along transitions: di( 01, 01. ri,s))
(22 r2; 86))
ds(s1, 82) + d'{ 02) + dr (r1, r2) + ds(s',, $22) Given Assumption 2.1 we consider the OT be- tween local distributions, where traces are drawn from episodic RL processes or infinite interactions (we show in Appendix A that considering the OT between trace-based distributions in the limit amounts to reasoning about stationary distributions). Our goal is to minimize Wa (Ex, Eo) so that
Eo (,(),,,,,') = -S JSXAMS Po 5, de (5, 5',
(2)
where P& is a transition decoder and En, denotes the stationary distribution of the latent model Mg. As proved by [Bousquet et al. (2017), this model allows to derive simpler form of the OT: instead of finding the optimal coupling of (i) the stationary distribution & of MM, and (ii) the behavioral model Eo. in the primal definition of W WJ(Cr, En). it sufficient to find an encoder q whose marginal is given by Qs,,
s') a,)) following Theorem, yielding a particular case of Wasserstein-autoer foolsssiinndeooltiiiiiiinnn et al. (2018):
ss ( a, s' |ss s, s'), and identical to & This is summarized in the Tolstikhin
Theorem 3.1. Let Ea and Pg be respectively behavioral model and transition decoder as defined in 2 Go: S - S be state-wise decoder, and tho be an action embedding function. Assume Po is deterministic with Dirac function Go(s, a, 8') (9o(s), (() (0 (8, a), Ro(s, a),9o( Go(8)), then
W WT(Cn,E))
inf q: q:Q-Eng s,,,,,-- 5, -,8,8'(((10,,8)
E
E
((,,,,,,,,)) Go(5,a,
5 and o: 3 A- (A) as parameterized state and action encoders with (, ā, 5' (s.a,8) 8') = (s) (ā s, a) (o)--s' and define the marginal encoder as Qu = Ex d.(. s, a, s'). Training the model components can be achieved via the objective:
Henceforth, fix o. S 5
min
E 8,0.T,8' En s.a.s' da.s"=o.(1ssse (1s,,,)))
E
d((ss a. r.s,,),,.......) Go(s,a,3)) + 3 DD(QQ,, Ex.).
where D is an arbitrary discrepancy metric and 3 > 0 hyperparameter. Intuitively, the encoder e, can be learned by enforcing its marginal distribution Q. to match EN, through this discrepancy. Remark 2. If M has a discrete action space, then learning A is not necessary. We can set Ā = A using identity functions for the action encoder and decoder (details Appendix (A.)).
When π is executed in M, observe that its parallel execution in MM is enabled by the action encoder Φd: given original state S, π first prescribes the action x(- s), which is then embedded in the latent space via a ~ 6. di (s), a) (cf. Fig. Ib). This parallel execution, along with setting D to W; yield an upper bound on the latent regularization, compliant with the bisimulation bounds. A two-fold regularizer is obtained thereby, defining the foundations of our objective function:
Lemma 3.2. Define T(8, a. 8') = Ex,a-s. [1o. ( 8, a). Po(s' 8, a)) as the distribution of drawing state-action pairs from interacting with M, embedding them the latent spaces, and finally letting them transition to their successor state in Mg. Then, Wx Wa(Qu. Exo) & W2(Ero T) + L
We therefore define the W Waaseesseein-Wessesstteee auto-encoded MDP) objective as:
min 4.0
E s, inE 5.a.5 La.s"~o. (11,,,3)))
[ds(s, G3) + da, vo(3, + ds Go))))))))) + L + 3 3-W,, + +Lx),
5

---

Published as a conference paper at ICLR 2023
Algorithm 1: Wasserstein" Auto-Encoded MDP
Input: batch size N, max. step T, no. of regularizer updates m, penalty coefficient 0 > 0 for f 1 to do
for i = 1 to N do
Sample transition S;; Qe, ri, s from the original environment via Ex Embed the transition into the latent space by drawing 8 a2 5 from $.(-
86; 04s
Make the latent space model transition to the next latent state: 5 ~ Po( 5,, a.) Sample a latent transition from Exgi : zi ~ Eng: de ~ To(- z;), and 2/K ~ Pal Ris di)
W Ei=1 (5,, 0,, 5 (2,, ay z/) ( az 5t a4 5/) at, St, 5;) P - Σ Gp(pE (5,, a,, st) (2,, d,, <f) + Gp( - (s,, 5,, ā, x), 3) 51)
Update the Lipschitz networks parameters w by ascending I/N (BW - P) it't mod m 0 then
- Σ (s( Go():)) + da(a,, (5 a:)) + dR (rs. Ro(s; d;)) ds ds(sl, go( Update the latent space model parameters (1,0) by descending I/N (L + B3 W)
function GP((pw: 2: y) U(0,
Gradient penalty for you : IR" R and at y E R > random noise; straight lines between æ and y
(1-c)y return (L) - 112
where We. (T ER.) and L.ST are respectively called steady-stafe and transition regularizers. The former allows to quantify the distance between the stationary distributions respectively induced by x in .M and Fy in Mg. further enabling the distillation. The latter allows to learn the latent dynamics. Note that L and L set over & instead of ÊT. - are not sufficient to ensure the bisimulation bounds (Eq. D running x in Me depends on the parallel execution of in the original model, which does not permit its (conventional) verification. Breaking this dependency is enabled by learning the distillation Ty through We_. as shown in Fig. Ib minimizing Wc, allows to make & and ETp closer together, further bridging the gap of the discrepancy between π and ae. At any time, recovering the local losses along with the linked bisimulation bounds in the objective function of the WPAE-MDP1 is allowed by considering the latent policy resulting from this distillation:
Theorem 3.3. Assume that truces are generated by running latent policy क E Il in the original environment and let dR be the usual Euclidean distance, then the W-AE-MDP objective is
min 1,8 8.6'6
ds (olo 1TT + B 3-WW,, +L)).
Optimizing the regularizers is enabled by the dual form of the OT: we introduce two parameterized networks, & and pP. constrained to be 1-Lipschitz and trained to attain the supremum of the dual:
We. (w) max
(())(s))
o (2,
w
5,a-Er T~ 1 (2), - (s),))
z,a ='
L.S (w) = max
p (s,,a,
E
as,
9,0, 8, (.la,..))
Details derive this tractable form of L (w) are in Appendix A.5 The networks constrained via the gradient penalty approach of Gulrajani et al. (2017). leveraging that any differentiable function is 1-Lipschitz iff it has gradients with norm at most everywhere (we show in Appendix this is still valid for relaxations of discrete spaces). The final learning process is presented in Algorithm |
3.2 DISCRETE LATENT SPACES
To enable the verification of latent models supported by the bisimulation guarantees of Eq. I we focus on the special case of discrete latent space models. Our approach relies continuous relaxation of discrete random variables, regulated by some temperature parameter(s) X: discrete random variables retrieved as À 0, which amounts applying rounding operator. For training, the

---

Published as a conference paper at ICLR 2023
000
Jogastie sompling
Edurry sncoding
Gombel
Go
coftman
oneskat enooding
R
+4
96
innfonttoon
P
training
6/%
odftma
logistle
Figure 2: W'AE-MDP architecture. Distances are depicted by red dotted lines.
lemperaturr--contollled relaxations to differentiate the objective and let the gradient flow through the network. When we deploy the latent policy in the environment and formally check the latent model, the zero-temperature limit is used. An overview of the approach is depticted in Fig. 2
State encoder. We work with binary representation of the latent states. First, this induces compact networks, able to deal with large discrete space via tractable number of parameter variables. But most importantly, this ensures that Assumption 2.1 is satisfied: let n log, [S], we reserve |AP] bits in S and each time s is passed to é., n API bits are produced and concatenated with {(s). ensuring a perfect reconstruction of the labels and further bisimulation bounds. To produce Bernoulli variables, o, deterministically maps s latent code z, passed the Heaviside (z) 12200 We train d. by using the smooth approximation (x a (2*/A), satisfying H limn Hy.
Latent distributions. Besides the discontinuity of their latent image space, a major challenge of optimizing over discrete distributions is sampling, required to be differentiable operation. We circumvent this by using concrete distributions (Jang et al. [017] Maddison et al. 2017): the idea is to sample reparameterizable random variables from X- p-rrameterizeed distributions. and applying differentiable, nonlinear operator in downstream. We use the Gumbel softmax trick to sample from distributions over (one-hot encoded) latent actions (o,", Rg). For binary distributions (Pa. Sn,), each relaxed Bernoulli with logit is retrieved by drawing logistic random variable located in 9/12 and scaled to 1/2, then applying a sigmoid downstream. We emphasize that this trick alone (as used by Comeil et al. 2018 Delgrange et al. 2022) is not sufficient: it yields independent Bernoullis, being too restrictive in general, which prevents from learning sound transition dynamics (cf. Example
Example Let M be the discrete MC of Fig.
{goal} 1/2
one-hot, AP {goal: (1, 0), unsafe <0. 1)}. We as- sume that bits are used for the (binary) state space, with S (5o (0,0, 0) 31 (1,0,0) 52 {0, 1), 5g: 1,) (the two first bits are reserved for the labels). Considering each bit as being independent is not suffi- cient to learn P: the optimal estimation ( 50) is in that case represented by the independent Bernoulli vector b (1/2, 1/2, 1/4), giving the probability to go from &o to each bit independently. This yields a poor estimation of the actual transition function: P Pa-(5g 56) (1 (1- (1-b)))))) P.. (8 55) by
funsafet 83
1/4
3/4
82
{unsafe)
Figure 3: Markov Chain with four states; labels are drawn next to their state.
b2) (1))-(111))) (1-bl)))))) by) = P (&2l 5o) = bi) b2 102-(11bb))) (1-ba) = 3/16, P (5g 50) = (1 (1-b,) - -b) by = 1/16.
We consider instead relaxed multivariate Bernoulli distributions by decomposing P A
product of conditionals: ((s) I[ P(s: 81: 81: where 8; is the 7th entry (bit) of We learn
7

---

Published as a conference paper at ICLR 2023
CIn - CurPakk
un Mouiaa Oar
any Auseut
Larerkander
un Perdukmn
Loss
My (4l)
0.91
0.25
0.75
1/00
0.00
0.25
8.83
0.30
2.25
0.99
116
0.00
0.25
0.50
1e5
(a) ADD objective: reconstruction loss, transition and steady-state regularizers
enn CarFaee
MouraacCar
ov Acrobot
eww i LurerLander
ow Ferdiker
WAR--10 100
WAR-MD0-
vAg 440 PP)
Local loos
1.5
3.00
0.2
3.53
3.75
1.00
3.00
0.2
0.53
3.75
1.50
0.00
325
0.5
0.75
1.00
0.00
0.2
355
075
(h) PAC local losses approximation for an emror of at most 10 and probability confidence 0.955
CarPon
Azzther
Perdiuu
100 adee
200
WAG MDD
-130
-05
2.5
179
500
2.5
25
-000
0.5
10
cogun $1000
enges
(c) Episode return obtained when executing the distilled policy in the original MDP (averaged over 30 episodes)
Figure 4: For each environment, we trained five different instances of the models with different random seeds: the solid line is the median and the shaded interval the interquartile range.
such distributions by introducing a masked autoregressive flow (MAF, Papamakarios et al. 2017) for relaxed Bernoullis via the recursion: 5; (l:+to//x), where 4 Logistic(O. 1), 0, f: 1-1). and f is a MADE (Germain et al., 2015), feedforward network implementing the conditional output dependency on the inputs via a mask that only keeps the necessary connections to enforce the conditional property. We use this MAF to model Pe and the dynamics related to the labels in ET..
We fix the logits of the remaining n APP| bits to O to allow for a fairly distributed latent space.
4 EXPERIMENTS
We evaluate the quality of latent space models learned and policies distilled through W'AE-MDPS To do so, we first trained deep-RL policies (DQN, Mnih et 2013 on discrete, and SAC, Haarnoja et al. 2018 on continuous action spaces) for various OpenAl benchmarks (Brockman et al. 2016),
which we then distill via approach (Figure 4. We thus evaluate (a) the WDAE-MDP training metrics, (b) the abstraction and representation quality via PAC local losses upper bounds (Delgrange) al. 2022), and (c) the distilled policy performance when deployed in the original environment. The
confidence metrics and performance are compared with those of VAE-MDPs. Finally, we formally verify properties in the latent model. The exact setting reproduce our results in Appendix B
Learning metrics. The objective (Fig. 4a is a weighted sum of the reconstruction loss and the two Wasserstein regularizers. The choice of defines the optimization direction. In contrast VAEs (cf. Appendix C) WAES indeed naturally avoid posterior collapse (Tolstikhin et al.) 2018), indicating that the latent space is consistently distributed. Optimizing the objective (Fig. 4aj effectively allows minimizing the local losses (Fig. 4b) and recovering the performance of the original policy (Fig. Fc).
E.e and Eoo via Local losses. For V- and WAEs, we formally evaluate PAC upper bounds on L
the algorithm of Delgrange et al. 2202)) (Fig 4b). The lower the local losses, the closer M and Mg are in terms of behaviors induced by To (ct. Eq. D. In VAEs, the losses are evaluated on transition function P obtained via frequency estimation of the latent transition dynamics (Delgrange et al. 2022), by reconstructing the transition model a posteriori and collecting data to estimate the transition probabilities (e.g., (Bzzile et al. 2020) Corneil et al. 2018). We thus also report the metrics for P. Our bounds quickly converge to close values in general for P, and P. whereas for VAEs, the convergence is slow and unstable, with P offering better bounds. We emphasize that WAEs do not require this additional reconstruction step to obtain losses that can be leveraged to assess the
8

---

Published as a conference paper at ICLR 2023
Table Formal Verification of distilled policies.
Values are computed for > = 0.99 (lower better).
Enviroament
step (10P)
A
[S]
[4]
L (PAC)
1 (PAC)
[V+,1
(sr(
CariPole
1.2
[
512
0.00499665
0.399636
3.71213
0,0316655
MountainCar
2.32
32
[122
1024
0,0141763
0.362323
2.83714
0
Acrobot
4.3
12
112,33
8192
0.0347698
0.649478
2.22006
0.0021911
LunarLander
3.2
-1.1
16384
0.0207205
0.131357
0,0372883
0,0702039
Pendulum
3.7
R89
[-2
8192
0.0266745
0.539508
4.33006
0.03484492
quality of the model, in contrast to VAEs, where learning Pg was performed via overly restrictive distributions, leading to poor estimation in general (ef. Ex. DT. Finally, when the distilled policies offer comparable performance (Fig. 4c our bounds are either close to or better than those of VAEs.
Distillation. The bisimulation guarantees [ are only valid for Tg. the policy under which formal properties can be verified. It crucial that mo achieves performance close to æ,, the original one, when deployed in the RL environment. We evaluate the performance of Te via the undiscounted episode return R., obtained by running Ta in the original model M. We observe that R approaches faster the original performance R,, for W- than VAEs: WAEs converge in a few steps for all environments, whereas the full learning budget is sometimes necessary with VAEs. The success in recovering the original performance emphasizes the representation quality guarantees (Eq. I induced by WAEs: when local losses are minimized, all original states that are embedded to the same representation are bisimilarly close. Distilling the policy over the representation, albeit discrete and hence coarser, still achieves effective performance since o, keeps only what is important to preserve behaviors, and thus values. Furthermore, the distillation can remove some non-robustness obtained during RL: Re prescribes the same actions for bisimilarly close states, whereas this is not necessarily the case for æ.
Formal verification. To formally verify Me. we implemented value iteration (VI) engine, handling the neural network encoding of the latent space for discounted properties, which is one of the most popular algorithms for checking property probabilities in MDPs (e.g.. (Baier& & Katoen 2008 Hensel et al. 2221] K Kwiatkowska et al. (022). We verify time-to-failure properties so, often used to check the failure rate of a system (Pnueli 1977p by measuring whether the agent fails before the end of the episode. Although simple, such properties highlight the upplicability of our approach on reachability events, which are building blocks to verify MDPs (Baier & Katoen 2008; cf. Appendix [.7]. In particular, we checked whether the agent reaches an unsafe position or angle (CartPole, LunarI does not reach its goal position (MountainCar, Acrobot), and does not reach and stay in a safe region of the system (Pendulum). Results are in Table [ for each environment, we select the distilled policy which gives the best trade-off between performance (episode return) and abstraction quality (local losses). As extra confidence metric, we report the value difference ||V,, IV (57) V V=. (5r) obtained by executing ny in M and Mo (Vz, (-) averaged while V me (.) is formally computed).
CONCLUSION
We presented WAE-MDPs, a framework for learning formally verifiable distillations of RL policiis with bisimulation guarantees. The latter, along with the learned abstraction of the unknown continuous environment to a discrete model, enables the verification. Our method overcomes the limitations of VAE-MDPs and our results show that it outperforms the latter in terms of learning speed, model quality, and performance, in addition being supported by stronger learning guarantees. As mentioned by Delgrange et al. (2022), distillation failure reveals the lack of robustness of original RL policies. In particular, we found that distilling highly sensiisee RL policies (such as robotics sinulations, e.g., Todorov et al. 2012) is laborious, even though the result remains formally verifiable.
We demonstrated the feasibility of our approach through the verification of reachability objectives, which are building blocks for stochastic model-checking (Baier & Katoen 2008). Besides the scope of this work, the verification of general discounted w-regular properties is theoretically allowed in our model via the rechability to components of standard constructions based on automata products
(e.g.. Baier et al. 2016 Sickert et al. 2016). and discounted games algorithms (Chatterjee et al. 2010). results, supported by Thm. suggest that WAE-MDP can be used as
Beyond distillation, general latent space learner for RL, further opening possibilities to combine RL and formal methods online when formal model is priori known, and address this way safety in RL. with guarantees.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We referenced in the main text the Appendix parts presenting the proofs additional details of every claim, Assumption, Lemma, and Theorem occurring in the paper. In addition, Appendix B dedicated to the presentation of the setup, hyperparameters, and other extra details required for reproducing the results of Section 4 We provide the source code of the implementation of our approach Supplementary material and we also provide the models saved during training that we used for model checking (i.e., reproducing the results of Table Additionally, we present in a notebook (eva 1 uat on htm²) videos demonstrating how our distilled policies behave in each environment, and code snippets showing how we formally verified the policies.
ACKNOWLEDGMENTS
This research received funding from the Flemish Government (AI Research Program) and was supported by the DESCARTES iBOF project. G.A. Perez is also supported by the Belgian FWO "SAlLor" project (G030020N). We thank Raphael Avalos for his valuable feedback during the preparation of this manuseript.
REFERENCES
Parand Alizadeh Alamdari, Guy Avni, Thomas A. Henzinger, and Anna Lukina. Formal methods with a touch of magic. In 2020 Formal Methods in Computer Aided Design, FMCAD 2020, Haifa, Israel, September 21-24, 2020, pp. 138-147. IEEE, 2020. doi: 0ttp:/77222220000mm07----------- sin.. 978-9-98--- 8.5 4444444888
21. URL http::tp//.mm. doi. wwoooooooo. 10. . 34
Alexander A. Alemi, Ben Poole, lan Fischer, Joshua Dillon, Rif A. Saurous, and Kevin Murphy, Fixing broken ELBO. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässann Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 159-168. PMLR, 2018. URL /tttt::/pprcce
ngs. m. r press /ttess/v8// w80/alllll...o emi 18a..1h.. 18a.html
Mohammed Alshickh, Roderick Bloem, Rüdiger Ehlers, Bettina Könighofer, Scott Niekum, and Ufuk Topcu. Safe reinforcement learning via shielding. In Sheila A. Mcllraith and Kilian Q. Weinberger (eds.), Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 2669-2678. AAAI Press, 2018 URL https: www. aaai org /t// index php/ CAPP/AAAI/A
Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial net- works. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International
Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 214-223. PMLR, 2017. URL httt:://rooceet ngs. ml /ttpss//////rrroossyy////
html
Edoardo Bacci and David Parker. Probabilistic guarantees for safe deep reinforcement learning. In Nathalie Bertrand and Nils Jansen (eds.), Formal Modeling and Analysis of Timed Systems 18th International Conference, FORMATS 2020, Vienna, Austria, September 1-3, 2020, Proceedings, volume 12288 of LNCS, pp. 238-248. Springer, 2020. doi: 10. 10000077988-3-00-5576663 14. URL
htttp:://do..orr
Christel Baier and Joost-Pieter Katoen. Principles of model checking. MIT Press, 2008 ISBN 978-0-262-02649-9,
Christel Baier, Stefan Kiefer, Joachim Klein, Sascha Klüppelholz, David Müller, and James Worrell. Markov chains and unambiguous büchi automata. Swarat Chaudhuri and Azadeh Farzan (eds.), Computer Aided Verification 28th International Conference, CAV 2016, Toronto, ON, Canada, July 17-23, 2016, Proceedings, Part I, volume 9779 of Lecture Notes in Computer Science, pp. 23 23-42. Springer, 2016. doi: 10.10077988-33319444558.
/tt.ttps///o.../////1 org/10
1007/978-33339--1528--8
10

---

Published as a conference paper at ICLR 2023
Hugo Bazille, Blaise Genest, Cyrille Jégourel, and Jun Sun. Global PAC bounds for learning discrete time markov chains. In Shuvendu K. Lahiri and Chao Wang (eds.), Computer Aided Verification 32nd International Conference, CAV 2020, Los Angeles. CA, USA, July 21-24, 2020, Proceedings, Part 11, volume 12225 of Lecture Notes in Computer Science, pp. 304-326. Springer, 2020. doi: (0007/978-33000-5329999 17. URL poittps://do... htt.//o.ooggg/// 1..0/7.0007/ 978 3- 030 53291-8 17
O. Bousquet, S. Gelly, I. Tolstikhin, Carl-Johann Simon-Gabriel, and B. Schölkopf. From optimal transport to generative modeling: the vegan cookbook. arXiv: Machine Learning, 2017.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, (bss/660..01550 2016. URL htt.tttpp///xii..orrttt.
abs/ 6ess/666..co 01540
Steven Carr, Nils Jansen, and Ufuk Topeu. Verifiable ran-based policies for pomdps under temporal logic constraints. In Christian Bessiere (ed.). Proceedings of the Twenty-Ninth International
Joint Conference on Artificial Intelligence, IJCAI 2020, pp. bttt:tt:: http://ii.oggg./// 10. 2tt9://....4/3/....a.....o 2020/570
ijcai.org, 2020. doi:
00.24996ijeiii2220/5777
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. Mico: Improved representations via sampling-basse based state similarity for markov decision processes. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 30113 30113-30126, 2021. URL http:: httpps::/pprcceeedngggg/////// wtw.rocceednngg.nuu..pp.... nwwrips....... wccpppper/22///aaeeeepeepp Edo 20000880002ee1112249611110110000
Abst html
Glenn Ceusters, Román Cantú Rodríguez, Alberte Bouso Gareía, Rüdiger Franke, Geert Deconinck, Lieve Helsen, Ann Nowé, Maarten Messagie, and Luis Ramirez Camargo. Model-predictive control and reinforcement learning in multi-energy system case studies. Applied Energy, 303; 1 11634, 2021. ISSN 0306-2619. doi: wttp:///iicnnerrrrcrc//////////0 // www. www.ecceeeeeeeeeccccccceeeceeeee wttprrrrreecccrcrerrrrrreeeeeeee
1764. URL https
Krishnendu Chatterjee, Luca de Alfaro, Rupak Majumdar, and Vishwanath Raman. Algorithms for game metrics (full version). Log. Methods Comput. Sci., 6(3), 2010. URL http:///t///ai
exq/abs/ 0309b43266
Dane S. Corneil, Wulfram Gerstner, and Johanni Brea. Efficient modelbased deep reinforcement learning with variational state tabulation. In Jennifer G. Dy and Andreas Krause (eds.), ), Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stockhoimsmüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1057-1066. PMLR, 2018. URL htttt::/prrcceeed////t ngs. mlr press v80/ wtwp://eee////eeoeeoc////ooooo.c corneilrneiooooocooco html
Florent Delgrange, Ann Nowé, and Guillermo A. Pérez. Distillation of rl policies with formal guarantees via variational abstraction of markov decision processes. Proceedings of the AAAI Conference on Artificial Intelligence, 3666)6697--6555, Jun. 2022. doi: 10. 10.1669//aai. v36i6.20602.
URL poattpss//ojs aaai. aetaiiorg/indeee. phppAAAI http:///AAAaaatiiiiaceeee htt.aa/.wwg....n......aaaaaaaaaa /ttp///AAA//.riiieeeeeeee//e2222
Josée Desharnais, Vineet Gupta, Radha Jagadeesan, and Prakash Panangaden. Metrics for labelled
markov processes. Theor. Comput. Sci., 318(3):3333554,, 2004. doi: 10. tos.006/.....2003.0.....
URL httpttt://....oeeetee.soe/ http:///.....2111.11111/111/1///
2003. 09.03.99..111.
Jiri Fajtl, Vasileios Argyriou, Dorothy Monekosso, and Paolo Remagnino. Latent bernoulli autoen- coder. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13- 18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 2964- 2974. PMLR, 2020. URL htth:tt::// httpt::::/pocceeiiioin
mwr. wttpss//ll//aaaaaaaeeeppeeeeeeee
html
Norm Ferns, Doina Precup, and Sophia Knight. Bisimulation for markov decision processes through families of functional expressions. In Franck van Breugel, Elham Kashefi, Catuscia Palamidessi, and Jan Rutten (eds.), Horizons of the Mind. A Tribute to Prakash Panangaden Essays Dedicated to Prakash Panangaden on the Occasion of His 60th Birthday, volume 8464 of LNCS, pp. 312-342.
11

---

Published as a conference paper at ICLR 2023
Springer, 2014. doi: (0.10007/978-3-11--06338 17. URL 9788-3319-00880-0
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Reseurch, pp. 2170-2179. PMLR, 2019, URL http:::/pprcceediiniit mlr. w.w.press/vve v http://aaaaaaaaa.aeaaaaa html
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. MADE: masked autoencoder for distribution estimation. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning. ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 881-889, JMLR.org, 2015. URL
http::/pproced:/ ngs mlm: press / v37 ge reaanmainnnn 11m...nn.......mmmm html
Robert Givan, Thomas I.. Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artif. Intell., 147(1-2) 2003. doi: 100111/0004037000000000 00376-4. URL htttt:t::://co wwg.:/io..ooooooooooo 100 1t///11.11011 1016/50004 3702 (C2) 00376-4
Ishaan Gulrajani, Faruk Ahmed, Martín Arjovsky, Vincent Dumoulin, and Aaron C. Courville. Improved training of wasserstein gans. In Isabelle Guyon, Ulrike von Luxburg, Samy Ben- gio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5767-5777, 2017. URL httttsss/ppocceedinnggs neuriinnggn...uu..eeeee.e.uu.ree neurips wttppppe//0111111 20tr//ah/////h////// / F32c3blc6dcc55233eee htt.e/dd222ee222cccooof066663333
Abstract html
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Jennifer G. Dy and Andreas Krause (cds.), Proceedings of the 35th International Conference on Machine Learning, 2CML2018 Stockholmsmmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1856-1865 1865. PMLR, 2018. URL http::tttp//oreeeeeeee mlr. w.r.ppess//v00 v80 hww./////aaaaaaaaaaaaaaa 8b. html
Mohammadhosein Hasanbeig, Natasha Yogananda Jeppu, Alessandro Abate, Tom Melham, and Daniel Kroening. Deepsynth: Automata synthesis for automatic task segmentation in deep reinforcement learning. In Thirty-Fifth AAAI Conference
Artificial Intelligence, AAAI 2021, Thirty-Third Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, aaa. www.aaaai....................... /rg
February 2-9, 2021, pp. 7647 AAAI Press, 2021. URL o
index. . php/ AAA.AAAAAAAAA. http///AA/aattcee/eeeeeeeeeeeeee 11633
Christian Hensel, Sehastian Junges, Joost-Pieter Katoen, Tim Quatmann, and Matthias Volk. The probabilistic model checker storm. International Journal on Software Tools for Technology doi:
Transfer, 2021. ISSN 1433-2787,
0rg/1001007/s10000 00000000000200000000000000000000 021-00633- 00633-2
Matthew D. Hoffman, David M. Blei, Chong Wang, and John W. Paisley. Stochastic variational infer- ence. J. Mach. Learn. Res., 14(1):1303-1347, 2013. RtLhhtt:// //l. acm. www.a//..o....aaaaaaaeeoe org/ciiaiionn
c cfm?id=2502622
Bojun Huang. Steady state analysis of episodic reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- -ien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Confer- ence on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12,
2020, virtual, 2020. URL httll/ttt////occciiiiiiie neurips wtt.://tt...ppppeee//eaeaps cc/papptr/222//aaall Abst html
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In 5th International Conference Learning Representations. ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL http:::://pppeeeeeeesss
12

---

Published as a conference paper at ICLR 2023
Nils Jansen, Bettina Könighofer, Sebastian Junges, Alex Serban, and Roderick Bloem. Safe Re- inforcement Learning Using Probabilistic Shields (Invited Paper). In Igor Konnov and Laura Kovács (eds.), 31st International Conference on Concurrency Theory (CONCUR 2020), vol- ume 171 of Leibniz International Proceedings in Informatics (IJPles), pp. 3:1-3:16, Dagstuhl, für Informatik. ISBN 978-3-95977-160- dwgpps.dagtuuuie www.aaasusuuu.oo...uuoooe
Germany, 2020. Schloss Bogstuul-LLeibnizZZaee
3. doi: http:///LLPP..CONCCRRR vrllttexte/2020/222
URL
Sebastian Junges, Nils Jansen, Christian Dehnert, Ufuk Topeu, and Joost-Pieted Katoen. Safety- constrained reinforcement learning for mdps. In Marsha Chechik and Jean-François Raskin (eds.). Tools and Algorithms for the Construction and Analysis of Systems 22nd International Conference, TACAS 2016, Eindhoven, The Netherlands, April 2-8, 2016, Proceedings, volume 9636 of LNCS, pp. 130-146 Springer, 2016. doi: 10.1107/978-3 9-11999788--662--477--99999 URLhttpps //doi www.i.ogg/...... 1007/ 000000777777-0 9988/9783333 1778---662249-7---
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http:
13112 6114
Marta Kwiatkowska, Gethin Norman, and David Parker. Probabilistic model checking and autonomy. Annual Review of Control, Robotics, and Autonomous Systems, 5(1):385 2022. doi: 10.1 11466mmmree--mmtoo00400000000000
URL.https///oo..rr/////
Kim Guldstrand Larsen and Arne Skou, Bisimulation through probabilistic testing. In Conference Record of the Sixteenth Annual ACM Symposium Principles of Programming Languages, Austin, Texas, USA, Jamuary 11-13, 1989, pp. 344-352, ACM Press, 1989. doi: 10.1145/75271.75007
URL https://doi.orgg: 10. 11..0....1117... 75307
Pieter K. Libin, Arno Moonens, Timothy Verstraeten, Fabian Perez-Sanjines, Niel Hens, Philippe Lemey, and Ann Nowé. Deep reinforcement learning for large-scale epidemic control. In Yux- iao Dong. Georgiana Ifrim, Dunja Mladenic, Craig Saunders, and Sofie Van Hoecke (eds.), Machine Learning and Knowledge Discovery in Databuses. Applied Data Science and Demo Truck European Conference, ECML PKDD 2020, Ghent, Belgium, September 14-18, 2020, Proceedings, Part V. volume 12461 of Lecture Notes in Computer Science, pp. 155-170. Springer, 2020. doi: 0-000797883-330---767---44 998-
hRtpps::: http://ii.orgg/..... org /ttp:://do..ooog10....0000
Michael L. Littman, Ufuk Topeu, Jie Fu, Charles Lee Isbell Jr., Min Wen, and James MacGlashan. Environmeent-indeeeedden task specifications via GLTL. CoRR, abs/1 (oss/704404341,, 2017. URL
http://arxiv..ogg/
04341
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relax- ation of discrete random variables. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.net, 2017. URL BRLLtttts////teiiie
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep re- inforcement learning. Nat., 518(7540):529-573 2015.
doi: 10.1038/nature142366 URL
htttps://doi.org/00 1002
Ann Nowe. Synthesis of "safe" fuzzy controllers based on reinforcemení learning. PhD thesis, Vrije Universiteit Brussel, 1994.
George Papamakarios, lain Murray, and Theo Pavlakou. Masked autoregressive flow for density estimation. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Gamett (eds.), Advances in Neural Information Processing Sys- tems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017,
13

---

Published as a conference paper at ICLR 2023
Long Beach, CA. USA, pp. 2338 2347,
http:///ww.....ooooo/////o/oooo. cc/
.html
Amir Pnueli. The temporal logic of programs. In 18th Annual Symposium on Foundations of
Computer Science, Providence, Rhode Island, USA, 31 October 1 November 1977, pp. 46-57. IEEE Computer Society, 1977. doi: 10.1 109/SFCS.1977.32. URL /tt.pstss:///...rrr
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming.
Wiley Series in Probability and Statistics. Wiley, 1994, ISBN 978-0-47161977-2. doi: 10.1002/
9780470316887. URL http:/tttt::.....tto org 000.1000000977071110066000
Tao Ren, Jianwei Niu, Jiahe Cui, Zhenchao Ouyang, and Xuefeng Liu. An application of multi-
objective reinforcement learning for efficient model-free control of canals deployed with iot net- works. Journal of Network and Computer Applications, 182:103049, 2021. ISSN 1084-8045. doi:
https://doi.or//10 10. 1.100.011111 2022 103049. URL https: www. scissciiiene encedi rect com/
Sg1
Salomon Sickert, Javier Esparza, Stefan Jaax, and Jan Kretínský. Limit-deterministic büchi au- tomata for linear temporal logic. In Swarat Chaudhuri and Azadeh Farzan (eds.), Computer
Aided Verification 28th International Conference, CAV 2016, Toronto, ON, Canada, July 17-23, 2016, Proceedings, Purt 11, volume 9780 of Lecture Notes in Computer Science, pp. 312-332. Springer, 2016. doi: 10.1007/978- 3-1-778---31--1550000 17. URL https:://oolii http:////oo.ogg/...000 org. 1007/ 11--3-319-1140---1 17.
Thiago D. Simão, Nils Jansen, and Matthijs T. J. Spaan. Alwayssafe: Reinforcement learning without safety constraint violations during training. In Frank Dignum, Alessio Lomuscio, Ulle Endriss, and Ann Nowé (eds.), AAMAS '21: 20th International Conference on Autonomous Agents and Multiagent Systems, Virtual Event, United Kingdom, May 3-7, 2021, pp. 1226-1235. ACM, 2021. URL httphhttts//ddddde. wtt.aaacooogooooooo//ooo 00t5//10..5555/ 5555/ 952. 3464094
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Ilya O.
Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schölkopf. Wasserstein auto- encoders. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL htttps://cpenreeee
net. wtt.uummmmmHkk1111111111
John N. Tsitsiklis. Asynchronous stochastic approximation and g-lammnng. Mach. Learn., 165 202, 1994. doi: 10. 10.1007/BF009933066 URL htRt:sttp:://dee wtt..///1110000//0000000000
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N.
Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 6306-6315, 2017. URL http://t/p//pperr http:://ppprrsssss..... cc.psper/pppr/rr/oo//pp/
Cédric Villani. Optimal Transport: Old and New. Springer Berlin Heidelberg, Berlin, Heidelberg, 2009. ISBN 978-3-540-71050-9,
doi: 10.10077978-3-5500 9007/978---540-7105009999 6. URL https://doil
org/10 10. 1007/90077888 34007777778--00-150000000
Andrew M. Wells, Morteza Lahijanian, Lydia E. Kavraki, and Moshe Y. Vardi. Lilf synthesis on probabilistic systems. In Jean-François Raskin and Davide Bresolin (eds.), Proceedings 11th International Symposium on Games, Automata, Logics, and Formal Verification, GandALF 2020, Brussels, Belgium, September 21-22, 2020, volume 326 of EPTCS, pp. 166-181, 2020. doi: www.iooooo..oo..oooooooooo.ooooe 4.0...222//PSSS 326.11
10.4400//PPYCS.3222
URL
Hongyu Zang, Xin Li, and Mingzhong Wang. Simsr: Simple distance-based state representations for deep reinforcement learning. Proceedings of the AAAÍ Conference on Artificial Intelligence, 36 (8):8997-9005, Jun. 2022. doi: 10. (0.1609/aaai.v. 20888. URL ltttps::///// aaai. /rg/
Index.phpp/AA//artiiccc www 20883
14

---

Published as a conference paper at ICLR 2023
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
net//formm i d= 2FCWDKRREU
Shunkang Zhang, Yuan Gao, Yuling Jiao, Jin Liu, Yang Wang, and Can Yang. Wasserstein- wasserstein auto-encoders. CoRR, wwwp/192..09333... 2019. URL http::/arxxi.. www.awwwwwwwaaaaaaaaaeaaooerrese org/ 1902. 0932309333
15

---

Published as a conference paper at ICLR 2023
APPENDIX
A THEORETICAL DETAILS ON WAE-MDPS
A.1 THE DISCREPANCY MEASURE
We show that reasoning about discrepancy measures between stationary distributions is sound in the context of infinite interaction and episodic RL processes. Let Po be a parameterized behavioral model that generate finite traces from the original environment (i.e., finite sequences of state, actions, and rewards of the form (Sorry do. Tor-->), do.--1;7or--1) our goal is find the best parameter which offers the most accurate reconstruction of the original traces issued from the original model M operating under æ. We demonstrate that, in the limit, considering the OT between trace-based distributions is equivalent to considering the OT between the stationary distribution of M, and the one of the behavioral model.
Let us first formally recall the definition of the metric on the transitions of the MDP.
Raw transition distance. Assume that S, A. and (m) are respectively equipped with metric ds. d.A. and dR, let us define the raw transition distance metric over transitions M. i.e. tuples of the form (s, a, s'), as d: S x A x Im(R) x S,
(51. 81 (52, a2 r2, <b2)) = ds (11, $2) + dd(ai, 02) + dR(r1. + ds (s($),5))
In a nutshell, d consists of the sum of the distance of all the transition components. Note that is well defined distance metric since the sum of distances preserves the identity of indiscernible, symmetry, and triangle inequality.
Trace-based distributions. The raw distance d allows to reason about transitions, we thus consider the distribution over transitions which occur along traces of length T to compare the dynamics of the original and behavioral models:
Dx [T] (s,
(s 81) T(als) P(s' (s,)) 1, 1 and
T
(s,T)ss,a,,s)))
+
E Po[[]
1 10
9ost 1. Po(]]
where Po[T] denotes the distribution over traces of length T' generated from Po. Intuitively, I/T- ET (s $1) can be seen the fraction of the time spent in along traces of length T, starting from the initial state Kulkarni 1995). Therefore, drawing (8, a, r²) ~ D [7] trivially follows:
is equivalent to drawing from 1/T Σ -111 et 8r). then respectively and s' from π(- s) and minimize the Wasserstein
P(. 8, a), to finally obtain r = R(s, a). Given N, our objective distance between those distributions: W (D, [T[T). The following Lemma enables optimizing the Wasserstein distance between the original MDP and the behavioral model when traces are drawn from episodic RL processes infinite interactions (Huang 2020).
Lemma A.1. Assume the existence of a stationary behavioral model Eg = limy- Po[TT then
lim Wa (D,, [T], Po[TT) = Wa(En. T--30
Proof. First, note that 3/r E 3y The result follows then from Villani 2009 Corollary 6.9).
weakly converges to & as T goes to op Kulkarni 1995
A.2 DEALING WITH DISCRETE ACTIONS
When the policy N executed in M already produces discrete actions, learning a latent action space is, in many cases, not necessary. We thus make the following assumptions:
Assumption A.2. Let n: - A(A") be the policy executed in M and assume that A (tractable) finite set. Then, we take A = A and oe as the identity function, i.e., 6/: S x A A", (8, a".
16

---

Published as a conference paper at ICLR 2023
Assumption A.3. Assume that the action space of the original environment M is a (tractable) finite set. Then, we take the as the identity function, i.e., 1/0 = 0,"
Concretely, the premise of Assumption A.2 typically occurs when π is a latent policy (see Rem. or when M has already a discrete action space. In the latter case, Assumption A.2 and A..3 amount to setting A = A and ignoring the action encoder and embedding function. Note that if a discrete action space is too large, or if the user explicitly aims for a coarser space, then the former is not considered as tractable, these assumptions do not hold, and the action space is abstracted to smaller set of discrete actions.
A.3 PROOF OF LEMMA
Notation. From now on, we write a 8, - (o)=== Φf 8,
Lemma 3.2. Define ((3, [, ā. 8 = Ee L (s)-)))))-)1111 5, a,a).Po($ 5, a) as the distribution of drawing state-action pairs from interacting with M, embedding them the latent spaces, and finally
letting them transition to their successor state Ma. Then, Wa (Q.: Ex.) W T) L.p
Proof. Wasserstein is compliant with the triangular inequality (Villani 2009), which gives us: W Ex.) & (Q(Q T) + Waz (T, (=,)
where
Wa
(note that Wa is reflexive (Villani 2009
sup E 5, 2. (PPo(a)
E
E S-E3y a- as,
E
E
I f a. 8"), and
(
Wa(Qu: T)
sup E
E
(8,0,5))
E
E
E
(s,a,s)
(3)
f Ee s.
(.
e 8.a (11.0))) 5 (F.
< E
E
sup
E 13, (s))) -
E
((,
(4)
SEF P
s'-Po(-(s.a)
E
E
E feFag s' P(1)a,))
sup
f(8')
E
f(8')
(5)
s,a
(s))
s'-Pot-e)()).a))
E
E
Wdz (.P s, a), Pol. (s))
5,a0fa a (s).a)
We pass from Eq. notice that
to Eq.
4 by the Jensen's 's inequality. To see how we pass from Eq. 4 to Eq. 5
F8 ff: f a1 51 - (52, 02. $ < a5, āI 5). o2 22
01, $1) - (32, a2 $2) W dg($1, 52) + d (2) + (3)) FJ= {f: (1,,
Observe now that and are fixed in the supremum computation of Eq. 4 all functions f considered and taken from F are of the form f(8. .). It is thus sufficient to consider the supremum over functions from the following subset of Fz:
(f: (5,,,5)) f(5 5,) & ds(5, da + ds (5)
(for 8, drawn from (.)
= =ff - a dg 52)
f(,
=f: f(5)) - (()) & dz (5)
=Fag-
Given a state s E S in the original model, the (parallel) execution of i Me is enabled through T(a, @ s) x(a s) (ā di(s), a) (cf. Fig. lb.. The local transition loss resulting from this
interaction is:
L = E
Waz (6. s, a), P(. ())
<a
= E
R
Waz (o.P(. s, Po (s),
e,a
17

---

Published as a conference paper at ICLR 2023
which finally yields the result.
A.4 PROOF OF THEBRREM 3.3
Before proving Theorem 3.3 let us introduce the following Lemma, that explicitly demonstrates the link between the transition regularizer the W2AE-MDP objectivv and the local transition loss
required to obtain the guarantees related to the bisimulation bounds of Eq. [
Lemma A.4. Assume that traces are generated by running क E II in the original environment, then
5,a ~ of (+),a)) Was (o.P(: a²) Pa(- (.(s), a)) L.
Proof. Since the latent policy म generates latent actions, Assumption A.2 holds, which means:
E 8,0" - a~ (H())
E
P. Waz (o. s,a"), Po( (- (s).
E
Was (c. 18.18,0)) a) Pol (.())
Theorem 3.3. Assume that traces are generated by running latent policy II in the original environment and let dR be the usual Euclidean distance, then the W2AE-MDP objective is
min 1,8 8,8'06x
ds(s', Go(l) 15 3-3--WW: L.)).
Proof. We distinguish two cases: (i) the case where the original and latent models share the same discrete action space, i.c., A = A. and (ii) the case where the two have a different action space (e.g., when the original action space is continuous), i.e., A. In both cases, the local losses term follows by definition L and Lemma A.4 When dR the Euclidean distance (or even the L1 distance since rewards are scalar values), the expected reward distance occurring in the expected trace-distance term d in the W2AE-MIPP objective directly translates the local loss L Concerning the local transition loss, in case (i), the result naturally follows from Assumption A.2 and A.3 In case (ii), only Assumption A.2 holds, meaning the action encoder term of the W MDP objective is ignored, but not the action embedding term appearing in Gg. Given ~ &F, recall that executing in M amounts embedding the produced latent actions
#(- (s() back to the original environment via ( (s), a) (cf. Rem. and Fig. [a]. Therefore, the projection of d( r, 8') Golo. (s), (s())) the action space A is ( (s), a), (o) for r = R(s, a) and s' P( s,
A.5 OPTIMIZING THE TRANSITION REGULARIZER
In the following, we detail how we derive a tractable form of our transition regularizer L (w). Opti- mizing the ground Kantorovich-Rubinte 1-Lipschitz network yP that need to be trained to attain the supremum of the dual:
duality is enabled via the introduction of parameterized,
L (w)
P (5))
max 8,0 5,5 (...0.---...) WE FFe of P.
P
Under this form, optimizing I.S (w) is intractable due to the expectation over the maximum. The following Lemma allows rewriting I.S- to make the optimization tractable through Monte Carlo estimation.
Lemma A.5. Let X.Y be two measurable sets, & E A(2), P: X - A()), - A(Y), and d: x [0, oo be metric Y. Then,
E Wa 2)) Q(-1z))= sup 1: X-
E
((x)(11)
E
p(x)(12)
2~5
((-aa)
92 -2~((1z))
18

---

Published as a conference paper at ICLR 2023
Proof. Our objective is show that
E
sup feFa
E P((a))
((v))()))
E
(()))(z)
(6)
(11(((
E 6: F- 2n Li ((-la)
sup
(1))
(1x)
02-2(11)
We start with (6
Construct 4p" X
Fa by setting for all x X
('(x)
arg sup fePa ((-1(()))
E
f(v)) Q(12)) f122).
This gives
E
sup
1(11)
(122
g^g
fET a1 P((l))
(2-((((2)
0 (z)
E
10 (x2)
g~ (-(r)())
()~((1z))
<
sup E
E
p(x)(1))
E
(12)
1:: x Fu an
Vi P(-Ix)
remains show that
Take
arg
sup
E
E 31 ((-2))
(())
E ((()
#: X
X-F
2
92~Q(112)
Then, for all x E X, we have 40" (x(())
E Fa which means:
E
(v1)
E (2-Q((1m)
4" (x)
21 >1~P(11 (1)
< sup
E feF P((
f(y1)
E
f(22)
(2-((112)
This finally yields
E P(-a)
(1())
E
40 (x))))()))
2x
(2~((11)
<
E
sup JeFayi-P(iB
E
f(y1)
1 x2-Q(-1a))
f(y2)
Corollary A.5.1. Let &x be a stationary distribution of Ma and X Ax8x A, then
-5-
sup E ii X' N
E
(((.
E
a(s,a,5,0)(5
800 Ee s,a (11,0)))
8'-Po(1(,.))
Consequently, we rewrite I.ST (w) as tractable maximization:
L (w) (w)
max wi PEffy s.a. Eo F. 5. ((..a,,,a,0.,(1)))
E
&e ((,,,,,,0,5))
s'-Po((s,a)
A.6 THE LATENT METRIC
In the following, we show that considering the Euclidean distance for and dg in the latent space for optimizing the regularizers We, and Lf is Lipschitz equivalent to considering continuous X-relaxation of the discrete metric 1 (xx y) = 1 Consequently, this also means it is consistently sufficient to enforce 1-Lispchitzness via the gradient penalty approach of Gulrajani et al. (2017) during training to maintain the guarantees linked to the regularizers in the aero-temperature limit, when the spaces are discrete.
19

---

Published as a conference paper at ICLR 2023
Lemma A.6. Let d be the usual Euclidean distance and dx: [0, [0, [0, 1[, (x, y) d(n.y) for E ]0, 1] and n N, then dy is a distance metric.
X+d(a.y)
Proof. The function dx is metric iff it satisfies the following axioms:
1. Identity of indiscernibles: If æ y, then (æx y) d(a.y) 0 0 since d is distance metric. Assume now that (æ,, y) 0 and take d(æ; y), for any a; y. Thus, or E [0, + and 0 = is only achieved in O, which only occurs whenever x: since d is distance metric.
X+dlæ.y] X+U
2. Symmelry:
dx(2, y) = d(x, y)
À + d(x y) d(y. x))
(d is a distance metric)
+ d(y, x))
dx .r)
3. Triangle inequality: Let æ. y, [0. 1]", the triangle inequality holds iff
dx y) + dx
> dx x)
(8)
d(x,y) + x
d(y.z)
d(x z)
x+
+ d(xx 2)
)d(xx +
d(x, 2)
X2 + Ad(x. + Ad(y) d(x,y)dd(v..
+ d(x,
Pd(x, y) + X 22d yd Ad(z,y)d(x; +
+ 2d(x x. y dl 2){(x, z)
>)d(0z + Ad(x, y),( + (y 2) + d(æ, y)d(y, z) z)((æ, 2)
(cross-product, with X and Im(d) [0, co
X²d(x, 3) + Xdd(y. 2) + 2xd(2 y) d(xx y)d(y, 2 z)d(x, X²d(x, z)
Since d is a distance metric, we have
Xd((x,)) y) X²d 2) Xdd(x,
(10)
and Im(d) [0, oo[, meaning
2Ad( y)d z) + d(xx, y)d(y. 2)d(x 2)
(11)
By Eq. 10]and II the incquality of Eq. holds. Furthermore, the fact that Eq. g and 9 yields the result.
are equivalent
Lemma A.7. Let d, d, as defined above, then (i) dx 2 1, and (ti) d, d, are Lipschitz-equivalent
Proof. Part (i) is straightforward by definition of d,. Distances d and dx are Lispchitz equivalent if and only 3a, > such that Væ, y e [0, 1 dx (z. y) d( [ y)
t.d(x, 4)
b d(z,y)
=a d(x, y)
b d(x,
X + d(x y)
b b < X + d(x y)
Taking Xtva and b Y yields the result.
20

---

Published as a conference paper at ICLR 2023
Corollary A.7.1. For all 3 > 1/2, SE SE a A, se and ā A, we have
1. Wax (T. & B. Wa (T,
2. Wa, ( ((. s, Po( a)) & B. Wa (o. 8, Pol. (.)))
Proof. By Lipschitz equivalence, taking 3 > 1/2 ensures that N. Væ. y E [0, 1]", dx (r.y) B dl x. y). Moreover, for any distributions P, Q. Wa (P,Q) & B Wa Q) (cf., e.g., Gelada ct al. 2019 Lemma A.4 for details).
In practice taking the hyperparameter 3 1/1 in the W'AE-MDP ensures that minimizing the 3-scaled regularizers w.r.t. also minimizes the regularizers w.r.t. the X-relaxation d,, being the discrete distribution in the zero-temperature limit. Note that optimizing over two different 31, 82 instead of a unique scale factor / is also good practice to interpolate between the two regularizers.
B EXPERIMENT DETAILS
The code for conducting and replicating experiments is available at httttps://ggthub. com/ www.nndegggaggeeeeeeoeooo mdp
B.1 SETUP
We used TENSORFLOW 2. 7.7.0 (Abadi et al. 2013) to implement the neural network architecture of W3AE-MDP, TEnsorFlow PROBABILITY 0. ....0 (Dillon et al. 2017) to handle the
probabilistic components of the latent model (e.g., latent distributions with reparameterization tricks, masked autoregressive flows, etc.), well TG-AGENTS 0111. (Guadarrama al. 2018 to handle the RL parts of the framework.
Models have been trained on a cluster running under CentOS Linux (Core) composed of mix of nodes containing Intel processors with the following CPU microarchitectures: (i) 10-core INTEL 2.5-2680v2 (ii) 14-core INTEL E5-266 v6044, and (iii) 20-core INTEL Xeon Gold 6148. We used 8 cores and 32 GB of memory for each run.
B.2 STATIONARY DISTRIBUTION
To sample from the stationary distribution &x of episodic learning environments operating under т E II, we implemented the recursive e-perturbation trick of Huang (2020). In a nutshell, the reset of the environment is explicitly added to the state space of M, which is entered the end of each episode and left with probability 1 to start a new one. We also added a special atomic proposition reset into AP to label this reset state and reason about episodic behaviors. For instance, this allows verifying whether the agent behaves safely during the entire episode, or it is able to reach a goal before the end of the episode.
B.3 ENVIRONMENTS WITH INITIAL DISTRIBUTION
Many environments do not necessarily have single initial state, but rather an initial distribution over states d, (SS). In that case, the results presented in this paper remain unchanged: it suffices add a dummy state
the state space J {s") so that 8 P(s s* a) = ds(s') for any action e A. Therefore, each time the reset of the environment is triggered, we make the MDP entering the initial state s", then transitioning to s' according to d,.
with the transition dynamics
B.4 LATENT SPACE DISTRIBUTION
As pointed out Sect. posterior collapse naturally avoided when optimizing WDAE-MDP. To illustrate that, report the distribution of latent states produced plots reveal that the latent space generated by mapping original states drawn from & during training to S via σ, is fairly distributed, for each environment.
by during training (Fig. 5). The
21

---

Published as a conference paper at ICLR 2023
CariPola
MountannCar
Acrobot
612
1024
8152
760
0144
Ialent state
Ialent stale
batont
258
6.12
4096
0
Ze5
405
6e5
Be5
2w5
4e5
6e5
Be5
2w5
4e5
646
Be5
slop
Lundor
Pandulum
8152
6144
Iatent 52800
4090
2048
stop
stop
Figure 5: Latent space distribution along training steps. The intensity of the blue hue corresponds to the frequency of latent states produced by de during training.
environment =s CartPole
environment MourtainCar
environment Acrabot
envtronment Lunart.nder
environment = Pendulum
difereros
15
0.0
0.5 sleg
1.0
0.0
0.5 step
1.0
0.0
0.5 stap
1.0
0.0
0.5
200000 slep
400000
105
slep
Figure
Absolute value difference
|Va.l reported along training steps.
B.5 DISTANCE METRICS: STATE, ACTION, AND REWARD RECONSTRUCTION
The choice of the distance functions ds, d and dR. plays role in the success of our approach. The usual Euclidean distance is often a good choice for all the transition components, but the scale, dimensionality, and nature of the inputs sometimes require using scaled, normalized, or other kinds of distances to allow the network to reconstruct each component. While we did not observe such requirements in our experiments (where we simply used the Euclidean distance), high dimensional observations (e.g., images) are example of data which could require tuning the state-distance function in such way, to make sure that the optimization of the reward or action reconstruction will not be disfavored compared to that of the states.
B.6 VALUE DIFFERENCE
In addition reporting the quality guarantees of the model along training steps through local losses (cf. Figure 4b). our experiments revealed that the absolute value difference V between the original and latent models operating under the latent policy quickly decreases and tends to converge to values in the same range (Figure This is consistent with the fact that minimizing local losses lead to close behaviors (cf. Eq. and that the value function is Lipschitz-continuous w.r.t. d=s (cf. Section
B.7 REMARK ON FORMAL VERIFICATION
Recall that our bisimulation guarantees come by construction of the latent space. Essentially, our learning algorithm spits out a distilled policy and a latent state space which already yields a guaranteed bisimulation distance between the original MDP and the latent MDP. This is the of how enable verification techniques like model checking. In particular, bisimulation guarantees mean that reachability probabilities in the latent MDP compared those in the original one close.
22

---

Published as a conference paper at ICLR 2023
Furthermore, the value difference of (omega-regular) properties (formulated through mu-calculus) obtained in the two models is bounded by this distance (cf. Sect. 2 and Chatterjee et al. 2010).
Reachabilityy is the key ingredient model-check MDPs. Model-checking properties is in most cases performed by reduction to the reachability of components or regions of the MDP: it either consists of (i) iteratively checking the reachability of the parts of the state space satisfying path formulae that comprise the specification, through a tree-like decomposition of the latter (e.g.. for (P.R-)CTL. properties, cf. BBierr & Katoen 2008), or (ii) checking the reachability to the part of the state space of product of the MDP with memory structure or an automaton that embeds the omega-regular property - e.g., for LTL (Baier et al., 2016: Sickert et al. 2016), LTLf (Wells) et al. 2020), or GLTL (Littman et al. 2017). among other specification formalisms. The choice of specification formalism is up to the user and depends on the case study. The scope of this work is focusing on learning to distill RL policies with bisimulation guarantees so that model checking can be applied, in order to reason about the behaviors of the agent. That being said, reachability is all we need to show that model checking can be applied.
B.8 HYPERPARAMETERS
W-AE-MDP parameters. All components (c.g., functions or distribution locations and scales, see Fig.P 2 are represented and inferred by neural networks (multilayer perceptrons). All the networks share the same architecture (i.e., number of layers and neurons per layer). We use a simple uniform experience replay of size 10 to store the transitions and sample them. The training starts when the agent has collected 10a transitions in MM. We used minibatches of size 128 to optimize the objective and we applied a minibatch update every time the agent executing т has performed 16 steps in M. We use the recursive e-perturbation trick of Huang 20200) with 3/4: when episode ends, it restarts from the initial state with probability 1/2; before re-starting an episode, the time spent in the reset state labeled with reset follows then the geometric distribution with expectation 5/1 3. We chose the same latent state-action space size than Delgrange et al. (2022). except for Lunarl.ander that we decreased log2 S| 14 and |
3 improve the scalability of the verification.
VAE-MDPs parameters. For the comparison of Sect. we used the exact same VAE-MDP hyperpa- rameter set as prescribed by |Delgrange et al. (2022). except for the state-action space of LunarLander that we also changed for scalability and fair comparison purpose
Hyperparameter search. To evaluate our W3AE-MDP, we realized a search in the parameter space defined in Table The best parameters found (in terms of trade-off between performance and latent quality) are reported in Table We used two different optimizers for minimizing the loss (referred as the minimizer) and computing the Wasserstein terms (reffered to the maximizer). We used ADAM (Kiggma& & Ba 2015) for the two, but we allow for different learning rates ADAM, and exponential decays ADAM,, ADAM2,. We also found that polynomial decay for ADAMa (c.g.. to 10-5 for 4 10 steps) is a good practice to stabilize the experiment learning curves, but is not necessary to obtain high-quality and performing distillation. Concerning the continuous relaxation of discrete distributions, we used different temperature for each distribution, as Maddison al. (2017) pointed out that doing so is valuable to improve the results. We further followed the guidelines of Maddison al. (2017) choose the interval of temperatures and did not schedule any annealing scheme (in contrast to VAE-MDPs). Essentially, the search reveals that the regularizer scale factors 3. (defining the optimization direction) as well as the encoder and latent transition temperatures are important to improve the performance of distilled policies. For the encoder temperature, we found nice spot in As = 4/s, which provides the best performance in general, whereas the choice of Xp, and B. are (latent-) environment dependent. The importance of the temperature parameters for the continuous relaxation of discrete distributions is consistent with the results of (Maddison et al., 2017). revealing that the success of the relaxation depends on the choice of the temperature for the different latent space sizes.
Labeling functions. We used the same labeling functions as those described by Delgrange al. (2022). For completeness, we recall the labeling function used for each environment in Table
TThe code for conducting the VAF-MDPs experiments is available at https: cewiithub.coom com l aaraang//ae_mm@e (GNU General Public License v3.0).
Florentée
23

---

Published as a conference paper at ICLR 2023
Table 2: Hyperparameter search. Xx refers the temperature used for W'AE-MDP component X.
Parameter
Range
ADAMa (minimizer)
[0.0001, 0.0002, 0.0003, 0.001 10.0001, 0.0002, 0.0003, 0.001
ADAMa (maximizer)
ADAM,,
0, 0.5, 0.9)
ADAM8,
0.9, 0.999
neurons per layer number of hidden layers
64, 128, 256, 512]
1, Leaky ReLU, tanh, solppus((o+2) (smooth ELU)
activation
3wt.
10, 25, 50, 75, 100
10, 25, 50, 75, 100
m
5, 10, 15, 10,
6
use s-mimic (cf. Delgrange 2022
True, False True, decay rate of 10-5 is used)
dp.
0.1, 1/a, 1/2, 2/, 2/5, 0.99
Xa
0.1, 1/a, 1/2, 2/, 3/5,
Ano
1/ -1, -1
Agd
11 4(44)
Table 3: Final hyperparameters used to evaluate W'AE-MDPs in Sect.
CartPole
MountainCar
Acrobot
LunarLander
Pendulum
loga S
9
10
13
14
13
[ activation
2 |A|
|A|
= |A|
3
3
tanh
ReLU
Leaky Relu
ReLU
ReLU
layers
[64, 64, 64]
[512, 512]
[512, 512]
[256]
[256, 256, 256]
ADAM, (minimizer)
0.0002
0.0001
0.0002
0.0003
0.0003
ADAM, (maximizer)
0.0002
0.0001
0.0001
0.0003
0.0003
ADAM,,
0.5
0.5
ADAM 52
0.999
0.999
0.999
0.999
0.999
B
25
10
50
25
Byee
75
100
10
100
25
5
20
20
15
5
6
20
20
10
e
0
0
0.5
Avo
1/3
1/3
0.1
0.75
3/3
No.
1/3
2/3
2/3
2/3
2/3
Aeo
2/3
1/3
0.5
0.5
0.5
1/3
1/3
24

---

Published as a conference paper at ICLR 2023
Environment
S
Description, for
((s)
(pi (p ...... Pray Preset
81 cart position
$2 cart velocity
11.5: unsafe cart position
CartPole
R44
pole angle (rad)
P2 e5300.15 unsafe pole angle
pole velocity tip
1 1.5-1.5 target pusition
81 position
MountainCar
R2
= 1 1aa-1ps right-hand side of the mountain
#2: velocity
Ps 1 s2 car going forward
Let 0,, B2 E 2n] be the angles of the two rotational joints,
1 1-er-aus 1- + RL agent target
P2 1 1:000 [0, v [35/2, 2n]
81 cos(0,)
Ps 10230: 01 [0,
#2 sin(0,) (0))
Acrobot
R66
P4 = 1a,30: 8g [0, */2| (35/2, 2n]
83 cos(
P5 = 1,,00 8g [0,
84 sin (0))
= 1,00: positive angular velocity (1)
ss angular velocity
= 10 positive angular velocity (2)
Set angular velocity 2
Let 0 [0. 2x] be the joint angle
1 3000((/)))) safe joint angle
$1 cos(O)
P2 a:30 A s/2| [35/2, 2n]
Pendulum
R
# sin
Ps a0000 [0, aT
#3: angular velucity
P4 1 positive angular velocity
horizontal coordinates
$2 vertical coordinates
Pt unsafe angle
Rg horizontal speed
P2 leg ground contact
Sy vertical speed
Py: lands rapidly
Lundrl.ander
R
R5: ship angle
P4: left inclination
8g angular speed
Ps: right inclination
sy: left leg contact
Pei motors shut down
8x: right leg contact
Table 4: Labeling functions for the OpenAl environments considered in our experiments (Delgrange et al., 2022). We provide a short description of the state space and the meaning of each atomic proposition. Recall that labels binary encoded, for [AP| (one bit is reserved for reset) and Preset iff s is a reset state (cf. Appendix B.2).
Time to failure properties. Based the labeling described in Table 4 formally detail the time to failure properties checked in Sect. 4 whose results are listed in Table I for each environment. Let Reset { reset] (O,.... ..... 1> (we assume here that the last bit indicates whether the current state is a reset state or not) and define + L A L2 iff L and 8 f L2 for any E S, then
CartPole: ş Reset 2/ Unsafe, where Unsafe (1, 1,
MountainCar: y Goal Reset, where Goal (1,0,0.0)
Acrobot: - 24 Reset, where Goal = (1, 0, ...
LunarLander: = -SafeLandingl4 Reset, where SafeLanding GroundContact ^ MotorsOff, GroundContact = (0, 0, 0, 0), and MotorsOff 00,0,0,0.0.1,00
Pendulum: 0 Safe ^ Reset), where Safe = (1, 11,0.00 0. 0, 0) QT - and 8; - TT iff $;+1 T, for any T AP, A aaa Traj. Intuitively, ip denotes the event of ending episode in an unsafe state, just before resetting the environment, which means that either the agent never reached the safe region or it reached and left at some point. Formally, &
So.co Gp:se
3i e N, s, # Safe $ } Reset
Traj.
C ON THE CURSE OF VARIATIONAL. MODELING
Posterior collapse is a well known issue occurring in variational models (see, e.g., e.g., Alemi al. 2018 Tolstikhin et al. 2018 He et al. 2019 Dong et al. 2020) which intuitively results a degenerate local optimum where the model learns ignore the latent space and only the reconstruction functions (i.e., the decoding distribution) to optimize the objective. VAE-MDPs are no exception, as pointed out in the original paper (Delgrange al. 2022 Section 4.3 and Appendix C.2).
25

---

Published as a conference paper at ICLR 2023
VAE-MOP (verilla)
VAE-MDP
512
512
384
384
Variational model
lafort 2bil
WAE-MPP
256
258
WIE MDP (vanllo)
128
128
Rate
205
4e5
G65
Be5
205
465
6e5
5e5
(a) Latent space distribution along training steps. The intensity of the blue hue corresponds to the frequeney of latent states produced from d, during training. The vanilla model collapses to single state.
0.2
(b) Rate of the variational model.
Vil MEP (anila)
Wil MGP
200
uhe
Varterionall shi vgD
10
15
-00
0.5
(c) Distortion of the variational model.
(d) Average point-wise entropy of 6. (.(+)s), for 8 E S drawn from the interaction with the original environ- ment.
(c) Performance of the resulting distilled policy To when deployed in the original environment (aver- aged over 30 episodes).
Figure 7: Comparison of the VAE-MDP in the CartPole environment (i) when the distortion and the rate are minimized as is (vanilla model) and (ii) when it makes use of annealing schemes, entropy regularization, and prioritized experience replay to avoid posterior collapse (cf. Delgrange et al. 2022). While the former clearly fails to learn a useful latent representation, the later does so meticulously and smoothly in two distinguishable phases: first, Φh focuses on fairly distributing the latent space, setting up the stage to the concrete optimization occurring from step 4- 10S, where the entropy of o is lowered, which allows to get the rate of the variational model away from zero. Five instances of the models are trained with different random seeds, with the same hyperparameters than in Sect. 4
Formally, VAE- and WAE-MDPs optimize their objective by minimizing two losses: reconstruction cost plus a regularizer term which penalizes a discrepancy between the encoding distribution and the dynamics of the latent space model. In VAE-MDPs, the former corresponds to the the distortion, and the later to the rate of the variational model (further details are given in Allmmi et al. 2018 Delgrange et al. 2022). while in our WAE-MDPs, the former corresponds to the raw transition distance and the later to both the steady-state and transition regularizers. Notably, the rate minimization of VAE- MDPs involves regularizing a stochastic embedding function (. s) point-wise, i.e., for all different input states 8 E S drawn from the interaction with the original environment. In contrast, the latent space regularization of the WAE-MDP involves the marginal embedding distribution Q, where the embedding function Ø. is not required to be stochastic. Alemi et al. (2018) showed that posterior collapse occurs in VAEs when the rate of the variational model is close to zero. leading to low-quality representation.
Posterior collapse in VAE-MDPs. We illustrate the sensitivity of VAE-MDPs to the posterior collapse problem in Fig. 7, through the CartPole environmen
P minimizing the distortion and the rate as is yields an embedding function which maps deterministically every input state to the same sink latent state (cf. Fig. 7aj. Precisely, there is latent state & S so that 6,(8) s) & I and A. This is a form of posterior collapse,
Po(8 8, a) & 1 whatever the state E S and action the resulting rate quickly drops to zero (cf. Fig 7b), and the resulting latent representation yields no information at all. This phenomenon is handled in VAE-MDPs by using (i) prioritized replay buffers that allow to focus on inputs that led to bad representation, and (ii) modifying the objective
In fact, the phenomenon of collapsing to few state occurs for all the environments considered in this paper when their prioritized experience replay not used, illustrated in Delgrange al. 2022 Appendix C.2.
26

---

Published as a conference paper at ICLR 2023
function for learning the latent space model the so-called evidence lower bound (Hoffman et al. 2013 Kingma& & Welling 2014), or ELBO for short and set up annealing schemes to eventually recover the ELBO at the end of the training process. Consequently, the resulting learning procedure focuses primarily on fairly distributing the latent space, to avoid it to collapse to a single latent state, to the detriment of learning the dynamics of the environment and the distillation of the RI. policy. Then, the annealing scheme allows to make the model learn to finally smoothly use the latent space maximize the ELBO, and achieve consequently a lower distortion at the "price" of a higher rate.
Impact of the resulting learning procedure. The aforementioned annealing process, used to avoid that every state collapses the same representation, possibly induces high entropy embedding function (Fig. 7d). which further complicates the learning of the model dynamics and the distillation in the first stage of the training process. fact, in this particular case, one can observe that the entropy reaches its maximal value, which yields fully random state embedding function. Recall that the VAE-MDP latent space is learned through independent Bernoulli distributions. Fig. 7d reports values centered around 1.188 in the first training phase, which corresponds to the entropy of the state embedding function when o. (- s) is uniformly distributed over S for any state E S: Joga -
Pi log p; - (1- p.) log(1 - p.) 4.188, where p; 1/2 for all i. The rate (Fig. 7b) drops to zero since the divergence pulls the latent dynamics towards this high entropy (yet another form of posterior collapse), which hinders the latent space model to learn useful representation. However, the annealing scheme increases the rate importance along training steps, which enables the optimization to eventually leave this local optimum (here around 10 training steps). This allows the learning procedure to leave the zero-rate spot, reduce the distortion (Fig. Tcp. and finally distill the original policy (Fig. Tcp.
H (. [s)]
As result, the whole engineering required to mitigate posterior collapse slows down the training procedure. This phenomenon reflected in Fig. [ VAB-MDPs need several steps to stabilize and set up the stage to the concrete optimization, whereas WAE-MDPs have no such requirements since they naturally do not suffer from collapsing issues (cf. Fig. 5). and are consequently faster to train.
Lack of representation guarantees. On the theoretical side, since VAE-MDPs are optimized via the ELBO and the local losses via the related variational proxies, VAE-MDPs do not leverage the representation quality guarantees induced by local losses (Eq. [ during the learning procedure (as explicitly pointed out by Delgrange et al., 2022] Sect. 1.): contrast to WAF-MDPs, when two original states are embedded to the same latent, abstract state, the former are not guaranteed to be bisimilarly close (i.c., the agent is not guaranteed to behave the same way from those two states by executing the policy), meaning those proxies do not prevent original states having distant values collapsing together to the same latent representation.
INDEX OF NOTATIONS
[cond]
indicator function: 1 the statement [cond] true, and otherwise Set of 1-Lipschitz functions w.r.t. the distance metric d Sigmoid function, with (Ix) 1/1+ex)(=a)
Fa
fo
A function fo X R modeled by neural network, parameterized by 6, where X is any measurable set
Latent Space Model
M (S, 87, Latent MDP with state space S, action space A, reward function R. labeling function é, atomic proposition space AP, and initial state 8y.
M, o, vb)
Latent space model of M
Latent action in A
Latent policy 5 A: can be executed in M via di 7(- ((())
d3
Distance metric over S
State embedding function, from S to S
27

---

Published as a conference paper at ICLR 2023
Action embedding function, from 8 to
oP
Distribution of drawing s' ~ P( 8, a), then embedding s' = ((s'), for any state E S and action a E
RR
Local reward loss under distribution & Local transition loss under distribution
Il
Set of (memoryless) latent policies
Latent state in S
V
Latent value function
Markoy Decision Processes
M
A,P, C, sP, MDP M with state space S, action space A, transition function P. labeling function l, atomic proposition space AP, and initial state 87.
Action in A
d.
Bisimulation pseudometric
Discount factor in [0, 1]
dA
Metric over the action space
dR
Metric over Im)
ds
Metric the state space
et
Limiting distribution of the MDP defined as (5 -PM, Succe Ouce any source state S
88 = 5 }), for
11
Set of memoryless policies of M Memoryless policy æ: S - A(A)
p.
Unique probability measure induced by the policy in M on the Borel o-algebra over measurable subsets of Traj
CUT
Constrained reachability event
M,
MDP obtained by replacing the initial state of M by s S
State in S
&a
Stationary distribution of M induced by the policy
d
Raw transition distance, i.c., metric over S A Im(R) S
Traj
Set of infinite trajectories of M
(So.T, Qur-1)
Trajectory
V
Value function for the policy æ
Probability Measure Theory
Discrepancy measure; D(P,Q) is the discrepancy between distributions P,Q A(X) A(X) Set of measures over complete, separable metric space X'
D
Logistic(H, s) Logistic distribution with location parameter and scale parameter
Wa
Wasserstein distance w.r.t. the metric d; Wa (P,Q) is the Wasserstein distance between distributions P,Q ((X)
Wasserstein Auto-encoded MDP
Ea
Behavioral model: distribution over S x A Im x S Mapping (5, a, 5") - (Go(3), (5, a), Ra(3, (3, a), Go(5))
Gg
6,4
Action encoder mapping 5 x A to A (A)
Go
State-wise decoder, from S to S
28

---

Published as a conference paper at ICLR 2023
Q.
Marginal encoding distribution over 8 Ax x 3: Es,a,''-f, d.(- s')
t
Stationary distribution of the latent model Mo, parameterized by
We.
Steady-state regularizer
p
Steady-state Lipschitz network
Temperature parameter
T
Distribution of drawing state-action pairs from interacting with M, embedding them to the latent spaces, and finally letting them transition to their successor state in Ma, in (S S)
Transition Lipschitz network
ADDITIONAL REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jía, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015, URL https: www. www.eernsrriio.o flw.. org/ Software available from tensorflow.org.
Olah,
Alexander A. Alemi, Ben Poole, Ian Fischer, Joshua V. Dillon, Rif A. Saurous, and Kevin Murphy, Fixing a broken ELBO. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässsn, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 169-168. PMLR, 2018 URL http::/ppoocedd.gg........ wtwp://www.......eeeeeeeeeeoeeee alemimlmmmiii 18a. html
Joshua V. Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A. Saurous. Tensorflow distributions, 2017.
Zhe Dong, Bryan A. Seybold, Kevin Murphy, and Hung H. Bui. Collapsed amortized variational inference for switching nonlinear dynamical systems. In Proceedings of the 37th International Conference Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 2638-2647. PMLR, 2020. URL http
//proceedings.oo mlr prw.sppespp............ / www.g////////..e.eeoo/ html
Sergio Guadarrama, Anoop Korattikara, Oscar Ramirez, Pablo Castro, Ethan Holly, Sam Fishman, Ke Wang, Ekaterina Gonina, Neal Wu, Efi Kokiopoulou, Luciano Sbaiz, Jamie Smith, Gábor Bartók, Jesse Berent, Chris Harris, Vincent Vanhoucke, and Eugene Brevdo. TF-Agents: A library for reinforcement learning in tensorflow. httph:tttt/t/. tttpt//ttttuuuu. 2018. URL httt:s/s//gittuu..oo/eeeeeetee agents [Online; accessed 25-June- 2019].
hwwph/wwwwww...tt/eeaeeeeeeeeeee www.or/rwww//////t.../s/ agent.s
Junxian He, Daniel Spokoyny, Graham Neubig, and Taylor Berg-Kirkpatrice Lagging inference networks and posterior collapse in variational autoencoders. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URLLtttp::/
wttp::wwwwwwteeteeteeee id dtrprrmmmpmdddyyyyypcpp
Matthew D. Hoffman, David M. Blei, Chong Wang, and John W. Paisley, Stochastic variational infer- ence. J. Mach. Learn. Res., 14(1):1003-1377 2013. htll:t acm. org/citation
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: 7/arxi wttg/ba////aaa////// 6980
29

---

Published as a conference paper at ICLR 2023
Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. In Yoshua Bengio and
Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014. Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http:
//arxiv.org/ abs / 1312. 6114
Vidyadhar G. Kulkarni. Modeling and Analysis of Stochastic Systems. Chapman & Hall, Ltd., GBR, 1995. ISBN 0412049910.
lya O. Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schölkopf. Wasserstein auto- encoders. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Truck Proceedings. OpenReview.net, 2018, URL
https://epenrevalt net/ forum? 1d=HkL 7
30


Reviews:
[{"review_id": "sHO2eW6MihM", "reviewer": "Reviewer_cw5C", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on improving the performance of deep reinforcement learning. The authors proposed Wasserstein auto-encoded Markov Decision Processes (WAE-MDPs) for solving this problem. Experimental results show that WAE-MDPs outperform VAE-MDPs. ", "summary_of_the_review": "The authors propose a novel method called WAE-MDP. Experimental results show that WAE-MDP outperforms VAE-MDP", "strength_and_weaknesses": "Strengths\n\nThe authors proposed a novel method called WAE-MDP.\nThe authors provided theoretical guarantees for WAE-MDP.\nExperimental results show that WAE-MDP outperforms VAE-MDP\n\nWeaknesses\n\nThe authors claim that one of the reasons why VAE-MDP does not work well is due to the mode collapse of VAE. However, limited evidence is provided that VAE suffers from mode collapse.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "n3MR6Aey37A", "reviewer": "Reviewer_MegB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper is devoted to verifying deep reinforcement learning (DRL) policies and proposes a latent space model (Wasserstein auto-encoded MDP, WAE-MDP) that addresses a number of issues pertaining to earlier latent space models by minimising the optimal transport between the behaviours of the agent executing the original policy and the distilled policy. Here, formal guarantees for the distilled policy are provided. Experimental results demonstrate the proposed approach to be faster in distilling policies and better in terms of the overall result quality.", "summary_of_the_review": "As I feel completely lost in the text, I am afraid I cannot give this paper a positive score. However, as I mentioned above, I could have overlooked important merits of the paper, hence, my score is \"borderline reject\".", "strength_and_weaknesses": "I should admit that I am no expert in this particular area and so I may have overlooked some of the important details of the work.\n\nHaving said that, I should bring the authors' attention to the fact that the paper is not helping much in this regard. Although it has a section on \"background\", in my opinion, it still does extremely little when introducing preliminary knowledge - the authors cut to the chase and start throwing bulky \"two-storey\" formulas, equations, and theorems with no motivation whatsoever. It may be okay for an expert in the field, e.g. the authors themselves, but an average reader like me will get lost immediately. The algorithm description provided follows the same style, which makes it really hard to appreciate the presentation provided.\n\nIn my view, \"formal verification\" means logical reasoning about a system and formally proving certain properties of the system. This paper does not make it clear what kind of reasoning engine is applied when verifying the properties. I should also say that it does not discuss the practicality of the proposed approach (say, from the scalability perspective).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TEduvWgqgek", "reviewer": "Reviewer_JqQE", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces the Wasserstein auto-encoded MDP (WAE-MDP), which is a latent space model that aims to overcome some of the shortcomings of a prior latent space model (VAE-MDP). Their method learns a (small) discrete representation of the state-action space. Their approach lends itself to formal guarantees on behaviour (based on the theory of bisimulation). In addition to enabling formal guarantees, the authors demonstrate their method is 10 times faster than VAE-MDP in distilling policies.", "summary_of_the_review": "A nice paper that provides a theoretically-motivated methodology for learning a latent space model on which formal verification can be performed.\n\nI support accepting the paper, provided the clarity points above are addressed.", "strength_and_weaknesses": "# Strengths\nThe paper is well written, well motivated, and properly evaluated. The theoretical claims are presented clearly and the proofs are correct (as far as I could tell).\n\nI wanted to commend the authors on Figures 1 and 2, which really nicely complement the text and help clarify the methods discussed.\n\n# Weaknesses\nThe main weakness for me is the complexity of the method presented, although that is difficult to resolve. The paper is quite clear, but there are a few points where it could be made clearer:\n1. There is a _lot_ of notation. It would be useful to provide an index of notation inthe appendix, as it can be hard to keep track of all of them.\n1. In the top of page 4, it is not clear whether $\\bar{\\mathcal{R}}$ and $\\bar{\\bf{P}}$ are learned or not. I think they are, but this should be made explicit. In particular, they should include the parameterization $\\theta$ (which is present in Figure 2) to make it clearer.\n1. In section 3.1, shouldn't there be an assumption on the scales of the metrics $d_S$, $d_A$, and $d_R$ being similar? Otherwise simply adding them doesn't make sense.\n1. Below Theorem 3.1 there is an objective presented. If I understand correctly that objective is the loss version of the equation in Thm 3.1. If so, what's the motivation for the last ($\\beta$-weighted) term?\n\n# Questions / suggestions\n1. In section 2 MDPs are introduced as having a single initial state $s_I$. Do the results not hold if one has an initial state distribution instead?\n1. In the bottom of page 2 it says $s\\models \\top$ if $\\ell(s)\\cap \\top \\ne\\emptyset$. Should it not be $\\top\\subseteq\\ell(s)$ instead? In the form presented $s\\models \\top$ as long as at least _one_ **AP** in $\\top$ is in $\\ell(s)$, which I don't think is what is meant.\n1. In the **Bisimulation** section of page 4, what is meant by \"the largest bisimulation relation\"?\n1. In the paragraph above Lemma 3.2 it says \"along with setting $D$ to $W_{\\overrightarrow{d}}$, yield ...\". This suggests your loss is somehow computing $(1 + \\beta)W_{\\overrightarrow{d}}$, or am I missing something?\n1. In Figure 4(c), why are no learning curves shown for RL policy? Where are the lines for original (DQN) and original (SAC)?\n1. In the paragraph below Algorithm 1 it appears the indicators \"former\" and \"latter\" are swapped.\n1. The authors use the term \"Gumble\", but it should be \"Gumble\".\n1. In the **Latent distributions** paragraph on page 7, it should read \"We emphasi**ze** that this trick alone...\"\n1. In the first paragraph of the conclusion it should say \"Our method overcome**s** the limitations...\"", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BjHbp8pWb5f", "reviewer": "Reviewer_67Kn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper learns a discrete abstraction of the state-action space of the MDP underlying an RL setup. This abstraction has smaller latent space and the distilled RL policy is more tractable for formal verification approaches such as model checking of bisimulation guarantees. The paper uses Waserstien autoencoders to overcome several limitations with the use of VAEs and builds on recent work by Delgrange et al in AAAI'22.", "summary_of_the_review": "This paper might be more suitable for COLT, SODA or STOCS. It is not clear if this paper has enough novelty and relevance to a machine learning audience. The paper is mostly about MDP abstraction. In addition to clarifying the weaknesses identified in this review, the authors are requested to draw a stronger connection to RL. The mismatch in theory and experiments in the paper suggest that this is a work in progress. ", "strength_and_weaknesses": "Strengths:\n\n+ The paper presents a framework for learning RL policy distillations with bisimulation guarantees. This enables formal verification over learned discrete abstraction of continuous environment. This approach improves upon VAE-MDP in terms of learning speed and model performance. \n\nWeaknesses: \n\n- While the paper talks about formal verification, it fails to clearly identify the kind of formal guarantees it aims to establish beyond bisimulation and reachability guarantees. While it makes a theoretical claim of verifying general discounted omega-regular properties, the experimental evaluation fall far short of it. \n\n- The experimental results are limited to toy examples - cartpole, pendulum, mountain car, acrobot and lunarlander. It would be helpful to consider more complex RL environment such as highD control examples.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster