Original: W ASSERSTEIN A UTO - ENCODED MDP S
Revised: W ASSERSTEIN A UTO - ENCODED MDP S

Original: F ORMAL V ERIFICATION OF E FFICIENTLY D ISTILLED RL P OLICIES WITH M ANY - SIDED G UARANTEES
Revised: F ORMAL V ERIFICATION OF E FFICIENTLY D ISTILLED RL P OLICIES WITH M ANY - SIDED G UARANTEES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Although deep reinforcement learning (DRL) has many success stories, the largescale deployment of policies learned through these advanced techniques in safetycritical scenarios is hindered by their lack of formal guarantees.
Revised: Although deep reinforcement learning (DRL) has many success stories, the largescale deployment of policies learned through these advanced techniques in safetycritical scenarios is hindered by their lack of formal guarantees.

Original: Variational Markov
Revised: Variational Markov

Original: Decision Processes (VAE-MDPs) are discrete latent space models that provide a reliable framework for distilling formally verifiable controllers from any RL policy.
Revised: Decision Processes (VAE-MDPs) are discrete latent space models that provide a reliable framework for distilling formally verifiable controllers from any RL policy.

Original: 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 (mode collapse, slow learning speed, poor dynamics estimates), primarily due to the absence of abstraction and representation guarantees to support latent optimization.
Revised: 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.

Original: 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.
Revised: 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.

Original: Our approach yields bisimulation guarantees while learning the distilled policy, allowing concrete optimization of the abstraction and representation model quality.
Revised: Our approach yields bisimulation guarantees while learning the distilled policy, allowing concrete optimization of the abstraction and representation model quality.

Original: Our experiments show that, besides distilling policies up to 10 times faster, the latent model quality is indeed better in general.
Revised: Our experiments show that, besides distilling policies up to 10 times faster, the latent model quality is indeed better in general.

Original: Moreover, we present experiments from a simple time-to-failure verification algorithm on the latent space.
Revised: Moreover, we present experiments from a simple time-to-failure verification algorithm on the latent space.

Original: The fact that our approach enables such simple verification techniques highlights its applicability.
Revised: The fact that our approach enables such simple verification techniques highlights its applicability.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: This work enables the formal verification of any such policies, learned by agents interacting with unknown, continuous environments modeled as Markov decision processes (MDPs).
Revised: This work enables the formal verification of any such policies, learned by agents interacting with unknown, continuous environments modeled as Markov decision processes (MDPs).

Original: 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).
Revised: 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).

Original: The latter are supported by bisimulation guarantees : intuitively, the agent behaves similarly in the original and latent models.
Revised: The latter are supported by bisimulation guarantees : intuitively, the agent behaves similarly in the original and latent models.

Original: 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.
Revised: 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.

Original: Variational MDPs (VAE-MDPs, Delgrange et al. 2022) offer a valuable framework for doing so.
Revised: Variational MDPs (VAE-MDPs, Delgrange et al. 2022) offer a valuable framework for doing so.

Original: 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 ).
Revised: 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 ).

Original: Whilst the bounds offer a confidence metric that enables the verification of performance and safety properties, VAE-MDPs suffer from several learning flaws.
Revised: Whilst the bounds offer a confidence metric that enables the verification of performance and safety properties, VAE-MDPs suffer from several learning flaws.

Original: 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.
Revised: 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.

Original: Second, variational autoencoders (VAEs) (Kingma & Welling, 2014; Hoffman et al., 2013) are known to suffer from mode collapse (e.g., Alemi et al. 2018) resulting in a deterministic mapping to a unique latent state in VAE-MDPs.
Revised: 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.

Original: 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.
Revised: 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.

Original: This requires extra regularizers, setting up annealing schemes and learning phases, and defining prioritized replay buffers to store transitions.
Revised: This requires extra regularizers, setting up annealing schemes and learning phases, and defining prioritized replay buffers to store transitions.

Original: Distillation through VAE-MDPs is thus a meticulous task, requiring a large step budget and tuning many hyperparameters.
Revised: Distillation through VAE-MDPs is thus a meticulous task, requiring a large step budget and tuning many hyperparameters.

Original: Building upon Wasserstein autoencoders (Tolstikhin et al., 2018) instead of VAEs, we introduce
Revised: Building upon Wasserstein autoencoders (Tolstikhin et al., 2018) instead of VAEs, we introduce

Original: Wasserstein auto-encoded MDPs (WAE-MDPs), which overcome those limitations.
Revised: Wasserstein auto-encoded MDPs (WAE-MDPs), which overcome those limitations.

Original: 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.
Revised: 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.

Original: In contrast to VAEs which rely on variational proxies, we derive a novel objective that directly incorporate the bisimulation bounds.
Revised: In contrast to VAEs which rely on variational proxies, we derive a novel objective that directly incorporate the bisimulation bounds.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Finally, we formally verified time-to-failure properties (e.g., Pnueli 1977) to emphasize the applicability of our approach.
Revised: Finally, we formally verified time-to-failure properties (e.g., Pnueli 1977) to emphasize the applicability of our approach.

Original: Other Related Work.
Revised: Other Related Work.

Original: Complementary works approach safe RL via formal methods (Junges et al., 2016; Alshiekh et al., 2018;
Revised: Complementary works approach safe RL via formal methods (Junges et al., 2016; Alshiekh et al., 2018;

Original: 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.
Revised: 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.

Original: They also include the work of Alamdari et al.
Revised: They also include the work of Alamdari et al.

Original: (2020), applying synthesis and model checking on policies distilled from RL, without quality guarantees.
Revised: (2020), applying synthesis and model checking on policies distilled from RL, without quality guarantees.

Original: 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).
Revised: 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).

Original: Finally,
Revised: Finally,

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: These two are, in contrast, optimized via the dual formulation of Wasserstein, as in Wassertein-GANs Arjovsky et al.
Revised: These two are, in contrast, optimized via the dual formulation of Wasserstein, as in Wassertein-GANs Arjovsky et al.

Original: 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.
Revised: 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.

Original: Works on representation learning for RL (Gelada et al., 2019; Castro et al., 2021; Zhang et al., 2021;
Revised: Works on representation learning for RL (Gelada et al., 2019; Castro et al., 2021; Zhang et al., 2021;

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In the following, we write ∆ p X q for the set of measures over (complete, separable metric space) X .
Revised: In the following, we write ∆ p X q for the set of measures over (complete, separable metric space) X .

Original: Markov decision processes (MDPs) are tuples M “ x S , A , P , R , ℓ, AP , s I y where S is a set of states ; A , a set of actions ; P : S ˆ A Ñ ∆ p S q , a probability transition function that maps the current state and action to a distribution over the next states; R : S ˆ A Ñ R , a reward function ; ℓ : S Ñ 2 AP , a labeling function over a set of atomic propositions AP ; and s I P S , the initial state .
Revised: Markov decision processes (MDPs) are tuples M “ x S , A , P , R , ℓ, AP , s I y where S is a set of states ; A , a set of actions ; P : S ˆ A Ñ ∆ p S q , a probability transition function that maps the current state and action to a distribution over the next states; R : S ˆ A Ñ R , a reward function ; ℓ : S Ñ 2 AP , a labeling function over a set of atomic propositions AP ; and s I P S , the initial state .

Original: If | A | “ 1 , M is a fully stochastic process called a Markov chain (MC).
Revised: If | A | “ 1 , M is a fully stochastic process called a Markov chain (MC).

Original: We write M s for the MDP obtained when replacing the initial state of M by s P S .
Revised: We write M s for the MDP obtained when replacing the initial state of M by s P S .

Original: An agent interacting in M produces trajectories , i.e., sequences of states and actions τ “ x s 0: T , a 0: T ´ 1 y where s 0 “ s I and s t ` 1 „ P p¨ | s t , a t q for t ă T .
Revised: An agent interacting in M produces trajectories , i.e., sequences of states and actions τ “ x s 0: T , a 0: T ´ 1 y where s 0 “ s I and s t ` 1 „ P p¨ | s t , a t q for t ă T .

Original: The set of infinite trajectories of M is Traj .
Revised: The set of infinite trajectories of M is Traj .

Original: We assume AP and labels being respectively one-hot and binary encoded.
Revised: We assume AP and labels being respectively one-hot and binary encoded.

Original: Given T Ď AP , we write s |ù T if s is labeled with T , i.e., ℓ p s q X T ‰ H , and s |ù ␣ T for s ­|ù T .
Revised: Given T Ď AP , we write s |ù T if s is labeled with T , i.e., ℓ p s q X T ‰ H , and s |ù ␣ T for s ­|ù T .

Original: We refer to MDPs with continuous state or s
Revised: We refer to MDPs with continuous state or s

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Original: [Equation]
Revised: [Equation]

Original: ¯ s ′ ϕ ψ
Revised: ¯ s ′ ϕ ψ

Original: M
Revised: M

Original: M
Revised: M

Original: Lξ ¯ π R
Revised: Lξ ¯ π R

Original: Lξ ¯
Revised: Lξ ¯

Original: [Figure]
Revised: [Figure]

Original: P
Revised: P

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: (¯ s
Revised: (¯ s

Original: [Figure]
Revised: [Figure]

Original: a )
Revised: a )

Original: [Figure]
Revised: [Figure]

Original: ¯ π
Revised: ¯ π

Original: [Figure]
Revised: [Figure]

Original: s ¯ s ϕ ¯ a π a ¯ s
Revised: s ¯ s ϕ ¯ a π a ¯ s

Original: ⋆ s ′ ¯ P P ¯ s ′ ϕ
Revised: ⋆ s ′ ¯ P P ¯ s ′ ϕ

Original: ϕ
Revised: ϕ

Original: A
Revised: A

Original: M M
Revised: M M

Original: Lξπ R
Revised: Lξπ R

Original: Lξπ P
Revised: Lξπ P

Original: R ( s,a )
Revised: R ( s,a )

Original: R (¯ s, ¯ a ) z ¯ a ′ z
Revised: R (¯ s, ¯ a ) z ¯ a ′ z

Original: ′ ¯
Revised: ′ ¯

Original: P ¯ π
Revised: P ¯ π

Original: W ξπ ξπ ¯ ξ ¯ π
Revised: W ξπ ξπ ¯ ξ ¯ π

Original: [Figure]
Revised: [Figure]

Original: action spaces as continuous MDPs .
Revised: action spaces as continuous MDPs .

Original: In that case, we assume S and A are complete separable metric spaces equipped with a Borel σ -algebra, and ℓ ´ 1 p T q is Borel-measurable for any T Ď AP .
Revised: In that case, we assume S and A are complete separable metric spaces equipped with a Borel σ -algebra, and ℓ ´ 1 p T q is Borel-measurable for any T Ď AP .

Original: Policies and stationary distributions.
Revised: Policies and stationary distributions.

Original: A (memoryless) policy π : S Ñ ∆ p
Revised: A (memoryless) policy π : S Ñ ∆ p

Original: A q prescribes which action to choose at each step of the interaction.
Revised: A q prescribes which action to choose at each step of the interaction.

Original: The set of memoryless policies of M is Π .
Revised: The set of memoryless policies of M is Π .

Original: The
Revised: The

Original: MDP M and π P Π induce an MC M π with unique probability measure P M π on the Borel σ -algebra over measurable subsets φ Ď Traj (Puterman, 1994).
Revised: MDP M and π P Π induce an MC M π with unique probability measure P M π on the Borel σ -algebra over measurable subsets φ Ď Traj (Puterman, 1994).

Original: We drop the superscript when the context is clear.
Revised: We drop the superscript when the context is clear.

Original: Define ξ tπ p s 1 | s q “ P M s π pt s 0: 8 , a 0: 8 | s t “ s 1 uq as the distribution giving the probability of being in each state of M s after t steps.
Revised: Define ξ tπ p s 1 | s q “ P M s π pt s 0: 8 , a 0: 8 | s t “ s 1 uq as the distribution giving the probability of being in each state of M s after t steps.

Original: B Ď S is a bottom strongly connected component (BSCC) of
Revised: B Ď S is a bottom strongly connected component (BSCC) of

Original: M π if (i) B is a maximal subset satisfying ξ tπ p s 1 | s q ą 0 for any s, s 1 P B and some t ě 0 , and (ii) E a „ π p¨| s q P p B | s, a q “ 1 for all s P S .
Revised: M π if (i) B is a maximal subset satisfying ξ tπ p s 1 | s q ą 0 for any s, s 1 P B and some t ě 0 , and (ii) E a „ π p¨| s q P p B | s, a q “ 1 for all s P S .

Original: The unique stationary distribution of B is ξ π P ∆ p B q .
Revised: The unique stationary distribution of B is ξ π P ∆ p B q .

Original: We write s, a „ ξ π for sampling s from ξ π then a from π .
Revised: We write s, a „ ξ π for sampling s from ξ π then a from π .

Original: An MDP M is ergodic if for all π P Π , the state space of M π consists of a unique aperiodic BSCC with ξ π “ lim t Ñ8 ξ tπ p¨ | s q for all s P S .
Revised: An MDP M is ergodic if for all π P Π , the state space of M π consists of a unique aperiodic BSCC with ξ π “ lim t Ñ8 ξ tπ p¨ | s q for all s P S .

Original: Value objectives.
Revised: Value objectives.

Original: Given π P Π , the value of a state s P S is the expected value of a random variable obtained by running π from s .
Revised: Given π P Π , the value of a state s P S is the expected value of a random variable obtained by running π from s .

Original: For a discount factor γ P r 0 , 1 s , we consider the following objectives.
Revised: For a discount factor γ P r 0 , 1 s , we consider the following objectives.

Original: (i) Discounted return : we write V π p s q “ E M s π
Revised: (i) Discounted return : we write V π p s q “ E M s π

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: ‰ for the expected discounted rewards accumulated along trajectories.
Revised: ‰ for the expected discounted rewards accumulated along trajectories.

Original: The typical goal of an RL agent is to learn a policy π ‹ that maximizes V π ‹ p s I q through interactions with the (unknown) MDP; (ii) Reachability : let C , T Ď AP , the (constrained) reachability event is C U T “ t s 0: 8 , a 0: 8 | D i P N , @
Revised: The typical goal of an RL agent is to learn a policy π ‹ that maximizes V π ‹ p s I q through interactions with the (unknown) MDP; (ii) Reachability : let C , T Ď AP , the (constrained) reachability event is C U T “ t s 0: 8 , a 0: 8 | D i P N , @

Original: j ă i, s j |ù C ^ s i |ù T u Ď
Revised: j ă i, s j |ù C ^ s i |ù T u Ď

Original: Traj .
Revised: Traj .

Original: We write V φπ p s q “ E M s π “ γ t ‹ 1 x s 0: 8 ,a 0: 8 y P φ ‰ for the discounted probability of satisfying φ “ C U T , where t ‹ is the length of the shortest trajectory prefix that allows satisfying φ .
Revised: We write V φπ p s q “ E M s π “ γ t ‹ 1 x s 0: 8 ,a 0: 8 y P φ ‰ for the discounted probability of satisfying φ “ C U T , where t ‹ is the length of the shortest trajectory prefix that allows satisfying φ .

Original: Intuitively, this denotes the discounted return of remaining in a region of the MDP where states are labeled with
Revised: Intuitively, this denotes the discounted return of remaining in a region of the MDP where states are labeled with

Original: C , until visiting for the first time a goal state labeled with T , and the return is the binary reward signal capturing this event.
Revised: C , until visiting for the first time a goal state labeled with T , and the return is the binary reward signal capturing this event.

Original: Safety w.r.t.
Revised: Safety w.r.t.

Original: failure states C can be expressed as the safety-constrained reachability to a destination T through ␣ C U T .
Revised: failure states C can be expressed as the safety-constrained reachability to a destination T through ␣ C U T .

Original: Notice that V φπ p s q “ P M s π p
Revised: Notice that V φπ p s q “ P M s π p

Original: φ q when γ “ 1 .
Revised: φ q when γ “ 1 .

Original: Latent MDP.
Revised: Latent MDP.

Original: Given the original (continuous, possibly unknown) environment model M , a latent space model is another (smaller, explicit)
Revised: Given the original (continuous, possibly unknown) environment model M , a latent space model is another (smaller, explicit)

Original: MDP M “ @
Revised: MDP M “ @

Original: S , A , P , R , ℓ, AP , s I
Revised: S , A , P , R , ℓ, AP , s I

Original: D with state-action space linked to the original one via state and action embedding functions : ϕ : S Ñ S and ψ : S ˆ A Ñ A .
Revised: D with state-action space linked to the original one via state and action embedding functions : ϕ : S Ñ S and ψ : S ˆ A Ñ A .

Original: Our goal is to learn the latter by optimizing an equivalence criterion between the two models.
Revised: Our goal is to learn the latter by optimizing an equivalence criterion between the two models.

Original: We refer to @
Revised: We refer to @

Original: M , ϕ, ψ
Revised: M , ϕ, ψ

Original: D as a latent space model of M and M as its latent MDP .
Revised: D as a latent space model of M and M as its latent MDP .

Original: We assume that d S is a metric on S , and write Π for the set of policies of M and V π for the values of running π P Π in M .
Revised: We assume that d S is a metric on S , and write Π for the set of policies of M and V π for the values of running π P Π in M .

Original: Remark 1 (Latent flow) .
Revised: Remark 1 (Latent flow) .

Original: The latent policy π can be seen as a policy in M (cf.
Revised: The latent policy π can be seen as a policy in M (cf.

Original: Fig.
Revised: Fig.

Original: 1a): states passed to π are first embedded with ϕ to the latent space, then the actions produced by π are executed via ψ in the original environment.
Revised: 1a): states passed to π are first embedded with ϕ to the latent space, then the actions produced by π are executed via ψ in the original environment.

Original: Let s P S , we write a „ π p¨ | s q for π p¨ | ϕ p s qq , then the reward and next state are respectively given by R p s, a q “ R p s, ψ p ϕ p s q , a qq and s 1 „ P p¨ | s, a q “ P p¨ | s, ψ p ϕ p s q , a qq .
Revised: Let s P S , we write a „ π p¨ | s q for π p¨ | ϕ p s qq , then the reward and next state are respectively given by R p s, a q “ R p s, ψ p ϕ p s q , a qq and s 1 „ P p¨ | s, a q “ P p¨ | s, ψ p ϕ p s q , a qq .

Original: 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 ξ P ∆
Revised: 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 ξ P ∆

Original: [Equation]
Revised: [Equation]

Original: where ϕ P p¨ | s, a q is the distribution of drawing s 1 „ P p¨ | s, a q then embedding s 1 “ ϕ p s  q , and
Revised: where ϕ P p¨ | s, a q is the distribution of drawing s 1 „ P p¨ | s, a q then embedding s 1 “ ϕ p s 1 q , and

Original: D is a discrepancy measure.
Revised: D is a discrepancy measure.

Original: Fig 1a depicts the losses when states and actions are drawn from a stationary distribution ξ π resulting from running π P Π in M .
Revised: Fig 1a depicts the losses when states and actions are drawn from a stationary distribution ξ π resulting from running π P Π in M .

Original: In this work, we focus on the case where D is the Wasserstein distance W d S : given two distributions P, Q over a measurable set X equipped with a metric d , W d is the solution of the optimal transport (OT) from P to Q , i.e., the minimum cost of changing P into Q (Villani, 2009): W d p P, Q q “ inf λ P Λ p P,Q q E x,y „ λ
Revised: In this work, we focus on the case where D is the Wasserstein distance W d S : given two distributions P, Q over a measurable set X equipped with a metric d , W d is the solution of the optimal transport (OT) from P to Q , i.e., the minimum cost of changing P into Q (Villani, 2009): W d p P, Q q “ inf λ P Λ p P,Q q E x,y „ λ

Original: d p x, y q , Λ p P, Q q being the set of all couplings of P and Q .
Revised: d p x, y q , Λ p P, Q q being the set of all couplings of P and Q .

Original: The Kantorovich duality yields W d p P, Q q “ sup f P F d E x „ P
Revised: The Kantorovich duality yields W d p P, Q q “ sup f P F d E x „ P

Original: f p x q ´ E x „ Q f p y q where F d is the set of 1-Lipschiz functions.
Revised: f p x q ´ E x „ Q f p y q where F d is the set of 1-Lipschiz functions.

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

Original: Bisimulation.
Revised: Bisimulation.

Original: A bisimulation B on M is a behavioral equivalence between states s 1 , s 2 P S so that, s 1 B s 2 iff (i) P p T | s 1 , a q “ P p T | s 2 , a q , (ii) ℓ p s 1 q “ ℓ p s 2 q , and (iii) R p s 1 , a q “ R p s 2 , a q for each action a PA
Revised: A bisimulation B on M is a behavioral equivalence between states s 1 , s 2 P S so that, s 1 B s 2 iff (i) P p T | s 1 , a q “ P p T | s 2 , a q , (ii) ℓ p s 1 q “ ℓ p s 2 q , and (iii) R p s 1 , a q “ R p s 2 , a q for each action a P

Original: Requirements (ii) and (iii) can be respectively relaxed depending on whether we focus only on behaviors formalized through AP or rewards.
Revised: Requirements (ii) and (iii) can be respectively relaxed depending on whether we focus only on behaviors formalized through AP or rewards.

Original: The relation can be extended to compare two MDPs (e.g., M and M ) by considering the disjoint union of their state space.
Revised: The relation can be extended to compare two MDPs (e.g., M and M ) by considering the disjoint union of their state space.

Original: We denote the largest bisimulation relation by „ .
Revised: We denote the largest bisimulation relation by „ .

Original: Characterized by a logical family of functional expressions derived from a logic L , bisimulation pseudometrics (Desharnais et al., 2004) generalize the notion of bisimilariy.
Revised: Characterized by a logical family of functional expressions derived from a logic L , bisimulation pseudometrics (Desharnais et al., 2004) generalize the notion of bisimilariy.

Original: More specifically, given a policy π P Π , we consider a family F of real-valued functions parameterized by a discount factor γ and defining the semantics of L in M π .
Revised: More specifically, given a policy π P Π , we consider a family F of real-valued functions parameterized by a discount factor γ and defining the semantics of L in M π .

Original: Such functional expressions allow to formalize discounted properties such as reachability, safety, as well as general ω -regular specifications (Chatterjee et al., 2008) and may include rewards as well (Ferns et al., 2014).
Revised: Such functional expressions allow to formalize discounted properties such as reachability, safety, as well as general ω -regular specifications (Chatterjee et al., 2010) and may include rewards as well (Ferns et al., 2014).

Original: The pseudometric d „ π is defined as the largest behavioral difference d „ π p s 1 , s 2 q “ sup f P F | f p s 1 q ´ f p s 2 q| , and its kernel is bisimilarity :
Revised: The pseudometric d „ π is defined as the largest behavioral difference d „ π p s 1 , s 2 q “ sup f P F | f p s 1 q ´ f p s 2 q| , and its kernel is bisimilarity :

Original: d „ π p s 1 , s 2 q “ 0 iff s 1 „ s 2 .
Revised: d „ π p s 1 , s 2 q “ 0 iff s 1 „ s 2 .

Original: In particular, value functions are Lipschitz-continuous w.r.t.
Revised: In particular, value functions are Lipschitz-continuous w.r.t.

Original: d „ π : | V ¨ π p s 1 q ´ V ¨ π p s 2 q| ď K d „ π p s 1 , s 2 q , where K is 1 { p 1 ´ γ q if rewards are included in F and 1 otherwise.
Revised: d „ π : | V ¨ π p s 1 q ´ V ¨ π p s 2 q| ď K d „ π p s 1 , s 2 q , where K is 1 { p 1 ´ γ q if rewards are included in F and 1 otherwise.

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

Original: Assumption 2.1.
Revised: Assumption 2.1.

Original: MDP M is ergodic, Im p R q is a bounded space scaled in r ´ 1 { 2 , 1 { 2 s , and the embedding function preserves the labels, i.e., ϕ p s q “ s ùñ ℓ p s q “ ℓ p s q for s P S , s P S .
Revised: MDP M is ergodic, Im p R q is a bounded space scaled in r ´ 1 { 2 , 1 { 2 s , and the embedding function preserves the labels, i.e., ϕ p s q “ s ùñ ℓ p s q “ ℓ p s q for s P S , s P S .

Original: Note that the ergodicity assumption is compliant with episodic RL and a wide range of continuous learning tasks (see Huang 2020; Delgrange et al.
Revised: Note that the ergodicity assumption is compliant with episodic RL and a wide range of continuous learning tasks (see Huang 2020; Delgrange et al.

Original: 2022 for detailed discussions on this setting).
Revised: 2022 for detailed discussions on this setting).

Original: Bisimulation bounds (Delgrange et al., 2022).
Revised: Bisimulation bounds (Delgrange et al., 2022).

Original: M being set over continuous spaces with possibly unknown dynamics, evaluating d „ can turn out to be particularly arduous, if not intractable.
Revised: M being set over continuous spaces with possibly unknown dynamics, evaluating d „ can turn out to be particularly arduous, if not intractable.

Original: A solution is to evaluate the original and latent model bisimilarity via local losses: fix π P Π , assume M is discrete, then given the induced stationary distribution ξ π in M , let s 1 , s 2 P S with ϕ p s 1 q “ ϕ p s 2 q :
Revised: A solution is to evaluate the original and latent model bisimilarity via local losses: fix π P Π , assume M is discrete, then given the induced stationary distribution ξ π in M , let s 1 , s 2 P S with ϕ p s 1 q “ ϕ p s 2 q :

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: The local losses and related bounds can be efficiently PAC-estimated.
Revised: The local losses and related bounds can be efficiently PAC-estimated.

Original: Our goal is to learn a latent model where the behaviors of the agent executing π can be formally verified, and the bounds offer a confidence metric allowing to lift the guarantees obtained this way back to the original model M , when the latter operates under π .
Revised: Our goal is to learn a latent model where the behaviors of the agent executing π can be formally verified, and the bounds offer a confidence metric allowing to lift the guarantees obtained this way back to the original model M , when the latter operates under π .

Original: We show in the following how to learn a latent space model by optimizing the aforementioned bounds, and distill policies π P Π obtained via any RL technique to a latent policy π P Π .
Revised: We show in the following how to learn a latent space model by optimizing the aforementioned bounds, and distill policies π P Π obtained via any RL technique to a latent policy π P Π .

Original:  W ASSERSTEIN A UTO - ENCODED MDP S
Revised:  W ASSERSTEIN A UTO - ENCODED MDP S

Original: Fix
Revised: Fix

Original: [Equation]
Revised: [Equation]

Original: D and @
Revised: D and @

Revised: M

Original: [Equation]
Revised: [Equation]

Original: , ϕ
Revised: , ϕ

Original: [Equation]
Revised: [Equation]

Original: , ψ
Revised: , ψ

Original: [Equation]
Revised: [Equation]

Original: D as a latent space model of M parameterized by ι and θ .
Revised: D as a latent space model of M parameterized by ι and θ .

Original: Our method relies on learning a behavioral model ξ θ of M from which we can retrieve the latent space model and distill π .
Revised: Our method relies on learning a behavioral model ξ θ of M from which we can retrieve the latent space model and distill π .

Original: This can be achieved via the minimization of a suitable discrepancy between ξ θ and M π .
Revised: This can be achieved via the minimization of a suitable discrepancy between ξ θ and M π .

Original: VAE-MDPs optimize a lower bound on the likelihood of the dynamics of M π using the Kullback-Leibler divergence , yielding (i) M θ , (ii) a distillation π θ of π , and (iii) ϕ ι and ψ θ .
Revised: VAE-MDPs optimize a lower bound on the likelihood of the dynamics of M π using the Kullback-Leibler divergence , yielding (i) M θ , (ii) a distillation π θ of π , and (iii) ϕ ι and ψ θ .

Original: Local losses are not directly minimized, but rather variational proxies that do not offer theoretical guarantees during the learning process.
Revised: Local losses are not directly minimized, but rather variational proxies that do not offer theoretical guarantees during the learning process.

Original: 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.
Revised: 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.

Original: Proofs of the claims made in this Section are provided in Appendix A.
Revised: Proofs of the claims made in this Section are provided in Appendix A.

Original: 3. T HE O BJECTIVE F UNCTION
Revised: 3. T HE O BJECTIVE F UNCTION

Original: Assume that S , A , and Im p R q are respectively equipped with metrics d S , d A , and d R , we define the raw transition distance metric ⃗d as the component-wise sum of distances between states, actions, and rewards occurring of along transitions: ⃗d px s 1 , a 1 , r 1 , s 1 1 y , x s 2 , a 2 , r 2 , s 1 2 yq “ d S 
Revised: Assume that S , A , and Im p R q are respectively equipped with metrics d S , d A , and d R , we define the raw transition distance metric ⃗d as the component-wise sum of distances between states, actions, and rewards occurring of along transitions: ⃗d px s 1 , a 1 , r 1 , s 1 1 y , x s 2 , a 2 , r 2 , s 1 2 yq “ d S [Equation] s

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: R p
Revised: R p

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Revised: 1 , r 2 q

Original: [Equation]
Revised: [Equation]

Original:  S p s 1 1 , s 1 2 q .
Revised: d S p s 1 1 , s 1 2 q .

Original: Given Assumption 2.1, we consider the OT between local distributions, where traces are drawn from episodic RL processes or infinite interactions (we show in Appendix A.1 that considering the OT between trace-based distributions in the limit amounts to reasoning about stationary distributions).
Revised: Given Assumption 2.1, we consider the OT between local distributions, where traces are drawn from episodic RL processes or infinite interactions (we show in Appendix A.1 that considering the OT between trace-based distributions in the limit amounts to reasoning about stationary distributions).

Original: Out goal is to minimize W ⃗d p ξ π , ξ θ q so that
Revised: Our goal is to minimize W ⃗d p ξ π , ξ θ q so that

Original: [Equation]
Revised: [Equation]

Original: where P θ is a transition decoder and ¯ ξ π θ denotes the stationary distribution of the latent model M θ .
Revised: where P θ is a transition decoder and ¯ ξ π θ denotes the stationary distribution of the latent model M θ .

Original: As proved by Bousquet et al.
Revised: As proved by Bousquet et al.

Original: (2017), this model allows to derive a simpler form of the OT: instead of finding the optimal coupling of (i) the stationary distribution ξ π of M π and (ii) the behavioral model ξ θ , in the primal definition of W ⃗d p ξ π , ξ θ q , it is sufficient to find an encoder q whose marginal is given by Q p s, a, s 1 q “ E s,a,s 1 „ ξ π q p s, a, s 1 | s, a, s 1 q and identical to ξ π .
Revised: (2017), this model allows to derive a simpler form of the OT: instead of finding the optimal coupling of (i) the stationary distribution ξ π of M π and (ii) the behavioral model ξ θ , in the primal definition of W ⃗d p ξ π , ξ θ q , it is sufficient to find an encoder q whose marginal is given by Q p s, a, s 1 q “ E s,a,s 1 „ ξ π q p s, a, s 1 | s, a, s 1 q and identical to ξ π .

Original: This is summarized in the following Theorem, yielding a particular case of Wasserstein-autoencoder Tolstikhin et al.
Revised: This is summarized in the following Theorem, yielding a particular case of Wasserstein-autoencoder Tolstikhin et al.

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: Let ξ θ and P θ be respectively a behavioral model and transition decoder as defined in
Revised: Let ξ θ and P θ be respectively a behavioral model and transition decoder as defined in

Original: Eq. 2, G θ : S Ñ S be a state-wise decoder, and ψ θ be an action embedding function.
Revised: Eq. 2, G θ : S Ñ S be a state-wise decoder, and ψ θ be an action embedding function.

Original: Assume P θ is deterministic with Dirac function G θ p s, a, s 1 q
Revised: Assume P θ is deterministic with Dirac function G θ p s, a, s 1 q

Original: [Equation]
Revised: [Equation]

Original: , then
Revised: , then

Original: [Equation]
Revised: [Equation]

Original: Henceforth, fix ϕ ι : S Ñ S and ϕ
Revised: Henceforth, fix ϕ ι : S Ñ S and ϕ

Original: A ι : S ˆ A Ñ ∆ `
Revised: A ι : S ˆ A Ñ ∆ `

Original: A ˘ as parameterized state and action encoders with ϕ ι p s, a, s 1 | s, a, s 1 q “ 1 ϕ ι p s q“ s ¨ ϕ
Revised: A ˘ as parameterized state and action encoders with ϕ ι p s, a, s 1 | s, a, s 1 q “ 1 ϕ ι p s q“ s ¨ ϕ

Original: A ι p a | s, a q ¨ 1 ϕ ι p s 1 q“ s 1 , and define the marginal encoder as
Revised: A ι p a | s, a q ¨ 1 ϕ ι p s 1 q“ s 1 , and define the marginal encoder as

Original: Q ι “ E s,a,s 1 „ ξ π ϕ ι p¨ | s, a, s 1 q .
Revised: Q ι “ E s,a,s 1 „ ξ π ϕ ι p¨ | s, a, s 1 q .

Original: Training the model components can be achieved via the objective:
Revised: Training the model components can be achieved via the objective:

Original: [Equation]
Revised: [Equation]

Original: where D is an arbitrary discrepancy metric and β ą 0 a hyperparameter.
Revised: where D is an arbitrary discrepancy metric and β ą 0 a hyperparameter.

Original: Intuitively, the encoder ϕ
Revised: Intuitively, the encoder ϕ

Original: ι can be learned by enforcing its marginal distribution Q ι to match ¯ ξ π θ through this discrepancy.
Revised: ι can be learned by enforcing its marginal distribution Q ι to match ¯ ξ π θ through this discrepancy.

Original: Remark 2 .
Revised: Remark 2 .

Original: If M has a discrete action space, then learning A is not necessary.
Revised: If M has a discrete action space, then learning A is not necessary.

Original: We can set A “ A using identity functions for the action encoder and decoder (details in Appendix A.2).
Revised: We can set A “ A using identity functions for the action encoder and decoder (details in Appendix A.2).

Original: When π is executed in M , observe that its parallel execution in M θ is enabled by the action encoder ϕ
Revised: When π is executed in M , observe that its parallel execution in M θ is enabled by the action encoder ϕ

Original: A ι : given an original state s P S , π first prescribes the action a „ π p¨ | s q , which is then embedded in the latent space via a „ ϕ
Revised: A ι : given an original state s P S , π first prescribes the action a „ π p¨ | s q , which is then embedded in the latent space via a „ ϕ

Original: A ι p¨ | ϕ ι p s q , a q (cf.
Revised: A ι p¨ | ϕ ι p s q , a q (cf.

Original: Fig.
Revised: Fig.

Original: 1b).
Revised: 1b).

Original: This parallel execution, along with setting
Revised: This parallel execution, along with setting

Original: D to W ⃗d , yield an upper bound on the latent regularization, compliant with the bisimulation bounds.
Revised: D to W ⃗d , yield an upper bound on the latent regularization, compliant with the bisimulation bounds.

Original: A two-fold regularizer is obtained thereby, defining the foundations of our objective function:
Revised: A two-fold regularizer is obtained thereby, defining the foundations of our objective function:

Original: Lemma 3.2.
Revised: Lemma 3.2.

Original: Define T p s, a, s
Revised: Define T p s, a, s

Original: [Equation]
Revised: [Equation]

Original: q
Revised: q

Original: [Equation]
Revised: [Equation]

Original: s q“ s ¨ ϕ
Revised: s q“ s ¨ ϕ

Original: A ι p a | s, a q¨ P θ p s 1 | s, a qs as the distribution of drawing state-action pairs from interacting with M , embedding them to the latent spaces, and finally letting them transition to their successor state in M θ .
Revised: A ι p a | s, a q¨ P θ p s 1 | s, a qs as the distribution of drawing state-action pairs from interacting with M , embedding them to the latent spaces, and finally letting them transition to their successor state in M θ .

Original: Then, W ⃗d
Revised: Then, W ⃗d

Original: [Equation]
Revised: [Equation]

Original: We therefore define the W 2 AE-MDP ( Wasserstein-Wasserstein auto-encoded MDP ) objective as: min
Revised: We therefore define the W 2 AE-MDP ( Wasserstein-Wasserstein auto-encoded MDP ) objective as: min

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Input: batch size N , max.
Revised: Input: batch size N , max.

Original: step T , no. of regularizer updates m , penalty coefficient δ ąfor t “ 1 to T do for i “ 1 to N do
Revised: step T , no. of regularizer updates m , penalty coefficient δ ąfor t “ 1 to T do for i “ 1 to N do

Original: Sample a transition s i , a i , r i , s 1 i from the original environment via ξ π
Revised: Sample a transition s i , a i , r i , s 1 i from the original environment via ξ π

Original: Embed the transition into the latent space by drawing s i , a i , s 1 i from ϕ ι p¨ | s i , a i , s 1 i q
Revised: Embed the transition into the latent space by drawing s i , a i , s 1 i from ϕ ι p¨ | s i , a i , s 1 i q

Original: Make the latent space model transition to the next latent state: s ‹ i „ P θ p¨ | s i , a i q
Revised: Make the latent space model transition to the next latent state: s ‹ i „ P θ p¨ | s i , a i q

Original: Sample a latent transition from ¯ ξ π θ : z i „ ¯ ξ π θ , a 1 i „
Revised: Sample a latent transition from ¯ ξ π θ : z i „ ¯ ξ π θ , a 1 i „

Original: π θ p¨ |
Revised: π θ p¨ |

Original: z i q , and z 1 i „ P θ p¨ | z
Revised: z i q , and z 1 i „ P θ p¨ | z

Original: i , a 1 i q
Revised: i , a 1 i q

Original: W Ð ř Ni “ 1 φ ξω p s i , a i , s ‹ i q ´ φ
Revised: W Ð ř Ni “ 1 φ ξω p s i , a i , s ‹ i q ´ φ

Original: ξω p
Revised: ξω p

Original: z i , a 1 i , z 1 i q
Revised: z i , a 1 i , z 1 i q

Original: ` φ P ω p
Revised: ` φ P ω p

Original: s i , a i , s i , a i , s
Revised: s i , a i , s i , a i , s

Original: 1 i q ´ φ P ω p s i , a i , s
Revised: 1 i q ´ φ P ω p s i , a i , s

Original: i , a i , s ‹
Revised: i , a i , s ‹

Original: i q
Revised: i q

Original: P Ð ř Ni “ 1 G P ` φ ξω , x s i , a i , s ‹
Revised: P Ð ř Ni “ 1 G P ` φ ξω , x s i , a i , s ‹

Original: i y , x z i , a 1 i , z 1 i y ˘ ` G P ` x ÞÑ φ P ω p s i , a i , s
Revised: i y , x z i , a 1 i , z 1 i y ˘ ` G P ` x ÞÑ φ P ω p s i , a i , s

Original: i , a i , x q , s 1 i , s ‹ i ˘
Revised: i , a i , x q , s 1 i , s ‹ i ˘

Original: Update the Lipschitz networks parameters ω by ascending 1 { N ¨ p β W ` δ P q if t mod m “ 0 then
Revised: Update the Lipschitz networks parameters ω by ascending 1 { N ¨ p β W ` δ P q if t mod m “ 0 then

Original: L Ð ř Ni “ 1 d S p s i , G θ p
Revised: L Ð ř Ni “ 1 d S p s i , G θ p

Original: s i qq
Revised: s i qq

Original: [Equation]
Revised: [Equation]

Original: d
Revised: d

Original: [Equation]
Revised: [Equation]

Original: p a i , ψ θ p s i , a i qq ` d R ` r i , R θ p s i , a i q ˘ ` d S p s 1 i , G θ p s 1 i qq
Revised: p a i , ψ θ p s i , a i qq ` d R ` r i , R θ p s i , a i q ˘ ` d S p s 1 i , G θ p s 1 i qq

Original: Update the latent space model parameters x ι, θ y by descending 1 { N ¨ p L ` β W q function G P p φ ω , x , y q Ź Gradient penalty for n ϵ „ U p 0 , 1 q
Revised: Update the latent space model parameters x ι, θ y by descending 1 { N ¨ p L ` β W q function G P p φ ω , x , y q Ź Gradient penalty for n ϵ „ U p 0 , 1 q

Original: ˜ x Ð ϵ x ` p 1 ´ ϵ q y where W ξ π “ W ⃗d `
Revised: ˜ x Ð ϵ x ` p 1 ´ ϵ q y where W ξ π “ W ⃗d `

Original: T , ¯ ξ π θ ˘ and L ξ π P are respectively called steady-state and transition regularizers.
Revised: T , ¯ ξ π θ ˘ and L ξ π P are respectively called steady-state and transition regularizers.

Original: The former allows to quantify the distance between the stationary distributions respectively induced by π in M and π θ in M θ , further enabling the distillation.
Revised: The former allows to quantify the distance between the stationary distributions respectively induced by π in M and π θ in M θ , further enabling the distillation.

Original: The latter allows to learn the latent dynamics.
Revised: The latter allows to learn the latent dynamics.

Original: Note that L ξ π R and L ξ π P — set over ξ π instead of ξ π θ — are not sufficient to ensure the bisimulation bounds (Eq. 1): running π in M θ depends on the parallel execution of π in the original model, which does not permit its (conventional) verification.
Revised: Note that L ξ π R and L ξ π P — set over ξ π instead of ξ π θ — are not sufficient to ensure the bisimulation bounds (Eq. 1): running π in M θ depends on the parallel execution of π in the original model, which does not permit its (conventional) verification.

Original: Breaking this dependency is enabled by learning the distillation π θ through W ξ π , as shown in Fig.
Revised: Breaking this dependency is enabled by learning the distillation π θ through W ξ π , as shown in Fig.

Original: 1b: minimizing W ξ π allows to make ξ π and ¯ ξ π θ closer together, further bridging the gap of the discrepancy between π and π θ .
Revised: 1b: minimizing W ξ π allows to make ξ π and ¯ ξ π θ closer together, further bridging the gap of the discrepancy between π and π θ .

Original: At any time, recovering the local losses along with the linked bisimulation bounds in the objective function of the
Revised: At any time, recovering the local losses along with the linked bisimulation bounds in the objective function of the

Original: W 2 AE-MDP is allowed by considering the latent policy resulting from this distillation:
Revised: W 2 AE-MDP is allowed by considering the latent policy resulting from this distillation:

Original: Theorem 3.3.
Revised: Theorem 3.3.

Original: Assume that traces are generated by running a latent policy π P Π in the original environment and let d R be the usual Euclidean distance, then the W 2 AE-MDP objective is
Revised: Assume that traces are generated by running a latent policy π P Π in the original environment and let d R be the usual Euclidean distance, then the W 2 AE-MDP objective is

Original: [Equation]
Revised: [Equation]

Original: Optimizing the regularizers is enabled by the dual form of the OT: we introduce two parameterized networks, φ ξω and φ P ω , constrained to be 1 -Lipschitz and trained to attain the supremum of the dual:
Revised: Optimizing the regularizers is enabled by the dual form of the OT: we introduce two parameterized networks, φ ξω and φ P ω , constrained to be 1 -Lipschitz and trained to attain the supremum of the dual:

Original: [Equation]
Revised: [Equation]

Original: ,a „ ξ π
Revised: ,a „ ξ π

Original: E a „ ϕ A ι p¨| ϕ
Revised: E a „ ϕ A ι p¨| ϕ

Original: ι p s q ,a q
Revised: ι p s q ,a q

Original: E s ‹ „ P θ p¨| ϕ
Revised: E s ‹ „ P θ p¨| ϕ

Original: ι p s q ,a q φ ξω p ϕ ι p s q , a, s ‹ q ´
Revised: ι p s q ,a q φ ξω p ϕ ι p s q , a, s ‹ q ´

Original: E z,a 1 ,z 1 „ ¯ ξ πθ φ ξω ` z, a 1 , z
Revised: E z,a 1 ,z 1 „ ¯ ξ πθ φ ξω ` z, a 1 , z

Original: 1 ˘
Revised: 1 ˘

Original: L ξ π
Revised: L ξ π

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: q “ max ω
Revised: q “ max ω

Original: E s
Revised: E s

Original: [Equation]
Revised: [Equation]

Original: a,s 1 „
Revised: a,s 1 „

Original: [Equation]
Revised: [Equation]

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: ,a,s
Revised: ,a,s

Original: [Equation]
Revised: [Equation]

Original: s,a,
Revised: s,a,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ` s, a,
Revised: ` s, a,

Original: [Equation]
Revised: [Equation]

Original: , a,
Revised: , a,

Original: [Equation]
Revised: [Equation]

Original: E s
Revised: E s

Original: [Equation]
Revised: [Equation]

Original: „ P
Revised: „ P

Original: [Equation]
Revised: [Equation]

Original: p¨
Revised: p¨

Original: [Equation]
Revised: [Equation]

Original: ,a q φ P
Revised: ,a q φ P

Original: [Equation]
Revised: [Equation]

Original: p s, a, s, a, s
Revised: p s, a, s, a, s

Original: [Equation]
Revised: [Equation]

Original: q ı
Revised: q ı

Original: Details to derive this tractable form of L ξ π P p ω q are in Appendix A.5.
Revised: Details to derive this tractable form of L ξ π P p ω q are in Appendix A.5.

Original: The networks are constrained via the gradient penalty approach of Gulrajani et al.
Revised: The networks are constrained via the gradient penalty approach of Gulrajani et al.

Original: (2017), leveraging that any differentiable function is 1 -Lipschitz
Revised: (2017), leveraging that any differentiable function is 1 -Lipschitz

Original: iff it has gradients with norm at most 1 everywhere (we show in Appendix A.6 this is still valid for relaxations of discrete spaces).
Revised: iff it has gradients with norm at most 1 everywhere (we show in Appendix A.6 this is still valid for relaxations of discrete spaces).

Original: The final learning process is presented in Algorithm 1.
Revised: The final learning process is presented in Algorithm 1.

Original: 3. D ISCRETE L ATENT S PACES
Revised: 3. D ISCRETE L ATENT S PACES

Original: To enable the verification of latent models supported by the bisimulation guarantees of Eq.
Revised: To enable the verification of latent models supported by the bisimulation guarantees of Eq.

Original: 1, we focus on the special case of discrete latent space models .
Revised: 1, we focus on the special case of discrete latent space models .

Original: Our approach relies on continuous relaxation of discrete random variables, regulated by some temperature parameter(s) λ : discrete random variables are retrieved as λ Ñ 0 , which amounts to applying a rounding operator.
Revised: Our approach relies on continuous relaxation of discrete random variables, regulated by some temperature parameter(s) λ : discrete random variables are retrieved as λ Ñ 0 , which amounts to applying a rounding operator.

Original: For training, we use the ϕ ι s,s ′ →
Revised: For training, we use the ϕ ι s,s ′ →

Original: G θ ˆ s, ˆ s ′
Revised: G θ ˆ s, ˆ s ′

Original: P θ
Revised: P θ

Original: R θ ¯ s → a → ˆ r ¯ s, ¯ s ′
Revised: R θ ¯ s → a → ˆ r ¯ s, ¯ s ′

Original: →
Revised: →

Original: ϕ A ι ¯ s → a → a ... one-hot encoding ¯ s → a → ψ θ ˆ a ...
Revised: ϕ A ι ¯ s → a → a ... one-hot encoding ¯ s → a → ψ θ ˆ a ...

Original: 11 eval.
Revised: 11 eval.

Original: training λ → 0 binary encoding λ →. 2 ...¯ s, z
Revised: training λ → 0 binary encoding λ →. 2 ...¯ s, z

Original: → a, a ′ → π θ z → evaluation training 0 { ⇒ ℓ ( s )=
Revised: → a, a ′ → π θ z → evaluation training 0 { ⇒ ℓ ( s )=

Original: ℓ (¯ s ) ⇒ ℓ ( s ′ )=
Revised: ℓ (¯ s ) ⇒ ℓ ( s ′ )=

Original: ℓ (¯ s ′ )
Revised: ℓ (¯ s ′ )

Original: l, l ′ λ →
Revised: l, l ′ λ →

Original: H λ φ P ω φ P ω
Revised: H λ φ P ω φ P ω

Original: L ξ π P r L ξ π R 0 ...
Revised: L ξ π P r L ξ π R 0 ...

Original: ¯ s ′ → ¯ s ⋆ → ¯ ξ ¯ π θ 0 0 ... 

Revised: ¯ s ′ → ¯ s ⋆ → ¯ ξ ¯ π θ 0 0 ...

Original: MAF φ ξω φ ξω
Revised: MAF φ ξω φ ξω

Original: W ξ π z labellogits ¯ s ⋆ z ′ a ′ ¯ s → a → ¯ s ⋆ →
Revised: W ξ π z labellogits ¯ s ⋆ z ′ a ′ ¯ s → a → ¯ s ⋆ →

Original: MAF ¯ s, ¯ s ′
Revised: MAF ¯ s, ¯ s ′

Original: Gumblesoftmax logistic sampling logistic sampling
Revised: Gumblesoftmax logistic sampling logistic sampling

Original: Gumblesoftmax ∼ ∼ ∼ ∼
Revised: Gumblesoftmax ∼ ∼ ∼ ∼

Original: [Figure]
Revised: [Figure]

Original: → temperature-controlled relaxations to differentiate the objective and let the gradient flow through the network.
Revised: → temperature-controlled relaxations to differentiate the objective and let the gradient flow through the network.

Original: When we deploy the latent policy in the environment and formally check the latent model, the zero-temperature limit is used.
Revised: When we deploy the latent policy in the environment and formally check the latent model, the zero-temperature limit is used.

Original: An overview of the approach is depticted in Fig.
Revised: An overview of the approach is depticted in Fig.

Original: State encoder.
Revised: State encoder.

Original: We work with a binary representation of the latent states.
Revised: We work with a binary representation of the latent states.

Original: First, this induces compact networks, able to deal with a large discrete space via a tractable number of parameter variables.
Revised: First, this induces compact networks, able to deal with a large discrete space via a tractable number of parameter variables.

Original: But most importantly, this ensures that Assumption 2.1 is satisfied: let n “ log 2 | S | , we reserve | AP | bits in S and each time s P S is passed to ϕ ι , n ´
Revised: But most importantly, this ensures that Assumption 2.1 is satisfied: let n “ log 2 | S | , we reserve | AP | bits in S and each time s P S is passed to ϕ ι , n ´

Original: | AP | bits are produced and concatenated with ℓ p s q , ensuring a perfect reconstruction of the labels and further bisimulation bounds.
Revised: | AP | bits are produced and concatenated with ℓ p s q , ensuring a perfect reconstruction of the labels and further bisimulation bounds.

Original: To produce Bernoulli variables, ϕ ι deterministically maps s to a latent code z , passed to the Heaviside H p
Revised: To produce Bernoulli variables, ϕ ι deterministically maps s to a latent code z , passed to the Heaviside H p

Original: z q “ 1 z ą 0 .
Revised: z q “ 1 z ą 0 .

Original: We train ϕ
Revised: We train ϕ

Original: ι
Revised: ι

Original: by using the smooth approximation
Revised: by using the smooth approximation

Original: H λ p
Revised: H λ p

Original: z q “ σ p 2 z { λ q , satisfying H “ lim λ Ñ 0 H λ .
Revised: z q “ σ p 2 z { λ q , satisfying H “ lim λ Ñ 0 H λ .

Original: Latent distributions.
Revised: Latent distributions.

Original: Besides the discontinuity of their latent image space, a major challenge of optimizing over discrete distributions is sampling , required to be a differentiable operation.
Revised: Besides the discontinuity of their latent image space, a major challenge of optimizing over discrete distributions is sampling , required to be a differentiable operation.

Original: We circumvent this by using concrete distributions (Jang et al., 2017; Maddison et al., 2017): the idea is to sample reparameterizable random variables from λ -parameterized distributions, and applying a differentiable, nonlinear operator in downstream.
Revised: We circumvent this by using concrete distributions (Jang et al., 2017; Maddison et al., 2017): the idea is to sample reparameterizable random variables from λ -parameterized distributions, and applying a differentiable, nonlinear operator in downstream.

Original: We use the Gumble softmax trick to sample from distributions over (one-hot encoded) latent actions ( ϕ A ι , π θ ).
Revised: We use the Gumble softmax trick to sample from distributions over (one-hot encoded) latent actions ( ϕ A ι , π θ ).

Original: For binary distributions ( P θ , ¯ ξ π θ ), each relaxed Bernoulli with logit α is retrieved by drawing a logistic random variable located in α { λ and scaled to 1 { λ , then applying a sigmoid in downstream.
Revised: For binary distributions ( P θ , ¯ ξ π θ ), each relaxed Bernoulli with logit α is retrieved by drawing a logistic random variable located in α { λ and scaled to 1 { λ , then applying a sigmoid in downstream.

Original: We emphasis that this trick alone (as used by Corneil et al.
Revised: We emphasis that this trick alone (as used by Corneil et al.

Original: ; Delgrange et al. 2022) is not sufficient: it yields independent Bernoullis, being too restrictive in general, which prevents from learning sound transition dynamics (cf.
Revised: ; Delgrange et al. 2022) is not sufficient: it yields independent Bernoullis, being too restrictive in general, which prevents from learning sound transition dynamics (cf.

Original: Ex. 1).
Revised: Ex. 1).

Original: ss
Revised: ss

Original: [Figure]
Revised: [Figure]

Original: goal u
Revised: goal u

Original: [Figure]
Revised: [Figure]

Original: u
Revised: u

Original: [Figure]
Revised: [Figure]

Original: u 1 {1 {1 {
Revised: u 1 {1 {1 {

Original: Example 1 .
Revised: Example 1 .

Original: Let M be the discrete MC of Fig.
Revised: Let M be the discrete MC of Fig.

Original: In one-hot, AP “ t goal : x 1 , 0 y , unsafe : x 0 , 1 yu .
Revised: In one-hot, AP “ t goal : x 1 , 0 y , unsafe : x 0 , 1 yu .

Original: We assume that 3 bits are used for the (binary) state space, with S “ t s 0 : x 0 , 0 , 0 y , s 1 : x 1 , 0 , 

Revised: We assume that 3 bits are used for the (binary) state space, with S “ t s 0 : x 0 , 0 , 0 y , s 1 : x 1 , 0 ,

Original: y , s 2 : x 0 , 1 , 

Revised: y , s 2 : x 0 , 1 ,

Original: y , s 3 : x 0 , 1 , 1 yu (the two first bits are reserved for the labels).
Revised: y , s 3 : x 0 , 1 , 1 yu (the two first bits are reserved for the labels).

Original: Considering each bit as being independent is not sufficient to learn P : the optimal estimation P θ ‹ p¨ | s 0 q is in that case represented by the independent Bernoulli vector b “ x 1 { 2 , 1 { 2 , 1 { 4 y , giving the probability to go from s 0 to each bit independently .
Revised: Considering each bit as being independent is not sufficient to learn P : the optimal estimation P θ ‹ p¨ | s 0 q is in that case represented by the independent Bernoulli vector b “ x 1 { 2 , 1 { 2 , 1 { 4 y , giving the probability to go from s 0 to each bit independently .

Original: This yields a poor estimation of the actual transition function: P
Revised: This yields a poor estimation of the actual transition function: P

Original: [Equation]
Revised: [Equation]

Original: We consider instead relaxed multivariate Bernoulli distributions by decomposing P P ∆ `
Revised: We consider instead relaxed multivariate Bernoulli distributions by decomposing P P ∆ `

Original: S ˘ as a product of conditionals: P p s q “ ś ni “ 1 P p s i | s 1: i ´ 1 q where s i is the i th entry (bit) of s .
Revised: S ˘ as a product of conditionals: P p s q “ ś ni “ 1 P p s i | s 1: i ´ 1 q where s i is the i th entry (bit) of s .

Original: We learn
Revised: We learn

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: such distributions by introducing a masked autoregressive flow (MAF, Papamakarios et al. 2017) for relaxed Bernoullis via the recursion: s i “ σ p l i ` α i { λ q , where l i „ Logistic p 0 , 1 q , α i “ f i p s 1: i ´ 1 q , 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.
Revised: such distributions by introducing a masked autoregressive flow (MAF, Papamakarios et al. 2017) for relaxed Bernoullis via the recursion: s i “ σ p l i ` α i { λ q , where l i „ Logistic p 0 , 1 q , α i “ f i p s 1: i ´ 1 q , 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.

Original: We use this MAF to model P θ and the dynamics related to the labels in ¯ ξ π θ .
Revised: We use this MAF to model P θ and the dynamics related to the labels in ¯ ξ π θ .

Original: We fix the logits of the remaining n ´ | AP | bits to 0 to allow for a fairly distributed latent space.
Revised: We fix the logits of the remaining n ´ | AP | bits to 0 to allow for a fairly distributed latent space.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate the quality of latent space models learned and policies distilled through W 2 AE-MDPs .
Revised: We evaluate the quality of latent space models learned and policies distilled through W 2 AE-MDPs .

Original: To do so, we first trained deep-RL policies (DQN, Mnih et al.
Revised: To do so, we first trained deep-RL policies (DQN, Mnih et al.

Original: 2015 on discrete, and SAC, Haarnoja et al.
Revised: 2015 on discrete, and SAC, Haarnoja et al.

Original: 2018 on continuous action spaces) for various OpenAI benchmarks (Brockman et al., 2016), which we then distill via our approach (Figure 4).
Revised: 2018 on continuous action spaces) for various OpenAI benchmarks (Brockman et al., 2016), which we then distill via our approach (Figure 4).

Original: We thus evaluate (a) the W 2 AE-MDP 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.
Revised: We thus evaluate (a) the W 2 AE-MDP 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.

Original: The confidence metrics and performance are compared with those of VAE-MDPs.
Revised: The confidence metrics and performance are compared with those of VAE-MDPs.

Original: Finally, we formally verify properties in the latent model.
Revised: Finally, we formally verify properties in the latent model.

Original: The exact setting to reproduce our results is in Appendix B.
Revised: The exact setting to reproduce our results is in Appendix B.

Original: Learning metrics.
Revised: Learning metrics.

Original: The objective (Fig. 4a) is a weighted sum of the reconstruction loss and the two Wasserstein regularizers.
Revised: The objective (Fig. 4a) is a weighted sum of the reconstruction loss and the two Wasserstein regularizers.

Original: The choice of β defines the optimization direction.
Revised: The choice of β defines the optimization direction.

Original: Mode collapse is not observed, naturally avoided in WAEs (Tolstikhin et al., 2018), reflecting that the latent space is consistently distributed.
Revised: Posterior collapse is not observed, naturally avoided in WAEs (Tolstikhin et al., 2018), reflecting that the latent space is consistently distributed.

Original: Optimizing the objective (Fig. 4a) effectively allows minimizing the local losses (Fig. 4b) and recovering the performance of the original policy (Fig. 4c).
Revised: Optimizing the objective (Fig. 4a) effectively allows minimizing the local losses (Fig. 4b) and recovering the performance of the original policy (Fig. 4c).

Original: Local losses.
Revised: Local losses.

Original: For V- and WAEs, we formally evaluate PAC upper bounds on L ξ πθ R and L ξ πθ P via the algorithm of Delgrange et al.
Revised: For V- and WAEs, we formally evaluate PAC upper bounds on L ξ πθ R and L ξ πθ P via the algorithm of Delgrange et al.

Original: (2022) (Fig 4b).
Revised: (2022) (Fig 4b).

Original: The lower the local losses, the closer M and
Revised: The lower the local losses, the closer M and

Original: M θ are in terms of behaviors induced by π θ (cf.
Revised: M θ are in terms of behaviors induced by π θ (cf.

Original: Eq. 1).
Revised: Eq. 1).

Original: In VAEs, the losses are evaluated on a transition function ˆ P obtained via frequency estimation of the latent transition dynamics (Delgrange et al., 2022), by reconstructing the transition model a posteriori and collecting data to estimate the transition probabilities (e.g., Bazille et al. 2020; Corneil et al. 2018).
Revised: In VAEs, the losses are evaluated on a transition function ˆ P obtained via frequency estimation of the latent transition dynamics (Delgrange et al., 2022), by reconstructing the transition model a posteriori and collecting data to estimate the transition probabilities (e.g., Bazille et al. 2020; Corneil et al. 2018).

Original: We thus also report the metrics for ˆ P .
Revised: We thus also report the metrics for ˆ P .

Original: Our bounds quickly converge to close values in general for P θ and ˆ P , whereas for VAEs, the convergence is slow and unstable, with ˆ P offering better bounds.
Revised: Our bounds quickly converge to close values in general for P θ and ˆ P , whereas for VAEs, the convergence is slow and unstable, with ˆ P offering better bounds.

Original: We emphasize that WAEs do not require this additional reconstruction step to obtain losses that can be leveraged to assess the
Revised: We emphasize that WAEs do not require this additional reconstruction step to obtain losses that can be leveraged to assess the

Original: [Table]
Revised: [Table]

Original: quality of the model, in contrast to VAEs, where learning P θ was performed via overly restrictive distributions, leading to poor estimation in general (cf.
Revised: quality of the model, in contrast to VAEs, where learning P θ was performed via overly restrictive distributions, leading to poor estimation in general (cf.

Original: Ex. 1).
Revised: Ex. 1).

Original: Finally, when the distilled policies offer comparable performance (Fig. 4c), our bounds are either close to or better than those of VAEs.
Revised: Finally, when the distilled policies offer comparable performance (Fig. 4c), our bounds are either close to or better than those of VAEs.

Original: Distillation.
Revised: Distillation.

Original: The bisimulation guarantees (Eq. 1) are only valid for π θ , the policy under which formal properties can be verified.
Revised: The bisimulation guarantees (Eq. 1) are only valid for π θ , the policy under which formal properties can be verified.

Original: It is crucial that π θ achieves performance close to π , the original one, when deployed in the RL environment.
Revised: It is crucial that π θ achieves performance close to π , the original one, when deployed in the RL environment.

Original: We evaluate the performance of π θ via the undiscounted episode return R π θ obtained by running π θ in the original model M .
Revised: We evaluate the performance of π θ via the undiscounted episode return R π θ obtained by running π θ in the original model M .

Original: We observe that R π θ approaches faster the original performance R π for W- than VAEs: WAEs converge in a few steps for all environments, whereas the full learning budget is sometimes necessary with VAEs.
Revised: We observe that R π θ approaches faster the original performance R π for W- than VAEs: WAEs converge in a few steps for all environments, whereas the full learning budget is sometimes necessary with VAEs.

Original: The success in recovering the original performance emphasizes the representation quality guarantees (Eq. 1) induced by WAEs: when local losses are minimized, all original states that are embedded to the same representation are bisimilarly close.
Revised: The success in recovering the original performance emphasizes the representation quality guarantees (Eq. 1) induced by WAEs: when local losses are minimized, all original states that are embedded to the same representation are bisimilarly close.

Original: 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.
Revised: 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.

Original: Furthermore, the distillation can remove some non-robustness obtained during RL: π θ prescribes the same actions for bisimilarly close states, whereas this is not necessarily the case for π .
Revised: Furthermore, the distillation can remove some non-robustness obtained during RL: π θ prescribes the same actions for bisimilarly close states, whereas this is not necessarily the case for π .

Original: Formal verification.
Revised: Formal verification.

Original: To formally verify M θ , 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).
Revised: To formally verify M θ , 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).

Original: We verify time-to-failure properties φ , often used to check the failure rate of a system (Pnueli, 1977) by measuring whether the agent fails before the end of the episode .
Revised: We verify time-to-failure properties φ , often used to check the failure rate of a system (Pnueli, 1977) by measuring whether the agent fails before the end of the episode .

Original: Although simple, such properties highlight the applicability of our approach on reachability events, which are building blocks to verify MDPs (Baier & Katoen, 2008).
Revised: Although simple, such properties highlight the applicability of our approach on reachability events, which are building blocks to verify MDPs (Baier & Katoen, 2008).

Original: 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).
Revised: 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).

Original: Results are in Table 1: for each environment, we select the distilled policy which gives the best trade-off between performance (episode return) and abstraction quality (local losses).
Revised: Results are in Table 1: for each environment, we select the distilled policy which gives the best trade-off between performance (episode return) and abstraction quality (local losses).

Original: As extra confidence metric, we report the value difference } V π θ } “ | V π θ p
Revised: As extra confidence metric, we report the value difference } V π θ } “ | V π θ p

Original: s I q
Revised: s I q

Original: ´ V π
Revised: ´ V π

Original: θ p s I q| obtained by executing π θ in M and M θ ( V π θ p¨q is averaged while V π θ p¨q is formally computed).
Revised: θ p s I q| obtained by executing π θ in M and M θ ( V π θ p¨q is averaged while V π θ p¨q is formally computed).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented WAE-MDPs, a framework for learning formally verifiable distillations of RL policies with bisimulation guarantees.
Revised: We presented WAE-MDPs, a framework for learning formally verifiable distillations of RL policies with bisimulation guarantees.

Original: The latter, along with the learned abstraction of the unknown continuous environment to a discrete model, enables the verification.
Revised: The latter, along with the learned abstraction of the unknown continuous environment to a discrete model, enables the verification.

Original: Our method overcome the limitations of
Revised: Our method overcome the limitations of

Original: 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.
Revised: 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.

Original: As mentioned by Delgrange et al.
Revised: As mentioned by Delgrange et al.

Original: , distillation failure reveals the lack of robustness of original RL policies.
Revised: , distillation failure reveals the lack of robustness of original RL policies.

Original: 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.
Revised: 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.

Original: We demonstrated the feasibility of our approach through the verification of reachability objectives, which are building blocks for stochastic model-checking (Baier & Katoen, 2008).
Revised: We demonstrated the feasibility of our approach through the verification of reachability objectives, which are building blocks for stochastic model-checking (Baier & Katoen, 2008).

Original: Besides the scope of this work, the verification of general discounted ω -regular properties is theoretically allowed in our model via the rechability to components of standard constructions based on automata products (e.g., Baier et al. 2016; Sickert et al. 2016), and discounted games algorithms (Chatterjee et al., 2008).
Revised: Besides the scope of this work, the verification of general discounted ω -regular properties is theoretically allowed in our model via the rechability to components of standard constructions based on automata products (e.g., Baier et al. 2016; Sickert et al. 2016), and discounted games algorithms (Chatterjee et al., 2010).

Original: Beyond distillation, our results, supported by Thm. 3.3, 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.
Revised: Beyond distillation, our results, supported by Thm. 3.3, 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.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: 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.
Revised: 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.

Original: In addition, Appendix B is dedicated to the presentation of the setup, hyperparameters, and other extra details required for reproducing the results of Section 4.
Revised: In addition, Appendix B is dedicated to the presentation of the setup, hyperparameters, and other extra details required for reproducing the results of Section 4.

Original: We provide the source code of the implementation of our approach in Supplementary material and we also provide the models saved during training that we used for model checking (i e., reproducing the results of Table 1).
Revised: We provide the source code of the implementation of our approach in Supplementary material and we also provide the models saved during training that we used for model checking (i e., reproducing the results of Table 1).

Original: Additionally, we present in a notebook ( code/evaluation.html ) videos demonstrating how our distilled policies behave in each environment, and code snippets showing how we formally verified the policies
Revised: Additionally, we present in a notebook ( code/evaluation.html ) videos demonstrating how our distilled policies behave in each environment, and code snippets showing how we formally verified the policies

Original: A PPENDIX A T HEORETICAL D ETAILS ON WAE-MDP S A. T HE D ISCREPANCY M EASURE
Revised: A PPENDIX A T HEORETICAL D ETAILS ON WAE-MDP S A. T HE D ISCREPANCY M EASURE

Original: We show that reasoning about discrepancy measures between stationary distributions is sound in the context of infinite interaction and episodic RL processes.
Revised: We show that reasoning about discrepancy measures between stationary distributions is sound in the context of infinite interaction and episodic RL processes.

Original: Let P θ 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 x s 0: T , a 0: T ´ 1 , r 0: T ´ 1 y ), our goal is to find the best parameter θ which offers the most accurate reconstruction of the original traces issued from the original model M operating under π .
Revised: Let P θ 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 x s 0: T , a 0: T ´ 1 , r 0: T ´ 1 y ), our goal is to find the best parameter θ which offers the most accurate reconstruction of the original traces issued from the original model M operating under π .

Original: We demonstrate that, in the limit, considering the OT between trace-based distributions is equivalent to considering the OT between the stationary distribution of M π and the one of the behavioral model.
Revised: We demonstrate that, in the limit, considering the OT between trace-based distributions is equivalent to considering the OT between the stationary distribution of M π and the one of the behavioral model.

Original: Trace-based distributions.
Revised: Trace-based distributions.

Original: The raw distance ⃗d allows to reason about transitions , we thus consider the distribution over transitions which occur along traces of length T to compare the dynamics of the original and behavioral models:
Revised: The raw distance ⃗d allows to reason about transitions , we thus consider the distribution over transitions which occur along traces of length T to compare the dynamics of the original and behavioral models:

Original: [Equation]
Revised: [Equation]

Original: where P θ r T s denotes the distribution over traces of length T , generated from P θ .
Revised: where P θ r T s denotes the distribution over traces of length T , generated from P θ .

Original: Intuitively, 1 { T ¨ ř
Revised: Intuitively, 1 { T ¨ ř

Original: Tt “ 1 ξ tπ p s | s I q can be seen as the fraction of the time spent in s along traces of length T , starting from the initial state Kulkarni (1995).
Revised: Tt “ 1 ξ tπ p s | s I q can be seen as the fraction of the time spent in s along traces of length T , starting from the initial state Kulkarni (1995).

Original: Therefore, drawing x s, a, r, s 1 y „ D π r T s trivially follows: it is equivalent to drawing s from 1 { T ¨ ř
Revised: Therefore, drawing x s, a, r, s 1 y „ D π r T s trivially follows: it is equivalent to drawing s from 1 { T ¨ ř

Original: Tt “ 1 ξ tπ p¨ | s I q , then respectively a and s 1 from π p¨ | s q and
Revised: Tt “ 1 ξ tπ p¨ | s I q , then respectively a and s 1 from π p¨ | s q and

Original: P p¨ | s, a q , to finally obtain r “ R p s, a q .
Revised: P p¨ | s, a q , to finally obtain r “ R p s, a q .

Original: Given T P N , our objective is to minimize the Wasserstein distance between those distributions: W ⃗d p D π r T s , P θ r T sq .
Revised: Given T P N , our objective is to minimize the Wasserstein distance between those distributions: W ⃗d p D π r T s , P θ r T sq .

Original: 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).
Revised: 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).

Original: Lemma A.1.
Revised: Lemma A.1.

Original: Assume the existence of a stationary behavioral model ξ θ
Revised: Assume the existence of a stationary behavioral model ξ θ

Original: “ lim T Ñ8 P θ r T s , then
Revised: “ lim T Ñ8 P θ r T s , then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: First, note that 1 { T ¨ ř Tt “ 1 ξ tπ p¨ | s I q weakly converges to ξ π as T goes to 8 Kulkarni (1995).
Revised: First, note that 1 { T ¨ ř Tt “ 1 ξ tπ p¨ | s I q weakly converges to ξ π as T goes to 8 Kulkarni (1995).

Original: The result follows then from (Villani, 2009, Corollary 6.9).
Revised: The result follows then from (Villani, 2009, Corollary 6.9).

Original: A. D EALING WITH D ISCRETE A CTIONS
Revised: A. D EALING WITH D ISCRETE A CTIONS

Original: When the policy π executed in M already produces discrete actions, learning a latent action space is, in many cases, not necessary.
Revised: When the policy π executed in M already produces discrete actions, learning a latent action space is, in many cases, not necessary.

Original: We thus make the following assumptions:
Revised: We thus make the following assumptions:

Original: Assumption A.2.
Revised: Assumption A.2.

Original: Let π : S Ñ ∆ p
Revised: Let π : S Ñ ∆ p

Original: A ‹ q be the policy executed in M and assume that A ‹ is a (tractable) finite set.
Revised: A ‹ q be the policy executed in M and assume that A ‹ is a (tractable) finite set.

Original: Then, we take A “ A ‹ and ϕ
Revised: Then, we take A “ A ‹ and ϕ

Original: A ι as the identity function, i.e., ϕ
Revised: A ι as the identity function, i.e., ϕ

Original: A ι : S ˆ A ‹ Ñ
Revised: A ι : S ˆ A ‹ Ñ

Original: A ‹ , x s, a ‹ y ÞÑ a ‹ .
Revised: A ‹ , x s, a ‹ y ÞÑ a ‹ .

Original: Assumption A.3.
Revised: Assumption A.3.

Original: Assume that the action space of the original environment M is a (tractable) finite set.
Revised: Assume that the action space of the original environment M is a (tractable) finite set.

Original: Then, we take ψ θ as the identity function, i.e., ψ θ “ ϕ A ι .
Revised: Then, we take ψ θ as the identity function, i.e., ψ θ “ ϕ A ι .

Original: Concretely, the premise of Assumption A.2 typically occurs when π is a latent policy (see Rem. 1) or when M has already a discrete action space.
Revised: Concretely, the premise of Assumption A.2 typically occurs when π is a latent policy (see Rem. 1) or when M has already a discrete action space.

Original: In the latter case, Assumption A.2 and A.3 amount to setting A “ A and ignoring the action encoder and embedding function.
Revised: In the latter case, Assumption A.2 and A.3 amount to setting A “ A and ignoring the action encoder and embedding function.

Original: 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.
Revised: 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.

Original: A. P ROOF OF L EMMA 3. 
Revised: A. P ROOF OF L EMMA 3.

Original: Notation.
Revised: Notation.

Original: From now on, we write ϕ ι p s, a | s, a q “ 1 ϕ ι p s q“ s ¨ ϕ
Revised: From now on, we write ϕ ι p s, a | s, a q “ 1 ϕ ι p s q“ s ¨ ϕ

Original: A ι p a | s, a q .
Revised: A ι p a | s, a q .

Original: Lemma 3.2.
Revised: Lemma 3.2.

Original: Define T p s, a, s 1 q “ E s,a „ ξ π r 1 ϕ ι p s q“ s ¨ ϕ
Revised: Define T p s, a, s 1 q “ E s,a „ ξ π r 1 ϕ ι p s q“ s ¨ ϕ

Original: A ι p a | s, a q¨ P θ p s 1 | s, a qs as the distribution of drawing state-action pairs from interacting with M , embedding them to the latent spaces, and finally letting them transition to their successor state in M θ .
Revised: A ι p a | s, a q¨ P θ p s 1 | s, a qs as the distribution of drawing state-action pairs from interacting with M , embedding them to the latent spaces, and finally letting them transition to their successor state in M θ .

Original: Then, W ⃗d
Revised: Then, W ⃗d

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Wasserstein is compliant with the triangular inequality (Villani, 2009), which gives us:
Revised: Wasserstein is compliant with the triangular inequality (Villani, 2009), which gives us:

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: (note that W ⃗d is reflexive (Villani, 2009))
Revised: (note that W ⃗d is reflexive (Villani, 2009))

Original: [Equation]
Revised: [Equation]

Original: We pass from Eq.
Revised: We pass from Eq.

Original: 3 to Eq.
Revised: 3 to Eq.

Original: 4 by the Jensen’s inequality.
Revised: 4 by the Jensen’s inequality.

Original: To see how we pass from Eq.
Revised: To see how we pass from Eq.

Original: 4 to Eq.
Revised: 4 to Eq.

Original: 5, notice that
Revised: 5, notice that

Original: [Equation]
Revised: [Equation]

Original: Observe now that s and a are fixed in the supremum computation of Eq. 4: all functions f considered and taken from F ⃗d are of the form f p s, a, ¨q .
Revised: Observe now that s and a are fixed in the supremum computation of Eq. 4: all functions f considered and taken from F ⃗d are of the form f p s, a, ¨q .

Original: It is thus sufficient to consider the supremum over functions from the following subset of F ⃗d :
Revised: It is thus sufficient to consider the supremum over functions from the following subset of F ⃗d :

Original: [Equation]
Revised: [Equation]

Original: a, s 1˘ ´ f ` s, a, s 1˘ ď d
Revised: a, s 1˘ ´ f ` s, a, s 1˘ ď d

Original: [Equation]
Revised: [Equation]

Original: p s, s q
Revised: p s, s q

Original: [Equation]
Revised: [Equation]

Original: d A p a, a q ` d S ` s 1 1 , s
Revised: d A p a, a q ` d S ` s 1 1 , s

Original: [Equation]
Revised: [Equation]

Original: ˘ u (for s , a drawn from ϕ ι ) “ t f : f ` s
Revised: ˘ u (for s , a drawn from ϕ ι ) “ t f : f ` s

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Original: [Equation]
Revised: [Equation]

Original: ˘
Revised: ˘

Original: [Equation]
Revised: [Equation]

Original: f `
Revised: f `

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ˘
Revised: ˘

Original: [Equation]
Revised: [Equation]

Original: `
Revised: `

Original: [Equation]
Revised: [Equation]

Original: ˘ u
Revised: ˘ u

Original: [Equation]
Revised: [Equation]

Original: `
Revised: `

Original: [Equation]
Revised: [Equation]

Original: u
Revised: u

Original: [Equation]
Revised: [Equation]

Original: Given a state s P S in the original model, the (parallel) execution of π in M θ is enabled through π p a, a | s q “ π p a | s q ¨ ϕ
Revised: Given a state s P S in the original model, the (parallel) execution of π in M θ is enabled through π p a, a | s q “ π p a | s q ¨ ϕ

Original: A ι p a | ϕ ι p s q , a q (cf.
Revised: A ι p a | ϕ ι p s q , a q (cf.

Original: Fig.
Revised: Fig.

Original: 1b).
Revised: 1b).

Original: The local transition loss resulting from this interaction is:
Revised: The local transition loss resulting from this interaction is:

Original: [Equation]
Revised: [Equation]

Original: which finally yields the result.
Revised: which finally yields the result.

Original: A. P ROOF OF T HEOREM 3. 
Revised: A. P ROOF OF T HEOREM 3.

Original: Before proving Theorem 3.3, let us introduce the following Lemma, that explicitly demonstrates the link between the transition regularizer of the W 2 AE-MDP objective and the local transition loss required to obtain the guarantees related to the bisimulation bounds of Eq. 1.
Revised: Before proving Theorem 3.3, let us introduce the following Lemma, that explicitly demonstrates the link between the transition regularizer of the W 2 AE-MDP objective and the local transition loss required to obtain the guarantees related to the bisimulation bounds of Eq. 1.

Original: Lemma A.4.
Revised: Lemma A.4.

Original: Assume that traces are generated by running π P Π in the original environment, then
Revised: Assume that traces are generated by running π P Π in the original environment, then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Since the latent policy π generates latent actions, Assumption A.2 holds, which means:
Revised: Since the latent policy π generates latent actions, Assumption A.2 holds, which means:

Original: [Equation]
Revised: [Equation]

Original: Theorem 3.3.
Revised: Theorem 3.3.

Original: Assume that traces are generated by running a latent policy π P Π in the original environment and let d R be the usual Euclidean distance, then the W 2 AE-MDP objective is
Revised: Assume that traces are generated by running a latent policy π P Π in the original environment and let d R be the usual Euclidean distance, then the W 2 AE-MDP objective is

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We distinguish two cases: (i) the case where the original and latent models share the same discrete action space, i.e., A “ A , and (ii) the case where the two have a different action space (e.g., when the original action space is continuous), i.e., A ‰ A .
Revised: We distinguish two cases: (i) the case where the original and latent models share the same discrete action space, i.e., A “ A , and (ii) the case where the two have a different action space (e.g., when the original action space is continuous), i.e., A ‰ A .

Original: In both cases, the local losses term follows by definition of L ξ π R and Lemma A.4.
Revised: In both cases, the local losses term follows by definition of L ξ π R and Lemma A.4.

Original: When d R 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 ⃗d in the W 2 AE-MDP objective directly translates to the local loss L ξ π R .
Revised: When d R 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 ⃗d in the W 2 AE-MDP objective directly translates to the local loss L ξ π R .

Original: Concerning the local transition loss, in case (i), the result naturally follows from Assumption A. and A.3.
Revised: Concerning the local transition loss, in case (i), the result naturally follows from Assumption A. and A.3.

Original: In case (ii), only Assumption A.2 holds, meaning the action encoder term of the W 2 AEMDP objective is ignored, but not the action embedding term appearing in G θ .
Revised: In case (ii), only Assumption A.2 holds, meaning the action encoder term of the W 2 AEMDP objective is ignored, but not the action embedding term appearing in G θ .

Original: Given s „ ξ π , recall that executing π in M amounts to embedding the produced latent actions a „ π p¨ | ϕ ι p s qq back to the original environment via a “ ψ θ p ϕ ι p s q , a q (cf.
Revised: Given s „ ξ π , recall that executing π in M amounts to embedding the produced latent actions a „ π p¨ | ϕ ι p s qq back to the original environment via a “ ψ θ p ϕ ι p s q , a q (cf.

Original: Rem. 1 and Fig.
Revised: Rem. 1 and Fig.

Original: 1a).
Revised: 1a).

Original: Therefore, the projection of ⃗d px s, a, r, s 1 y , G θ p ϕ ι p s q , a, ϕ ι p s 1 qqq on the action space
Revised: Therefore, the projection of ⃗d px s, a, r, s 1 y , G θ p ϕ ι p s q , a, ϕ ι p s 1 qqq on the action space

Original: A is d
Revised: A is d

Original: A p ψ θ p ϕ ι p s q , a q , ψ θ p ϕ ι p s q , a qq “ 0 , for r “ R p s, a q and s 1 „ P p¨ | s, a q .
Revised: A p ψ θ p ϕ ι p s q , a q , ψ θ p ϕ ι p s q , a qq “ 0 , for r “ R p s, a q and s 1 „ P p¨ | s, a q .

Original: A. O PTIMIZING THE T RANSITION R EGULARIZER
Revised: A. O PTIMIZING THE T RANSITION R EGULARIZER

Original: In the following, we detail how we derive a tractable form of our transition regularizer L ξ π P p ω q .
Revised: In the following, we detail how we derive a tractable form of our transition regularizer L ξ π P p ω q .

Original: Optimizing the ground Kantorovich-Rubinstein duality is enabled via the introduction of a parameterized, 1 -Lipschitz network φ P ω , that need to be trained to attain the supremum of the dual:
Revised: Optimizing the ground Kantorovich-Rubinstein duality is enabled via the introduction of a parameterized, 1 -Lipschitz network φ P ω , that need to be trained to attain the supremum of the dual:

Original: [Equation]
Revised: [Equation]

Original: Under this form, optimizing L ξ π P p ω q is intractable due to the expectation over the maximum.
Revised: Under this form, optimizing L ξ π P p ω q is intractable due to the expectation over the maximum.

Original: The following Lemma allows us rewriting L ξ π P to make the optimization tractable through Monte Carlo estimation.
Revised: The following Lemma allows us rewriting L ξ π P to make the optimization tractable through Monte Carlo estimation.

Original: Lemma A.5.
Revised: Lemma A.5.

Original: Let X , Y be two measurable sets, ξ P ∆ p X q , P : X Ñ ∆ p Y q , Q : X Ñ ∆ p Y q , and d : Y ˆ Y Ñ r 0 , `8r be a metric on Y .
Revised: Let X , Y be two measurable sets, ξ P ∆ p X q , P : X Ñ ∆ p Y q , Q : X Ñ ∆ p Y q , and d : Y ˆ Y Ñ r 0 , `8r be a metric on Y .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Our objective is to show that
Revised: Our objective is to show that

Original: [Equation]
Revised: [Equation]

Original: We start with (6) ď (7).
Revised: We start with (6) ď (7).

Original: Construct φ ‹ : X Ñ F d by setting for all x P X
Revised: Construct φ ‹ : X Ñ F d by setting for all x P X

Original: [Equation]
Revised: [Equation]

Original: This gives us
Revised: This gives us

Original: [Equation]
Revised: [Equation]

Original: It remains to show that (6) ě (7).
Revised: It remains to show that (6) ě (7).

Original: Take
Revised: Take

Original: [Equation]
Revised: [Equation]

Original: Then, for all x P X
Revised: Then, for all x P X

Original: [Equation]
Revised: [Equation]

Original: we have φ ‹ p x q P F d which means:
Revised: we have φ ‹ p x q P F d which means:

Original: [Equation]
Revised: [Equation]

Original: This finally yields
Revised: This finally yields

Original: [Equation]
Revised: [Equation]

Original: Corollary A.5.1.
Revised: Corollary A.5.1.

Original: Let ξ π be a stationary distribution of M π and X “ S ˆ
Revised: Let ξ π be a stationary distribution of M π and X “ S ˆ

Original: A ˆ S ˆ A , then
Revised: A ˆ S ˆ A , then

Original: [Equation]
Revised: [Equation]

Original: ˘ff
Revised: ˘ff

Original: Consequently, we rewrite L ξ π P p ω q as a tractable maximization:
Revised: Consequently, we rewrite L ξ π P p ω q as a tractable maximization:

Original: [Equation]
Revised: [Equation]

Original: A. T HE L ATENT M ETRIC
Revised: A. T HE L ATENT M ETRIC

Original: In the following
Revised: In the following

Original: [Equation]
Revised: [Equation]

Original: we show that considering the Euclidean distance for ⃗d and d S in the latent space for optimizing the regularizers W ξ π and L ξ π P is Lipschitz equivalent to considering a continuous λ -relaxation of the discrete metric 1 ‰ p x , y q “ 1 x ‰ y .
Revised: we show that considering the Euclidean distance for ⃗d and d S in the latent space for optimizing the regularizers W ξ π and L ξ π P is Lipschitz equivalent to considering a continuous λ -relaxation of the discrete metric 1 ‰ p x , y q “ 1 x ‰ y .

Original: Consequently, this also means it is consistently sufficient to enforce 1 -Lispchitzness via the gradient penalty approach of Gulrajani et al.
Revised: Consequently, this also means it is consistently sufficient to enforce 1 -Lispchitzness via the gradient penalty approach of Gulrajani et al.

Original: (2017) during training to maintain the guarantees linked to the regularizers in the zero-temperature limit, when the spaces are discrete.
Revised: (2017) during training to maintain the guarantees linked to the regularizers in the zero-temperature limit, when the spaces are discrete.

Original: Lemma A.6.
Revised: Lemma A.6.

Original: Let d be the usual Euclidean distance and d λ : r 0 , 1 s n ˆ r 0 , 1 s n Ñ r 0 , 1 r , x x , y y ÞÑ d p x , y q λ ` d p x , y q for λ P s 0 , 1 s and
Revised: Let d be the usual Euclidean distance and d λ : r 0 , 1 s n ˆ r 0 , 1 s n Ñ r 0 , 1 r , x x , y y ÞÑ d p x , y q λ ` d p x , y q for λ P s 0 , 1 s and

Original: n P N , then d λ is a distance metric
Revised: n P N , then d λ is a distance metric

Original: Proof.
Revised: Proof.

Original: The function d λ is a metric iff it satisfies the following axioms:
Revised: The function d λ is a metric iff it satisfies the following axioms:

Original: Identity of indiscernibles : If x “ y , then d λ p x , y q “
Revised: Identity of indiscernibles : If x “ y , then d λ p x , y q “

Original: [Equation]
Revised: [Equation]

Original: 0 since d is a distance metric.
Revised: 0 since d is a distance metric.

Original: Assume now that d λ p x , y q “ 0 and take α “ d p x , y q , for any x , y .
Revised: Assume now that d λ p x , y q “ 0 and take α “ d p x , y q , for any x , y .

Original: Thus, α P r 0 , `8r and 0 “ αλ ` α is only achieved in α “ 0 , which only occurs whenever x “ y since d is a distance metric.
Revised: Thus, α P r 0 , `8r and 0 “ αλ ` α is only achieved in α “ 0 , which only occurs whenever x “ y since d is a distance metric.

Original: Symmetry :
Revised: Symmetry :

Original: [Equation]
Revised: [Equation]

Original: Triangle inequality : Let x , y ,
Revised: Triangle inequality : Let x , y ,

Original: z P r 0 , 1 s n , the triangle inequality holds iff
Revised: z P r 0 , 1 s n , the triangle inequality holds iff

Original: [Equation]
Revised: [Equation]

Original: (cross-product, with λ ą 

Revised: (cross-product, with λ ą

Original: and Im p
Revised: and Im p

Original: d q P r 0 , 8r )
Revised: d q P r 0 , 8r )

Original: [Equation]
Revised: [Equation]

Original: Since d is a distance metric, we have
Revised: Since d is a distance metric, we have

Original: [Equation]
Revised: [Equation]

Original: and Im p
Revised: and Im p

Original: d q P r 0 , 8r , meaning
Revised: d q P r 0 , 8r , meaning

Original: [Equation]
Revised: [Equation]

Original: By Eq.
Revised: By Eq.

Original: 10 and 11, the inequality of Eq.
Revised: 10 and 11, the inequality of Eq.

Original: 9 holds.
Revised: 9 holds.

Original: Furthermore, the fact that Eq.
Revised: Furthermore, the fact that Eq.

Original: and 9 are equivalent yields the result.
Revised: and 9 are equivalent yields the result.

Original: Lemma A.7.
Revised: Lemma A.7.

Original: Let d , d λ as defined above, then (i) d λ ÝÝÝÑ λ Ñ 0 1 ‰ and (ii) d, d λ are Lipschitz-equivalent.
Revised: Let d , d λ as defined above, then (i) d λ ÝÝÝÑ λ Ñ 0 1 ‰ and (ii) d, d λ are Lipschitz-equivalent.

Original: Proof.
Revised: Proof.

Original: Part (i) is straightforward by definition of d λ .
Revised: Part (i) is straightforward by definition of d λ .

Original: Distances d and d λ are Lispchitz equivalent if and only if D a, b ą 0 such that @ x , y P r 0 , 1 s n ,
Revised: Distances d and d λ are Lispchitz equivalent if and only if D a, b ą 0 such that @ x , y P r 0 , 1 s n ,

Original: [Equation]
Revised: [Equation]

Original: Taking a “ 1 λ `? n and b “ 1 λ yields the result.
Revised: Taking a “ 1 λ `? n and b “ 1 λ yields the result.

Original: Corollary A.7.1.
Revised: Corollary A.7.1.

Original: For all β ě 1 { λ , s P S , a P A , s P S , and a P A , we have
Revised: For all β ě 1 { λ , s P S , a P A , s P S , and a P A , we have

Original: W d λ
Revised: W d λ

Original: [Equation]
Revised: [Equation]

Original: W d λ ` ϕ ι P p¨ | s, a q , P θ p¨ | s, a q ˘ ď β ¨ W d ` ϕ ι P p¨ | s, a q , P θ p¨ | s, a q ˘
Revised: W d λ ` ϕ ι P p¨ | s, a q , P θ p¨ | s, a q ˘ ď β ¨ W d ` ϕ ι P p¨ | s, a q , P θ p¨ | s, a q ˘

Original: [Figure]
Revised: [Figure]

Original: Proof.
Revised: Proof.

Original: By Lipschitz equivalence, taking β ě 1 { λ ensures that @ n P N , @ x , y P r 0 , 1 s n , d λ p x , y q ď β ¨ d p x , y q .
Revised: By Lipschitz equivalence, taking β ě 1 { λ ensures that @ n P N , @ x , y P r 0 , 1 s n , d λ p x , y q ď β ¨ d p x , y q .

Original: Moreover, for any distributions P, Q , W d λ p P, Q q ď β ¨ W d p P, Q q (cf., e.g., Gelada et al. 2019, Lemma A.4 for details).
Revised: Moreover, for any distributions P, Q , W d λ p P, Q q ď β ¨ W d p P, Q q (cf., e.g., Gelada et al. 2019, Lemma A.4 for details).

Original: In practice, taking the hyperparameter β ě 1 { λ in the W 2 AE-MDP ensures that minimizing the β -scaled regularizers w.r.t .
Revised: In practice, taking the hyperparameter β ě 1 { λ in the W 2 AE-MDP ensures that minimizing the β -scaled regularizers w.r.t .

Original: d also minimizes the regularizers w.r.t .
Revised: d also minimizes the regularizers w.r.t .

Original: the λ -relaxation d λ , being the discrete distribution in the zero-temperature limit.
Revised: the λ -relaxation d λ , being the discrete distribution in the zero-temperature limit.

Original: Note that optimizing over two different β 1 , βinstead of a unique scale factor β is also a good practice to interpolate between the two regularizers.
Revised: Note that optimizing over two different β 1 , βinstead of a unique scale factor β is also a good practice to interpolate between the two regularizers.

Original: B E XPERIMENT D ETAILS B. S ETUP
Revised: B E XPERIMENT D ETAILS B. S ETUP

Original: We used T ENSOR F LOW 2.7.0 (Abadi et al., 2015) to implement the neural network architecture of our W 2 AE-MDP , T ENSOR F LOW P ROBABILITY 0.15.0 (Dillon et al., 2017) to handle the probabilistic components of the latent model (e.g., latent distributions with reparameterization tricks, masked autoregressive flows, etc.), as well as TF-A GENTS 0.11.0 (Guadarrama et al., 2018) to handle the RL parts of the framework.
Revised: We used T ENSOR F LOW 2.7.0 (Abadi et al., 2015) to implement the neural network architecture of our W 2 AE-MDP , T ENSOR F LOW P ROBABILITY 0.15.0 (Dillon et al., 2017) to handle the probabilistic components of the latent model (e.g., latent distributions with reparameterization tricks, masked autoregressive flows, etc.), as well as TF-A GENTS 0.11.0 (Guadarrama et al., 2018) to handle the RL parts of the framework.

Original: Models have been trained on a cluster running under CentOS Linux 7 (Core) composed of a mix of nodes containing Intel processors with the following CPU microarchitectures: (i) 10-core
Revised: Models have been trained on a cluster running under CentOS Linux 7 (Core) composed of a mix of nodes containing Intel processors with the following CPU microarchitectures: (i) 10-core

Original: INTEL E5-2680v2 , (ii) 14-core INTEL E5-2680v4 , and (iii) 20-core INTEL Xeon
Revised: INTEL E5-2680v2 , (ii) 14-core INTEL E5-2680v4 , and (iii) 20-core INTEL Xeon

Original: Gold 6148 .
Revised: Gold 6148 .

Original: We used 8 cores and 32 GB of memory for each run.
Revised: We used 8 cores and 32 GB of memory for each run.

Original: B. S TATIONARY D ISTRIBUTION
Revised: B. S TATIONARY D ISTRIBUTION

Original: To sample from the stationary distribution ξ π of episodic learning environments operating under π P Π , we implemented the recursive ϵ -perturbation trick of Huang (2020).
Revised: To sample from the stationary distribution ξ π of episodic learning environments operating under π P Π , we implemented the recursive ϵ -perturbation trick of Huang (2020).

Original: In a nutshell, the reset of the environment is explicitly added to the state space of M , which is entered at the end of each episode and left with probability 1 ´ ϵ to start a new one.
Revised: In a nutshell, the reset of the environment is explicitly added to the state space of M , which is entered at the end of each episode and left with probability 1 ´ ϵ to start a new one.

Original: We also added a special atomic proposition reset into AP to label this reset state and reason about episodic behaviors.
Revised: We also added a special atomic proposition reset into AP to label this reset state and reason about episodic behaviors.

Original: 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.
Revised: 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.

Original: B. L ATENT SPACE DISTRIBUTION
Revised: B. L ATENT SPACE DISTRIBUTION

Original: As pointed out in Sect.
Revised: As pointed out in Sect.

Original: 4, mode collapse is naturally avoided when optimizing W 2 AE-MDP .
Revised: 4, posterior collapse is naturally avoided when optimizing W 2 AE-MDP .

Original: To illustrate that, we report the distribution of latent states produced by ϕ ι during training (Fig. 5).
Revised: To illustrate that, we report the distribution of latent states produced by ϕ ι during training (Fig. 5).

Original: The plots reveal that the latent space generated by mapping original states drawn from ξ π during training to S via ϕ ι is fairly distributed, for each environment.
Revised: The plots reveal that the latent space generated by mapping original states drawn from ξ π during training to S via ϕ ι is fairly distributed, for each environment.

Original: B. V ALUE DIFFERENCE
Revised: B. V ALUE DIFFERENCE

Original: In addition to reporting the quality guarantees of the model along training steps through local losses (cf.
Revised: In addition to reporting the quality guarantees of the model along training steps through local losses (cf.

Original: Figure 4b)
Revised: Figure 4b)

Original: , our experiments revealed that the absolute value difference } V π θ } between the original and latent models operating under the latent policy quickly decreases and tends to converge to values in the same range (Figure 6).
Revised: , our experiments revealed that the absolute value difference } V π θ } between the original and latent models operating under the latent policy quickly decreases and tends to converge to values in the same range (Figure 6).

Original: This is consistent with the fact that minimizing local losses lead to close behaviors (cf.
Revised: This is consistent with the fact that minimizing local losses lead to close behaviors (cf.

Original: Eq. 1) and that the value function is Lipschitz-continuous w.r.t.
Revised: Eq. 1) and that the value function is Lipschitz-continuous w.r.t.

Original: d „ π θ (cf.
Revised: d „ π θ (cf.

Original: Section 2).
Revised: Section 2).

Original: B. H YPERPARAMETERS
Revised: B. H YPERPARAMETERS

Original: W 2 AE-MDP parameters.
Revised: W 2 AE-MDP parameters.

Original: All components (e.g., functions or distribution locations and scales, see
Revised: All components (e.g., functions or distribution locations and scales, see

Original: Fig. 2) are represented and inferred by neural networks (multilayer perceptrons).
Revised: Fig. 2) are represented and inferred by neural networks (multilayer perceptrons).

Original: All the networks share the same architecture (i.e., number of layers and neurons per layer).
Revised: All the networks share the same architecture (i.e., number of layers and neurons per layer).

Original: We use a simple uniform experience replay of size 10 6 to store the transitions and sample them.
Revised: We use a simple uniform experience replay of size 10 6 to store the transitions and sample them.

Original: The training starts when the agent has collected 10 4 transitions in M .
Revised: The training starts when the agent has collected 10 4 transitions in M .

Original: We used minibatches of size 128 to optimize the objective and we applied a minibatch update every time the agent executing π has performed 16 steps in M .
Revised: We used minibatches of size 128 to optimize the objective and we applied a minibatch update every time the agent executing π has performed 16 steps in M .

Original: We use the recursive ϵ -perturbation trick of Huang (2020) with ϵ “ 3 { 4 : when an episode ends, it restarts from the initial state with probability 1 { 4 ; before re-starting an episode, the time spent in the reset state labeled with reset follows then the geometric distribution with expectation ϵ { 1 ´ ϵ “ 3 .
Revised: We use the recursive ϵ -perturbation trick of Huang (2020) with ϵ “ 3 { 4 : when an episode ends, it restarts from the initial state with probability 1 { 4 ; before re-starting an episode, the time spent in the reset state labeled with reset follows then the geometric distribution with expectation ϵ { 1 ´ ϵ “ 3 .

Original: We chose the same latent state-action space size than Delgrange et al.
Revised: We chose the same latent state-action space size than Delgrange et al.

Original: (2022), except for LunarLander that we decreased to log 2 ˇˇ S ˇˇ “ 14 and ˇˇ A ˇˇ “ 3 to improve the scalability of the verification.
Revised: (2022), except for LunarLander that we decreased to log 2 ˇˇ S ˇˇ “ 14 and ˇˇ A ˇˇ “ 3 to improve the scalability of the verification.

Original: VAE-MDPs parameters.
Revised: VAE-MDPs parameters.

Original: For the comparison of Sect.
Revised: For the comparison of Sect.

Original: 4, we used the exact same VAE-MDP hyperparameter set as prescribed by Delgrange et al.
Revised: 4, we used the exact same VAE-MDP hyperparameter set as prescribed by Delgrange et al.

Original: (2022), except for the state-action space of LunarLander that we also changed for scalability and fair comparison purpose.
Revised: (2022), except for the state-action space of LunarLander that we also changed for scalability and fair comparison purpose.

Original: Hyperparameter search.
Revised: Hyperparameter search.

Original: To evaluate our W 2 AE-MDP , we realized a search in the parameter space defined in Table 2.
Revised: To evaluate our W 2 AE-MDP , we realized a search in the parameter space defined in Table 2.

Original: The best parameters found (in terms of trade-off between performance and latent quality) are reported in Table 3.
Revised: The best parameters found (in terms of trade-off between performance and latent quality) are reported in Table 3.

Original: We used two different optimizers for minimizing the loss (referred to as the minimizer) and computing the Wasserstein terms (reffered to as the maximizer).
Revised: We used two different optimizers for minimizing the loss (referred to as the minimizer) and computing the Wasserstein terms (reffered to as the maximizer).

Original: We used
Revised: We used

Original: A DAM (Kingma & Ba, 2015) for the two, but we allow for different learning rates A DAM α and exponential decays A DAM β 1 , A DAM β 2 .
Revised: A DAM (Kingma & Ba, 2015) for the two, but we allow for different learning rates A DAM α and exponential decays A DAM β 1 , A DAM β 2 .

Original: We also found that polynomial decay for A DAM
Revised: We also found that polynomial decay for A DAM

Original: α
Revised: α

Original: (e.g., to 10 ´ 5 for 4 ¨ 10 5 steps) is a good practice to stabilize the experiment learning curves, but is not necessary to obtain high-quality and performing distillation.
Revised: (e.g., to 10 ´ 5 for 4 ¨ 10 5 steps) is a good practice to stabilize the experiment learning curves, but is not necessary to obtain high-quality and performing distillation.

Original: Concerning the continuous relaxation of discrete distributions, we used a different temperature for each distribution, as Maddison et al.
Revised: Concerning the continuous relaxation of discrete distributions, we used a different temperature for each distribution, as Maddison et al.

Original: (2017) pointed out that doing so is valuable to improve the results.
Revised: (2017) pointed out that doing so is valuable to improve the results.

Original: We further followed the guidelines of
Revised: We further followed the guidelines of

Original: Maddison et al.
Revised: Maddison et al.

Original: (2017) to choose the interval of temperatures and did not schedule any annealing scheme (in contrast to VAE-MDPs).
Revised: (2017) to choose the interval of temperatures and did not schedule any annealing scheme (in contrast to VAE-MDPs).

Original: Essentially, the search reveals that the regularizer scale factors β ¨ (defining the optimization direction) as well as the encoder and latent transition temperatures are
Revised: Essentially, the search reveals that the regularizer scale factors β ¨ (defining the optimization direction) as well as the encoder and latent transition temperatures are

Original: [Table] [Table]
Revised: [Table] [Table]

Original: important to improve the performance of distilled policies.
Revised: important to improve the performance of distilled policies.

Original: For the encoder temperature, we found a nice spot in λ ϕι “ 2 { 3 , which provides the best performance in general, whereas the choice of λ P θ and β ¨ are (latent-) environment dependent.
Revised: For the encoder temperature, we found a nice spot in λ ϕι “ 2 { 3 , which provides the best performance in general, whereas the choice of λ P θ and β ¨ are (latent-) environment dependent.

Original: The importance of the temperature parameters for the continuous relaxation of discrete distributions is consistent with the results of (Maddison et al., 2017), revealing that the success of the relaxation depends on the choice of the temperature for the different latent space sizes.
Revised: The importance of the temperature parameters for the continuous relaxation of discrete distributions is consistent with the results of (Maddison et al., 2017), revealing that the success of the relaxation depends on the choice of the temperature for the different latent space sizes.

Original: Labeling functions.
Revised: Labeling functions.

Original: We used the same labeling functions as those described by Delgrange et al.
Revised: We used the same labeling functions as those described by Delgrange et al.

Original: For completeness, we recall the labeling function used for each environment in Table 4.
Revised: For completeness, we recall the labeling function used for each environment in Table 4.

Original: Time to failure properties.
Revised: Time to failure properties.

Original: Based on the labeling described in Table 4, we formally detail the time to failure properties checked in Sect.
Revised: Based on the labeling described in Table 4, we formally detail the time to failure properties checked in Sect.

Original: 4 whose results are listed in Table 1 for each environment.
Revised: 4 whose results are listed in Table 1 for each environment.

Original: Let
Revised: Let

Original: Reset “ t reset u “ x 0 , . .
Revised: Reset “ t reset u “ x 0 , . .

Original: , 1 y (we assume here that the last bit indicates whether the current state is a reset state or not) and define s |ù L 1 ^ L 2 iff s |ù L 1 and s |ù L 2 for any s P S , then
Revised: , 1 y (we assume here that the last bit indicates whether the current state is a reset state or not) and define s |ù L 1 ^ L 2 iff s |ù L 1 and s |ù L 2 for any s P S , then

Original: • CartPole : φ “ ␣ Reset U Unsafe , where Unsafe “ x 1 , 1 , 0 y
Revised: • CartPole : φ “ ␣ Reset U Unsafe , where Unsafe “ x 1 , 1 , 0 y

Original: [Table]
Revised: [Table]

Original: • MountainCar : φ “ ␣ Goal U Reset , where Goal “ x 1 , 0 , 0 , 0 y • Acrobot : φ “ ␣ Goal U Reset , where Goal “ x 1 , 0 , . . .
Revised: • MountainCar : φ “ ␣ Goal U Reset , where Goal “ x 1 , 0 , 0 , 0 y • Acrobot : φ “ ␣ Goal U Reset , where Goal “ x 1 , 0 , . . .

Original: , 0 y • LunarLander : φ “ ␣ SafeLanding U Reset , where SafeLanding “ GroundContact ^ MotorsOff , GroundContact “ x 0 , 1 , 0 , 0 , 0 , 0 , 0 y , and MotorsOff “ x 0 , 0 , 0 , 0 , 0 , 1 , 0 y •
Revised: , 0 y • LunarLander : φ “ ␣ SafeLanding U Reset , where SafeLanding “ GroundContact ^ MotorsOff , GroundContact “ x 0 , 1 , 0 , 0 , 0 , 0 , 0 y , and MotorsOff “ x 0 , 0 , 0 , 0 , 0 , 1 , 0 y •

Original: Pendulum : φ “ ♢ p␣ Safe ^ ⃝ Reset q , where Safe “ x 1 , 0 , 0 , 0 , 

Revised: Pendulum : φ “ ♢ p␣ Safe ^ ⃝ Reset q , where Safe “ x 1 , 0 , 0 , 0 ,

Original: y , ♢ T “ ␣H U T , and s i |ù ⃝ T iff s i ` 1 |ù T , for any T Ď AP , s i : 8 , a i : 8 P Traj .
Revised: y , ♢ T “ ␣H U T , and s i |ù ⃝ T iff s i ` 1 |ù T , for any T Ď AP , s i : 8 , a i : 8 P Traj .

Original: Intuitively, φ denotes the event of ending an episode in an unsafe state, just before resetting the environment,which means that either the agent never reached the safe region or it reached and left it atsome point.
Revised: Intuitively, φ denotes the event of ending an episode in an unsafe state, just before resetting the environment, means that either the agent never reached the safe region or it reached and left it at point.

Original: Formally, φ “ t s 0: 8 , a 0: 8 | D i P N , s i ­|ù
Revised: Formally, φ “ u Ď Traj N THE CURSE OF V ARIATIONAL M ODELING

Revised: Posterior collapse is a well known issue occurring in variational models (see, e.g., Alemi et al.

Original: Safe ^ s i ` 1 |ù Reset u Ď Traj .
Revised: Tolstikhin et al. He et al. 2019; Dong et al. 2020) which intuitively results in a degenerate local optimum where the model learns to ignore the latent space and use only the reconstruction functions (i.e., the decoding distribution) to optimize the objective. VAE-MDPs are no exception, as pointed out in the original paper (Delgrange et al., 2022, Section 4.3 and Appendix C.2). Formally, VAE- and WAE-MDPs optimize their objective by minimizing two losses: a reconstruction cost plus a regularizer term which penalizes a discrepancy between the encoding distribution and the dynamics of the latent space model.

Revised: In VAE-MDPs, the former corresponds to the the distortion , and the later to the rate of the variational model (further details are given in Alemi et al. 2018; Delgrange et al. 2022), while in our WAE-MDPs, the former corresponds to the raw transition distance and

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: the CartPole environment (i) when the distortion and the rate are minimized as is ( vanilla model ) and (ii) when it makes use of annealing schemes, entropy regularization, and prioritized experience replay to avoid posterior collapse (cf.

Revised: Delgrange et al. 2022).

Revised: While the former clearly fails to learn a useful latent representation, the later does so meticulously and smoothly in two distinguishable phases: first, ϕ ι focuses on fairly distributing the latent space, setting up the stage to the concrete optimization occurring from step 4 ¨ 10 5 , where the entropy of ϕ ι is lowered, which allows to get the rate of the variational model away from zero.

Revised: Five instances of the models are trained with different random seeds, with the same hyperparameters than in Sect.

Revised: the later to both the steady-state and transition regularizers.

Revised: Notably, the rate minimization of VAEMDPs involves regularizing a stochastic embedding function ϕ

Revised: ι p¨ | s q point-wise , i.e., for all different input states s P S drawn from the interaction with the original environment.

Revised: In contrast, the latent space regularization of the WAE-MDP involves the marginal embedding distribution Q ι where the embedding function ϕ

Revised: ι is not required to be stochastic.

Revised: Alemi et al.

Revised: (2018) showed that posterior collapse occurs in VAEs when the rate of the variational model is close to zero, leading to low-quality representation.

Revised: Posterior collapse in VAE-MDPs.

Revised: We illustrate the sensitivity of VAE-MDPs to the posterior collapse problem in Fig.

Revised: 7, through the CartPole environment 2 : minimizing the distortion and the rate as is yields an embedding function which maps deterministically every input state to the same sink latent state (cf.

Revised: Fig.

Revised: 7a).

Revised: Precisely, there is a latent state s P S so that ϕ ι p s | s q « 1 and

Revised: P θ p s | s, a q « 1 whatever the state s P S and action a P A .

Revised: This is a form of posterior collapse, the resulting rate quickly drops to zero (cf.

Revised: Fig 7b), and the resulting latent representation yields no information at all.

Revised: This phenomenon is handled in VAE-MDPs by using (i) prioritized replay buffers that allow to focus on inputs that led to bad representation, and (ii) modifying the objective function for learning the latent space model — the so-called evidence lower bound (Hoffman et al., 2013; Kingma & Welling, 2014), or ELBO for short — and set up annealing schemes to eventually recover the ELBO at the end of the training process.

Revised: Consequently, the resulting learning procedure focuses primarily on fairly distributing the latent space, to avoid it to collapse to a single latent state, to the detriment of learning the dynamics of the environment and the distillation of the RL policy.

Revised: Then, the annealing scheme allows to make the model learn to finally smoothly use the latent space to maximize the ELBO, and achieve consequently a lower distortion at the “price” of a higher rate.

Revised: Impact of the resulting learning procedure.

Revised: The aforementioned annealing process, used to avoid that every state collapses to the same representation, possibly induces a high entropy embedding function (Fig. 7d), which further complicates the learning of the model dynamics and the distillation in the first stage of the training process.

Revised: In fact, in this particular case, one can observe that the entropy reaches its maximal value, which yields a fully random state embedding function.

Revised: Recall that the VAE-MDP latent space is learned through independent Bernoulli distributions.

Revised: Fig.

Revised: 7d reports values centered around 4 .

Revised: 188 in the first training phase, which corresponds to the entropy of the state embedding function when ϕ ι p¨ | s q is uniformly distributed over S for any state s P S :

Revised: [Equation]

Revised: p i q

Revised: [Equation]

Revised: 188 , where p i “ 1 { 2 for all i .

Revised: The rate (Fig.

Revised: 7b) drops to zero since the divergence pulls the latent dynamics towards this high entropy (yet another form of posterior collapse), which hinders the latent space model to learn a useful representation.

Revised: However, the annealing scheme increases the rate importance along training steps, which enables the optimization to eventually leave this local optimum (here around 4 ¨ 10training steps).

Revised: This allows the learning procedure to leave the zero-rate spot, reduce the distortion (Fig. 7c), and finally distill the original policy (Fig. 7e).

Revised: As a result, the whole engineering required to mitigate posterior collapse slows down the training procedure.

Revised: This phenomenon is reflected in Fig.

Revised: 4: VAE-MDPs need several steps to stabilize and set up the stage to the concrete optimization, whereas WAE-MDPs have no such requirements since they naturally do not suffer from collapsing issues (cf.

Revised: Fig. 5), and are consequently faster to train.

Revised: Lack of representation guarantees.

Revised: On the theoretical side, since VAE-MDPs are optimized via the ELBO and the local losses via the related variational proxies, VAE-MDPs do not leverage the representation quality guarantees induced by local losses (Eq. 1) during the learning procedure (as explicitly pointed out by Delgrange et al., 2022, Sect. 4.1.): in contrast to WAE-MDPs, when two original states are embedded to the same latent, abstract state, the former are not guaranteed to be bisimilarly close (i.e., the agent is not guaranteed to behave the same way from those two states by executing the policy), meaning those proxies do not prevent original states having distant values collapsing together to the same latent representation.

Revised: A DDITIONAL R