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.

§ INTRODUCTION

Reinforcement learning (RL) is 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., 2021], or effective canal control  [Ren et al., 2021].

RL enables learning high performance controllers by introducing general nonlinear function approximators (such as 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 is 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 al., 2022) offer a valuable framework for doing so.
The distillation is provided with PAC-verifiable 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 on variational proxies to the bisimulation bounds, meaning there is no learning guarantee on the quality of the latent model via its optimization.

Second, variational autoencoders (VAEs) [Kingma & Welling, 2014, Hoffman et al., 2013]

are known to suffer from posterior collapse (e.g., Alemi et al., 2018)
resulting in a deterministic mapping to a 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 a meticulous task, requiring a 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 reconstructed from the latent model operating under the distilled policy.

In contrast to VAEs 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-processing operations.

Those theoretical claims are reflected 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 et 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 deep-RL policies [Bacci & Parker, 2020, Carr et al., 2020] but rely on a known environment model, among other assumptions (e.g., deterministic or discrete environment).
Finally, DeepSynth [Hasanbeig et al., 2021] allows learning a formal model from execution traces, with the different purpose of guiding the agent towards sparse and non-Markovian rewards.

On the latent space training side, WWAEs [Zhang 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 in Wassertein-GANs [Arjovsky et al., 2017].
Similarly to VQ-VAEs [van den Oord et al., 2017] 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, Zang et al., 2022] consider bisimulation metrics to 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 a deterministic MDP and without verifiable confidence measurement.

§ BACKGROUND

In the following, we write $\distributions{\measurableset}$ for the set of measures over (complete, separable metric space) $\measurableset$.
The index of all the notations introduced along the paper is available at the end of the Appendix.

Markov decision processes

(MDPs) are tuples $\mdp = \mdptuple$ where
$\states$ is a set of states; $\actions$, a set of actions; $\probtransitions \colon \states \times \actions \to \distributions{\states}$, a probability transition function that maps the current state and action to a distribution over the next states; $\rewards \colon \states \times \actions \to \R$, a reward function; $\labels \colon \states \to 2^\atomicprops$, a labeling function over a set of atomic propositions $\atomicprops$; and $\sinit \in \states$, the initial state.

If $|\actions| = 1$, $\mdp$ is a fully stochastic process called a Markov chain (MC).
We write $\mdp_\state$ for the MDP obtained when replacing the initial state of $\mdp$ by $\state \in \states$.

An agent interacting in $\mdp$ produces trajectories, i.e., sequences of states and actions $\trajectory = \trajectorytuple{\state}{\action}{T}$
where $\state_0 = \sinit$ and $\state_{t + 1} \sim \probtransitions\fun{\sampledot \mid \state_t, \action_t}$ for $t < T$.
The set of infinite trajectories of $\mdp$ is .

We assume $\atomicprops$ and labels being respectively one-hot and binary encoded.
Given $\labelset{T} \subseteq \atomicprops$, we write $\state \models \labelset{T}$ if $\state$ is labeled with $\labelset{T}$, i.e., $\labels\fun{\state}\cap \labelset{T} \neq \emptyset$, and $\state \models \neg \labelset{T}$ for $\state \not\models \labelset{T}$.

We refer to MDPs with continuous state or action spaces as continuous MDPs.
In that case, we assume $\states$ and $\actions$ are complete separable metric spaces equipped with a Borel $\sigma$-algebra,
and $\labels^{-1}\fun{\labelset{T}}$ is Borel-measurable for any $\labelset{T} \subseteq \atomicprops$.

Policies and stationary distributions. A (memoryless) policy $\policy \colon \states \to \distributions{\actions}$
prescribes which action to choose at each step of the interaction.

The set of memoryless policies of $\mdp$ is $\mpolicies{\mdp}$.
The MDP $\mdp$ and $\policy \in \mpolicies{\mdp}$
induce an MC $\mdp_\policy$

with unique probability measure $\Prob^\mdp_\policy$ on the Borel $\sigma$-algebra over measurable subsets $\varphi \subseteq \inftrajectories{\mdp}$ [Puterman, 1994].
We drop the superscript when the context is clear.

Define $\stationary{\policy}^{t}\fun{\state' \mid \state} =  \Prob^{\mdp_\state}_\policy(\{{\seq{\state}{\infty}, \seq{\action}{\infty}}
\mid \state_t = \state'\})$ as the distribution giving the probability of being in each state of $\mdp_\state$ after $t$ steps.

$B\subseteq \states$ is a bottom strongly connected component (BSCC) of $\mdp_\policy$ if (i) $B$ is a maximal subset satisfying $\stationary{\policy}^t\fun{\state'\mid \state} > 0$  for any  $\state, \state' \in B$ and some $t \geq 0$, and (ii)
$\expectedsymbol{\action \sim \policy\fun{\sampledot \mid \state}} \probtransitions\fun{B \mid \state, \action}=1$ for all $\state \in \states$.
The unique stationary distribution of $B$ is
$\stationary{\policy} \in \distributions{B}$.
We write $\state, \action \sim \stationary{\policy}$ for sampling $\state$ from $\stationary{\policy}$ then $\action$ from $\policy$.
An MDP $\mdp$ is ergodic if for all $\policy \in \mpolicies{\mdp}$, the state space of $\mdp_\policy$ consists of a unique aperiodic BSCC
with $\stationary{\policy} = \lim_{t \to \infty} \stationary{\policy}^t\fun{\sampledot \mid \state}$ for all $\state \in \states$.

Value objectives.
Given $\policy \in \mpolicies{\mdp}$, the value of a state $\state \in \states$ is the expected value of a random variable obtained by running $\policy$ from $\state$.
For a discount factor $\discount \in \mathopen[0, 1\mathclose]$, we consider the following objectives.
(i) Discounted return: we write $\values{\policy}{}{\state} = \expectedsymbol{\policy}^{\mdp_{\state}}\left[ \sum_{t = 0}^{\infty} \discount^t \rewards\fun{\state_t, \action_t} \right]$ for the expected discounted rewards accumulated along trajectories.
The typical goal of an RL agent is to learn a policy $\policy^{\star}$ that maximizes $\values{\policy^{\star}}{}{\sinit}$ through interactions with the (unknown) MDP;
(ii) Reachability: let $\labelset{C}, \labelset{T} \subseteq \atomicprops$, the (constrained) reachability
event is $\until{\labelset{C}}{\labelset{T}} =
\{\, \seq{\state}{\infty}, \seq{\action}{\infty} \,|\, \exists i \in \N, \forall j < i, \state_j \models {\labelset{C}} \wedge \state_i \models {\labelset{T}}\, \} \subseteq \inftrajectories{\mdp}$.

We write $\values{\policy}{\varphi}{\state} = \expectedsymbol{\policy}^{\mdp_{\state}}\left[ \discount^{t^{\star}} \condition{\tuple{\seq{\state}{\infty}, \seq{\action}{\infty}} \, \in \, \varphi} \right]$ for the discounted probability of satisfying $\varphi = \until{\labelset{C}}{\labelset{T}}$,
where $t^{\star}$ is the length of the shortest trajectory prefix that allows satisfying $\varphi$.
this denotes the discounted return of remaining in a region of the MDP where states are labeled with $\labelset{C}$, until visiting for the first time a goal state labeled with $\labelset{T}$, and the return is the binary reward signal capturing this event.
Safety w.r.t. failure states $\labelset{C}$ can be expressed as the safety-constrained reachability to a destination $\labelset{T}$ through $\until{\neg \labelset{C}}{\labelset{T}}$.
Notice that $\values{\policy}{\varphi}{\state} = \Prob_{\policy}^{\mdp_{\state}}\fun{\varphi}$ when $\discount = 1$.

Latent MDP. Given the original (continuous, possibly unknown) environment model $\mdp$, a latent space model is another (smaller, explicit) MDP $\latentmdp = \latentmdptuple$ with state-action space linked to the original one via state and action embedding functions:  $\embed\colon \states \to \latentstates$ and $\embeda\colon \latentstates \times \latentactions \to \actions$.

We refer to $\tuple{\latentmdp, \embed, \embeda}$ as a latent space model of $\mdp$ and $\latentmdp$ as its latent MDP.
Our goal is to learn $\tuple{\latentmdp, \embed, \embeda}$ by optimizing an equivalence criterion between the two models.

We assume that $\distance_{\latentstates}$ is a metric on $\latentstates$, and

write  $\latentmpolicies$ for the set of policies of $\latentmdp$
and ${\latentvaluessymbol{\latentpolicy}{}}$ for the values of running $\latentpolicy \in \latentmpolicies$ in $\latentmdp$.

The latent policy $\latentpolicy$ can be seen as a policy in $\mdp$ (cf. Fig. <ref>): states passed to $\latentpolicy$ are first embedded with $\embed$ to the latent space, then the actions produced by $\latentpolicy$ are executed via $\embeda$ in the original environment.
Let $\state \in \states$, we write
$\latentaction \sim \latentpolicy\fun{\sampledot \mid \state}$ for
$\latentpolicy\fun{\sampledot \mid \embed\fun{\state}}$, then
the reward and next state are respectively given by $\rewards\fun{\state, \latentaction} = \rewards\fun{\state, \embeda\fun{\embed\fun{\state}, \latentaction}}$ and $ \state' \sim \probtransitions\fun{\sampledot \mid \state, \latentaction} = \probtransitions\fun{\sampledot \mid \state, \embeda\fun{\embed\fun{\state}, \latentaction}}$.

Execution of the latent policy $\latentpolicy$ in the original and latent MDPs, and local losses.

Parallel execution of the original RL policy $\policy$ 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 in red.
Distilling $\policy$ into $\latentpolicy$ via flow (b) by minimizing

allows closing the gap between flows (a) and (b).

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 $\stationary{} \in \distributions{\states \times \latentactions}$:

\begin{align*}
\localrewardloss{\stationary{}} &= \expectedsymbol{\state,  \latentaction \sim \stationary{}}\left| \rewards\fun{\state, \latentaction} - \latentrewards\fun{\embed\fun{\state}, \latentaction} \right|, %\\
\localtransitionloss{\stationary{}} &= \expectedsymbol{\state, \latentaction \sim \stationary{}}
%\wassersteindist{\distance_{\latentstates}}{\embed\probtransitions\fun{\sampledot \mid \state, \latentaction}}{\latentprobtransitions\fun{\sampledot \mid \embed\fun{\state}, \latentaction}}
D\fun{\embed\probtransitions\fun{\sampledot \mid \state, \latentaction}, \latentprobtransitions\fun{\sampledot \mid \embed\fun{\state}, \latentaction}}
\end{align*}

where $\embed\probtransitions\fun{\sampledot \mid \state, \latentaction}$ is
the distribution
of drawing
$\state' \sim \probtransitions\fun{\sampledot \mid \state, \latentaction}$ then

embedding $\latentstate' = \embed\fun{\state'}$, and $D$ is a discrepancy measure.
Fig <ref> depicts the losses when states and actions are drawn from a stationary distribution $\stationary{\latentpolicy}$

resulting from running $\latentpolicy \in \latentmpolicies$ in $\mdp$.

In this work, we focus on the case where $D$ is the Wasserstein distance $\wassersteinsymbol{\distance_{\latentstates}}$:

given two distributions $P, Q$ over a measurable set $\measurableset$
equipped with a metric $\distance$,

$\wassersteinsymbol{\distance}$ is the solution of the optimal transport (OT) from $P$ to $Q$, i.e.,

the minimum cost of changing $P$ into $Q$ [Villani, 2009]:
\wassersteindist{\distance}{P}{Q} = \inf_{\coupling \in \couplings{P}{Q}} \expectedsymbol{x, y \sim \coupling} \distance\fun{x, y},

$\couplings{P}{Q}$ being the set of all couplings of $P$ and $Q$.

The Kantorovich duality yields

%\wassersteindist{\distance}{P}{Q} = \sup_{f \in \Lipschf{\distance}} \int f\fun{x} \, dP\fun{x} - \int f\fun{y} \, dQ\fun{y},
\wassersteindist{\distance}{P}{Q} = \sup_{f \in \Lipschf{\distance}} \expectedsymbol{x \sim P} f\fun{x}  - \expectedsymbol{x \sim Q} f\fun{y}
%\wassersteindist{\distance}{P}{Q} = \sup_{f \in \Lipschf{\distance}}\, \expectedsymbol{x \sim P} f\fun{x} - \expectedsymbol{y \sim Q} f\fun{y} %\wassersteindist{\distance}{P}{Q} = \sup_{f \in \Lipschf{\distance}}\; \expectedsymbol{x \sim P} f\fun{x}  - \expectedsymbol{x \sim Q} f\fun{y}
\label{eq:wasserstein-dual}

where $\Lipschf{\distance}$ is the set of 1-Lipschiz functions.

Local losses are related to a well-established behavioral equivalence between transition systems, called bisimulation.

A bisimulation $\bisimulation$
on $\mdp$ is a behavioral equivalence between states $\state_1, \state_2 \in \states$

so that, $\state_1 \, \bisimulation \, \state_2$ iff
(i) $\probtransitions(T \mid \state_1, \action) = \probtransitions(T \mid \state_2, \action)$,
(ii) $\labels(\state_1) = \labels(\state_2)$,

and (iii) $\rewards(\state_1, \action) = \rewards(\state_2, \action)$

for each action $\action \in \actions$ and (Borel measurable) equivalence class $T \in \states / \bisimulation$.
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 $\atomicprops$ or rewards.
The relation can be extended to compare two MDPs (e.g., $\mdp$ and $\latentmdp$) by considering the disjoint union of their state space.
We denote the largest bisimulation relation by $\sim$.

Characterized by a logical family of functional expressions

derived from a logic $\logic$,
bisimulation pseudometrics [Desharnais et al., 2004] generalize the notion of bisimilariy.

More specifically, given a policy $\policy \in \mpolicies{\mdp}$, we consider a family $\functionalexpr$ of real-valued functions parameterized by a discount factor $\discount$ and defining the semantics of $\logic$ in $\mdp_\policy$.
Such functional expressions allow to formalize discounted properties such as reachability, safety, as well as general $\omega$-regular specifications

[Chatterjee et al., 2010]
and may include rewards as well [Ferns et al., 2014].

The pseudometric $\bidistance_{\policy} $
is defined as the largest behavioral difference $\bidistance_\policy\fun{\state_1, \state_2} = \sup_{f \in \functionalexpr} \left| f\fun{\state_1} - f\fun{\state_2} \right|$,

and its kernel is bisimilarity: $\bidistance_{\policy}\fun{\state_1, \state_2} = 0$ iff $\state_1 \sim \state_2$.

In particular, value functions are Lipschitz-continuous w.r.t. $\bidistance_{\policy}$:

$| \values{\policy}{\scalebox{.95}{$\cdot$}}{\state_1} - \values{\policy}{\scalebox{.95}{$\cdot$}}{\state_2}| \leq
%\nicefrac{1}{(1-\discount)^{K_{\varphi}}} \cdot \bidistance^{\varphi}_{\policy}\fun{\state_1, \state_2}$,
K \bidistance_{\policy}\fun{\state_1, \state_2}$,

where $K$ is $\nicefrac{1}{\fun{1 - \discount}}$ if rewards are included in $\functionalexpr$ and $1$ otherwise.

To ensure the upcoming bisimulation guarantees, we make the following assumptions:

MDP $\mdp$ is ergodic, $\images{\rewards}$ is a bounded space scaled in $\left[\nicefrac{-1}{2}, \nicefrac{1}{2} \right]$, and
the embedding function preserves the labels, i.e., $\embed\fun{\state} = \latentstate \implies \labels\fun{\state} = \latentlabels\fun{\latentstate}$ for $\state \in \states$, $\latentstate \in \latentstates$.

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 on this setting).

Bisimulation bounds [Delgrange et al., 2022].

$\mdp$ being set over continuous spaces with possibly unknown dynamics, evaluating $\bidistance$ can turn out to be particularly arduous, if not intractable.
A solution is to evaluate the original and latent model bisimilarity via local losses:
$\latentpolicy \in \latentmpolicies$, assume $\latentmdp$ is discrete, then given the induced stationary distribution $\stationary{\latentpolicy}$ in $\mdp$, let $\state_1, \state_2 \in \states$ with $\embed\fun{\state_1} = \embed\fun{\state_2}$:

\begin{align}
\expectedsymbol{\state \sim \stationary{\latentpolicy}} \bidistance_{{\latentpolicy}}\fun{\state, \embed\fun{\state}} &\leq %K_{\varphi}
\frac{\localrewardloss{\stationary{\latentpolicy}} + \discount \localtransitionloss{\stationary{\latentpolicy}}}{1 - \discount}, &
\bidistance_{\latentpolicy}\fun{\state_1, \state_2} &\leq \Big( \frac{\localrewardloss{\stationary{\latentpolicy}} +\discount \localtransitionloss{\stationary{\latentpolicy}}}{1 - \discount}\Big) \fun{\stationary{\latentpolicy}^{-1}\fun{\state_1} + \stationary{\latentpolicy}^{-1}\fun{\state_2}}.
\label{eq:bidistance-bound}
\end{align}

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 a latent model where the behaviors of the agent executing $\latentpolicy$ can be formally verified, and the bounds offer a confidence metric allowing to lift the guarantees obtained this way back to the original model $\mdp$, when the latter operates under $\latentpolicy$. We show in the following how to learn a latent space model by optimizing the aforementioned bounds, and distill policies $\policy \in \mpolicies{\mdp}$ obtained via any RL technique to a latent policy $\latentpolicy \in \latentmpolicies$.

§ WASSERSTEIN AUTO-ENCODED MDPS

Fix $ \latentmdp_{\decoderparameter} = \tuple{\latentstates, \latentactions, \latentprobtransitions_{\decoderparameter}, \latentrewards_{\decoderparameter}, \latentlabels, \atomicprops, \zinit}$ and $\tuple{\latentmdp_{\decoderparameter}, \embed_{\encoderparameter}, \embeda_{\decoderparameter}}$ as a latent space model of $\mdp$ parameterized by $\encoderparameter$ and $\decoderparameter$.
Our method relies on learning a behavioral model

of $\mdp$ from which we can retrieve the latent space model

and distill $\pi$.
This can be achieved via the minimization of a suitable discrepancy between $\stationarydecoder$ and

VAE-MDPs optimize a lower bound on the likelihood of the dynamics of $\mdp_\policy$ using the Kullback-Leibler divergence,

yielding (i) $\latentmdp_{\decoderparameter}$, (ii) a distillation  $\latentpolicy_{\decoderparameter}$ of $\policy$, and (iii) $\embed_{\encoderparameter}$ and $\embeda_{\decoderparameter}$.

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

§.§ The Objective Function

Assume that $\states$, $\actions$, and $\images{\rewards}$ are respectively equipped with metrics $\distance_{\states}$, $\distance_{\actions}$, and $\distance_{\rewards}$,

we define the  raw transition distance metric $\transitiondistance$ as the component-wise sum of distances between states, actions, and rewards occurring of along transitions:
\tracedistance\fun{\tuple{\state_1, \action_1, \reward_1, \state'_1}, \tuple{\state_2, \action_2, \reward_2, \state'_2}} = \distance_\states\fun{\state_1, \state_2} + \distance_{\actions}\fun{\action_1, \action_2} + \distance_{\rewards}\fun{\reward_1, \reward_2} + \distance_{\states}\fun{\state_1', \state'_2}.\notag

Trace distributions.
The raw distance $\tracedistance$ allows to reason about transitions, we thus consider the distribution over transitions which occur along trajectories of length $T$ to compare the dynamics of the original and behavioral models:

\begin{align*}
\mathcal{D}_\policy\left[ T \right]\fun{\state, \action, \reward, \state'} &= \frac{1}{T} \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\state \mid \sinit} \cdot \policy\fun{\action \mid \state} \cdot \probtransitions\fun{\state' \mid \state, \action} \cdot \condition{=}\fun{\reward, \rewards\fun{\state, \action}}, \text{ and} \\
\mathcal{P}_\decoderparameter[T]\fun{\state, \action, \reward, \state'} &= \frac{1}{T} \sum_{t = 1}^{T} \expectedsymbol{\seq{\state}{t}, \seq{\action}{t - 1}, \seq{\reward}{t - 1} \sim \decoder[t]}{\condition{=}\fun{\tuple{\indexof{\state}{t - 1}, \indexof{\action}{t - 1}, \indexof{\reward}{t - 1}, \indexof{\state}{t}}, \tuple{\state, \action, \reward, \state'}}},
\end{align*}

where $\decoder[T]$ denotes the distribution over trajectories of length $T$ that include rewards, generated from $\decoder$.
Intuitively, $\nicefrac{1}{T} \cdot \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\state \mid \sinit}$

can be seen as the fraction of the time spent in $\state$
along traces of length $T$, starting from the initial state <cit.>.
Therefore, drawing $\tuple{\state, \action, \reward, \state'} \sim \mathcal{D}_\policy\left[ T \right]$ trivially follows:
it is equivalent to drawing $\state$ from $\nicefrac{1}{T} \cdot \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\cdot \mid \sinit}$, then respectively $\action$ and $\state'$ from $\policy\fun{\cdot \mid \state}$ and $\probtransitions\fun{\cdot \mid \state, \action}$, to finally obtain $\reward = \rewards\fun{\state, \action}$.

Given $T \in \N$, our objective is to minimize the Wasserstein between those distributions:

The following Lemma enables optimizing the Wasserstein distance between the original MDP and the behavioral model
when traces are drawn from episodic RL processes or infinite interactions [Huang, 2020].

Assume the existence of a stationary behavioral model $\stationarydecoder = \lim_{T \to \infty} \mathcal{P}_{\decoderparameter}[T]$, then

\begin{equation*}
\lim_{T \to \infty} \wassersteindist{\tracedistance}{\mathcal{D}_{\policy}[T]}{\mathcal{P}_{\decoderparameter}[T]} = \wassersteindist{\tracedistance}{\stationary{\policy}}{\stationarydecoder}.
\end{equation*}

First, note that $\nicefrac{1}{T} \cdot \sum_{t = 1}^T \stationary{\policy}^t\fun{\cdot \mid \sinit}$ weakly converges to $ \stationary{\policy}$ as $T$  goes to $\infty$ <cit.>. The result follows then from .

Remove the following Lemma? (it doesn't add anything)
Limitations. The following Lemma states that optimizing the Wasserstein between these stationary distributions only yields a lower bound on the Wasserstein distance between distributions over traces generated from the input MDP and those generated from the behavioral model, normalized by their size at the limit.

Assume the existence of a stationary behavioral model $\stationarydecoder$ defined as $\lim_{T \to \infty} \mathcal{P}[T]$.

\begin{equation*}
\lim_{T \to \infty}\, \frac{1}{T}\, \wassersteindist{\tracedistance}{\mdp_{\policy}[T]}{\decoder[T]} \geq \wassersteindist{\tracedistance}{\stationary{\policy}}{\stationarydecoder}
\end{equation*}

where $\mdp_{\policy}[T]$ denotes the distribution over $\traces{\mdp_\policy}\fun{T}$, the latter being the set of traces in $\mdp_{\policy}$ of size $T$.

Given $\trace \in \traces{\mdp_\policy}\fun{T}$, denote by $\xrightarrow{t}_{\trace} = \tuple{\state_t, \action_t, \reward_t, \state_{t + 1}, \labeling_{t + 1}}$ the transition occurring at time $t < T$ along $\trace$.
Observe that for any $f \in \Lipschf{\tracedistance}$, the function $f'\colon \traces{\mdp_{\policy}} \to \R, \, \trace \mapsto \sum_{t = 0}^{T - 1} f(\xrightarrow{t}_{\trace})$ is also $1$-Lipschitz
since $f(\xrightarrow{t}_{\trace_1}) - f({\xrightarrow{t}_{\trace_2}}) \leq \tracedistance({\xrightarrow{t}_{\trace_1}, \xrightarrow{t}_{\trace_2}})$ for all $\trace_1, \trace_2 \in \bigcup_{\state \in \states} \traces{\mdp_{\state, \policy}}$, which implies $f'\fun{\trace_1} - f'\fun{\trace_2} \leq \tracedistance\fun{\trace_1, \trace_2}$ by definition of $\tracedistance$.
We prove the result by considering the set $\mathcal{F}_{\Sigma} = \set{f \in \Lipschf{\tracedistance} \colon f\fun{\trace} = \sum_{t = 0}^{T - 1} f(\xrightarrow{t}_{\trace})}$ instead of $\Lipschf{\tracedistance}$:

\begin{align*}
& \frac{1}{T}\, \wassersteindist{\tracedistance}{\mdp_{\policy}[T]}{\decoder[T]}\\
= & \sup_{f \in \Lipschf{\tracedistance}} \, \frac{1}{T} \expectedsymbol{\trace \sim \mdp_{\policy}[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} - \frac{1}{T} \expectedsymbol{\trace \sim \decoder[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} \\
\geq & \sup_{f \in \mathcal{F}_{\Sigma}} \, \frac{1}{T} \expectedsymbol{\trace \sim \mdp_{\policy}[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} - \frac{1}{T} \expectedsymbol{\trace \sim \decoder[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} \tag{since $\mathcal{F}_{\Sigma} \subseteq \Lipschf{\tracedistance}$} \\
= & \sup_{f \in \Lipschf{\tracedistance}} \frac{1}{T} \expected{\trace \sim \mdp_{\policy}[T]}{\sum_{t = 0}^{T - 1} f\fun{\state_t, \action_t, \reward_t, \state_{t + 1}, \labeling_{t + 1}}} - \frac{1}{T} \expected{\trace \sim \decoder[T]}{\sum_{t = 0}^{T - 1} f\fun{\state_t, \action_t, \reward_t, \state_{t + 1}, \labeling_{t + 1}}} \\
= & \sup_{f \in \Lipschf{\tracedistance}} \,  \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \mathcal{D}_{\policy}[T]}{f\fun{\state, \action, \reward, \state', \labeling'}} - \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \mathcal{P}_{\decoderparameter}[T]}{f\fun{\state, \action, \reward, \state', \labeling'}}\\
= & \wassersteindist{\tracedistance}{\mathcal{D}_{\policy}[T]}{\mathcal{P}_{\decoderparameter}[T]}
\end{align*}

We obtain the last line by compactness of $\states, \actions$, and $\rewards$ as well as $\atomicprops$ being finite.
The same development than the proof of Lemma [VAE-MDP, Appendix] and applying the Portmanteau's Theorem allows us to pass from an expectation over $\mdp_{\policy}[T]$ to $\mathcal{D}_\policy[T]$ and $\decoder[T]$ to $\mathcal{P}_{\decoderparameter}[T]$.
Applying Lemma <ref> finally yields the result.

Let $\stationary{\policy}$ be the stationary distribution of $\mdp_\policy$ and $\stationarydecoder$ be the stationary behavioral model defined above, then

\begin{align*}
\lim_{T \to \infty} \frac{1}{T} \, \wassersteindist{\tracedistance}{\mdp_\policy[T]}{\decoder[T]}
%\lim_{T \to \infty} \sup_{f \in \Lipschf{\tracedistance}} \frac{1}{T} \left| \expectedsymbol{\trace \sim \mdp_{\policy}[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} - \expectedsymbol{\trace \sim \decoder} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}}\right|\\
&=\sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \frac{1}{T} \fun{ \expectedsymbol{\trace \sim \mdp_{\policy}[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} - \expectedsymbol{\trace \sim \decoder[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} }.
\end{align*}

Observe that

\begin{equation*}
\lim_{T \to \infty} \frac{1}{T}\, \wassersteindist{\tracedistance}{\mdp_\policy[T]}{\decoder[T]} = \lim_{T \to \infty} \sup_{f \in \Lipschf{\tracedistance}} \frac{1}{T} \fun{ \expectedsymbol{\trace \sim \mdp_{\policy}[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} - \expectedsymbol{\trace \sim \decoder[T]} f\fun{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}, \seq{\labeling}{T}} }.
\end{equation*}

Take $f^\star \in \Lipschf{\tracedistance}\,$ and $T \in \N_0$, we have

\begin{align*}
& &\expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f^\star\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f^\star\fun{\trace}}&\leq \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}} \\
& \text{then,} &\lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f^\star\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f^\star\fun{\trace}} & \leq \liminf_{T \to \infty} \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}}. \tag{since this holds for any $T \in \N_{0}$}
\end{align*}

Observe that $f^\star$ is arbitrary in $\Lipschf{\tracedistance}$.
Thus, taking $f^\star = \arg \sup_{f \in \Lipschf{\tracedistance}} \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}}$ yields in particular

\begin{align}
\sup_{f \in \Lipschf{\tracedistance}} \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}} & \leq \liminf_{T \to \infty} \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}}. \label{eq:limsup-suplim-leq}
\end{align}

Now, let $\delta > 0$. Then, there is a step number $T_0 \in \N_0$ from which, for all $t^\star \geq T_0$,

\begin{equation}
\left| \expected{\trace \sim \mdp_{\policy}[t^\star]}{\frac{1}{t^\star} \, f\fun{\trace}} - \expected{\trace \sim \decoder[t^\star]}{\frac{1}{t^\star} \, f\fun{\trace}}  - \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationary{\policy}}{f\fun{\state, \action, \reward, \state', \labeling'}} + \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationarydecoder}{f\fun{\state, \action, \reward, \state', \labeling'}}\right| \leq \delta \label{eq:limit-delta}
\end{equation}

whenever $f \in \Lipschf{\tracedistance}$, by definition of $\lim_{T \to \infty} \frac{1}{T} \fun{\expectedsymbol{\trace \sim \mdp_{\policy}[T]}\, f\fun{\trace} - \expectedsymbol{\trace \sim \decoder[T]}\, f\fun{\trace}}$

and Lem. <ref>.
For all $n \geq T_0$, there is a $1$-Lipschitz function $f_n \in \Lipschf{\tracedistance}$ with

\begin{align*}
& \sup_{f \in \Lipschf{\tracedistance}} \, \expected{\trace \sim \mdp_{\policy}[n]}{\frac{1}{n} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n]}{\frac{1}{n} \, f\fun{\trace}} - \delta \\
\leq & \expected{\trace \sim \mdp_{\policy}[n]}{\frac{1}{n} \, f_n\fun{\trace}} - \expected{\trace \sim \decoder[n]}{\frac{1}{n} \, f_n\fun{\trace}}
\tag{e.g., take $f_n = \arg \sup_{f \in \Lipschf{\tracedistance}} \, \expected{\trace \sim \mdp_{\policy}[n]}{\frac{1}{n} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n]}{\frac{1}{n} \, f\fun{\trace}}$}\\
\leq & \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationary{\policy}}{ f_n\fun{\state, \action, \reward, \state', \labeling'}} - \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationarydecoder}{f_n\fun{\state, \action, \reward, \state', \labeling'}} + \delta \tag{by Eq.~\ref{eq:limit-delta}, because $n \geq T_0$} \\
= & \lim_{T \to \infty} \, \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T}\, f_n\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f_n\fun{\trace}} + \delta \\
\leq & \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace\sim \mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace\sim \decoder[T]}{\frac{1}{T}\, f\fun{\trace}} + \delta \tag{by definition of $\sup_{f \in \Lipschf{\tracedistance}}$}
\end{align*}

In particular, take $n^\delta = T_0 + \lceil \frac{1}{\delta} \rceil$, we have

\begin{align*}
&& \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}} - \delta
&\leq \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f\fun{\trace}} + \delta\\
&\equiv & \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}} \phantom{- \delta}
&\leq \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f\fun{\trace}} + 2 \delta
\end{align*}

Since this holds for any arbitrary $\delta > 0$, we have

\begin{align*}
&& \limsup_{\delta \to 0} \, \sup_{f \in \Lipschf{\tracedistance}} \, \expected{\trace \sim \mdp_{\policy}[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n^\delta]}{\frac{1}{n^{\delta}} \, f\fun{\trace}}
&\leq \lim_{\delta \to 0} \, \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{ \frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{ \frac{1}{T} \, f\fun{\trace}} + 2 \delta
%&\equiv    & \limsup_{T \to \infty} \, \sup_{f \in \Lipschf{\tracedistance}} \, \expected{\trace \sim \mdp_{\policy}[n]}{\frac{1}{n} \, f\fun{\trace}} - \expected{\trace \sim \decoder[n]}{\frac{1}{n} \, f\fun{\trace}  } &\leq \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace\sim\mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace\sim\decoder[T]}{\frac{1}{T}\, f\fun{\trace}} \tag{because $n^{\delta} \to \infty$ when $\delta \to 0$}
\end{align*}

Notice that $n^{\delta} \to \infty$ when $\delta \to 0$. Therefore, take $T = n^\delta$

\begin{align}
&& \limsup_{T \to \infty} \, \sup_{f \in \Lipschf{\tracedistance}} \, \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f\fun{\trace}  } &\leq \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace\sim\mdp_{\policy}[T]}{\frac{1}{T}\, f\fun{\trace}} - \expected{\trace\sim\decoder[T]}{\frac{1}{T}\, f\fun{\trace}} \label{eq:limsup-to-suplim}
\end{align}

By Eq. <ref> and <ref> $\liminf_{T \to \infty}$ and $\limsup_{T \to \infty}$ coincides, which finally yields

\begin{equation*}
\lim_{T \to \infty} \sup_{f \in \Lipschf{\tracedistance}} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f\fun{\trace}} =  \sup_{f \in \Lipschf{\tracedistance}} \, \lim_{T \to \infty} \expected{\trace \sim \mdp_{\policy}[T]}{\frac{1}{T} \, f\fun{\trace}} - \expected{\trace \sim \decoder[T]}{\frac{1}{T} \, f\fun{\trace}}.
\end{equation*}

Let $\stationary{\policy}$ be the stationary distribution of $\mdp_\policy$ and $\stationarydecoder$ be the stationary behavioral model defined above, then

\begin{equation*}
\lim_{T \to \infty} \frac{1}{T} \, \wassersteindist{\tracedistance}{\mdp_\policy[T]}{\decoder[T]} = \wassersteindist{\tracedistance}{\stationary{\policy}}{\stationarydecoder}.% \sup_{f \in \Lipschf{\tracedistance}} \, \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationary{\policy}} f\fun{\state, \action, \reward, \state', \labeling'} - \expectedsymbol{\state, \action, \reward, \state', \labeling' \sim \stationarydecoder} f\fun{\state, \action, \reward, \state',  \labeling'}.
\end{equation*}

§.§ The Objective Function

Optimizing the regularizers

is enabled by the dual form of the OT: we introduce two parameterized networks, $\steadystatenetwork$ and $\transitionlossnetwork$, constrained to be $1$-Lipschitz and trained to attain the supremum of the dual:

\[
\steadystateregularizer{\policy}\fun{\wassersteinparameter} = \max_{\wassersteinparameter %\colon \steadystatenetwork \in \Lipschf{\tracedistance}
} \; \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentaction \sim \actionencoder\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action}}\expectedsymbol{\latentstate^{\star} \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}} \steadystatenetwork\fun{\embed_{\encoderparameter}\fun{\state}, \latentaction, \latentstate^{\star}}
- \expectedsymbol{\latentvariable, \latentaction', \latentvariable' \sim \latentstationaryprior}\steadystatenetwork\fun{\latentvariable, \latentaction', \latentvariable'}
\]

\[
\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter} = \max_{\wassersteinparameter
%\colon \transitionlossnetwork \in \Lipschf{\distance_{\scriptscriptstyle \latentstates}}
} \expectedsymbol{\state, \action, \state' \sim \stationary{\policy}} \expectedsymbol{\latentstate, \latentaction, \latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action, \state'}}\Big[{\transitionlossnetwork\fun{\state, \action, \latentstate, \latentaction, \latentstate'} - \! \expectedsymbol{\latentstate^{\star} \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \transitionlossnetwork\fun{\state, \action, \latentstate, \latentaction, \latentstate^{\star}}}\Big]
\]

Details to derive this tractable form of $\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter}$ are in Appendix <ref>.

The networks are 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 $1$ everywhere
(we show in Appendix <ref> this is still valid for relaxations of discrete spaces).

The final learning process is presented in Algorithm <ref>.

Distances are depicted by red dotted lines.

§.§ Discrete Latent Spaces

To enable the verification of latent models supported by the bisimulation guarantees of Eq. <ref>, we focus on the special case of discrete latent space models.
Our approach relies on continuous relaxation of discrete random variables, regulated by some temperature parameter(s) $\temperature$: discrete random variables are retrieved as $\temperature \to 0$, which amounts to applying a rounding operator.
For training, we use the temperature-controlled 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. <ref>.

State encoder.
We work with a binary representation of the latent states.
First, this induces compact networks, able to deal with a large discrete space via a tractable number of parameter variables.
But most importantly, this
ensures that Assumption <ref> is satisfied:
let $n = \log_2 | \latentstates|$,
we reserve $\left| \atomicprops \right|$ bits in $\latentstates$ and each time $\state \in \states$ is passed to $\embed_{\encoderparameter}$, $n- \left|\atomicprops\right|$ bits are produced and concatenated with $\labels\fun{\state}$, ensuring a perfect reconstruction of the labels and further bisimulation bounds.

To produce Bernoulli variables, $\embed_{\encoderparameter}$ deterministically maps $\state$ to a latent code $\vz$, passed to the Heaviside $H\fun{\vz} = \condition{\vz > 0}$.
We train $\embed_{\encoderparameter}$ by using the smooth approximation $H_{\temperature}\fun{\vz} = \sigma\fun{\nicefrac{2\vz}{\temperature}}$, satisfying $H = \lim_{\temperature \to 0} H_{\temperature}$.

Latent distributions.
Besides the discontinuity of their latent image space, a major challenge of optimizing over discrete distributions is sampling, required to be a differentiable operation.
We circumvent this by using concrete distributions [Jang et al., 2017, Maddison et al., 2017]:
the idea is to sample reparameterizable random variables from $\temperature$-parameterized distributions, and applying a differentiable, nonlinear operator in downstream.
We use the Gumbel softmax trick to sample from distributions over (one-hot encoded) latent actions ($\actionencoder$, $\latentpolicy_{\decoderparameter}$).
For binary distributions ($\latentprobtransitions_{\decoderparameter}$, $\latentstationaryprior$), each relaxed Bernoulli with logit $\alpha$ is retrieved by drawing a logistic random variable located in $\nicefrac{\alpha}{\temperature}$ and scaled to $\nicefrac{1}{\temperature}$, then applying a sigmoid in downstream.
We emphasize that this trick alone (as used by Corneil 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 <ref>).

[->,shorten >=1pt,auto,node distance=2.5cm,bend angle=45, scale=0.6, font=]
state=[draw,circle,text centered,minimum size=7mm,text width=4mm]
act=[fill,circle,inner sep=1pt,minimum size=1.5pt, node distance=1cm]
empty=[text centered, text width=15mm]
[state] (0) at (0, 0)     $\latentstate_0$;
[state] (1) at (2.5, .8)   $\latentstate_1$;
[empty] (l1) at (1.2, 1.4)  $\{\textit{\color{OliveGreen}goal}\}$;
[state] (2) at (2.5, -.8)  $\latentstate_2$;
[empty] (l2) at (1., -1.3) $\{\textit{\color{BrickRed}unsafe}\}$;
[state] (3) at (-2.5, 0)    $\latentstate_3$;
[empty] (l2) at (-2.5, 1)   $\{\textit{\color{BrickRed}unsafe}\}$;
[->] (0) – node[above]$\nicefrac{1}{2}$ (1);
[->] (0) – node[below]$\nicefrac{1}{4}$ (2);
[->] (0) – node[above]$\nicefrac{1}{4}$ (3);
[->] (1) edge[out=335,in=25,loop] node[left]$1$ (1);
[->] (2) edge[out=335,in=25,loop] node[left]$1$ (2);
[->] (3) edge[out=215,in=155,loop] node[right]$1$ (3);

Markov Chain with four states
; labels are drawn next to their state.

Let $\latentmdp$ be the discrete MC of Fig. <ref>
(the labels of $\latentmdp$ are drawn next to their state).
Let $\latentmdp$ be the discrete MC of Fig. <ref>.
In one-hot, $\atomicprops = \{\textit{goal}\!: \tuple{{\color{OliveGreen}1}, 0}, \textit{unsafe}\!: \tuple{0, {\color{BrickRed}1}}\}$.
We assume that $3$ bits are used for the (binary) state space, with $\latentstates = \{\latentstate_0 \!: \tuple{{0, 0}, 0}, \latentstate_1\!: \tuple{{\color{OliveGreen}1}, 0, 0}, \latentstate_2\!: \tuple{0, {\color{BrickRed}1}, 0}, \latentstate_3 \!: \tuple{0, {\color{BrickRed}1}, 1}\}$ (the two first bits are reserved for the labels).
Considering each bit as being independent is not sufficient to learn $\latentprobtransitions$:

the optimal estimation $\latentprobtransitions_{\decoderparameter^{\star}}\fun{\sampledot \mid \latentstate_0}$ is in that case represented by the independent Bernoulli vector $\mathbf{b} = \tuple{\nicefrac{1}{2}, \nicefrac{1}{2}, \nicefrac{1}{4}}$, giving the probability to go from $\latentstate_0$ to each bit independently.
This yields a poor estimation of the actual transition function:

\latentprobtransitions_{\decoderparameter^{\star}}\fun{\latentstate_0 \mid \latentstate_0} = (1 - \mathbf{b}_1) \cdot (1 - \mathbf{b}_2) \cdot (1 - \mathbf{b}_3) = %0.1875, \,
\latentprobtransitions_{\decoderparameter^{\star}}\fun{\latentstate_1 \mid \latentstate_0} = \mathbf{b}_1 \cdot (1 - \mathbf{b}_2) \cdot (1 - \mathbf{b}_3) = %0.1875, \,
\latentprobtransitions_{\decoderparameter^{\star}}\fun{\latentstate_2 \mid \latentstate_0} = (1 - \mathbf{b}_1) \cdot \mathbf{b}_2 \cdot (1 - \mathbf{b}_3) = %0.1875,\,
\nicefrac{3}{16}, \,
\latentprobtransitions_{\decoderparameter^{\star}}\fun{\latentstate_3 \mid \latentstate_0} = (1 - \mathbf{b}_1) \cdot \mathbf{b}_2 \cdot \mathbf{b}_3 = \nicefrac{1}{16}

We consider instead relaxed multivariate Bernoulli distributions by decomposing $P \in \distributions{\latentstates}$
as a product of conditionals:

P\fun{\latentstate} = \prod_{i=1}^{n} P\fun{\latentstate_{i} \mid \latentstate_{1 \colon i-1}}

where $\latentstate_i$ is the $i^\text{th}$ entry (bit) of $\latentstate$.
We learn such distributions by introducing

a masked autoregressive flow (MAF, Papamakarios et al., 2017) for relaxed Bernoullis via the recursion:

$\latentstate_i = \sigma\fun{\nicefrac{l_i + \alpha_i}{\temperature}}$, where $l_i \sim \logistic{0}{1}$, $\alpha_i = f_i\fun{\latentstate_{1\colon i - 1}}$, and $f$ is a MADE [Germain et al., 2015], a 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 $\latentprobtransitions_{\decoderparameter}$ and the dynamics
related to

the labels in $\latentstationaryprior$.

We fix the logits of the remaining $n - \left| \atomicprops \right|$ bits to $0$ to allow for a fairly distributed latent space.

§ EXPERIMENTS

objective: reconstruction loss, transition and steady-state regularizers

PAC local losses approximation for an error of at most $10^{-2}$ and probability confidence $0.955$

Episode return obtained when executing the distilled policy in the original MDP
(averaged over $30$ episodes)

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.

We evaluate the quality of latent space models learned and policies distilled through .
To do so, we first trained deep-RL policies (DQN, Mnih et al., 2015 on discrete, and SAC, Haarnoja et al., 2018 on continuous action spaces) for various OpenAI benchmarks [Brockman et al., 2016], which we then distill via our approach (Figure <ref>).

We thus evaluate

(a) the training metrics,

(b) the abstraction and representation quality via PAC local losses upper bounds [Delgrange et 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 to reproduce our results is in Appendix <ref>.

Learning metrics. The objective (Fig. <ref>) is a weighted sum of the reconstruction loss and the two Wasserstein regularizers.

The choice of $\beta$ defines the optimization direction.

Posterior collapse is not observed, naturally avoided in WAEs [Tolstikhin et al., 2018],
which reflects
that the latent space is consistently distributed (see Appendix <ref> for a discussion and a concrete illustration of collapsing issues occurring in VAE-MDPs).
In contrast to VAEs (cf. Appendix <ref>), WAEs indeed naturally avoid posterior collapse [Tolstikhin et al., 2018], indicating that the latent space is consistently distributed.

Optimizing the objective (Fig. <ref>) effectively allows minimizing the local losses (Fig. <ref>) and recovering the performance of the original policy (Fig. <ref>).

Local losses.

For V- and WAEs, we formally evaluate PAC upper bounds on $\localrewardloss{\stationary{\latentpolicy_{\decoderparameter}}}$ and $\localtransitionloss{\stationary{\latentpolicy_{\decoderparameter}}}$ via the algorithm of [Delgrange et al., 2022] (Fig <ref>).
The lower the local losses, the closer $\mdp$ and $\latentmdp_{\decoderparameter}$ are in terms of behaviors induced by $\latentpolicy_{\decoderparameter}$ (cf. Eq. <ref>).
In VAEs, the losses are evaluated on a transition function $\hat{\probtransitions}$ 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., Bazille et al., 2020, Corneil et al., 2018).
We thus also report the metrics for $\hat{\probtransitions}$.

Our bounds quickly converge to close values in general for $\latentprobtransitions_{\decoderparameter}$  and $\hat{\probtransitions}$, whereas for VAEs, the convergence is slow and unstable, with $\hat{\probtransitions}$ offering better bounds.
We emphasize that WAEs do not require this additional reconstruction step to obtain losses that can be leveraged to assess the quality of the model, in contrast to VAEs, where learning $\latentprobtransitions_{\decoderparameter}$ was performed via overly restrictive distributions, leading to poor estimation in general (cf. Ex. <ref>).

Finally, when the distilled policies offer comparable performance (Fig. <ref>), our bounds are either close to or better than those of VAEs.

The bisimulation guarantees (Eq. <ref>) are only valid for $\latentpolicy_{\decoderparameter}$, the policy under which formal properties can be verified.
It is crucial that $\latentpolicy_{\decoderparameter}$ achieves performance close to $\policy$, the original one, when deployed in the RL environment.
We evaluate the performance of $\latentpolicy_{\decoderparameter}$ via the undiscounted episode return $\episodereturn{\latentpolicy_{\decoderparameter}}$ obtained by running $\latentpolicy_{\decoderparameter}$ in the original model $\mdp$.

We observe that $\episodereturn{\latentpolicy_{\decoderparameter}}$ approaches faster the original performance $\episodereturn{\policy}$ 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. <ref>) 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 new representation, albeit discrete and hence coarser, still achieves effective performance since

keeps only what is important to preserve behaviors, and thus values.

Furthermore, the distillation can remove some non-robustness obtained during RL:

$\latentpolicy_{\decoderparameter}$ prescribes the same actions for bisimilarly close states, whereas this is not necessarily the case for $\policy$.

Formal verification.
To formally verify $\latentmdp_\decoderparameter$, we implemented a 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., 2021, Kwiatkowska et al., 2022).
We verify time-to-failure properties $\varphi$, often used to check the failure rate of a system [Pnueli, 1977] by measuring whether the agent fails before the end of the episode.
Although simple, such properties highlight the applicability of our approach on reachability events, which are building blocks to verify MDPs
(Baier & Katoen, 2008; cf. Appendix <ref> for a discussion).
(Baier & Katoen, 2008; cf. Appendix <ref>).
In particular, we checked whether the agent reaches an unsafe position or angle (CartPole, LunarLander), 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 <ref>: 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
$\norm{V_{\latentpolicy_{\decoderparameter}}} = |\values{\latentpolicy_{\decoderparameter}}{}{\sinit} - \latentvalues{\latentpolicy_{\decoderparameter}}{}{\zinit}|$ obtained by executing $\latentpolicy_{\decoderparameter}$ in $\mdp$ and $\latentmdp_{\decoderparameter}$

($\values{\latentpolicy_{\decoderparameter}}{}{\sampledot}$ is averaged while $\latentvalues{\latentpolicy_{\decoderparameter}}{}{\sampledot}$ is formally computed).

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 crux of how we enable verification techniques like model checking.
In particular, bisimulation guarantees mean that reachability probabilities in the latent MDP compared to those in the original one are close.
Furthermore, the value difference of (omega-regular) properties (formulated through mu-calculus) obtained in the two models is bounded by this distance (cf. Sect. <ref> and Chatterjee et al., 2010).

Reachability is the key ingredient to 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. Baier & Katoen, 2008), or (ii) checking the reachability to the part of the state space of a product of the MDP with a 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.

§ CONCLUSION

We presented  WAE-MDPs, a framework for learning formally verifiable distillations of RL policies 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 to 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 noise-sensitive RL policies (such as robotics simulations, 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 $\omega$-regular properties is theoretically allowed in our model via the rechability to components of standard constructions based on automata products (e.g., ), and discounted games algorithms [Chatterjee et al., 2010].

Beyond distillation, our results, supported by Thm. <ref>, suggest that our WAE-MDP can be used as a general latent space learner for RL, further opening possibilities to combine RL and formal methods online when no formal model is a priori known, and  address this way safety in  RL with guarantees.

§.§.§ Reproducibility Statement

We referenced in the main text the Appendix parts presenting the proofs or additional details of every claim, Assumption, Lemma, and Theorem occurring in the paper.
In addition, Appendix <ref> is dedicated to the presentation of the setup, hyperparameters, and other extra details required for reproducing the results of Section <ref>.
We provide the source code of the implementation of our approach in Supplementary material [available at <https://github.com/florentdelgrange/wae_mdp>], and we also provide the models saved during training that we used for model checking (i.e., reproducing the results of Table <ref>).
Additionally, we present in a notebook () 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 “SAILor” project (G030020N).
We thank Raphael Avalos for his valuable feedback during the preparation of this manuscript.

[Alamdari et al., 2020]
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.
URL <https://doi.org/10.34727/2020/isbn.978-3-85448-042-6_21>.

[Alemi et al., 2018]
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ässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of
Proceedings of Machine Learning Research, pp.  159–168. PMLR,
URL <http://proceedings.mlr.press/v80/alemi18a.html>.

[Alshiekh et al., 2018]
Mohammed Alshiekh, Roderick Bloem, Rüdiger Ehlers, Bettina
Könighofer, Scott Niekum, and Ufuk Topcu.
Safe reinforcement learning via shielding.
In Sheila A. McIlraith 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.

[Arjovsky et al., 2017]
Martín Arjovsky, Soumith Chintala, and Léon Bottou.
Wasserstein generative adversarial networks.
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 <http://proceedings.mlr.press/v70/arjovsky17a.html>.

[Bacci & Parker, 2020]
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.  231–248. Springer, 2020.
URL <https://doi.org/10.1007/978-3-030-57628-8_14>.

[Baier & Katoen, 2008]
Christel Baier and Joost-Pieter Katoen.
Principles of model checking.
MIT Press, 2008.
ISBN 978-0-262-02649-9.

[Baier et al., 2016]
Christel Baier, Stefan Kiefer, Joachim Klein, Sascha Klüppelholz, David
Müller, and James Worrell.
Markov chains and unambiguous büchi automata.
In 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–42. Springer, 2016.
URL <https://doi.org/10.1007/978-3-319-41528-4_2>.

[Bazille et al., 2020]
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 II, volume 12225 of Lecture
Notes in Computer Science, pp.  304–326. Springer, 2020.
URL <https://doi.org/10.1007/978-3-030-53291-8_17>.

[Bousquet et al., 2017]
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.

[Brockman et al., 2016]
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman,
Jie Tang, and Wojciech Zaremba.
Openai gym.
CoRR, abs/1606.01540, 2016.
URL <http://arxiv.org/abs/1606.01540>.

[Carr et al., 2020]
Steven Carr, Nils Jansen, and Ufuk Topcu.
Verifiable rnn-based policies for pomdps under temporal logic
In Christian Bessiere (ed.), Proceedings of the Twenty-Ninth
International Joint Conference on Artificial Intelligence, IJCAI 2020,
pp.  4121–4127. ijcai.org, 2020.
URL <https://doi.org/10.24963/ijcai.2020/570>.

[Castro et al., 2021]
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland.
Mico: Improved representations via sampling-based state similarity
for markov decision processes.
In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy
Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural
Information Processing Systems 34: Annual Conference on Neural Information
Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp.  30113–30126, 2021.

[Ceusters et al., 2021]
Glenn Ceusters, Román Cantú Rodríguez, Alberte Bouso Garcí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:0 117634, 2021.
ISSN 0306-2619.

[Chatterjee et al., 2010]
Krishnendu Chatterjee, Luca de Alfaro, Rupak Majumdar, and Vishwanath Raman.
Algorithms for game metrics (full version).
Log. Methods Comput. Sci., 60 (3), 2010.
URL <http://arxiv.org/abs/0809.4326>.

[Corneil et al., 2018]
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,
Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of
Proceedings of Machine Learning Research, pp.  1057–1066. PMLR,
URL <http://proceedings.mlr.press/v80/corneil18a.html>.

[Delgrange et al., 2022]
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,
360 (6):0 6497–6505, Jun. 2022.
URL <https://ojs.aaai.org/index.php/AAAI/article/view/20602>.

[Desharnais et al., 2004]
Josée Desharnais, Vineet Gupta, Radha Jagadeesan, and Prakash Panangaden.
Metrics for labelled markov processes.
Theor. Comput. Sci., 3180 (3):0 323–354,
URL <https://doi.org/10.1016/j.tcs.2003.09.013>.

[Fajtl et al., 2020]
Jiri Fajtl, Vasileios Argyriou, Dorothy Monekosso, and Paolo Remagnino.
Latent bernoulli autoencoder.
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,
URL <http://proceedings.mlr.press/v119/fajtl20a.html>.

[Ferns et al., 2014]
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.  319–342. Springer, 2014.
URL <https://doi.org/10.1007/978-3-319-06880-0_17>.

[Gelada et al., 2019]
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G.
Deepmdp: Learning continuous latent space models for representation
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 Research, pp.  2170–2179. PMLR,
URL <http://proceedings.mlr.press/v97/gelada19a.html>.

[Germain et al., 2015]
Mathieu Germain, Karol Gregor, Iain 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://proceedings.mlr.press/v37/germain15.html>.

[Givan et al., 2003]
Robert Givan, Thomas L. Dean, and Matthew Greig.
Equivalence notions and model minimization in markov decision
Artif. Intell., 1470 (1-2):0 163–223, 2003.
URL <https://doi.org/10.1016/S0004-3702(02)00376-4>.

[Gulrajani et al., 2017]
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 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, December 4-9, 2017, Long Beach, CA,
USA, pp.  5767–5777, 2017.

[Haarnoja et al., 2018]
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.
Soft actor-critic: Off-policy maximum entropy deep reinforcement
learning with a stochastic actor.
In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the
35th International Conference on Machine Learning, ICML 2018,
Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of
Proceedings of Machine Learning Research, pp.  1856–1865. PMLR,
URL <http://proceedings.mlr.press/v80/haarnoja18b.html>.

[Hasanbeig et al., 2021]
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 on Artificial Intelligence,
AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial
Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in
Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021,
pp.  7647–7656. AAAI Press, 2021.
URL <https://ojs.aaai.org/index.php/AAAI/article/view/16935>.

[Hensel et al., 2021]
Christian Hensel, Sebastian Junges, Joost-Pieter Katoen, Tim Quatmann, and
Matthias Volk.
The probabilistic model checker storm.
International Journal on Software Tools for Technology
Transfer, 2021.
ISSN 1433-2787.
URL <https://doi.org/10.1007/s10009-021-00633-z>.

[Hoffman et al., 2013]
Matthew D. Hoffman, David M. Blei, Chong Wang, and John W. Paisley.
Stochastic variational inference.
J. Mach. Learn. Res., 140 (1):0 1303–1347,
URL <http://dl.acm.org/citation.cfm?id=2502622>.

[Huang, 2020]
Bojun Huang.
Steady state analysis of episodic reinforcement learning.
In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell,
Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural
Information Processing Systems 33: Annual Conference on Neural Information
Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.

[Jang et al., 2017]
Eric Jang, Shixiang Gu, and Ben Poole.
Categorical reparameterization with gumbel-softmax.
In 5th International Conference on Learning Representations,
ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track
Proceedings. OpenReview.net, 2017.
URL <https://openreview.net/forum?id=rkE3y85ee>.

[Jansen et al., 2020]
Nils Jansen, Bettina Könighofer, Sebastian Junges, Alex Serban, and
Roderick Bloem.
Safe Reinforcement Learning Using Probabilistic Shields (Invited
In Igor Konnov and Laura Kovács (eds.), 31st International
Conference on Concurrency Theory (CONCUR 2020), volume 171 of Leibniz
International Proceedings in Informatics (LIPIcs), pp.  3:1–3:16,
Dagstuhl, Germany, 2020. Schloss Dagstuhl–Leibniz-Zentrum für
ISBN 978-3-95977-160-3.
URL <https://drops.dagstuhl.de/opus/volltexte/2020/12815>.

[Junges et al., 2016]
Sebastian Junges, Nils Jansen, Christian Dehnert, Ufuk Topcu, and
Joost-Pieter 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.
URL <https://doi.org/10.1007/978-3-662-49674-9_8>.

[Kingma & Welling, 2014]
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>.

[Kwiatkowska et al., 2022]
Marta Kwiatkowska, Gethin Norman, and David Parker.
Probabilistic model checking and autonomy.
Annual Review of Control, Robotics, and Autonomous Systems,
50 (1):0 385–410, 2022.
URL <https://doi.org/10.1146/annurev-control-042820-010947>.

[Larsen & Skou, 1989]
Kim Guldstrand Larsen and Arne Skou.
Bisimulation through probabilistic testing.
In Conference Record of the Sixteenth Annual ACM Symposium on
Principles of Programming Languages, Austin, Texas, USA, January 11-13,
1989, pp.  344–352. ACM Press, 1989.
URL <https://doi.org/10.1145/75277.75307>.

[Libin et al., 2020]
Pieter J. 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 Yuxiao Dong, Georgiana Ifrim, Dunja Mladenic, Craig Saunders, and
Sofie Van Hoecke (eds.), Machine Learning and Knowledge Discovery in
Databases. Applied Data Science and Demo Track - 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.
URL <https://doi.org/10.1007/978-3-030-67670-4_10>.

[Littman et al., 2017]
Michael L. Littman, Ufuk Topcu, Jie Fu, Charles Lee Isbell Jr., Min Wen, and
James MacGlashan.
Environment-independent task specifications via GLTL.
CoRR, abs/1704.04341, 2017.
URL <http://arxiv.org/abs/1704.04341>.

[Maddison et al., 2017]
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh.
The concrete distribution: A continuous relaxation of discrete
random variables.
In 5th International Conference on Learning Representations,
ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track
Proceedings. OpenReview.net, 2017.
URL <https://openreview.net/forum?id=S1jE5L5gl>.

[Mnih et al., 2015]
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness,
Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland,
Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis
Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and
Demis Hassabis.
Human-level control through deep reinforcement learning.
Nat., 5180 (7540):0 529–533, 2015.
URL <https://doi.org/10.1038/nature14236>.

[Nowe, 1994]
Ann Nowe.
Synthesis of “safe” fuzzy controllers based on reinforcement
PhD thesis, Vrije Universiteit Brussel, 1994.

[Papamakarios et al., 2017]
George Papamakarios, Iain 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 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.  2338–2347, 2017.

[Pnueli, 1977]
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.
URL <https://doi.org/10.1109/SFCS.1977.32>.

[Puterman, 1994]
Martin L. Puterman.
Markov Decision Processes: Discrete Stochastic Dynamic
Wiley Series in Probability and Statistics. Wiley, 1994.
ISBN 978-0-47161977-2.
URL <https://doi.org/10.1002/9780470316887>.

[Ren et al., 2021]
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 networks.
Journal of Network and Computer Applications, 182:0
103049, 2021.
ISSN 1084-8045.

[Sickert et al., 2016]
Salomon Sickert, Javier Esparza, Stefan Jaax, and Jan Kretínský.
Limit-deterministic büchi automata 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, Part II, volume 9780 of
Lecture Notes in Computer Science, pp.  312–332. Springer, 2016.
URL <https://doi.org/10.1007/978-3-319-41540-6_17>.

[Simão et al., 2021]
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 <https://dl.acm.org/doi/10.5555/3463952.3464094>.

[Todorov et al., 2012]
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.

[Tolstikhin et al., 2018]
Ilya O. Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard
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 <https://openreview.net/forum?id=HkL7n1-0b>.

[Tsitsiklis, 1994]
John N. Tsitsiklis.
Asynchronous stochastic approximation and q-learning.
Mach. Learn., 160 (3):0 185–202, 1994.
URL <https://doi.org/10.1007/BF00993306>.

[van den Oord et al., 2017]
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.

[Villani, 2009]
Cédric Villani.
Optimal Transport: Old and New.
Springer Berlin Heidelberg, Berlin, Heidelberg, 2009.
ISBN 978-3-540-71050-9.
URL <https://doi.org/10.1007/978-3-540-71050-9_6>.

[Wells et al., 2020]
Andrew M. Wells, Morteza Lahijanian, Lydia E. Kavraki, and Moshe Y. Vardi.
Ltlf 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.
URL <https://doi.org/10.4204/EPTCS.326.11>.

[Zang et al., 2022]
Hongyu Zang, Xin Li, and Mingzhong Wang.
Simsr: Simple distance-based state representations for deep
reinforcement learning.
Proceedings of the AAAI Conference on Artificial Intelligence,
360 (8):0 8997–9005, Jun. 2022.
URL <https://ojs.aaai.org/index.php/AAAI/article/view/20883>.

[Zhang et al., 2021]
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey
Learning invariant representations for reinforcement learning without
In 9th International Conference on Learning Representations,
ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
URL <https://openreview.net/forum?id=-2FCwDKRREu>.

[Zhang et al., 2019]
Shunkang Zhang, Yuan Gao, Yuling Jiao, Jin Liu, Yang Wang, and Can Yang.
Wasserstein-wasserstein auto-encoders.
CoRR, abs/1902.09323, 2019.
URL <http://arxiv.org/abs/1902.09323>.

§ APPENDIX

§ THEORETICAL DETAILS ON WAE-MDPS

§.§ 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 $\decoder$ 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 $\tuple{\seq{\state}{T}, \seq{\action}{T - 1}, \seq{\reward}{T - 1}}$), our goal is to find the best parameter $\decoderparameter$ which offers the most accurate reconstruction of the original traces issued from the original model $\mdp$ operating under $\policy$.
We demonstrate that, in the limit, considering the OT between trace-based distributions is equivalent to considering the OT between the stationary distribution of $\mdp_{\policy}$ 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 $\states$, $\actions$, and $\images{\rewards}$ are respectively equipped with metric $\distance_{\states}$, $\distance_{\actions}$, and $\distance_{\rewards}$,
let us define the raw transition distance metric over transitions of $\mdp$, i.e., tuples of the form $\tuple{\state, \action, \reward, \state'}$, as $\transitiondistance \colon \states \times \actions \times \images{\rewards} \times \states$,

\begin{equation}
\tracedistance\fun{\tuple{\state_1, \action_1, \reward_1, \state'_1}, \tuple{\state_2, \action_2, \reward_2, \state'_2}} = \distance_\states\fun{\state_1, \state_2} + \distance_{\actions}\fun{\action_1, \action_2} + \distance_{\rewards}\fun{\reward_1, \reward_2} + \distance_{\states}\fun{\state_1', \state'_2}.\notag
\end{equation}

In a nutshell, $\transitiondistance$ consists of the sum of the distance of all the transition components.
Note that it is a well defined distance metric since the sum of distances preserves the identity of indiscernible, symmetry, and triangle inequality.

Trace-based distributions.
The raw distance $\tracedistance$ 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:

\begin{align*}
\mathcal{D}_\policy\left[ T \right]\fun{\state, \action, \reward, \state'} &= \frac{1}{T} \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\state \mid \sinit} \cdot \policy\fun{\action \mid \state} \cdot \probtransitions\fun{\state' \mid \state, \action} \cdot \condition{\reward=\rewards\fun{\state, \action}}, \text{ and} \\
\mathcal{P}_\decoderparameter[T]\fun{\state, \action, \reward, \state'} &= \frac{1}{T} \sum_{t = 1}^{T} \expectedsymbol{\seq{\state}{t}, \seq{\action}{t - 1}, \seq{\reward}{t - 1} \sim \decoder[t]}{\condition{\tuple{\indexof{\state}{t - 1}, \indexof{\action}{t - 1}\indexof{\reward}{t - 1}, \indexof{\state}{t}}= \tuple{\state, \action, \reward, \state'}}},
\end{align*}

where $\decoder[T]$ denotes the distribution over traces of length $T$, generated from $\decoder$.
Intuitively, $\nicefrac{1}{T} \cdot \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\state \mid \sinit}$

can be seen as the fraction of the time spent in $\state$
along traces of length $T$, starting from the initial state 10.5555/280952.
Therefore, drawing $\tuple{\state, \action, \reward, \state'} \sim \mathcal{D}_\policy\left[ T \right]$ trivially follows:
it is equivalent to drawing $\state$ from $\nicefrac{1}{T} \cdot \sum_{t = 1}^{T} \stationary{\policy}^t\fun{\cdot \mid \sinit}$, then respectively $\action$ and $\state'$ from $\policy\fun{\cdot \mid \state}$ and $\probtransitions\fun{\cdot \mid \state, \action}$, to finally obtain $\reward = \rewards\fun{\state, \action}$.

Given $T \in \N$, our objective is to minimize the Wasserstein distance between those distributions:

The following Lemma enables optimizing the Wasserstein distance between the original MDP and the behavioral model
when traces are drawn from episodic RL processes or infinite interactions [Huang, 2020].

Assume the existence of a stationary behavioral model $\stationarydecoder = \lim_{T \to \infty} \mathcal{P}_{\decoderparameter}[T]$, then

\begin{equation*}
\lim_{T \to \infty} \wassersteindist{\tracedistance}{\mathcal{D}_{\policy}[T]}{\mathcal{P}_{\decoderparameter}[T]} = \wassersteindist{\tracedistance}{\stationary{\policy}}{\stationarydecoder}.
\end{equation*}

First, note that $\nicefrac{1}{T} \cdot \sum_{t = 1}^T \stationary{\policy}^t\fun{\cdot \mid \sinit}$ weakly converges to $ \stationary{\policy}$ as $T$  goes to $\infty$ 10.5555/280952. The result follows then from <cit.>.

§.§ Dealing with Discrete Actions

When the policy $\policy$ executed in $\mdp$ already produces discrete actions, learning a latent action space is, in many cases, not necessary.
We thus make the following assumptions:

Let $\policy \colon \states \to \distributions{\actions^{\star}}$ be the policy executed in $\mdp$ and assume that $\actions^{\star}$ is a (tractable) finite set.
Then, we take $\latentactions=\actions^{\star}$ and $\actionencoder$ as the

identity function,
i.e., $\actionencoder \colon \latentstates \times \actions^{\star} \to \actions^{\star}, \, \tuple{\latentstate, \action^{\star}} \mapsto \action^{\star}$.

Assume that the action space of the original environment $\mdp$ is a (tractable) finite set.
Then, we take $\embeda_{\decoderparameter}$ as the identity function, i.e., $\embeda_{\decoderparameter} = \actionencoder$.

Concretely, the premise of Assumption <ref> typically occurs when $\policy$ is a latent policy (see Rem. <ref>) or when $\mdp$ has already a discrete action space.
In the latter case, Assumption <ref> and <ref> amount to setting $\latentactions = \actions$ 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 a smaller set of discrete actions.

§.§ Proof of Lemma <ref>

From now on, we write $\embed_{\encoderparameter}\fun{\latentstate, \latentaction \mid \state, \action} = \condition{\embed_\encoderparameter\fun{\state}=\latentstate} \cdot \actionencoder\fun{\latentaction \mid \latentstate, \action}$.

Wasserstein is compliant with the triangular inequality [Villani, 2009], which gives us:

\begin{align*}
\wassersteindist{\transitiondistance}{\encoder}{ \latentstationaryprior}
\leq \wassersteindist{\transitiondistance}{\encoder}{\originaltolatentstationary{}} + \wassersteindist{\distance_{\latentstates}}{\originaltolatentstationary{}}{ \latentstationaryprior},
\end{align*}

\begin{align}
&\wassersteindist{\transitiondistance}{\originaltolatentstationary{}}{ \latentstationaryprior} \tag{note that $\wassersteinsymbol{\tracedistance}$ is reflexive [Villani, 2009]} \\
=& \sup_{f \in \Lipschf{\transitiondistance}} \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} f\fun{\latentstate, \latentaction, \latentstate'}
- \expectedsymbol{\latentstate \sim \latentstationaryprior}\expectedsymbol{\latentaction \sim \latentpolicy_{\decoderparameter}\fun{\sampledot \mid \latentstate}}\expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} f\fun{\latentstate, \latentaction, \latentstate'}\text{, and} \notag \\[.5em]
& \wassersteindist{\transitiondistance}{\encoder}{\originaltolatentstationary{}} \notag \\
\sup_{f \in \Lipschf{\transitiondistance}} \expectedsymbol{\state, \action, \state' \sim \stationary{\policy}}\expectedsymbol{\latentstate, \latentaction, \latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action, \state'}} f\fun{\latentstate, \latentaction, \latentstate'} - \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} f\fun{\latentstate, \latentaction, \latentstate'} \label{eq:proof-lemma-wdist-triangular-inequality-1} \\
\leq& \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}} \; \sup_{f \in \Lipschf{\transitiondistance}} \expectedsymbol{\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}} f\fun{\latentstate, \latentaction, \embed_{\encoderparameter}\fun{\state'}} - \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} f\fun{\latentstate, \latentaction, \latentstate'} \label{eq:proof-lemma-wdist-triangular-inequality-2} \\
=& \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\scriptscriptstyle \actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action}} \; \sup_{f \in \Lipschf{\distance_{\latentstates}}} \expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}} f\fun{\latentstate'} - \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}} f\fun{\latentstate'} \label{eq:proof-lemma-wdist-triangular-inequality-3} \\
=& \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\scriptscriptstyle \actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}}. \notag
\end{align}

We pass from Eq. <ref> to Eq. <ref> by the Jensen's inequality.
To see how we pass from Eq. <ref> to Eq. <ref>, notice

\begin{align*}
\Lipschf{\transitiondistance} &= \left\{f \colon f\fun{\latentstate_1, \latentaction_1, \latentstate_1'} - f\fun{\latentstate_2, \latentaction_2, \latentstate_2'} \leq \transitiondistance\fun{\tuple{\latentstate_1, \latentaction_1, \latentstate_1'}, \tuple{\latentstate_2, \latentaction_2, \latentstate_2'}}\right\} \\
\Lipschf{\transitiondistance} &= \set{f \colon f\fun{\latentstate_1, \latentaction_1, \latentstate_1'} - f\fun{\latentstate_2, \latentaction_2, \latentstate_2'} \leq \distance_{\latentstates}\fun{\latentstate_1, \latentstate_2} + \distance_{\latentactions}\fun{\latentaction_1, \latentaction_2} + \distance_{\latentstates}\fun{\latentstate_1', \latentstate_2'}}
\end{align*}

Observe now that $\latentstate$ and $\latentaction$ are fixed in the supremum computation of Eq. <ref>: all functions $f$ considered and taken from $\Lipschf{\transitiondistance}\,$ are of the form $f\fun{\latentstate, \latentaction, \sampledot}$.
It is thus sufficient to consider the supremum over functions from the following subset of $\Lipschf{\transitiondistance}\,$:

\begin{align*}
&\set{f \colon f\fun{\latentstate, \latentaction, \latentstate_1'} - f\fun{\latentstate, \latentaction, \latentstate_2'} \leq \distance_{\latentstates}\fun{\latentstate, \latentstate} + \distance_{\latentactions}\fun{\latentaction, \latentaction} +  \distance_{\latentstates}\fun{\latentstate_1', \latentstate_2'}} \\ \tag{for $\latentstate$, $\latentaction$ drawn from $\embed_{\encoderparameter}$} \\
=&\set{f \colon f\fun{\latentstate, \latentaction, \latentstate_1'} - f\fun{\latentstate, \latentaction, \latentstate_2'} \leq \distance_{\latentstates}\fun{\latentstate_1', \latentstate_2'}}\\
=&\set{f \colon f\fun{\latentstate_1'} - f\fun{\latentstate_2'} \leq \distance_{\latentstates}\fun{\latentstate_1', \latentstate_2'}} \\
=& \Lipschf{\distance_{\latentstates}}.
\end{align*}

Given a state $\state \in \states$ in the original model, the (parallel) execution of $\policy$ in $\latentmdp_{\decoderparameter}$ is enabled through $\policy\fun{\action, \latentaction \mid \state} = \policy\fun{\action \mid \state} \cdot \actionencoder\fun{\latentaction \mid \embed_{\encoderparameter}\fun{\state}, \action}$ (cf. Fig. <ref>).
The local transition loss resulting from this interaction is:

\begin{align*}
\localtransitionloss{\stationary{\policy}} &= \expectedsymbol{\state, \tuple{\action, \latentaction} \sim \stationary{\policy}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}}{\latentprobtransitions\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}}\\
&= \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\scriptscriptstyle \actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}},
\end{align*}

which finally yields the result.

§.§ Proof of Theorem <ref>

Before proving Theorem <ref>, let us introduce the following Lemma, that explicitly demonstrates the link between the transition regularizer of the objective and the local transition loss required to obtain the guarantees related to the bisimulation bounds of Eq. <ref>.

Assume that traces are generated by running $\latentpolicy \in \latentpolicies$ in the original environment, then

\begin{align*}
\expectedsymbol{\state, \action^\star \sim \stationary{\latentpolicy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action^{\star}}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action^{\star}}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}} = \localtransitionloss{\stationary{\latentpolicy}}.
\end{align*}

Since the latent policy $\latentpolicy$ generates latent actions, Assumption <ref> holds, which means:

\begin{align*}
& \expectedsymbol{\state, \action^{\star} \sim \stationary{\latentpolicy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action^{\star}}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action^{\star}}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}}\\
=& \expectedsymbol{\state, \latentaction \sim \stationary{\latentpolicy}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \latentaction}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}}\\
=& \, \localtransitionloss{\stationary{\latentpolicy}}.
\end{align*}

We distinguish two cases: (i) the case where the original and latent models share the same discrete action space, i.e., $\actions = \latentactions$, and (ii) the case where the two have a different action space (e.g., when the original action space is continuous), i.e., $\actions \neq \latentactions$.
In both cases, the local losses term follows by definition of $\localrewardloss{\stationary{\latentpolicy}}$ and Lemma <ref>.
When $\distance_{\rewards}$ is the Euclidean distance (or even the $L_1$ distance since rewards are scalar values), the expected reward distance occurring in the expected trace-distance term $\tracedistance$ in the objective directly translates to the local loss $\localrewardloss{\stationary{\latentpolicy}}$.

Concerning the local transition loss, in case (i), the result naturally follows from  Assumption <ref> and <ref>.
In case (ii), only Assumption <ref> holds, meaning the action encoder term of the objective is ignored, but not the action embedding term appearing in $G_{\decoderparameter}$.
Given $\state \sim \stationary{\latentpolicy}$, recall that executing $\latentpolicy$ in $\mdp$ amounts to embedding the produced latent actions $\latentaction \sim \latentpolicy\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}}$ back to the original environment via $\action = \embeda_{\decoderparameter}\fun{\embed_{\encoderparameter}\fun{\state}, \latentaction}$ (cf. Rem. <ref> and Fig. <ref>).
Therefore, the projection of $\tracedistance\fun{\tuple{\state, \action, \reward, \state'}, G_{\decoderparameter}\fun{\embed_{\encoderparameter}\fun{\state}, \latentaction, \embed_{\encoderparameter}\fun{\state'}}}$ on the action space $\actions$ is $\distance_{\actions}\fun{\embeda_{\decoderparameter}\fun{\embed_{\encoderparameter}\fun{\state}, \latentaction}, \embeda_{\decoderparameter}\fun{\embed_{\encoderparameter}\fun{\state}, \latentaction}} = 0$, for $\reward = \rewards\fun{\state, \action}$ and $\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}$.

§.§ Optimizing the Transition Regularizer

In the following, we detail how we derive a tractable form of
our transition regularizer $\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter}$.

Optimizing the ground Kantorovich-Rubinstein duality is enabled via the introduction of a parameterized, $1$-Lipschitz network
$\transitionlossnetwork$, that need to be trained to attain the supremum of the dual:

\begin{align*}
\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter} &= \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}} \; \max_{\wassersteinparameter \colon \transitionlossnetwork \in \Lipschf{\distance_{\latentstates}}} \; \expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}} \transitionlossnetwork\fun{\latentstate'} - \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \transitionlossnetwork\fun{\latentstate'}. %\label{eq:transition-regularizer}
%=& \expectedsymbol{\state, \action \sim \stationary{\policy}}\expectedsymbol{\latentaction \sim \embed_{\encoderparameter}^{\scriptscriptstyle \actions}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \action}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \embed_{\encoderparameter}\fun{\state}, \latentaction}}. \notag
\end{align*}

Under this form, optimizing $\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter}$ is intractable due to the expectation over the maximum.

The following Lemma allows us rewriting $\localtransitionloss{\stationary{\policy}}$ to make the optimization tractable through Monte Carlo estimation.

Let $\measurableset, \varmeasurableset$ be two measurable sets, $\stationary{} \in \distributions{\measurableset}$, $P \colon \measurableset \to \distributions{\varmeasurableset}, Q \colon \measurableset \to \distributions{\varmeasurableset}$, and $\distance \colon \varmeasurableset \times \varmeasurableset \to \mathopen[ 0, + \infty \mathclose[$ be a metric on $\varmeasurableset$. Then,

\begin{align*}
&\expectedsymbol{x \sim \stationary{}}{\wassersteindist{\distance}{P\fun{\sampledot \mid x}}{Q\fun{\sampledot \mid x}}} = \sup_{\varphi \colon \measurableset \to \Lipschf{\distance}} \expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_2}}
\end{align*}

Our objective is to show that

\begin{align}
&\expected{x \sim \stationary{}}{\sup_{f \in \Lipschf{\distance}} \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}}\varphi\fun{y_1}\fun{x} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi\fun{y_2}\fun{x}} \label{eq:expected-cond-wasserstein}\\
=& \sup_{\varphi \colon \measurableset \to \Lipschf{\distance}} \expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_2}} \label{eq:sup-expected-cond-wasserstein}
\end{align}

We start with (<ref>) $\leq$ (<ref>). Construct $\varphi^\star \colon \measurableset \to \Lipschf{\distance}$ by setting for all $x \in \measurableset$

\begin{equation*}
\varphi^\star\fun{x} = \arg\sup_{f \in \Lipschf{\distance}} \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} f\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} f\fun{y_2}.
\end{equation*}

This gives us

\begin{align*}
&\expected{x \sim \stationary{}}{\sup_{f \in \Lipschf{\distance}} \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} f\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} f\fun{y_2}} \\
=& \expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_2}} \\
\leq& \sup_{\varphi \colon \measurableset \to \Lipschf{\distance}} \expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_2}}.
\end{align*}

It remains to show that (<ref>) $\geq$ (<ref>).

\[
\varphi^{\star} = \arg\sup_{\varphi \colon \measurableset \to \Lipschf{\distance}} \expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi\fun{x}\fun{y_2}}.
\]

Then, for all $x \in \measurableset$, we have $\varphi^{\star}\fun{x} \in \Lipschf{\distance}$ which means:

\begin{align*}
& \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_2} \\
\leq & \sup_{f \in \Lipschf{\distance}} \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} f\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} f\fun{y_2}
\end{align*}

This finally yields

\begin{align*}
&\expected{x \sim \stationary{}}{\expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} \varphi^{\star}\fun{x}\fun{y_2}} \\
\leq& \expected{x \sim \stationary{}}{\sup_{f \in \Lipschf{\distance}} \expectedsymbol{y_1 \sim P\fun{\sampledot \mid x}} f\fun{y_1} - \expectedsymbol{y_2 \sim Q\fun{\sampledot \mid x}} f\fun{y_2}}.
\end{align*}

Let  $\stationary{\policy}$ be a stationary distribution of $\mdp_\policy$ and $\measurableset = \states \times \actions \times \latentstates \times \latentactions$, then

\begin{align*}
\localtransitionloss{\stationary{\policy}} = \sup_{\varphi \colon \measurableset \to \scriptscriptstyle \Lipschf{\distance_{\latentstates}}} \expectedsymbol{\state, \action, \state' \sim \stationary{\policy}} \expected{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}}{\varphi\fun{\state, \action, \latentstate, \latentaction}\fun{\embed_{\encoderparameter}\fun{\state'}} - \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}}{\varphi\fun{\state, \action, \latentstate, \latentaction}}\fun{\latentstate'}}
\end{align*}

Notice that

\begin{align*}
& \expectedsymbol{\state, \action \sim \stationary{\policy}} \wassersteindist{\distance_{\latentstates}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot \mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \\
=& \expectedsymbol{\state, \action \sim \stationary{\policy}}{\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left\{  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right\}}, \tag{e.g., Villani, 2009}
\end{align*}

our objective is therefore to show that

\begin{align}
& \expectedsymbol{\state, \action \sim \stationary{\policy}}{\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left\{  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right\}} \label{eq:wae-mdp-loss-proof-1} \\
=& \sup_{\varphi \colon \states \times \actions \to \Lipschf{\distance_{\latentstates}}} \, \expected{\state, \action, \state' \sim \stationary{\policy}}{\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi\tuple{\state, \action}\fun{\latentstate'}}. \label{eq:wae-mdp-loss-proof-2}
\end{align}

We start with the $(\ref{eq:wae-mdp-loss-proof-1}) \leq (\ref{eq:wae-mdp-loss-proof-2})$ direction. For all $\state \in \states, \action \in \actions$, construct $\varphi^{\star} \colon \states \times \actions \to \Lipschf{\distance_{\latentstates}}$ by setting

\begin{align*}
\varphi^\star{\tuple{\state, \action}} = \arg\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left[  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right].
\end{align*}

This gives us

\begin{align*}
& \expectedsymbol{\state, \action \sim \stationary{\policy}}{\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left\{  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right\}}\\
=& \expected{\state, \action \sim \stationary{\policy}}{\expectedsymbol{\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi^\star\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi^\star\tuple{\state, \action}\fun{\latentstate'}} \\
\leq&  \sup_{\varphi \colon \states \times \actions \to \Lipschf{\distance_{\latentstates}}} \, \expected{\state, \action \sim \stationary{\policy}}{\expectedsymbol{\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi\tuple{\state, \action}\fun{\latentstate'}}\\
=& \sup_{\varphi \colon \states \times \actions \to \Lipschf{\distance_{\latentstates}}} \, \expected{\state, \action, \state' \sim \stationary{\policy}}{\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi\tuple{\state, \action}\fun{\latentstate'}}.
\end{align*}

It remains to prove the other direction, i.e., $(\ref{eq:wae-mdp-loss-proof-1}) \geq (\ref{eq:wae-mdp-loss-proof-2})$.

\begin{align*}
\varphi^{\star} = \arg\sup_{\varphi \colon \states \times \actions \to \Lipschf{\distance_{\latentstates}}} \, \expected{\state, \action, \state' \sim \stationary{\policy}}{\expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi\tuple{\state, \action}\fun{\latentstate'}},
\end{align*}

then for all $\state \in \states$, $\action \in \actions$, $\varphi^\star\tuple{\state, \action} \in \Lipschf{\distance_{\latentstates}}$ and

\begin{align*}
&& &{\expectedsymbol{\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi^\star\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi^\star\tuple{\state, \action}\fun{\latentstate'}} \\
&&\leq& {\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left\{  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right\}} \\
&\equiv&&  \expected{\state, \action \sim \stationary{\policy}}{\expectedsymbol{\state' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}}\varphi^\star\tuple{\state, \action}\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} \varphi^\star\tuple{\state, \action}\fun{\latentstate'}} \\
&&\leq& \expectedsymbol{\state, \action \sim \stationary{\policy}}{\sup_{f \in \Lipschf{\distance_{\latentstates}}} \left\{  \expectedsymbol{\state ' \sim \probtransitions\fun{\sampledot \mid \state, \action}}\expectedsymbol{\latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state'}} f\fun{\latentstate'} - \expectedsymbol{\latentstate \sim \embed_{\encoderparameter}\fun{\sampledot\mid \state}}  \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \action}} f\fun{\latentstate'}\right\}},
\end{align*}

which yields the result.

Corollary <ref> allows us re-writing our objective as

\begin{multline*}
\min_{\encoderparameter, \decoderparameter}\max_{\omega} \, \expectedsymbol{\state, \action, \reward, \state' \sim \stationary{\policy}} \, \expectedsymbol{\latentstate, \latentaction_1, \latentstate' \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action, \state'}} \expectedsymbol{\latentstate_1' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \expectedsymbol{\latentstate_2 \sim \stationary{\latentpolicy_{\decoderparameter}}} \expectedsymbol{\latentaction_2 \sim \latentpolicy_{\decoderparameter}\fun{\sampledot \mid \latentstate_2}}\expectedsymbol{\latentstate_2' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate_2, \latentaction_2}} \\
\left[\tracedistance\fun{\tuple{\state, \action, \reward, \state'}, G_{\decoderparameter}\fun{\latentstate, \latentaction_1, \latentstate'}} + \beta \cdot \fun{ \varphi_{\omega}^{\stationary{}}\fun{\latentstate, \latentaction_1, \latentstate_1'} - \varphi^{\stationary{}}_{\omega}\fun{\latentstate_2, \latentaction_2, \latentstate_2'} + \varphi^{\probtransitions}_{\omega}\fun{\state, \action, \latentstate, \latentaction, \latentstate'} - \varphi^{\probtransitions}_{\omega}\fun{\state, \action, \latentstate, \latentaction, \latentstate_1'} } \right]
\end{multline*}

where we assume that $\varphi^{\stationary{}}_{\wassersteinparameter}$ and $\varphi^{\latentprobtransitions}_{\wassersteinparameter}$ are functions parameterized by $\wassersteinparameter$ compliant with the 1-Lipschitz constraint.
Consequently, we rewrite $\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter}$ as a tractable maximization:

\begin{equation*}
\localtransitionloss{\stationary{\policy}}\fun{\wassersteinparameter} = \max_{\wassersteinparameter \colon \transitionlossnetwork \in \Lipschf{\distance_{\scriptscriptstyle \latentstates}}} \; \expectedsymbol{\state, \action, \state' \sim \stationary{\policy}} \expected{\latentstate, \latentaction \sim \embed_{\encoderparameter}\fun{\sampledot \mid \state, \action}}{\transitionlossnetwork\fun{\state, \action, \latentstate, \latentaction, \embed_{\encoderparameter}\fun{\state'}} - \expectedsymbol{\latentstate' \sim \latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \transitionlossnetwork\fun{\state, \action, \latentstate, \latentaction, \latentstate'}}.
\end{equation*}

§.§ The Latent Metric

In the following, we show that considering the Euclidean distance for $\tracedistance$ and $\distance_{\latentstates}$ in the latent space for optimizing  the regularizers $\steadystateregularizer{{\policy}}$ and $\localtransitionloss{\stationary{\policy}}$ is Lipschitz equivalent to considering a continuous $\temperature$-relaxation of the discrete metric $\condition{\neq}\fun{\vx, \vy} = \condition{\vx \neq \vy}$.
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 zero-temperature limit, when the spaces are discrete.

Let $\distance$ be the usual Euclidean distance and $\distance_{\temperature}\colon \mathopen[0, 1 \mathclose]^n \times \mathopen[0, 1 \mathclose]^n \to \mathopen[0, 1\mathclose[, \, \tuple{\vx, \vy} \mapsto \frac{\distance\fun{\vx, \vy}}{\temperature + \distance\fun{\vx, \vy}}$ for $\temperature \in \mathopen]0, 1\mathclose]$ and $n \in \N$, then $\distance_{\temperature}$ is a distance metric.

The function $\distance_{\temperature}$ is a metric iff it satisfies the following axioms:

* Identity of indiscernibles:
If $\vx=\vy$, then $\distance_{\lambda}\fun{\vx, \vy} = \frac{\distance\fun{\vx, \vy}}{\temperature + \distance\fun{\vx, \vy}} = \frac{0}{\temperature + 0} = 0$ since $\distance$ is a distance metric.
Assume now that $\distance_\temperature\fun{\vx, \vy} = 0$ and take $\alpha = \distance\fun{\vx, \vy}$, for any $\vx, \vy$. Thus, $\alpha \in \left[0, +\infty\right[$ and $0 = \frac{\alpha}{\temperature + \alpha}$ is only achieved in $\alpha = 0$, which only occurs whenever $\vx = \vy$ since $\distance$ is a distance metric.

* Symmetry:

\begin{align*}
\distance_\temperature\fun{\vx, \vy} &= \frac{\distance\fun{\vx, \vy}}{\temperature + \distance\fun{\vx, \vy}}\\
&= \frac{\distance\fun{\vy, \vx}}{\temperature + \distance\fun{\vy, \vx}} \tag{$\distance$ is a distance metric} \\
&= \distance_\temperature\fun{\vy, \vx}
\end{align*}

* Triangle inequality:
Let $\vx, \vy, \vz \in \mathopen[0, 1\mathclose]^n$, the triangle inequality holds iff

\begin{align}
&&\distance_\temperature\fun{\vx, \vy} + \distance_{\temperature}\fun{\vy, \vz} &\geq \distance_{\temperature}\fun{\vx, \vz}& \label{eq:triangle-ineq-d-temperature} \\
&\equiv& \frac{\distance\fun{\vx, \vy}}{\temperature + \distance\fun{\vx, \vy}} + \frac{\distance\fun{\vy, \vz}}{\temperature + \distance\fun{\vy, \vz}} & \geq \frac{\distance\fun{\vx, \vz}}{\temperature + \distance\fun{\vx, \vz}}& \notag \\
&\equiv& \frac{\temperature \distance\fun{\vx, \vy} + \temperature \distance\fun{\vy, \vz} + 2\distance\fun{\vx, \vy} \distance\fun{\vy, \vz}}{\temperature^2 + \temperature \distance\fun{\vx, \vy} + \temperature \distance\fun{\vy, \vz} + \distance\fun{\vx, \vy}\distance\fun{\vy, \vz}} & \geq \frac{\distance\fun{\vx, \vz}}{\temperature + \distance\fun{\vx, \vz}}& \notag
\end{align}

\begin{align}
&\equiv& &\temperature^2 \distance\fun{\vx, \vy} + \temperature^2 \distance\fun{\vy, \vz} + 2 \temperature \distance\fun{\vx, \vy}\distance\fun{\vy, \vz} + \notag \\
&&&\temperature\distance\fun{\vx, \vy} \distance\fun{\vx, \vz} + \temperature\distance\fun{\vy, \vz} \distance\fun{\vx, \vz} + 2 \distance\fun{\vx, \vy} \distance\fun{\vy, \vz}\distance\fun{\vx, \vz} & \notag \\
&& \geq& \temperature^2 \distance\fun{\vx, \vz} + \temperature\distance\fun{\vx, \vy} \distance\fun{\vx, \vz} + \temperature\distance\fun{\vy, \vz} \distance\fun{\vx, \vz} + \distance\fun{\vx, \vy} \distance\fun{\vy, \vz} \distance\fun{\vx, \vz} \tag{cross-product, with $\temperature > 0$ and $\images{\distance} \in \mathopen[0, \infty\mathclose[$} \\
&\equiv& &\temperature^2 \distance\fun{\vx, \vy} + \temperature^2 \distance\fun{\vy, \vz} + 2 \temperature \distance\fun{\vx, \vy}\distance\fun{\vy, \vz} +
\distance\fun{\vx, \vy} \distance\fun{\vy, \vz}\distance\fun{\vx, \vz}
\geq \temperature^2 \distance\fun{\vx, \vz} \label{eq:triangle-ineq-relaxation}
\end{align}

Since $\distance$ is a distance metric, we have

\begin{equation}
\temperature^2 \distance\fun{\vx, \vy} +
\temperature^2 \distance\fun{\vy, \vz} \geq
\temperature^2 \distance\fun{\vx, \vz} \label{eq:triangle-ineq-scaled}
\end{equation}

and $\images{\distance}\in \mathopen[0, \infty\mathclose[$, meaning

\begin{equation}
2 \temperature \distance\fun{\vx, \vy}\distance\fun{\vy, \vz} +
\distance\fun{\vx, \vy} \distance\fun{\vy, \vz}\distance\fun{\vx, \vz} \geq 0 \label{eq:triangle-ineq-upper-bound}
\end{equation}

By Eq. <ref> and <ref>, the inequality of Eq. <ref> holds.
Furthermore, the fact that Eq. <ref> and <ref> are equivalent yields the result.

Let $\distance$, $\distance_{\temperature}$ as defined above, then
(i) $ \distance_{\temperature} \xrightarrow[\temperature \to 0]{} \condition{\neq}$

and (ii) $\distance, \distance_\temperature$ are Lipschitz-equivalent.

Part (i) is straightforward by definition of $\distance_{\temperature}$.
Distances $\distance$ and $\distance_{\temperature}$ are Lispchitz equivalent if and only if $\exists a, b > 0$ such that $\forall \vx, \vy \in \mathopen[ 0, 1 \mathclose]^n$,

a ·,     ≤   _,            ≤b ·,

≡    a ·,     ≤   , /+ ,     ≤b ·,

≡    a     ≤    1/+ ,     ≤b

Taking $a = \frac{1}{\temperature + \sqrt{n}}$ and $b = \frac{1}{\temperature}$ yields the result.

For all $\beta \geq \nicefrac{1}{\temperature}$, $\state \in \states$, $\action \in \actions$, $\latentstate \in \latentstates$, and $\latentaction \in \latentactions$, we have

* $\wassersteindist{\distance_{\temperature}}{\originaltolatentstationary{}}{\latentstationaryprior} \leq \beta \cdot \wassersteindist{\distance}{\originaltolatentstationary{}}{\latentstationaryprior}$

* $\wassersteindist{\distance_{\temperature}}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot\mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}} \leq \beta \cdot \wassersteindist{\distance}{\embed_{\encoderparameter}\probtransitions\fun{\sampledot\mid \state, \action}}{\latentprobtransitions_{\decoderparameter}\fun{\sampledot \mid \latentstate, \latentaction}}$

By Lipschitz equivalence, taking $\beta \geq \nicefrac{1}{\temperature}$ ensures that $\forall n \in \N$, $\forall \vx, \vy \in \mathopen[0, 1\mathclose]^n, \, \distance_{\temperature}\fun{\vx, \vy} \leq \beta \cdot \distance\fun{\vx, \vy}$.
Moreover, for any distributions $P, Q$, $\wassersteindist{\distance_{\temperature}}{P}{Q} \leq \beta \cdot \wassersteindist{\distance}{P}{Q}$

(cf., e.g.,  for details).

In practice, taking the hyperparameter $\beta \geq \nicefrac{1}{\temperature}$ in the ensures that minimizing the $\beta$-scaled regularizers w.r.t$.$ $\distance$ also minimizes the regularizers w.r.t$.$ the $\temperature$-relaxation $\distance_\temperature$, being the discrete distribution in the zero-temperature limit.

Note that optimizing over two different $\beta_1, \beta_2$ instead of a unique scale factor $\beta$ is also a good practice to interpolate between the two regularizers.

§ EXPERIMENT DETAILS

The code for conducting and replicating our experiments is available at <https://github.com/florentdelgrange/wae_mdp>.

§.§ Setup

We used TensorFlow  tensorflow2015-whitepaper to implement the neural network architecture of our , TensorFlow Probability  dillon2017tensorflow to handle the probabilistic components of the latent model (e.g., latent distributions with reparameterization tricks, masked autoregressive flows, etc.), as well as TF-Agents  TFAgents to handle the RL parts of the framework.

Models have been trained on a cluster running under  composed of a mix of nodes containing Intel processors with the following CPU microarchitectures:
(i) , (ii) , and (iii) .
We used $8$ cores and $32$ GB of memory for each run.

§.§ Stationary Distribution

To sample from the stationary distribution $\stationary{\policy}$ of episodic learning environments operating under $\policy \in \mpolicies{\mdp}$, we implemented the recursive $\epsilon$-perturbation trick of [Huang, 2020].
In a nutshell, the reset of the environment is explicitly added to the state space of $\mdp$, which is entered at the end of each episode and left with probability $1 - \epsilon$ to start a new one.
We also added a special atomic proposition $\labelset{reset}$ into $\atomicprops$ 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 if it is able to reach a goal before the end of the episode.

§.§ Environments with initial distribution

Many environments do not necessarily have a single initial state, but rather an initial distribution over states $d_I \in \distributions{\states}$.
In that case, the results presented in this paper remain unchanged: it suffices to add a dummy state $\state^{\star}$ to the state space $\states \cup \set{\state^{\star}}$ so that $\sinit = \state^{\star}$ with the transition dynamics $\probtransitions\fun{\state' \mid \state^{\star}, \action} = d_I\fun{\state'}$ for any action $\action \in \actions$.
Therefore, each time the reset of the environment is triggered, we make the MDP entering the initial state $\state^{\star}$, then transitioning to $\state'$ according to $d_I$.

§.§ Latent space distribution

Latent space distribution along training steps. The intensity of the blue hue corresponds to the frequency of latent states produced by $\embed_{\encoderparameter}$ during training.

As pointed out in Sect. <ref>, posterior collapse is naturally avoided when optimizing .
To illustrate that, we report the distribution of latent states produced by $\embed_{\encoderparameter}$ during training (Fig. <ref>).
The plots reveal that the latent space generated by mapping original states drawn from $\stationary{\policy}$ during training to $\latentstates$ via $\embed_{\encoderparameter}$ is fairly distributed, for each environment.

§.§ Distance Metrics: state, action, and reward reconstruction

The choice of the distance functions $\distance_{\states}$, $\distance_{\actions}$, and $\distance_{\rewards}$, plays a 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 an example of data which could require tuning the state-distance function in such a way, to make sure that the optimization of the reward or action reconstruction will not be disfavored compared to that of the states.

§.§ Value difference