Title: Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficiently
Venue: NeurIPS 2024
Accepted: yes
Keywords: Optimal transport

Markov chains

bisimulation metrics

Markov decision processes
Abstract: We propose a new framework for formulating optimal transport distances between Markov chains. Previously known formulations studied couplings between the entire joint distribution induced by the chains, and derived solutions via a reduction to dynamic programming (DP) in an appropriately defined Markov decision process. This formulation has, however, not led to particularly efficient algorithms so far, since computing the associated DP operators requires fully solving a static optimal transport problem, and these operators need to be applied numerous times during the overall optimization process. In this work, we develop an alternative perspective by considering couplings between a ``flattened'' version of the joint distributions that we call discounted occupancy couplings, and show that calculating optimal transport distances in the full space of joint distributions can be equivalently formulated as solving a linear program (LP) in this reduced space. This LP formulation formulation allows us to port several algorithmic ideas from other areas of optimal transport theory. In particular, our formulation makes it possible to introduce an appropriate notion of entropy regularization into the optimization problem, which in turn enables us to directly calculate optimal transport distances via a Sinkhorn-like method we call Sinkhorn Value Iteration (SVI). We show both theoretically and empirically that this method converges quickly to an optimal coupling, essentially at the same computational cost of running vanilla Sinkhorn in each pair of states. Along the way, we point out that our optimal transport distance exactly matches the common notion of bisimulation metrics between Markov chains, and thus our results also apply to computing such metrics, and in fact our algorithm turns out to be significantly more efficient than the best known methods developed so far for this purpose.

Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficiently

Abstract

We propose a new framework for formulating optimal transport distances between Markov chains. Previously known formulations studied couplings between the entire joint distribution induced by the chains, and derived solutions via a reduction to dynamic programming (DP) in an appropriately defined Markov decision process. This formulation has, however, not led to particularly efficient algorithms so far, since computing the associated DP operators requires fully solving a static optimal transport problem, and these operators need to be applied numerous times during the overall optimization process. In this work, we develop an alternative perspective by considering couplings between a 'flattened' version of the joint distributions that we call discounted occupancy couplings, and show that calculating optimal transport distances in the full space of joint distributions can be equivalently formulated as solving a linear program (LP) in this reduced space. This LP formulation allows us to port several algorithmic ideas from other areas of optimal transport theory. In particular, our formulation makes it possible to introduce an appropriate notion of entropy regularization into the optimization problem, which in turn enables us to directly calculate optimal transport distances via a Sinkhorn-like method we call Sinkhorn Value Iteration (SVI). We show both theoretically and empirically that this method converges quickly to an optimal coupling, essentially at the same computational cost of running vanilla Sinkhorn in each pair of states. Along the way, we point out that our optimal transport distance exactly matches the common notion of bisimulation metrics between Markov chains, and thus our results also apply to computing such metrics, and in fact our algorithm turns out to be significantly more efficient than the best known methods developed so far for this purpose.

1 Introduction

Measuring distances between structured objects and sequences is an important problem in a variety of areas of science. The more structured the objects become, the harder it gets to define appropriate notions of distances, as good notions of proximity need to take into account the possibly complex relationships between the constituent parts of each object. The possibility that the objects in question may be random further complicates the picture, and in such cases it becomes more natural to measure distances between the underlying joint probability distributions. Within the specific context of comparing stochastic processes, two natural notions of distance have emerged over the past decades: the notion of probabilistic bisimulation metrics that takes its root in modal logic and theoretical computer science [Sangiorgi, 2009, Abate, 2013], and the notion of optimal-transport distances that originates from probability theory [Villani, 2009, Peyré and Cuturi, 2019]. In this paper, we show that bisimulation metrics are in fact optimal-transport distances, and we make use of this observation to derive efficient algorithms for computing distances between stochastic processes.

The two distance notions have found strikingly different applications. Bisimulation emerged within the area of theoretical computer science as one of the most important important concepts in concurrency theory and formal verification of computer systems [Park, 1981, Milner, 1989], and has been extended to probabilistic transition systems by Larsen and Skou [1989]. Within machine learning, bisimulation metrics have become especially popular in the context of reinforcement learning (RL) due to the work of Ferns et al. [2004], and have become one of the few standard tools of representation learning [Jiang, 2018, 2024]. In particular, the work of Ferns et al. [2004] advocates for using bisimilarity as a basis for state aggregation, measuring similarities of states in terms of similarities of two chains M X and M Y that only differ in their initial state. While this approach has inspired numerous follow-up works [Gelada et al., 2019, Castro, 2020, Agarwal et al., 2021b, Zhang et al., 2021, Hansen-Estruch et al., 2022, Castro et al., 2022], ultimately this line of work has failed to discover efficient algorithms for computing bisimulation metrics and has largely resorted to heuristics for computing similarity metrics.

On the other hand, optimal transport (OT) has found numerous applications in areas as diverse as economics [Galichon, 2016], signal processing [Kolouri et al., 2017], or genomics [Schiebinger et al., 2019]. Within machine learning, it has been used for the similarly diverse areas of domain adaptation [Courty et al., 2016], generative modeling [Arjovsky et al., 2017, Song et al., 2020, Shi et al., 2024], representation learning [Courty et al., 2018], and, perhaps most relevant to our work, as a way of measuring distances between graphs [Titouan et al., 2019, Chen et al., 2022, Chuang and Jegelka, 2022]. The recent works of Yi et al. [2021], Brugère et al. [2024] propose to define graph distances via studying the behavior of random walks defined on the graph, thus reducing the problem of comparing graphs to comparing stochastic processes-exactly the subject of the present paper. Other applications of OT between stochastic processes include generative modeling for sequential data [Xu et al., 2020], pricing and hedging in mathematical finance [Backhoff-Veraguas et al., 2017], and analyzing multistage stochastic optimization problems [Pflug, 2010, Bartl and Wiesel, 2022]. It appears however that the literature on optimal transport for stochastic processes has apparently not yet discovered connections with bisimulation metrics and the rich intellectual history behind it. Also, applications of optimal transport for representation learning within the context of reinforcement learning appear to be nonexistent.

In this paper we observe that, despite their apparent differences, bisimulation metrics and optimal transport distances are one and the same. Furthermore, we provide a new perspective on both OT distances and bisimulation metrics by formulating the distance metric as the solution of a linear program (LP) in the space of 'occupancy couplings', a finite-dimensional projection of the infinitedimensional process laws. Building on tools from computational optimal transport [Peyré and Cuturi, 2019] and entropy-regularized Markov decision processes [Neu et al., 2017, Geist et al., 2019], we design an algorithm that effectively combines Sinkhorn's algorithm [Sinkhorn and Knopp, 1967, Cuturi, 2013] with an entropy-regularized version of the classic Value Iteration algorithm [Bellman, 1957, Neu et al., 2017]. Building on recent work on computational optimal transport [Altschuler et al., 2017, Ballu and Berthet, 2023], we provide theoretical guarantees for the resulting algorithm (called Sinkhorn Value Iteration) and perform numerical studies that demonstrate its effectiveness for computing distances between Markov chains.

Notations. For a finite set S , we use ∆ S to denote the set of all probability distributions over S . We will denote infinite sequences by x = ( x 0 , x 1 , . . . ) and the corresponding subsequences as x n = ( x 0 , x 1 , . . . , x n ) . For two sets X and Y , we will often write XY to abbreviate the directproduct notation X × Y , and for two indices x and y and a function f : XY → Z , we will often write f ( xy ) instead of f ( x, y ) to save space. Also, we will denote scalar products by ⟨· , ·⟩ and use ∥·∥ p to denote the ℓ p -norm.

2 Preliminaries

We study the problem of measuring distances between pairs of finite Markov chains. Specifically, we consider two stationary Markov processes M X = ( X , P X , ν 0 , X ) and M Y = ( Y , P Y , ν 0 , Y ) , where

X and Y are the finite state spaces with cardinalities m = |X| and n = | Y | ,

P X : X → ∆ X and P Y : Y → ∆ Y are the transition kernels that determine the evolution of the states as P X ( x ′ | x ) = P [ X t +1 = x ′ | X t = x ] and P Y ( y ′ | y ) = P [ Y t +1 = y ′ | Y t = y ] for all t , and

ν 0 , X and ν 0 , Y are the initial-state distributions with X 0 ∼ ν 0 , X and Y 0 ∼ ν 0 , Y .

Without significant loss of generality, we will suppose that the initial states are fixed almost surely as X 0 = x 0 and Y 0 = y 0 , and refer to their corresponding joint distribution as ν 0 = δ x 0 ,y 0 . These objects together define a sequence of joint distributions P [( X 0 , X 1 , . . . , X n ) = ( x 0 , x 1 , . . . , x n )] and P [( Y 0 , Y 1 , . . . , Y n ) = ( y 0 , y 1 , . . . , y n )] for each n , which together define respective the laws of the infinite sequences X = ( X 1 , X 2 , . . . ) and Y = ( Y 1 , Y 2 , . . . ) via Kolmogorov's extension theorem. With a slight abuse of notation, we will use M X and M Y to denote the corresponding measures that satisfy M X ( x n ) = P [ X n = x n ] and M Y ( y n ) = P [ Y n = y n ] for any x ∈ X ∞ , y ∈ Y ∞ and n . The corresponding conditional distributions are denoted as M X ( x n | x n -1 ) = P [ X n = x n | X n -1 = x n -1 ] and M Y ( y n | y n -1 ) = P [ Y n = y n | Y n -1 = y n -1 ] .

2.1 Optimal transport between Markov chains

Our main object of interest in this work is a notion of optimal transport distance between infinitehorizon Markov chains. Several previous works have studied such distances (which are discussed in detail in Appendix A), and our precise definition we give below is closest to Moulos [2021], O'Connor et al. [2022] and Brugère et al. [2024]. We consider Markov chains on state spaces where a 'ground metric' (or 'ground cost') c : X ×Y → R + is available to measure distances between any two individual states x ∈ X and y ∈ Y , with the distance denoted as c ( x, y ) . For any two sequences x = ( x 1 , x 2 , . . . ) and y = ( y 1 , y 2 , . . . ) , we define the discounted total cost

where γ ∈ (0 , 1) is the discount factor that expresses the preference that two sequences be considered further apart if they exhibit differences at earlier times in terms of the ground cost c . Following the optimal-transport literature, we will consider distances between the stochastic processes M X and M Y via the notion of couplings . To this end, we define a coupling of M X and M Y as a stochastic process evolving on the joint space XY , with its law defined for all n as M XY ( x n y n ) = P [ X n = x n , Y n = y n ] , required to satisfy ∑ y n ∈Y n M XY ( x n y n ) = M X ( x n ) and ∑ x n ∈X n M XY ( x n y n ) = M Y ( y n ) . We will define the set of all such couplings as Π .

The notion of couplings defined above does not respect the temporal structure of the Markov chains M X and M Y appropriately: while by definition the distribution of state X n may only be causally influenced by past states X k with k < n , the general notion of coupling above allows the state X n to be influenced by future states Y k with k ≥ n as well. To rule out this possibility (and following past works mentioned in the introduction), we will introduce the notion of bicausal couplings . A coupling M XY is called bicausal if and only if it satisfies

for all sequences x, y ∈ X ∞ ×Y ∞ and all n . Denoting the set of all bicausal couplings by Π bc, we define our optimal transport distance as

where we emphasize the dependence of the distance on x 0 , y 0 explicitly with our notation.

By noticing that the optimization problem outlined above can be reformulated as a Markov decision process (MDP), Moulos [2021] has shown that the infimum in (2) is achieved within the family of Markovian couplings that satisfy

for all sequences of state pairs and all values of n . Furthermore, it can be seen that Markovian couplings are fully specified in terms of transition couplings of the form π : XY → ∆ XY , with π ( x ′ y ′ | xy ) standing for P [ ( X t +1 , Y t +1 ) = ( x ′ , y ′ ) | ( X t , Y t ) = ( x t , y t )] under the law induced by the coupling. We say that a transition coupling is valid if it satisfies the marginal constraints ∑ y ′ π ( x ′ y ′ | xy ) = P X ( x ′ | x ) and ∑ x ′ π ( x ′ y ′ | xy ) = P Y ( y ′ | y ) . Defining the set of such valid transition couplings in state pair xy by Π xy = { p ∈ ∆ XY : ∑ y ′ p ( x ′ y ′ ) = P X ( x ′ | x ) , ∑ x ′ p ( x ′ y ′ ) =

P Y ( y ′ | y ) } , Moulos [2021] introduces an MDP M with an infinite action set corresponding to picking the joint next-state couplings in Π xy . An optimal transition coupling can be found by solving the following Bellman optimality equations of the MDP M :

The infimum on the right-hand side is achieved by an optimal transition coupling π ∗ ( ·| xy ) = arg min p ∈ Π xy ∑ x ′ y ′ p ( x ′ y ′ ) V ∗ ( x ′ y ′ ) . The solution V ∗ is unique and can be shown to satisfy V ∗ ( xy ) = W γ ( M X , M Y ; c, x, y ) for all xy ∈ XY . For completeness, we include the precise definition of the MDP M and the proofs of these results in Appendix B.

2.2 Bisimulation metrics

The notion of bisimulation metrics has been introduced by Desharnais et al. [1999, 2004] and van Breugel and Worrell [2001], with the purpose of defining distances between Markov chains, using a methodology rooted in modal logic that at first may appear entirely different from the optimal-transport framework described above. We only give a very high-level overview of the classic logic-based characterization here (as the fine details are irrelevant to the final conclusion that this section is headed to), and refer the reader to the additional discussion in Appendix A for further reading. Desharnais et al. [1999, 2004] considered labeled Markov chains where a labeling function r : X ∪ Y → R assigns labels to each state, and defined bisimulation metrics via

where F γ is a family of functional expressions generated by a certain grammar, and f M X : X → R and f M Y : Y → R are the respective 'interpretations' of each f ∈ F on the Markov chains M X and M Y . Building on this formulation, van Breugel and Worrell [2001] have shown that the distance metric can be equivalently characterized by the solution of a fixed-point equation, whose expression was subsequently used by Ferns et al. [2004] to define bisimulation metrics for Markov decision processes where r takes the role of a reward function, and the evolution of states may be influenced by actions. The case of having no actions available corresponds to our setting, where their definition of a bisimulation metric simplifies to the solution of the fixed-point equation

The solution to this system is unique and satisfies U ∗ ( xy ) = d γ ( M X , M Y ; r, x, y ) . Putting this result side-by-side with Equation (3), one can immediately realize that bisimulation metrics and our notion of optimal transport distances coincide when picking the ground cost function c ( xy ) = (1 -γ ) | r ( x ) -r ( y ) | . To our knowledge, this remarkable observation has not been publicly made anywhere in either the optimal-transport or the bisimulation-metric literature. This connection has several important implications, which we have already discussed at some length in Section 1. We relegate further discussion of these metrics in the light of this observation to Appendix A.

3 Optimal transport between Markov chains as a linear program

Optimal transport problems can typically be formulated as linear programs (LPs), since couplings can be characterized as joint distributions satisfying a set of easily-expressed linear constraints (see, e.g., Chapter 3 in Peyré and Cuturi 2019). Our problem is no exception, and in fact the original problem statement of Equation 2 can be expressed in this form: the constraints defining Π bc are all linear in M XY . However, M XY is an infinite-dimensional object and thus this formulation is not instructive for developing computationally tractable algorithms. We address this problem in this section, where we define an equivalent LP formulation that replaces the infinite-dimensional optimization variable with an appropriate low-dimensional projection. Our framework builds on the classic LP formulation of optimal control in MDPs first proposed in the 1960's [Manne, 1960, de Ghellinck, 1960, d'Epenoux, 1963, Denardo, 1970], and covered thoroughly in several standard textbooks (e.g., Section 6.9 of Puterman 1994).

In order to set things up, we need to start with some important definitions. We say that a transition coupling π : XY → ∆ XY generates a trajectory ( X 0 , Y 0 , X 1 , Y 1 , . . . ) if ( X 0 , Y 0 ) ∼ ν 0 and the

subsequent state-pairs are drawn independently from the transition coupling as ( X t +1 , Y t +1 ) ∼ π ( ·| X t , Y t ) for all t . Then, we define the occupancy coupling µ π associated with this process as a distribution over XY × XY , with each of its entries xy, x ′ y ′ defined as

where E π [ · ] emphasizes that the trajectory of state-pairs has been generated by π . In words, µ π ( xy, x ′ y ′ ) is the discounted number of times that the quadruple ( xy, x ′ y ′ ) is visited by the process. With this definition, it is easy to notice that the objective optimized in Equation (2) can be rewritten as a linear function of µ π . Indeed, suppose that M XY is the law of the process generated by π as described above, so that we can write

We say that an occupancy coupling µ is valid if it is generated by a valid transition coupling. It is easy to verify that every valid occupancy coupling µ ∈ ∆ XY×XY satisfies the following three constraints:

We refer to the first equality constraint as the flow constraint and the second and third ones as the transition coherence constraints for M X and M Y , respectively 1 . We show in the following lemma that the above conditions uniquely characterize the set of valid occupancy couplings.

Lemma 1. A distribution µ ∈ ∆ XY×XY is a valid occupancy coupling associated with some transition coupling π : XY → ∆ XY if and only if it satisfies Equations (6) -(8) .

One important consequence of this result is that for each transition coupling, one can associate a transition coupling π µ with µ π µ = µ , with entries satisfying ∑ x ′′ y ′′ µ ( xy, x ′′ y ′′ ) π µ ( x ′ y ′ | xy ) = µ ( xy, x ′ y ′ ) . A proof is provided in Appendix B.2. Having established in the previous section that stationary occupancy couplings are sufficient to achieve the supremum in the definition of the OT distance of Equation (2), this result immediately implies the following.

Theorem 1. Let B , S X and S Y respectively stand for the sets of all distributions µ ∈ ∆ XY×XY that satisfy equations (6) , (7) , and (8) . Then,

Since the constraints on µ in the above reformulation are all linear, the optimization problem stated above is obviously a linear program. Notably, the resulting LP is not the standard dual to the LP associated with the MDP formulation introduced in Section 2.1, which would result in an infinite-dimensional LP with one constraint associated with each of the continuously-valued actions. Such an infinite-dimensional reformulation was previously considered by Chen et al. [2012] in the context of computing bisimulation metrics, who used it as a tool for analysis rather than algorithm design. Notably, our formulation results in a finite-dimensional LP with |X| 2 |Y| 2 variables and |X| 2 + |X| |Y| + |Y| 2 constraints. Building on the celebrated result of Ye [2011] (and similarly to the work of Chen et al., 2012 mentioned above), one can show that our LP can be solved in strongly polynomial time via an appropriate adaptation of the simplex method or the classic policy iteration method of Howard [1960]. We propose an alternative methodology in the next section.

4 Sinkhorn Value Iteration

Solving optimal-transport problems via standard LP solvers (like variations of the simplex method or interior-point methods) is known to be empirically hard, and thus we seek alternatives to this approach towards optimizing our own LP defined in the previous section. In computational optimal transport, a paradigm shift was initiated by Cuturi [2013] who successfully applied entropic regularization to the classic LP objective of optimal transport, and solved the resulting optimization method through an iterative algorithm called the Sinkhorn-Knopp method (sometimes simply called Sinkhorn's algorithm, due to Sinkhorn and Knopp, 1967). This method is based on finding a feasible point in a two-constraint problem by alternately satisfying one and the other, and has resulted in practical algorithms that were orders of magnitude faster than all previously studied methods. Entropy regularization and Sinkhorn's algorithm have thus became the most important cornerstones of computational optimal transport. Drawing on the same principles as well as the theory of entropy-regularized Markov decision processes [Neu et al., 2017, Geist et al., 2019], we develop a computationally effective algorithm for computing optimal transport distances between Markov chains below.

4.1 Formal definition: Mirror Sinkhorn on the space of occupancy couplings

Our method is an adaptation of a version of Sinkhorn's algorithm called Mirror Sinkhorn, first proposed and analyzed by Ballu and Berthet [2023]. This method combines mirror-descent-style updates [Nemirovski and Yudin, 1983, Beck and Teboulle, 2003] with alternating projections to two convex sets whose intersection corresponds to the feasible set we seek to optimize over. In our adaptation, we choose the two sets as

that can be seen to be the set of distributions µ that satisfy both Equations (6) and (7), and

which is the set of distributions µ that satisfy both Equations (6) and (8). Naturally, the intersection of the two sets corresponds to valid occupancy couplings. It remains to define an appropriate notion of entropy for the purpose of regularization. Following Neu et al. [2017], we will use the conditional relative entropy defined between two joint distributions µ, µ ′ ∈ ∆ XYXY as

̸

It is an easy exercise to show that H is a Bregman divergence that is convex in its first argument (see, e.g., Appendix A.1 of Neu et al., 2017). Note however that H ( µ ∥ µ ′ ) can be zero even if µ = µ ′ and thus it is not strongly convex in µ .

With these ingredients, we symbolically define our algorithm as calculating the sequence of updates

for each k = 1 , . . . , K -1 , where µ 1 is the occupancy coupling associated to the trivial coupling π 1 ( ·| xy ) = P X ( ·| x ) ⊗ P Y ( ·| y ) for each state pair xy ∈ XY , η > 0 is a stepsize (or learning-rate) parameter and B k is chosen to be B X in odd rounds and B Y in even rounds. By adapting tools from the theory of entropy-regularized Markov decision processes, the updates can be computed in closed form by solving a system of equations closely resembling the regularized Bellman equations. In particular, we define the Bellman-Sinkhorn operators for a given transition coupling π as the operators

T π X : R XY×X → R XY×X and T π Y : R XY×Y → R XY×Y acting on functions V X ∈ R XY×X and V Y ∈ R XY×Y respectively as

and

Then, for odd rounds, the updates can be calculated by solving the fixed-point equations V k = T π k X V k , defining the shorthand Q k ( xy, x ′ y ′ ) = c ( xy ) + γ ∑ x ′′ P X ( x ′′ | x ′ ) V k ( x ′ y ′ , x ′′ ) , and subsequently updating the transition coupling π k multiplicatively as

It is easy to verify that this transition coupling satisfies ∑ y ′ π k +1 ( x ′ y ′ | xy ) = P ( x ′ | x ) . The updates for even rounds are computed analogously with the roles of X and Y swapped. We respectively refer to the fixed-point equations V k = T π k X V k and V k = T π k Y V k as the Bellman-Sinkhorn equations for M X and M Y , and the functions V k and Q k as value functions . The following proposition (proved in Appendix E.1) formally establishes the equivalence between the two update rules.

Proposition 1. Let µ k +1 and π k +1 be specified for each k as in Equations (9) and (10) , respectively. Then, µ k +1 = µ π k +1 holds for all k .

4.2 Practical implementation

The algorithm described above can be seen as performing online Mirror Sinkhorn updates in each state pair xy with a sequence of cost functions Q k , which are computed via solving the Bellman-Sinkhorn equations. Since T π X and T π Y are easily seen to be contractive with respect to the supremum norm with contraction factor γ (as shown by a standard calculation included in Appendix E.3), these equations can be solved by an adaptation of the classic Value Iteration method of Bellman [1957]. Concretely, we repeatedly apply the Bellman-Sinkhorn operators until the fixed point is reached up to sufficient precision (controlled by the number of update steps m ). We call the resulting method Sinkhorn Value Iteration (SVI), and provide its pseudocode as Algorithm 1.

Notably, while SVI is defined in its abstract form as a sequence of updates in the space of occupancy couplings µ k , its implementation only works with transition couplings π k . The final output of SVI is a transition coupling π out, obtained by computing the average µ K = 1 K ∑ K k =1 µ π k of all occupancy couplings, computing π µ K and then rounding the result to a valid transition coupling. In particular we apply a simple rounding procedure due to Altschuler et al. [2017] individually on π µ K ( ·| xy ) for each state-pair xy -for the full details, see Appendix E.2. Besides π out, SVI also outputs an estimate of V ∗ in the form of the value function V π out , as defined in Equation (11) in Appendix B.1. This function can be computed efficiently by solving the linear system of Bellman equations V π out ( xy ) = c ( xy ) + γ ∑ x ′ y ′ π out ( x ′ y ′ | xy ) V π out ( x ′ y ′ ) .

A number of small simplifying steps can be made to make the algorithm easier to implement. First, instead of obtaining π µ K via the computationally expensive procedure described above, one can simply run the rounding procedure on the final transition coupling π K and output the result. Second,

while theoretical analysis suggests setting m = ∞ in order to make sure that all projection steps are perfect, such exact computation may be unnecessary and inefficient in practice, and thus (much) smaller values can be used instead. Third, for small values of η the softmax function used in the definition of the Bellman-Sinkhorn operator can be accurately approximated by an average with respect to π k ( x ′ y ′ | xy ) /P X ( x ′ | x ) , which suggests a simple alternative to the projection steps. This approximates SVI similarly as to how the Mirror Descent Modified Policy Iteration method of Geist et al. [2019] approximates the mirror descent method of Neu et al. [2017] (see also Azar et al., 2012). The resulting method (that we refer to as Sinkhorn Policy Iteration , or SPI) is presented in detail along with its theoretical analysis in Appendix D. We study effects of these implementation choices via a sequence of experiments in Section 5.

4.3 Convergence guarantees

The following theorem establishes a guarantee on the number of iterations necessary for V π out ( xy ) be an ε -accurate approximation of the transport cost W γ ( M X , M Y ; c, x, y ) for any x, y .

Theorem 2. Suppose that Sinkhorn Value Iteration is run for K steps with regularization parameter η = 1 4 ∥ c ∥ ∞ √ (1 -γ ) 3 log |X||Y| K , and initialized with the uniform coupling defined for each xy, x ′ y ′ as π 1 ( x ′ y ′ | xy ) = 1 |X||Y| . Then, for any x 0 y 0 ∈ XY , the output satisfies V π out ( x 0 y 0 ) ≤ W γ ( M X , M Y ; c, x 0 , y 0 ) + ε if the number of iterations is at least

The proof is relegated to Section C, and we present a similar performance guarantee for SPI in Appendix D. Importantly, these guarantees technically only hold when setting m = ∞ , which is a limitation we discuss in more detail in Section 6. The condition that π 1 is chosen as the uniform coupling is not necessary and simply made to make the statement easier to state. A more detailed statement of the bound is provided in Appendix C.4.

5 Experiments

We have conducted a range of experiments on some simple environments with the purpose of illustrating the numerical properties of our algorithms and some aspects of the distance metrics we studied. Due to space restrictions, we only report a very limited subsample of the results below, and refer the reader to Appendix F for the complete suite 2 .

One set of experiments we report here addresses the biggest open question left behind our theory: the effect of the choice of m on the quality of the updates. For this experiment, we use the classic '4-rooms' environment first studied by Sutton et al. [1999], and run both SVI (Algorithm 1) and SPI (Algorithm 2) for a range of different choices of m , and a fixed γ = 0 . 95 . The results of this study are shown in Figure 1. The plots indicate that the estimates produced by both algorithms converge towards the true distance at a rate that is basically unaffected by m , and in particular even a value of m = 1 remains competitive. This observation is consistent across all of our experiments. Also, the output of SPI appears to converge slightly more slowly towards the optimum in this experiment, but this observation is not entirely consistent and can be likely ascribed to the fact that the learning rate was not optimized to favor either algorithm in this experiment. In most experiments, the two algorithms performed very similarly, up to some small occasional differences.

We have also conducted a number of experiments to illustrate the potential of optimal-transport distances for comparing Markov chains of different sizes and transition functions. In the experiment we show here, we compare two Markov chains illustrated in Figure 2. The first Markov chain M X is a simple, nine-state 'gridworld' environment, which has its initial state in the upper left corner (denoted as s 0 ) and a reward of +1 in the lower left corner (shown in blue). The second, M Y , is an instance of the 4-rooms environment, where each room is a rotation of the aforementioned small grid. The transition kernels in both environments are uniform distributions over the adjacent cells in the four principal directions. The plot shows the distances between the two chains as a function of the initial state of M Y , revealing an intuitive pattern of similarities that captures the symmetries of M Y .

(a) SVI (Algorithm 1)

(b) SPI (Algorithm 2)

10-2

m:1

m:2

m:5

m: 10

m:25

10-41

100

101

102

103

Iterations

10-2

m: 1

m:2

m:5

m:10

m:25

10-4.

100

101

102

103

Iterations

V*(so,÷) =

+1

+1

So

七=

+1

+1

6 Discussion

We discuss some further aspects of our framework and results below.

Representation learning for reinforcement learning. Among the numerous applications listed in Section 1 and Appendix A, the most interesting for us is using our metrics for representation learning in RL. As mentioned earlier, bisimulation metrics have been extensively used for this purpose in the past. In particular, almost all such work uses bisimulation metrics to compare states within the same MDP and use the resulting similarity metrics for merging states that are at low distance (an approach called 'state aggregation'). As our results highlight, this is a rather narrow view of what bisimulation metrics are capable of: they can define similarity metrics between processes that live on potentially different state spaces, which in particular can be used to select representations by minimizing the distance between a high-dimensional process and a set of low-dimensional representations. Curiously, our LP formulation may allow differentiating the distances with respect to the transition kernels, which we believe will be an important property for future developments in representation learning for RL.

Limitations of the theory. In their current form, our theoretical guarantees in Theorems 2 and 6 only apply to perfect projection and evaluation steps, corresponding to setting m = ∞ . We conjecture that this limitation can be addressed with a more careful analysis, and results similar to those of Theorems 2 and 6 can be shown, potentially at the price of a worse dependence on the effective horizon 1 / (1 -γ ) [Scherrer et al., 2012, 2015], by making use of the techniques of Geist et al. [2019] and Moulin and Neu [2023] for analyzing regularized dynamic-programming algorithms.

From dynamic programming to learning from data. This paper focuses on computing distances between known Markov processes via dynamic-programming-style methods. In the most interesting applications however, the transition kernels are unknown, which requires the development of new tools. We are confident that our framework can serve as a solid basis for such developments, and in particular that one can port many ideas from the field of reinforcement learning that is essentially all about turning dynamic-programming methods into algorithms that can learn from interaction data. Additionally, we believe that our LP formulation in Section 3 makes it much easier to import further ideas from computational optimal transport, and in particular that stochastic optimization methods like those of Genevay et al. [2016] can be adapted to solving our linear programs.

Acknowledgments.

G. Neu would like to thank Marin Ballu and Quentin Berthet for clarifying a number of details of their analysis of Mirror Sinkhorn, Tristan Brugère for pointing us to their publicly available code, and Anna Korba and Stefan Schrott for kindly helping out with some references about optimal transport between stochastic processes. G. Neu was supported by the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (Grant agreement No. 950180). Anders Jonsson is partially supported by the EU ICT-48 2020 project TAILOR (No. 952215), AGAUR SGR, and the Spanish grant PID2019-108141GB-I00. Javier Segovia-Aguas is supported by MCIN/AEI /10.13039/501100011033 under the Maria de Maeztu Units of Excellence Programme (CEX2021-001195-M).

References

Abate. Approximation metrics based on probabilistic bisimulations for general state-space Markov processes: a survey. Electronic Notes in Theoretical Computer Science , 297:3-25, 2013.

Abbasi-Yadkori, P. Bartlett, K. Bhatia, N. Lazic, C. Szepesvari, and G. Weisz. POLITEX: Regret bounds for policy iteration using expert prediction. In International Conference on Machine Learning (ICML) , pages 3692-3702, 2019.

Aczel. Non-well-founded sets . CSLI lecture notes, no. 14, 1988.

Agarwal, S. M. Kakade, J. D. Lee, and G. Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research , 22(98): 1-76, 2021a.

Agarwal, M. C. Machado, P. S. Castro, and M. G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations (ICLR) , 2021b.

Altschuler, J. Niles-Weed, and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. Advances in Neural Information Processing Systems , 30, 2017.

Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML) , pages 214-223, 2017.

G. Azar, V. Gómez, and H. J. Kappen. Dynamic policy programming. Journal of Machine Learning Research , 13(Nov):3207-3245, 2012.

Backhoff-Veraguas, M. Beiglbock, Y. Lin, and A. Zalashko. Causal transport in discrete time and applications. SIAM Journal on Optimization , 27(4):2528-2562, 2017.

Backhoff-Veraguas, D. Bartl, M. Beiglböck, and M. Eder. All adapted topologies are equal. Probability Theory and Related Fields , 178:1125-1172, 2020.

Ballu and Q. Berthet. Mirror Sinkhorn: Fast online optimization on transport polytopes. In International Conference on Machine Learning (ICML) , 2023.

Bartl and J. Wiesel. Sensitivity of multiperiod optimization problems in adapted wasserstein distance. arXiv preprint arXiv:2208.05656 , 2022.

Bayraktar and B. Han. Fitted value iteration methods for bicausal optimal transport. arXiv preprint arXiv:2306.12658 , 2023.

Beck and M. Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters , 31(3):167-175, 2003.

Bellman. Dynamic Programming . Princeton University Press, Princeton, New Jersey, 1957.

P. Bertsekas. Monotone mappings with application in dynamic programming. SIAM Journal on Control and Optimization , 15(3):438-464, 1977. doi: 10.1137/0315031.

Bian and A. Abate. On the relationship between bisimulation and trace equivalence in an approximate probabilistic context. In Foundations of Software Science and Computation Structures (FoSSaCS) , pages 321-337, 2017.

Brugère, Z. Wan, and Y. Wang. Distances for Markov chains, and their differentiation. In International Conference on Algorithmic Learning Theory (ALT) , pages 282-336, 2024.

X.-R. Cao. Single sample path-based optimization of Markov chains. Journal of optimization theory and applications , 100:527-548, 1999.

S. Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In AAAI Conference on Artificial Intelligence (AAAI) , pages 10069-10076, 2020.

S. Castro, T. Kastner, P. Panangaden, and M. Rowland. A kernel perspective on behavioural metrics for Markov decision processes. Transactions on Machine Learning Research , 2022.

Chen, F. van Breugel, and J. Worrell. On the complexity of computing probabilistic bisimilarity. In Foundations of Software Science and Computation Structure (FoSSaCS) , pages 437-451, 2012.

Chen, S. Lim, F. Mémoli, Z. Wan, and Y. Wang. Weisfeiler-Lehman meets Gromov-Wasserstein. In International Conference on Machine Learning (ICML) , pages 3371-3416, 2022.

C.-Y. Chuang and S. Jegelka. Tree mover's distance: Bridging graph metrics and stability of graph neural networks. Advances in Neural Information Processing Systems , 35:2944-2957, 2022.

Courty, R. Flamary, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence , 39(9):1853-1865, 2016.

Courty, R. Flamary, and M. Ducoffe. Learning Wasserstein embeddings. In International Conference on Learning Representations (ICLR) , pages 1-13, 2018.

Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems , 26, 2013.

de Ghellinck. Les problèmes de décisions séquentielles. Cahiers du Centre d'Études de Recherche Opérationnelle , 2:161-179, 1960.

V. Denardo. On linear programming in a Markov decision problem. Management Science , 16(5): 281-288, 1970.

d'Epenoux. A probabilistic production and inventory problem. Management Science , 10(1): 98-108, 1963.

Desharnais, V. Gupta, R. Jagadeesan, and P. Panangaden. Metrics for labeled Markov systems. In International Conference on Concurrency Theory (CONCUR) , pages 258-273, 1999.

Desharnais, R. Jagadeesan, V. Gupta, and P. Panangaden. The metric analogue of weak bisimulation for probabilistic processes. In IEEE Symposium on Logic in Computer Science (LICS) , pages 413-422, 2002.

Desharnais, V. Gupta, R. Jagadeesan, and P. Panangaden. Metrics for labelled Markov processes. Theoretical Computer Science , 318(3):323-354, 2004.

Eckstein and G. Pammer. Computational methods for adapted optimal transport. The Annals of Applied Probability , 34(1A):675-713, 2024.

Even-Dar, S. M. Kakade, and Y. Mansour. Online Markov decision processes. Mathematics of Operations Research , 34(3):726-736, 2009.

Ferns, P. Panangaden, and D. Precup. Metrics for finite Markov decision processes. In Uncertainty in Artificial Intelligence (UAI) , pages 162-169, 2004.

Ferns, P. S. Castro, D. Precup, and P. Panangaden. Methods for computing state similarity in Markov decision processes. In Uncertainty in Artificial Intelligence (UAI) , pages 174-181, 2006.

Forti and F. Honsell. Set theory with free construction principles. Annali Scuola Normale Superiore, Pisa , Serie IV X(3):493-522, 1983.

Galichon. Optimal Transport Methods in Economics . Princeton University Press, 2016.

Geist, B. Scherrer, and O. Pietquin. A theory of regularized Markov decision processes. In International Conference on Machine Learning (ICML) , pages 2160-2169, 2019.

Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Learning continuous latent space models for representation learning. In International Conference on Machine Learning (ICML) , 2019.

Genevay, M. Cuturi, G. Peyré, and F. Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems , 29, 2016.

Giacalone, C.-C. Jou, and S. A. Smolka. Algebraic reasoning for probabilistic concurrent systems. In IFIP Conference on Programming Concepts and Methods , pages 443-458, 1990.

Givan, T. Dean, and M. Greig. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence , 147:163-223, 2003.

Hansen-Estruch, A. Zhang, A. Nair, P. Yin, and S. Levine. Bisimulation makes analogies in goalconditioned reinforcement learning. In International Conference on Machine Learning (ICML) , pages 8407-8426, 2022.

A. Howard. Dynamic programming and Markov processes . John Wiley, 1960.

Jiang. Notes on state abstractions, 2018.

Jiang. A note on loss functions and error compounding in model-based reinforcement learning. arXiv preprint arXiv:2404.09946 , 2024.

Jonsson and K. G. Larsen. Specification and refinement of probabilistic processes. In IEEE Symposium on Logic in Computer Science (LICS) , pages 167-183, 1991.

Kakade. A natural policy gradient. Advances in Neural Information Processing Systems , 14: 1531-1538, 2001.

Kakade and J. Langford. Approximately optimal approximate reinforcement learning. In International Conference on Machine Learning (ICML) , pages 267-274, 2002.

V. Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS) , volume 37, pages 199-201, 1942.

Kemertas and A. Jepson. Approximate policy iteration with bisimulation metrics. arXiv preprint arXiv:2202.02881 , 2022a.

Kemertas and A. Jepson. Approximate policy iteration with bisimulation metrics. Transactions of Machine Learning Research , 2022b.

Kolouri, S. R. Park, M. Thorpe, D. Slepcev, and G. K. Rohde. Optimal mass transport: signal processing and machine-learning applications. IEEE Signal Processing Magazine , 34(4):43-59, 2017.

B. Kruskal. Multidimensional scaling by optimizing goodness of fit to a nonmetric hypothesis. Psychometrika , 29(1):1-27, 1964.

G. Larsen and A. Skou. Bisimulation through probabilistic testing. In ACM Symposium on Principles of Programming Languages (POPL) , pages 344-352, 1989.

Lassalle. Causal Transport Plans and Their Monge-Kantorovich Problems . Taylor & Francis, 2018.

Lattimore and C. Szepesvári. Bandit algorithms . Cambridge University Press, 2020.

S. Manne. Linear programming and sequential decisions. Management Science , 6(3):259-267, 1960.

Milner. Communication and Concurrency . Prentice Hall, 1989.

Moulin and G. Neu. Optimistic planning via regularized dynamic programming. In International Conference on Machine Learning (ICML) , 2023.

Moulos. Bicausal optimal transport for Markov chains via dynamic programming. In IEEE International Symposium on Information Theory (ISIT) , pages 1688-1693, 2021.

Nemirovski and D. Yudin. Problem Complexity and Method Efficiency in Optimization . Wiley Interscience, 1983.

Neu, A. Jonsson, and V. Gómez. A unified view of entropy-regularized Markov decision processes. arXiv preprint arXiv:1705.07798 , 2017.

O'Connor, K. McGoff, and A. B. Nobel. Optimal transport for stationary Markov chains via policy iteration. Journal of Machine Learning Research , 23(1):2175-2226, 2022.

M. R. Park. Concurrency and automata on infinite sequences. In GI Symposium on Theoretical Computer Science , volume 104 of Lecture Notes in Computer Science , pages 167-183. Springer, 1981.

Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning , 11(5-6):355-607, 2019.

C. Pflug. Version-independence and nested distributions in multistage stochastic optimization. SIAM Journal on Optimization , 20(3):1406-1420, 2010.

C. Pflug and A. Pichler. A distance for multistage stochastic optimization models. SIAM Journal on Optimization , 22(1):1-23, 2012.

L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming . WileyInterscience, April 1994.

Sangiorgi. On the origins of bisimulation and coinduction. ACM Transactions on Programming Languages and Systems (TOPLAS) , 31(4):1-41, 2009.

Scherrer. On the performance bounds of some policy search dynamic programming algorithms. arXiv preprint arXiv:1306.0539 , 2013.

Scherrer, V. Gabillon, M. Ghavamzadeh, and M. Geist. Approximate modified policy iteration. In International Conference on Machine Learning (ICML) , pages 1207-1214, 2012.

Scherrer, M. Ghavamzadeh, V. Gabillon, B. Lesner, and M. Geist. Approximate modified policy iteration and its application to the game of Tetris. Journal of Machine Learning Research , 16: 1629-1676, 2015.

Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, L. Lee, J. Chen, J. Brumbaug, P. Rigollet, K. Hochedlinger, R. Jaenisch, A. Regev, and E. S. Lander. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell , 176(4):928-943, 2019.

Shi, V. De Bortoli, A. Campbell, and A. Doucet. Diffusion Schrödinger bridge matching. Advances in Neural Information Processing Systems , 36, 2024.

Sinkhorn and P. Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics , 21(2):343-348, 1967.

Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

S. Sutton, D. Precup, and S. Singh. Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial intelligence , 112(1-2):181-211, 1999.

Titouan, N. Courty, R. Tavenard, and R. Flamary. Optimal transport for structured data with application on graphs. In International Conference on Machine Learning (ICML) , pages 62756284, 2019.

van Benthem. Modal Logic and Classical Logic . Bibliopolis, 1983.

van Breugel and J. Worrell. An algorithm for quantitative verification of probabilistic transition systems. In International Conference on Concurrency Theory (CONCUR) , pages 336-350, 2001.

Villani. Optimal transport: old and new , volume 338. Springer, 2009.

Xu, L. K. Wenliang, M. Munn, and B. Acciaio. COT-GAN: Generating sequential data via causal optimal transport. Advances in neural information processing systems , 33:8798-8809, 2020.

Ye. The simplex and policy-iteration methods are strongly polynomial for the Markov decision problem with a fixed discount rate. Mathematics of Operations Research , 36(4):593-603, 2011.

Yi, K. O'Connor, K. McGoff, and A. B. Nobel. Alignment and comparison of directed networks via transition couplings of random walks. arXiv preprint arXiv:2106.07106 , 2021.

Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations (ICLR) , 2021.

A Extended discussion of related work

In this appendix we include a discussion of related work that could not be accommodated in the main text due to space limitations.

A.1 Bisimulation

The concept of bisimulation originated independently in modal logic [van Benthem, 1983], computer science [Park, 1981, Milner, 1989] and set theory [Forti and Honsell, 1983, Aczel, 1988], with firm roots in fixed-point theory. Bisimulation was originally devised as a tool for determining whether or not two processes are behaviorally equivalent in the sense that no test can distinguish between the labels they generate. Being a much less demanding notion of equivalence than isomorphism (which is generally NP-hard to verify), the notion of bisimulation has had significant impact in concurrency theory and formal verification of computer systems, and has become a standard tool for model checking. We refer the reader to the very enjoyable paper of Sangiorgi [2009] for a detailed history of bisimulation and related concepts.

Larsen and Skou [1989] developed a theory of probabilistic bisimulation between stochastic processes. Their approach is based on comparing interpretations of logical formulas on stochastic processes, roughly saying that two processes are probabilistically bisimilar if all formulas acting on the sequence of labels encountered along the corresponding random trajectories follow the same probability distribution. Their definition can be most simply presented when the two processes live on the same state space and follow the same transition kernel, but are initialized at two different states. In this setup, bisimulation reduces to a relation between individual states, which can be described formally as follows. Given a stationary Markov process M X = ( X , P X , ν 0 , X ) as defined in Section 2 and a label function L : X → R , a probabilistic bisimulation R is a relation on X × X that satisfies the following property: two states x and x ′ are bisimilar (denoted xRx ′ ) if and only if L ( x ) = L ( x ′ ) and for each subset C in the partition X \ R induced by R , it holds that

Jonsson and Larsen [1991] define another similarity notion called 'satisfaction relation' based on couplings, and prove that probabilistic bisimilarity and satisfaction are equivalent notions of similarity. These works show that bisimulation is indeed an equivalence relation, and that when two processes are initialized from two states within the same equivalence class (i.e., they are bisimilar ), then they will not only transition to bisimilar states in the next step but will in fact continue to evolve in a way that is indistinguishable based on the labels (in the sense that they will produce the same distribution over sequences of labes).

Giacalone et al. [1990], Desharnais et al. [1999, 2004] and van Breugel and Worrell [2001] relax the restrictive notion of exact probabilistic bisimulation and introduce real-valued pseudometrics that measure the degree of bisimilarity between two states. These notions rely on real-valued labeling functions. Giacalone et al. [1990] gives a notion of ε -bisimulation which relaxes the condition L ( x ) = L ( x ′ ) in the definition of the hard bisimulation relation given above, and only requires equality to hold up to some ε > 0 . Desharnais et al. [1999, 2004] go further and define a genuinely real-valued extension of bisimulation relations by defining bisimulation metrics as described in the main text (and particularly Equation (4)). A fixed-point characterization of these bisimulation metrics was established by van Breugel and Worrell [2001] and Desharnais et al. [2002]. These works show that the resulting distance notion is in fact a pseudometric, and two processes are bisimilar if and only if they are at distance zero. This justifies seeing bisimilarity metrics as 'soft' extensions of the binary relation of bisimilarity.

Interestingly, the first bisimulation metrics all make use of concepts from optimal transport in some way or another: Desharnais et al. [1999, 2002] already note that their definition is inspired by the Wasserstein distance (which they call the 'Hutchinson metric'), and the fixed-point characterization of van Breugel and Worrell [2001] also make use of this distance to compare transition kernels (cf. Equation 5). Precisely, their definition that we recalled as Equation (4) is admittedly inspired by the Kantorovich dual representation of the Wasserstein distance between probability distributions over metric spaces. To our knowledge, this connection with optimal transport has not been explored further in the literature, and in particular no 'primal' counterpart based on couplings has been discovered so far.

Givan et al. [2003] adapt probabilistic bisimulation to Markov decision processes (MDPs), requiring states (or state-action pairs) to have identical rewards and transition probabilities for two MDPs to be bisimilar. Ferns et al. [2004, 2006] introduce bisimulation metrics for MDPs, essentially using the fixed-point characterization of van Breugel and Worrell [2001] as a starting point for their definition. Once again, their definition makes use of the Wasserstein distance between the transition kernels (called Kantorovich metric in the paper), but no deeper connection between bisimulation metrics and optimal transport is discussed.

Castro [2020] defines bisimulation metrics for MDPs with respect to a given policy π , which thus falls back to the standard definition of bisimulation metrics for Markov chains as studied in the works of Desharnais et al. [1999, 2004] and van Breugel and Worrell [2001]. Kemertas and Jepson [2022b] adjusts the definition of Ferns et al. [2004, 2006] by replacing the Wasserstein distance with the entropy-regularized Wasserstein distance proposed by Cuturi [2013], which allows them to apply the Sinkhorn algorithm to compute bisimulation metrics for MDPs. The resulting approach can be seen to be nearly identical to the methods proposed by O'Connor et al. [2022] and Brugère et al. [2024] for approximately solving the fixed-point equations (3) in the context of optimal transport-see Section A.2 for further discussion of these works.

Chen et al. [2012] study the computational complexity of computing bisimulation metrics. Similar to our work, the authors formulate a linear program that characterizes bisimulation metrics in an equivalent way to other common definitions, though the linear program has one constraint per next-state coupling, which makes their LP intractable as stated. They use their LP formulation as an analytic tool to show the existence of a polynomial-time algorithm to solve the fixed-point equations (5) (which essentially amounts to Bellman's value iteration algorithm implemented in the MDP we describe in Appendix B). Each step of the resulting algorithm solves one optimal-transport problem per state pair via the network simplex algorithm, which is known to be impractical for this purpose in comparison with Sinkhorn-style methods [Cuturi, 2013]. In the context of optimal transport, a closely related linear program has been discovered by Backhoff-Veraguas et al. [2017], whose framework is more general in that it mostly focuses on general (potentially non-Markovian) stochastic processes, but with the limitation that only finite horizons are considered. We discuss further developments on this topic in Section A.2.

Finally, Bian and Abate [2017] show that ε -bisimilar Markov processes generate distributions over finite-length trajectories that are close in total variation distance. Since this distance is a special case of the Wasserstein distance (with the Hamming metric over sequences as ground metric), this result can be seen to establish some relation between OT distances and bisimulation, but the link is rather weak in the sense that no equivalence is shown between the two notions. Indeed, these results only imply that nearly-bisimilar processes generate nearly-identical trajectory distributions, but the reverse implication is not shown to hold.

A.2 Optimal Transport

Optimal transport [Villani, 2009] studies the problem of transporting mass between two density functions p and q , given a cost function that measures the transport distance between any pair of points. In the classic formulation of Kantorovich [1942], the vehicle used to transport mass is a coupling, that is, a joint probability distribution whose marginals equal p and q . The problem of finding an optimal coupling that minimizes the total transport distance can be formulated as a linear program. Historically, the resulting LPs have been solved via standard solvers like the network simplex method or interior-point methods, which lead to algorithms with polynomial runtime guarantees but rather poor empirical performance. Cuturi [2013] successfully advocated for adding entropy regularization to the standard LP objective, which enabled algorithms that are orders of magnitude faster than previously proposed methods.

In this work we consider a problem of optimal transport between stochastic processes with a temporal dimension. This topic has recently started to receive attention in the OT literature, mostly focusing on stochastic processes with finite horizon [Pflug and Pichler, 2012, Backhoff-Veraguas et al., 2017, Lassalle, 2018]. In this setting (often called 'adapted transport', 'causal transport', or 'bicausal transport'), the problem is to transport mass between joint distributions of sequences of elements, which can be formulated as an optimization problem over the set of causal couplings (i.e., the set of couplings over joint distributions over sequences that respect the temporal order inherent in the process). Backhoff-Veraguas et al. [2017] have observed that, due to the linearity of the

causality constraints, this optimization problem can be phrased as a linear program, which however is infinite-dimensional and thus intractable to solve directly. They complement this view by providing dynamic-programming principles for characterizing the structure of the optimal coupling, which, in the special case of Markov processes, boils down to the finite-horizon version of the fixed point equations (3). This development essentially mirrors the LP formulation and dynamic-programming principles put forth by Chen et al. [2012] in the context of computing bisimulation metrics (cf. the discussion in Section A.1).

Still on the front of computing optimal transport distances, a notable contribution is due to Eckstein and Pammer [2024], who propose and analyze a version of Sinkhorn's algorithm for optimal transport on the space of stochastic processes. When specialized to Markov processes, their algorithm can be seen to be very closely related to ours, the technical explanation being that in finite-horizon Markov processes the entropy of path distributions that they use as regularization can be seen to be equal to the conditional entropy that our method uses for the same purpose. The resulting algorithm performs iterative Bregman projections via backward recursion over the finite time horizon, with computational steps that are essentially identical to applying our Bellman-Sinkhorn operators. That said, their analysis relies very heavily on the finite-horizon structure of the problem and as such it is not applicable in our considerably more challenging infinite-horizon problem setting.

The more recent works of Moulos [2021], O'Connor et al. [2022], Bayraktar and Han [2023] and Brugère et al. [2024] have investigated optimal-transport distances between infinite-horizon Markov chains. O'Connor et al. [2022] considered the undiscounted version of our problem and proposed to compute optimal transition couplings via an adaptation of approximate policy iteration (cf. Scherrer 2013) to an appropriately adjusted version of the MDP we describe in Appendix B. Their key algorithmic idea is approximating the greedy policy update steps by running Sinkhorn's algorithm for each pair of states. Essentially the same idea was used by Brugère et al. [2024] to solve the discounted problem that is the subject of the present paper, with the difference that their method takes approximate value iteration as its starting point. Both of these approaches are closely related to the alternative fixed-point definition of bisimulation metrics using Sinkhorn divergences due to Kemertas and Jepson [2022b], as mentioned in Section A.1. While these approaches are nearly as effective as our Sinkhorn Value Iteration method in practice, their black-box use of Sinkhorn's algorithm make them difficult to analyze theoretically, and difficult to build further theory on.

To wrap up, let us mention some results that in a sense have already foreshadowed our observation about the relation of OT distances and bisimulation metrics. First, we note that Yi et al. [2021], Brugère et al. [2024] proposed to study optimal transport distances of Markov chains defined over graphs as a means of studying the similarity of the underlying graphs. The purpose of these works was to define a notion of distance that is less demanding than isomorphism, but is still grounded in fundamental theory and can be computed effectively-which is precisely the reason that the notion of bisimulation was originally introduced in the 1980s in the context of formal verification by Park [1981] and Milner [1989]. Finally, the work of Backhoff-Veraguas et al. [2020] has established that 'all adapted topologies are equal' on the space of laws of stochastic processes, understood in the sense that a large number of topologies (including the one induced by optimal-transport metrics defined in terms of bicausal couplings) are in fact identical. While one may argue with their sweeping claim that all such topologies are equal, it may not be surprising in light of their results that the topology induced by bisimulation metrics is also identical to these well-studied topologies (as revealed to be true by our observations in this paper).

B Optimal Transport as a Markov Decision Process

Consider the two Markov processes M X and M Y with initial states x 0 and y 0 , respectively. To compute the optimal transport cost W γ ( M X , M Y ; c, x 0 , y 0 ) and the optimal transition coupling π ∗ , we can introduce a Markov decision process

M = ( XY , A , q, γ, c, ν 0 ) where:

XY is the state space, defined as the set of joint states of M X and M Y ,

A ( xy ) = Π xy is the set of applicable actions in state xy ∈ XY , corresponding to the set of valid couplings of P X ( ·| x ) and P Y ( ·| y ) ,

q ( ·| xy, a ) = a is the transition probability distribution, which is fully determined by the action a ∈ Π xy ,

γ is the discount factor,

c : XY → [0 , ∞ ) is the cost function that maps joint states to positive real numbers,

ν 0 = δ x 0 y 0 is the initial state distribution.

The objective of the agent in this MDP is to select its sequence of actions A 0 , A 1 , . . . in a way that minimizes the total discounted cost E [ ∑ ∞ t =0 γ t c ( X t , Y t )] , where each state pair is drawn according to the action taken by the agent as ( X t , Y t ) ∼ A t . The sequence of actions is generated by a sequence of history dependent policies π t ∈ Π HD : H t → ∆ A ( X t ,Y t ) where H t = ( X 0 , Y 0 , . . . , X t , Y t ) . A first remark is that we can restrict ourselves to deterministic policies. Indeed, the action set is convex at every time step and both the reward and the transition probability distributions are linear in the action. A second remark is that bicausal couplings correspond exactly to history-dependent deterministic policies. Indeed, by bicausality, the bicausal coupling M XY is generated by the policy M XY ( x n , y n | ¯ x n -1 ¯ y n -1 ) . We will denote this policy by π M XY .

Of special interest are stationary deterministic (or Markovian) policies of the form π : XY → A , mapping joint states ( X t , Y t ) to actions in A ( X t , Y t ) as A t = π ( X t , Y t ) . Such policies correspond exactly with transition couplings as defined in the main text as mappings π : XY → ∆ XY of the same type. Accordingly, we will sometimes write π ( ·| xy ) to refer to the distribution π ( xy ) ∈ A ( xy ) below.

B.1 Value functions, optimal policies, and sufficiency of transition couplings

Each policy π induces a value function V π , defined in each state xy ∈ XY as

where the expectation is taken with respect to the stochastic process induced by the policy π . Here, X t and Y t are random variables representing the state of the two processes at time t . In particular, we have that for any bicausal coupling M XY ∈ Π bc ,

And as a result we can relate the optimal transport cost to the optimal value function of the MDP

Building on classic results of MDP theory, it can be shown that there exists an optimal Markovian policy π ∗ whose value function V ∗ = V π ∗ satisfies V ∗ ( xy ) ≤ V π ( xy ) for all policies π and joint states xy , and said optimal value function V ∗ satisfies the Bellman optimality equations

where T is the Bellman operator acting on a function V ∈ R XY as

This is precisely the set of equations in Equation (3). Since V ∗ is optimal, V ∗ ( xy ) equals the optimal transport cost W γ ( M X , M Y ; c, x, y ) for each state xy . An optimal policy π ∗ achieves the infimum in each state xy , with associated value function V π ∗ = V ∗ . These claims are summarized in the following theorem, stated in nearly identical form by Moulos [2021].

Theorem 3 (cf. Theorem 1 in Moulos 2021) . Under the above conditions, the following hold:

There exists an optimal Markovian transition coupling π ∗ such that for any policy π ∈ Π HD and any joint states xy , V π ∗ ( xy ) ≤ V π ( xy ) .

The value function of π ∗ satisfies W γ ( M X , M Y ; c, x, y ) = V π ∗ ( xy ) .

There exists a unique solution to the Bellman optimality equation (3) denoted V ∗ ∈ R XY .

We have that V ∗ = V π ∗ .

The above theorem justifies considering Markovian transition couplings when computing the optimal transport cost W γ ( M X , M Y ; c, x 0 , y 0 ) .

Proof. One can straightforwardly check that our setting is an instance of optimal control problems with additive cost functional studied in Bertsekas [1977]. In particular, the contraction assumption (Assumption C in the paper mentioned above) is satisfied. We define pointwise V ∗ ( xy ) = inf π ∈ Π HD V π ( xy ) , and Proposition 1 of Bertsekas [1977] shows that V ∗ is the unique solution to the Bellman optimality equation.

Now, for the existence of an optimal stationary policy or transition coupling in our terminology, an additional technical condition on the action set must be carefully verified. For every xy ∈ XY , ℓ ∈ [0 , ∞ ) and k , we define the set

This set is compact as the intersection of the compact set A ( xy ) with a closed set, the preimage of a closed set by a continuous application. Knowing that U k ( xy, λ ) is compact, we can then apply Proposition 14 of Bertsekas [1977] and that gives us the existence of an optimal Markovian transition coupling π ∗ that satisfies V π ∗ = V ∗ . In particular, for any policy π ∈ Π HD and joint state xy , we have that V π ∗ ( xy ) = V ∗ ( xy ) ≤ V π ( xy ) by definition of V ∗ .

B.2 Occupancy measures and occupancy couplings

In a finite Markov decision process, occupancy measures express the discounted number of times that a given state and action are visited on expectation by the controlled stochastic process. This notion is not meaningfully applicable in the MDP formulation of our optimal-transport problem, given that the action space is infinite. However, the closely related notion of occupancy coupling can be seen to play a similar role in that it allows expressing the total-discounted-cost objective as a linear function, and that the set of valid occupancy couplings can be fully characterized in terms of a finite number of linear constraints. In what follows, we prove this latter key property of occupancy couplings, stated as Lemma 1 in the main text.

Proof of Lemma 1. We begin by showing that the occupancy coupling µ π associated with any transition coupling π ∈ Π XY satisfies the following system of equations (sometimes called the 'Bellman

flow equations'):

Indeed, this can be shown to follow from the definition of occupancy couplings as

where in the first step we used the stationarity of the transition coupling π , then the definition of ν 0 , followed by the law of total probability, and finally the stationarity of the Markov chain that allowed us to recognize µ π ( x ′′ y ′′ , xy ) in the last step. Now, summing both sides of Equation (12) for all x ′ y ′ , we can confirm that µ π indeed satisfies Equation (6). Furthermore, summing the two sides of Equation (12) over all x ′ , we get

where the first step in the second line follows from Equation 6, and the 4th line comes from the fact that π ( ·| xy ) is a coupling of P X ( ·| x ) and P Y ( ·| y ) . This verifies that µ π satisfies Equation (8), and the same reasoning can be used to verify that it also satisfies Equation 7.

Conversely, suppose that µ ∈ R XY×XY + satisfies Equations (6), (7), and (8). Define ν µ ( xy ) = ∑ x ′ y ′ µ ( xy, x ′ y ′ ) and let

̸

̸

We will verify that π µ defines a valid Markovian coupling and that µ is the state action occupancy measure of π µ . If ν µ ( xy ) = 0 , then ∑ x ′ π µ ( x ′ y ′ | xy ) = P Y ( y ′ | y ) ∑ x ′ P X ( x ′ | x ) = P Y ( y ′ | y ) . If ν µ ( xy ) = 0 , then we have

where we have used that µ satisfies the constraint of Equation (8). In any case, we have that ∑ x ′ π µ ( x ′ y ′ | xy ) = P Y ( y ′ | y ) . By symmetry, using that µ satisfies Equation (7), we also have

∑ y ′ µ ( x ′ y ′ | xy ) = P X ( x ′ | x ) . Hence we have demonstrated that π µ is a valid Markovian coupling of P X and P Y . To proceed, observe that the occupancy coupling associated with π µ satisfies

We will verify that µ π µ = µ by showing that this system of equations has a unique solution. In order to see this, let us recall the definition of ν µ and reorder Equation (6) as

Introducing the matrix Z ∈ R |X||Y|×|X||Y| with entries Z ( xy, x ′ y ′ ) = π µ ( x ′ y ′ | xy ) and representing the functions ν 0 and ν µ in matrix form, this system of equations can be written as

Now, thanks to the Perron-Frobenius theorem, the stochastic matrix Z has spectral radius 1 , and thus ( I -γZ ) is invertible, meaning that there is a unique solution ν µ = (1 -γ )( I -γZ ) -1 ν 0 . This in turn implies that µ = µ π µ , thus verifying that µ is indeed an occupancy coupling induced by a valid transition coupling π µ if it satisfies Equations (6)-(8). This concludes the proof.

C The proof of Theorem 2

The proof is composed of two main parts: showing a bound on the regret of the iterates µ 1 , µ 2 , . . . , µ K , and then accounting for the errors incurred when rounding the average iterate µ K = 1 K ∑ K k =1 µ k to a feasible occupancy coupling. We start by stating some general results that will be useful throughout the analysis, and then study the two sources of error mentioned above separately. For any µ ∈ R XY×XY , we define Eµ ( xy ) = ∑ x ′ ,y ′ µ ( x ′ y ′ , xy ) , E T µ ( xy ) = ∑ x ′ ,y ′ µ ( xy, x ′ y ′ ) and

̸

In particular, we always have that µ ( xy, x ′ y ′ ) = E T µ ( xy ) π µ ( x ′ y ′ | xy ) . Furthermore, for any π : XY → ∆ XY , we will denote by ˜ π = ρ ( π ) the rounded transition coupling obtained by adapting the rounding procedure of Altschuler et al. [2017], and we will let ρ ( µ ) denote the corresponding occupancy coupling µ ˜ π µ (cf. Section E.2 for the description and analysis of this method). In particular, ˜ π will always be a valid transition coupling associated with P X , P Y . Finally we define π out = ρ ( π µ K ) and µ out = µ π out , and we let µ ∗ be an optimal occupancy coupling achieving the minimum in the problem formulation of Theorem 1.

With these notations, we decompose the overall error of the output as

The first sum on the right-hand side corresponds to the rounding error, and the second one to the socalled regret of the sequence of iterates µ k . This latter sum can be controlled by adapting arguments from the classic analysis of mirror-descent methods [Nemirovski and Yudin, 1983, Beck and Teboulle, 2003], with some ideas adopted from the Mirror Sinkhorn analysis of Ballu and Berthet [2023] that will also come in handy for analyzing the rounding errors. We state these tools first below, and then analyze the two terms in the above decomposition separately.

C.1 General tools

We begin with a version of the classic 'three-point identity' for mirror-descent methods (e.g., Lemma 4.1 of Beck and Teboulle, 2003) adapted to our specific setting that involves alternating projections to the sets B X and B Y . The result is similar to Lemma A.3 of Ballu and Berthet [2023], which we reprove here with a more standard methodology (as used for proving, e.g., Theorem 28.4 of Lattimore and Szepesvári, 2020).

Lemma 2. Let µ ∗ ∈ B X ∩ B Y be arbitrary. Then,

Proof. We first recall that H is the Bregman divergence associated with the conditional entropy function C ( µ ) = ∑ xy,x ′ y ′ µ ( xy, x ′ y ′ ) log µ ( xy,x ′ y ′ ) ∑ x ′′ y ′′ µ ( xy,x ′′ y ′′ ) (cf. Appendix A.1 in Neu et al. 2017). For the actual proof, let us consider the case of k odd when µ k +1 ∈ B X . Then, by definition, we have that µ k +1 is the minimizer of Ψ k +1 ( µ ) = η ⟨ µ, c ⟩ + H ( µ ∥ µ k ) on this set. Noticing that the gradient of Ψ k +1 at µ is written as ∇ Ψ k +1 ( µ ) = ηc + ∇C ( µ ) -∇C ( µ k ) , the first-order optimality condition over the convex set B X implies that the following inequality holds for any µ ∈ B X :

In particular, using this result for µ = µ ∗ (which is indeed in B X ), the claim follows from using the standard three-point identity of Bregman divergences that states

Repeating the same argument for even rounds (and noticing that the comparator also satisfies µ ∗ ∈ B Y as needed for that case) completes the proof.

The following standard lemma will also be useful for studying various notions of distances between occupancy couplings: the total variation distance ∥ µ -µ ′ ∥ = ∑ ′ ′ | µ ( xy, x ′ y ′ ) -µ ′ ( xy, x ′ y ′ ) |

the relative entropy D ( µ ∥ µ ′ ) = ∑ xy,x ′ y ′ µ ( xy, x ′ y ′ ) log µ ′ ( xy,x ′ y ′ ) , and the conditional relative entropy introduced earlier.

1 xy,x y , µ ( xy,x ′ y ′ )

Lemma 3. For any two occupancy couplings µ and µ ′ , we have

Proof. The first inequality is Pinsker's. For proving the second inequality, we let ν = E T µ and ν ′ = E T µ ′ be the state occupancies associated with µ and µ ′ , respectively. Then, we write the following:

where the final step follows from using the data-processing inequality for the relative entropy. Reordering the terms concludes the proof.

C.2 Constraint violations

Let us begin with some definitions. First of all we introduce a quantity that measures the extent to which an occupancy coupling µ violates the transition coherence constraints. Specifically, we will measure the violations of the X -constraints by

and the violations of the Y -constraints by

and the overall constraint violations will be written as

Note that we have δ X ( µ k ) = 0 for odd rounds and δ Y ( µ k ) = 0 for even rounds by definition of the updates. We also define the rounding error associated with an occupancy coupling µ as the average total variation distance between the transition coupling π µ and its rounded counterpart ˜ π µ = ρ ( π µ ) :

The first statement establishes a link between the quality of an occupancy coupling and the occupancy coupling obtained after rounding the transition coupling to satisfy the constraints.

Lemma 4. For any µ satisfying Equation (6) , we have

Proof. Let ˜ π µ = ρ ( π µ ) , ˜ µ = ρ ( µ ) , ν µ = E T µ , and ν ˜ µ = E T ˜ µ . Furthermore, let us define the shorthand notation ν ◦ π to denote the composition of the state-pair distribution ν with the transition coupling π as ( ν ◦ π )( xy, x ′ y ′ ) = π ( x ′ y ′ | xy ) ν ( xy ) . Then, we have

where in the last step we used that E is non-expansive with respect to the ℓ 1 -norm. Reordering gives

and putting everything together proves the claim of the lemma.

The next result relates the rounding errors to the constraint violations.

Lemma 5. For any µ ∈ ∆ XY×XY , we have

The proof of this result builds on the error analysis of the rounding procedure of Altschuler et al. [2017], and can be found in Appendix E.2. Finally, the last technical lemma (inspired by Lemma A.4 of Ballu and Berthet, 2023) bounds the rounding errors in terms of the change rate of the occupancy couplings.

Lemma 6. For any k ≥ 1 ,

Proof. We study the case where the k th update is a B Y projection. For this proof, it will be convenient to introduce the following notation. We define I X : R XY×XY → R XY×X and (with some abuse of notation), P X : R XY×XY → R XY×X as the linear operators that respectively act on µ via the assignment ( I X µ )( xy, x ′ ) = ∑ y ′ µ ( xy, x ′ y ′ ) and ( P X µ )( xy, x ′ ) = ∑ x ′′ ,y ′′ µ ( xy, x ′′ y ′′ ) P X ( x ′ | x ) . This allows us to write δ X ( µ ) = ∥ ( I X -P X ) µ ∥ , so that we have the expression

1

where we have also used δ Y ( µ k ) = 0 that holds thanks to the fact that k is even. Moreover, the ( k +1) st update is a B X -projection and thus we have I X µ k +1 = P X µ k +1 . Hence,

where we have used the fact that both I X and P X are non-expansions for the ℓ 1 -norm in the last line, which follows from the data-processing inequality for the total variation distance. We then conclude the analysis for the even rounds by replacing µ k +1 with µ k -1 in the argument above, and repeating the same reasoning for odd rounds completes the overall proof.

Having established these elementary results, we now turn to addressing the main technical hurdle: bounding the cumulative rounding errors.

Theorem 4. The cumulative constraint violations of the iterates produced by Sinkhorn Value Iteration satisfy

Proof. We start by applying Lemma 6 to show δ ( µ k ) ≤ 2 ∥ µ k -µ k +1 ∥ 1 , which reduces our task to bounding ∑ K k =1 ∥ µ k -µ k +1 ∥ 1 . We do this as follows, for any fixed α > 0 :

where we have finally used that δ ( µ k +1 ) ≤ 2 ∥ µ k -µ k +1 ∥ 1 (Lemma 6). Now, we need to make sure that 4 η ∥ c ∥ ∞ α (1 -γ ) 2 ≤ 1 to turn this into a meaningful result. In particular, setting α = 8 η ∥ c ∥ ∞ (1 -γ ) 2 guarantees that the constant in question equals 1 2 , so that we can reorder the terms to obtain

This concludes the proof.

C.3 Regret analysis

In this section, we bound the regret of the iterates produced by Sinkhorn Value Iteration.

Theorem 5. The regret of the mirror Sinkhorn procedure satisfies

Proof. We first apply Lemma 2 to obtain the bound

Adding up both sides for all k = 1 , 2 , . . . , K , we get

This concludes the proof.

C.4 The proof of Theorem 2

The proof follows from applying the above results to bounding the rounding errors and the regret. The first of these is handled as follows:

where the first and second inequalities respectively come from Lemmas 4 and 5, the third one from the convexity of δ , and the last inequality comes from the application of Theorem 4.

The regret is then bounded using Theorem 5 as

Putting both bounds together, we obtain

Now, to prove the actual claim of the theorem, we now let π 1 be the uniform transition coupling and µ 1 be the associated occupancy coupling. In this case, the conditional relative entropy can be upper bounded as H ( µ ∗ ∥ µ 1 ) ≤ log |X||Y| , and we can further bound

where the second to last line is obtained by picking η = 1 4 ∥ c ∥ ∞ √ (1 -γ ) 3 log |X||Y| K and the last line is obtained by noticing that 1 K ≤ √ log |X||Y| K (1 -γ ) 3 holds whenever K ≥ 1 , |X| ≥ 2 , |Y| ≥ 2 and γ ∈ (0 , 1) . Finally, note that

Now, it can be directly verified that the right-hand side is indeed at most ε whenever K is greater than the expression given in the theorem, thus concluding the proof.

D Sinkhorn Policy Iteration

We describe here a simple alternative to Sinkhorn Value Iteration called Sinkhorn Policy Iteration (SPI). After introducing this method heuristically, we provide a formal performance analysis, and finally explain its relation to SVI.

The core concept underlying the definition of SPI is the notion of Q-functions , defined analogously to action-value functions in an MDP. The Q-function associated with a transition coupling π is a function Q π : XY × XY → R , with each of its entries defined as

Analogously to the results presented in Section B.1, it is possible to show that the Q-function of a given transition coupling π satisfies the Bellman equations

and that the Q-value function of the optimal transition coupling Q ∗ = Q π ∗ satisfies the Bellman optimality equations

These are respectively the fixed points of the Bellman operator T π : R XY×XY → R XY×XY defined via

and the Bellman optimality operator T : R XY×XY → R XY×XY defined via

The system of equations T Q ∗ = Q ∗ is essentially as hard to solve (or even harder) than the Bellman optimality equations for V ∗ stated earlier. However, one can develop an algorithmic approach toward finding an optimal transition coupling by drawing inspiration from the literature on regularized dynamic programming.

In particular, we develop below an analogue of an entropy-regularized policy iteration scheme that is known under many names in the RL literature: Natural Policy Gradients by Kakade [2001] and Agarwal et al. [2021a], MDP-Expert by Even-Dar et al. [2009], Mirror-Descent Policy Iteration by Geist et al. [2019], POLITEX by Abbasi-Yadkori et al. [2019], Policy Mirror Descent by Agarwal et al. [2021a], and the list goes on. This method can be directly adapted to our setting as follows. Starting from an arbitrary initial transition coupling π 1 , SPI performs the following sequence of updates for each k = 1 , 2 , . . . , K :

Round the transition coupling π k to ˜ π k = ρ ( π k ) ,

update the Q-function by solving the fixed-point equation Q k = T ˜ π k Q k ,

if k is odd, then update the transition coupling as

else if k is even, then update the transition coupling as

As in the case of SVI, it is easy to verify that these transition couplings satisfy the required marginal constraints. Furthermore, one can verify that the updates defined above exactly correspond to running an instance of Mirror Sinkhorn [Ballu and Berthet, 2023] in each state-pair xy with the sequence of cost functions Q 1 ( xy, · ) , Q 2 ( xy, · ) , . . . , Q K ( xy, · ) , similarly how the entropy-regularized policy iteration methods run an instance of entropic mirror descent in each state. We discuss various aspects of this algorithm below.

D.1 Practical implementation

Like SVI, this method can be seen as performing online Mirror Sinkhorn updates in each state pair xy with a sequence of cost functions Q k , which are computed via solving the linear system of Bellman equations Q k = T ˜ π k Q k . The occupancy couplings produced by SPI are denoted by ˜ µ k = µ ˜ π k , and the final output of the method is produced by averaging these occupancies as µ out = 1 K ∑ K k =1 ˜ µ k , and then extracting the transition coupling π out = π µ out . Notably, there is no need to round this transition coupling since each ˜ µ k satisfies all constraints by construction, and so does their average µ out due to the linearity of the constraints.

One advantage of SPI over SVI is that finding an exact solution for the fixed-point equation Q k = T π k Q k is easier than computing the fixed points required by SVI, thanks to the fact that this is a linear system of equations. A downside of the method is that it requires to run the rounding procedure after each update, at least for the theoretical guarantees to remain valid. The impact of these steps may however be negligible in practical implementations. Similarly to SVI, the ideal updates of SPI can be approximated by applying the Bellman operator to the Q-functions only a small number of times m , with m = ∞ corresponding to the ideal implementation analyzed above. We present a pseudocode for SPI as Algorithm 2.

D.2 Convergence guarantees

In what follows, we show the following performance guarantee for SPI.

Theorem 6. Suppose that Sinkhorn Policy Iteration is run for K steps with regularization parameter η = 1 -γ 3 ∥ c ∥ ∞ √ 8 log |X||Y| K , and initialized with the uniform coupling defined for each xy, x ′ y ′ as π 1 ( x ′ y ′ | xy ) = 1 |X||Y| . Then, for any x 0 y 0 ∈ XY , the output satisfies V π out ( x 0 y 0 ) ≤ W γ ( M X , M Y ; c, x 0 , y 0 ) + ε if the number of iterations is at least

As the analyses of all regularized policy iteration methods listed above, this one also starts with establishing the following claim that corresponds to the classic performance difference lemma (often attributed to Kakade and Langford, 2002, but proposed much earlier in works like Cao, 1999 and even Howard, 1960). To state the result, we define V k ( xy ) = ∑ x ′ y ′ ˜ π k ( x ′ y ′ | xy ) Q k ( xy, x ′ y ′ ) and the operators E + : R XY×XY → R XY and E -: R XY×XY → R XY with each element given by ( E -V )( xy, x ′ y ′ ) = V ( xy ) and ( E + V )( xy, x ′ y ′ ) = V ( x ′ y ′ ) . Then, the following bound holds on the instantaneous regret of SPI in round k .

Proof. The proof follows from elementary properties of occupancy couplings. First, note that by the definition of the value function V k , we have

Furthermore, by multiplying both sides of the Bellman equations Q k = c + γE + V k with µ ∗ and using the flow constraints E T + µ ∗ = γE T -µ ∗ +(1 -γ ) ν 0 , we obtain

The result follows after reordering the terms.

Given the above lemma, we can readily express the regret of SPI as follows:

where we can recognize the regret of Mirror Sinkhorn in each state pair xy ∈ XY . Thus, applying the bound of Theorem 3.1 of Ballu and Berthet [2023] to each of these terms (while noting that ∥ Q k ∥ ∞ ≤ ∥ c ∥ ∞ / (1 -γ ) holds for all k ) gives

and thus putting the bounds together we obtain

Now, setting π 1 as the uniform coupling, we can further upper bound the conditional relative entropy as H ( µ ∗ ∥ µ 1 ) ≤ log |X||Y| , and after setting η = 1 -γ 3 ∥ c ∥ ∞ √ 8 log |X||Y| K , the bound becomes

Finally, note that

Using a crude upper bound 36 / 8 ≤ 5 verifies the claim of Theorem 6.

D.3 Relation to Sinkhorn Value Iteration

While on the surface, SPI may seem only loosely related to SVI, a closer connection can be drawn by making the following observations. First, observe that the transition-coupling updates exactly match the updates of Sinkhorn Value Iteration, although there is an apparent difference in how the Q-functions are defined. To expose the similarity between the two methods better, let us consider an even round in which π k satisfies the X -marginal conditions ∑ y ′ π k ( x ′ y ′ | xy ) = P ( x ′ | x ) . Thus, introducing the notation V X ( xy, x ′ ) = ∑ y ′ π k ( x ′ y ′ | xy ) P ( x ′ | x ) Q k ( xy, x ′ y ′ ) , we can multiply the Bellman equations by π k ( x ′ y ′ | xy ) /P ( x ′ | x ) and sum them up to obtain

where the approximation is accurate as η approaches zero. Thus, for small values of η , the SinkhornBellman operator used by Sinkhorn Value Iteration is an accurate approximation of the Bellman operator used by Sinkhorn Policy Iteration, and thus one may reasonably expect their respective fixed points to be close as well (this intuition may, however easily fail and is not necessary for our analysis above).

E Auxiliary proofs and technical results

In this appendix we prove several technical results from the main text.

E.1 Proof of Proposition 1

We prove the claim by showing that the solution of the constrained optimization problem of Equation (9) is equivalent to the transition-coupling update rule specified in Equation (10). To this end, we study the Lagrangian of the optimization problem (9) corresponding to the update for odd rounds, B X , and note that the update rule for even rounds, B Y , can be worked out analogously. By introducing Lagrange multipliers V X ( xy, x ′ ) for each constraint in B X , we obtain the Lagrangian

A quick calculation (cf. Appendix A.1 of Neu et al., 2017) shows that the derivative of H ( µ ∥ µ k ) satisfies

where we have used that π µ k = π k holds by definition of µ k and π k . To proceed, for a fixed V X , we set the gradient of the Lagrangian to zero and solve for the transition coupling π k +1 , which gives

Then, the correct choice of V X has to be such that the constraint ∑ y ′ π k +1 ( x ′ y ′ | xy ) = P X ( x ′ | x ) is satisfied. To see this, suppose that this condition is indeed verified and that µ k +1 is the occupancy coupling associated with π k +1 . We need to show that µ k +1 indeed verifies the condition defining B X . To this end, notice that µ k +1 satisfies Equation (6), and thus the condition can be simply written as

which, after recalling the relation π k +1 ( x ′ y ′ | xy ) = µ ( xy,x ′ y ′ ) ∑ x ′′ y ′′ µ ( xy,x ′′ y ′′ ) can be indeed seen to hold if ∑ y ′ π k +1 ( x ′ y ′ | xy ) = P X ( x ′ | x ) is true.

Enforcing this constraint gives the following expression for V X ( xy, x ′ ) :

To conclude, we need to ensure that this system of equations has a unique solution. In order to do this, we recall the definition of the Bellman-Sinkhorn operator T π k X : R XY×X → R XY×X , acting on a function f as

With this notation, we can directly verify that V X satisfies T π k X V X = V X . Furthermore, it can be shown that the operator T π k X is a γ -contraction in supremum norm (cf. Lemma 8), and thus it has a unique fixed point by the Banach fixed-point theorem. We finally note that defining Q k ( xy, x ′ y ′ ) = ( c ( xy ) + γ ∑ x ′′ P X ( x ′′ | x ′ ) V X ( x ′ y ′ , x ′′ )) , the transition-coupling update derived above can be rewritten as

This concludes the proof.

E.2 Rounding procedure and proof of Lemma 5

We adapt the rounding procedure stated as Algorithm 2 of Altschuler et al. [2017] and reproduced here as Algorithm 3 (where / denotes element-wise division in the pseudocode). Formally, for two probability distributions p ∈ ∆( X ) and q ∈ ∆( Y ) , the set of valid couplings is U p,q = { P ∈ R XY + : P · 1 = p ; P T · 1 = q } . For a nonnegative matrix F ∈ R XY + , the rounding procedure outputs a valid coupling ρ ( F, p, q ) ∈ U p,q which, by Lemma 7 of Altschuler et al. [2017], satisfies

We will now define the rounding procedure for a (not necessarily valid) transition coupling π ∈ R XY×XY by using the aforementioned rounding procedure at each state pair as

With some abuse of notation, we will write the resulting transition coupling as ˜ π = ρ ( π ) and the associated occupancy coupling as ˜ µ = µ ˜ π = ρ ( µ ) . Because of the correctness of the original rounding procedure, this transition coupling is valid, and so is the associated occupancy coupling. We can now proceed to the proof of Lemma 5.

Proof of Lemma 5. Let µ ∈ R XY×XY + , and as before define ν µ ( xy ) = ∑ x ′ y ′ µ ( xy, x ′ y ′ ) and

̸

For arbitrary state pairs xy , we use Lemma 7 of Altschuler et al. [2017] to obtain that

Now, multiplying by ν µ ( xy ) and summing over xy , we get

where we used the fact that µ ( xy, x ′ y ′ ) = ν µ ( xy ) π µ ( x ′ y ′ | xy ) for any xy, x ′ y ′ , and the definitions δ X , δ Y and δ . Finally, we notice that ∆( µ ) = ∑ xy ν µ ( xy ) ∥ ˜ π µ ( ·| xy ) -π µ ( ·| xy ) ∥ 1 , which concludes the proof.

E.3 The contraction property of the Bellman-Sinkhorn operator

Lemma 8. Let π : XY → ∆ XY be arbitrary and consider the associated Bellman-Sinkhorn operator T π acting on a function f : XY × Y → R as

Then, T π is a γ -contraction for the supremum norm ∥·∥ ∞ , that is, for any two functions f 1 , f 2 : XY × X , we have

Proof. The claim easily follows from using the standard fact that the function g p ( z ) = log ∑ y ′ p ( y ′ ) e z ( y ′ ) is 1 -smooth with respect to the supremum norm, so that for any two vectors z , we have | g p ( z ) -g p ( z ) | ≤ ∥ z -z ′ ∥ ∞ . To apply this result, we define q ( y ′ | xy ) = π ( x ′ y ′ | xy ) /P ( x ′ | x ) and z 1 ( xy, x ′ y ′ ) = c ( xy ) + γ ∑ x ′′ P ( x ′′ | x ′ ) f 1 ( x ′ y ′ , x ′′ ) and z 2 ( xy, x ′ y ′ ) = c ( xy ) + γ ∑ x ′′ P ( x ′′ | x ′ ) f 2 ( x ′ y ′ , x ′′ ) , so that we can write

where the last step follows from the straightforward calculation

This concludes the proof.

F Additional experimental results

In this appendix we present the results of additional experiments not included in the main text.

F.1 Impact of the regularization parameter η

Besides the parameter m that we have already studied experimentally in Section 5, the only tuning parameter of SVI and SPI is the regularization parameter η , which takes the role of a learning rate. In this experiment, we study two random walks run on two separate 4-room environments [Sutton et al., 1999], with two separate reward functions r X and r Y that together define the ground cost function c ( x, y ) = | r X ( x ) -r Y ( y ) | for each state pair. The results of this study for K = 2 · 10 4 iterations are shown in Figure 3. The error is computed as the difference between the distance estimate produced by the algorithm and a near-optimal distance obtained by running Algorithm 1 for a very small value of η and a large number of iterations. We observe that higher values of η lead to faster error reduction in the initial steps, but eventually prevent convergence to the true solution. In contrast, choosing smaller learning rates enables convergence to better solutions, at the cost of making slower progress initially. An intuitive explanation for this is that for larger values of η , the iterates converge rapidly to the broad proximity of an optimal solution, but then reach a cycle where they continue to perform large updates to the transition coupling which results in large constraint violations, which necessitate large updates, and so on. This is formally verified by the fact that our bound on the constraint violation terms in Theorem 4 are increasing for large values of η . Notably, employing a time-dependent learning-rate schedule (inspired by Ballu and Berthet [2023]) with η k ∼ 1 / √ k leads to the best performance. This strategy leverages faster convergence initially and, for sufficiently large number of iterations, also achieves near-optimal solutions.

(a) Algorithm 1

(b) Algorithm 2

10-1

10-2

Error

n:10/vK

10-3

n:0.1

n:1

n:10

10-4

100

101

102

103

104

Iterations

10-1

10-2

Error

n:10/vK

10-3

n:0.1

n:1

n:10

10-4

100

101

102

103

104

Iterations

F.2 Comparison with alternative methods

We now turn to studying the computational complexity of our algorithms, and compare them empirically to some existing methods that have been proposed for computing optimal transport distances and bisimulation metrics between Markov chains. Specifically, we will focus here on two previous methods: the method of [O'Connor et al., 2022] that we refer to as 'EntripicOTC' and Algorithm 2 of [Brugère et al., 2024] that we call 'dWL'. We adapt both of these methods with some minor changes to our setting. First, we remove one scaling factor from the transport cost in the definition of the distance defined by Brugère et al. [2024] so that it matches ours. Second, the algorithm of O'Connor et al. [2022] is originally defined for the infinite-horizon average-cost case, and thus we made appropriate changes to adapt it to the discounted case by replacing their approximate policy evaluation step by T applications of the discounted Bellman evaluation operator. As pointed out in Appendix A.2, the resulting methods are closely related, and can be regarded as approximate dynamic programming methods for solving the MDP formulation of our optimal transport problem presented in Appendix B. We also recall that the algorithm proposed by Kemertas and Jepson [2022a] for the purpose of computing bisimulation metrics also falls into the same class of approximate dynamic programming methods, and nearly matches the method of [Brugère et al., 2024]. The comparison

Y=0.75

y =0.9

Y=0.99

1000.0

SVI

100.0

SPI

dWL

10.0

dOTC

1.0

0.1

0.01

5

1015

25

5

1015

25

5

1015

25

below is based on the original Python implementation 3 of Brugère et al. [2024] and our own Python adaptation of the MATLAB code 4 of O'Connor et al. [2022].

One difficulty that we had to face in these experiments is having to tune various hyperparameters of each method (such as number of iterations and regularization parameter), which can each influence the quality of the solution and the computation time. For a fair comparison between the methods, we have adopted the following procedure to obtain our results. First, we estimate a ground truth obtained by running one of the algorithms for a very high number of iterations and a very low level of regularization, and then use this ground truth as a comparator to adjust the hyperparameters of the algorithms so that they are close to this value in as little wall-clock time as possible. While all the algorithms perform similar operations, their total runtimes turn out to be rather different and heavily dependent on problem parameters such as the discount factor γ . The comparison between the resulting runtimes of each method is shown on Figure 4 as a function of the size of the Markov chains, and for a three different choices of the discount factor, for a set of randomly generated MDPs (following the setup described in Section 7.1 of O'Connor et al., 2022).

First, we observe that EntropicOTC is a policy-iteration-like method, and as such it needs fewer iterations than the rest of the algorithms we tested to converge to the true cost, but each iteration requires running an expensive policy evaluation subroutine until convergence. This computational cost eventually adds up in a way that this algorithm has always ended up being the slowest among all that we have tested, although its performance has proved notably robust to changes in the discount factor γ . Second, we note that the updates of dWL are much cheaper to compute, especially for large regularization parameters. However, the errors of this value-iteration-like method compound much more rapidly than in the case of EntropicOTC, which makes it especially hard to tune the hyperparameters of this method. This problem is especially pronounced when the discount factor is large, which is the most interesting regime as it leads to distances with much stronger discrimination power.

The plots shown on Figure 4 indicate that our methods find optimal couplings consistently more efficiently than the other methods in the regime we studied, leading to up to 10 times faster runtimes. For the case of sufficiently small values of γ , the algorithm of Brugère et al. [2024] sometimes performs competitively, but the hyperparameters leading to good performance are much harder to find than in the case of our methods. In our experience, the massive speedup achieved by our methods can be largely ascribed to maintaining the transition couplings π k between iterations, as opposed to computing these afresh by running Sinkhorn's algorithm from scratch for each update as done by all other competing methods. Adjusting these other algorithms by maintaining the couplings in memory and using them to warm-start the subsequent updates makes them competitive with our methods, and in fact doing so makes them quite similar to SVI and SPI. Our algorithms use such warm-starts as a primary design choice as opposed to an obscure implementation detail, which is ultimately responsible for the computational efficiency of all these dynamic-programming methods.

The computational time per iteration of each of these methods grows roughly at a rate of n 4 , with n being the number of states of the Markov chains. In the case of our methods, this is easily explained by noting that applying the Bellman-Sinkhorn operators and updating the transition couplings requires |X| 2 |Y| 2 operations in total. This matches the runtime necessary for running the regularized policy improvement subroutines employed by O'Connor et al. [2022] and Brugère et al. [2024], which consists of running an instance of Sinkhorn's algorithm in each pair of states. The computational cost of all these methods can be improved by leveraging the sparsity of transition kernels: in particular, if at most S states are reachable with positive probability in both of the chains, the complexity of the updates can be trivially improved to |X||Y| S 2 . We did not pursue this direction in our experiments as our goal was to compare the basic versions of each studied method, and we believe that our conclusions would not be altered if we were to implement this improvement for all methods.

F.3 Optimal transport distances as similarity metrics

We finally provide a range of experiments that illustrate how the optimal-transport distances we studied in this paper can be used to capture relationships and symmetries in groups of varying Markov chains. To this end, we have generated 35 different '4-room' instances and computed their pairwise distances. Each instance differs in its initial state and position of the obstacles (amounting to changes in the transition kernel), while maintaining a fixed reward function, with one reward located in each room except for the upper left room. Crossing a door between each room results in a negative reward. For each instance, we have studied the Markov chain induced by the corresponding optimal policy (which amounts to taking the shortest path toward the closest positive reward, modulo the additional randomness inherent to the transitions). Figure 5 presents a 3D visualization (where the z -dimension is represented by a color gradient) generated using Multidimensional Scaling (MDS) [Kruskal, 1964] based on the computed pairwise distances. One can observe a clear clustered structure, where instances with similar behaviors are grouped closely together. As the figure highlights, the resulting metrics capture the intuitive similarities and symmetries between each process, which indicates the potential usefulness of optimal transport distances and bisimulation metrics for comparing Markov chains under minimal structural assumptions made on the state spaces and the transition functions.

30

30

20

20

10

10

0

0

-10

-10

-20

-20

-30

-30

-30

-20

-10

0

10

20

30

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( ❤/a116/a116♣/a115✿✴✴♥✐♣/a115✳❝❝✴ ♣✉❜❧✐❝✴❣✉✐❞❡/a115✴❈♦❞❡❙✉❜♠✐/a115/a115✐♦♥/a80♦❧✐❝② ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( ❤/a116/a116♣/a115✿ ✴✴♥✐♣/a115✳❝❝✴♣✉❜❧✐❝✴❣✉✐❞❡/a115✴❈♦❞❡❙✉❜♠✐/a115/a115✐♦♥/a80♦❧✐❝② ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: There is no randomness in any of the experiments that could be meaningfully assessed via error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics ❤/a116/a116♣/a115✿✴✴♥❡✉/a114✐♣/a115✳❝❝✴♣✉❜❧✐❝✴❊/a116❤✐❝/a115●✉✐❞❡❧✐♥❡/a115 ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: There is no societal impact of the work performed.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, ♣❛♣❡/a114/a115✇✐/a116❤❝♦❞❡✳❝♦♠✴❞❛/a116❛/a115❡/a116/a115 has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: There are no new assets introduced in the paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe goal of this work is to study optimal transport (OT) distances between pairs of finite Markov chains, providing a novel relation between OT distances and probabilistic bisimulation metrics. The proposed linear program builds on ideas from optimal control in Markov decision processes, and the designed algorithm for solving the proposed linear program combines Sinkhorn's algorithm with an entropy-regularized version of the classic Value Iteration algorithm. Convergence guarantees and computational complexity analysis of the method are provided in the main text and in the appendix.\n\nSTRENGTHS:\nThe paper is well-organized, the topic of study is very interesting, and the contribution is novel. Previous works on OT and bisimulation metrics are well-referenced and up-to-date.\n\nWEAKNESSES:\nThe main theoretical contributions are presented in Section 3, where the authors introduce the definition of occupancy couplings $\\mu$. These are distributions over $\\mathcal{XY}\\times \\mathcal{XY}$. It is not entirely clear why they need to duplicate the variables $(x,y)$ and consider quadrupes $(xy,x'y')$.\n\nQUESTIONS:\n- In the abstract, the authors say \"In this work, we develop an alternative perspective by considering couplings between a “flattened” version of the joint distributions that we call discounted occupancy couplings [...]\" Could you expand on what do you mean by \"flatten\"? \n\n- In the introduction, the authors say that \"The possibility that the objects in question may be random further complicates the picture, and in such cases it becomes more natural to measure distances between the underlying joint probability distributions.\" Could you explain/expand?\n\n- Section 2.1: \n\nDoes $\\gamma$ depend on $t$? \n\nIs $\\gamma^t\\in (0,1)$?\n\n- Line 99, if Markov chains are considered, isn't it that $M_{\\mathcal{X}}$ $(x_n | \\bar{x}_{n-1})$ $=$ $M_{\\mathcal{X}}$ $(x_n|x_{n-1})$ ?\n\n- Line 148 and equation (4): Which is the dependence of $d_\\gamma$ on the labeling function $r$?\n \nLine 149: Which is the dependence of the set $\\mathcal{F}_\\gamma$ on $\\gamma$?\n\nIt is clear from eq. (5) that $U^*$ does explicitly depend on $r$ and $\\gamma$. Could you expand on the equality $U^*=d_\\gamma$? \n\n- Explain the pevious to last equality in (6). How do you make the variables $x',y'$ to apprear?\n\n- Explain the last equality in (6). Why do you have $\\langle \\mu^\\pi,c\\rangle$ if c depends on $(x,y)$ and $\\mu^\\pi$ depends on $x,y,x',y'$?\n\nComments:\n\n- The term \"discounted occupancy couplings\" mentioned in the abtract is not excatly used in the main text, the authors mainly use just \"occupancy couplings\". \n\n- $n$ is used for the size $|\\mathcal{Y}|$ (line 87) and as an index. \n\nIsn't it $|\\mathcal{X}|=m=|\\mathcal{Y}|$? In the experiments, why the choice of $n$ is not analyzed too?\n\nTypos:\n\n- Abstract: line 17, \"method that we call\", instead of \"method we call\".\n\n- Notations: line 79, \"the corresponding finite subsequences\", instead of \"the corresponding subsequences\"\n\n- Line 39 (Introduction): \"important\" is repeated twice. \n\n- Notations: Add the definitions of $\\mathcal{X}^\\infty$, $\\mathcal{X}^n$, $\\Delta_{\\mathcal{X}}$.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis submission \n\nNamely, the authors define a notion of optimal transport distance between Markov chains on state spaces with a ground metric. This notion of distance differs from standard optimal transport, as the set of couplings is restricted to the set of so-called bicausal couplings. Using the results of Moulos [2021], it is shown that this Markovian optimal transport can be characterized in terms of the solution of the Bellman optimality equations for a Markov decision process. \n\nUsing this equivalence, it is demonstrated that, in certain cases, bisimulation metrics can be though of, equivalently, as Markovian optimal transport with a specific cost function by exploiting their connection to solutions of certain fixed-point equations. \n\nNext, it is shown that the optimal value for the derived Markovian optimal transport problem can be computed by solving a finite-dimensional linear program, where the constraint set consists of the intersection of three sets. Rather than directly solving this linear program, the authors propose to regularize the problem using a conditional entropy (by analogy with entropic regularization of optimal transport distances which can then be solved efficiently using Sinkhorn iterations). To solve the entropy regularized problem, the Sinkhorn Value Iteration algorithm is proposed and the number of steps required to obtain a desired accuracy in estimating the Markovian optimal transport distance is provided. The authors also propose an analyze an alternative algorithm which they dub Sinkhorn Policy Iteration. \n\nThe paper concludes with some numerical experiments to illustrate the performance of the proposed algorithms.\n\nSTRENGTHS:\nIn my opinion, the submission is well-written and its contributions relative to the broader literature are clearly identified. \n\nThe main contributions of this work, identifying bisimulation metrics as a type of optimal transport problem and providing some new algorithms for estimating bisimulation metrics by using this connection is of interest and is, to my knowledge, novel.\n\nWEAKNESSES:\n1. It appears that the connection between the Markovian optimal transport problem and the finite-dimensional linear program provided in Theorem 1 only enables the computation of the optimal value for optimal transport problem, but does not allow the recovery of the optimal bicausal coupling. If this is the case, this should be further clarified in the text. \n\n2. While the analysis for the Sinkhorn Value Iteration is a nice addition, the guarantees are a bit confusing. Notably, if we wish to attain a precision of $\\epsilon$ in estimating the Markovian optimal transport one requires $K=O(1/\\epsilon^2)$ iterations and setting the regularization parameter to $\\eta =C/\\sqrt K=O(\\epsilon)$. However, if epsilon is sufficiently small, the objective in the update (10) will be dominated by the entropy which is minimized by $\\mu=\\mu_k$. Naively I would assume that $1/\\eta$ should be small so that minimization of the linear term dominates.\n\nQUESTIONS:\nApart from the points mentioned above, it would be helpful for the authors to state the complexity of the Sinkhorn Value Iteration algorithm in the main text (rather than only in the appendix).\n\nLIMITATIONS:\nThe authors address the limitations of their work in the discussion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work integrates optimal transport with Markov chains by proposing an alternative joint distribution between Markov processes, namely \"discounted occupancy couplings\". They show that optimal transport distances can be computed as a a linear program (LP) in reduced space. This improves the computational efficiency of OT between markov decision process (MDP) over the previous methods. The previous methods, as they reviewed, often requires complex dynamic programming algorithms. They showed that the new formulation can be extended to the well-known entropy regularization problem, employing Sinkhorn-like iterations to solve, making the computation scalable with large problem. This paper provides both theoretical and experimental supports for the new formulation. In the end, they discuss the potential applications on RL models, limitations and challenges.\n\nSTRENGTHS:\nThe main contribution of this work is demonstration of optimal transport shares the same formulation as probabilistic bisimulation metrics, which is popular in the practice of reinforcement learning (RL). By establishing that solving for optimal transport is equivalent to computing bisimulation metrics, the authors creates a novel link between RL and optimal transport theory.\n- The formulation of OT on MDP is novel. This work provides an new prospective of Bisimulation metric between MDPs.\n- Theoretical analysis looks good.\n- Two computational solutions (Sinkhorn-like iterations) were proposed and tested in experiments sections.\n\nWEAKNESSES:\nThe theoretical guarantees are based on the assumption of perfect projection steps $m=\\infty$, which is not practical. As stated in line 290, such exact computation is unnecessary in practice. This claim needs a further theoretical support.\n\nQUESTIONS:\n- What is the best practice of bisimulation metric computation? Could your provide the comparison with your methods?\n- What does the entropy means in the sense of the coupling between MDPs? Some explanations on this will be appreciated \n- How does this metric works with data contains noise as we know OT is sensitive to outliers?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}