Distances for Markov chains from sample streams
Authors: Sergio Calo
Anders Jonsson
Gergely Neu
Ludovic Schwartz
Javier Segovia-Aguas
Venue: NeurIPS 2025 poster ()
Keywords: Optimal Transport
Markov Chains
bisimulation
primal dual
stochastic methods
Primary area: theory

Abstract: Bisimulation metrics are powerful tools for measuring similarities between stochastic processes, and specifically 
Markov chains. Recent advances have uncovered that bisimulation metrics are, in fact, optimal-transport distances, 
which has enabled the development of fast algorithms for computing such metrics with provable accuracy and 
runtime guarantees. However, these recent methods, as well as all previously known methods, assume full knowledge of the transition 
dynamics. This is often an impractical assumption in most real-world scenarios, where typically only sample trajectories are available.
In this work, we propose a stochastic optimization method that addresses this limitation
and estimates bisimulation metrics based on sample access, without requiring explicit transition models.
Our approach is derived from a new linear programming (LP) formulation of bisimulation metrics, which we solve using a 
stochastic primal-dual optimization method. We provide theoretical guarantees on the sample complexity of the algorithm 
and validate its effectiveness through a series of empirical evaluations.

Distances for Markov chains from sample streams
Sergio Calo
Anders Jonsson
Ludovic Schwartz Javier Segovia-Aguas Universitat Pompeu Fabra, Barcelona, Spain
Gergely Neu
(sergio. calo, anders, jonason, gergely new, ludovio. javier segovia /uut.
Abstract
Bisimulation metrics are powerful tools for measuring similarities between stochas- tie processes, and specifically Markov chains. Recent advances have uncovered that bisimulation metrics are, in fact, optimal-transport distances, which has enabled the development of fast algorithms for computing such metrics with provable accuracy and runtime guarantees. However, these recent methods, well as all previously known methods, assume full knowledge of the transition dynamics. This is often impractical assumption in most real-world scenarios, where typically only sam- ple trajectories are available. In this work, we propose a stochastic optimization method that addresses this limitation and estimates bisimulation metrics based on sample access, without requiring explicit transition models, Our approach is derived from a new linear programming (LP) formulation of bisimulation metries, which we solve using stochastic primal-dual optimization method. We provide theoretical guarantees on the sample complexity of the algorithm and validate its effectiveness through a series of empirical evaluations.
Introduction
Computing similarity metries between stochastic processes an important mathematical problem with numerous promising use cases in diverse areas such as mathematical finance, computational neuroscience, biology, and computer science. Within machine learning, potential applications include representation learning for dynamical systems and reinforcement learning [Zhang et al.) [2021 Chen and Pan 2022). ). fitting and comparing sequence models [Xu al., 2020 Tao et al. [024] or prediction tasks graph-structured data [Titouan et al.) 2019 Brugère et al. [024]. While there exist several rigorous frameworks for defining such similarity metrics and studying their properties, computing them typically requires full knowledge of the probability law of the processes to compare, which is not available in just about any case of practical interest. In this paper, we address this problem by developing methods for estimating similarity metrics for family of stochastic processes, based only on sample streams and without requiring any prior information about the underlying process laws.
We focus on family of similarity metrics known as bisimulation metrics, originating from theoretical computer science for purposes of formal verification of computer programs [Park] 1981 Milner 1989 Desharnais et al. 1999. van Breugel and Worrell 2001]. This notion of process similarity has gained popularity within reinforcement learning (RL), where its potential for learning state representations has been recognized by the early works of Givan et al. (2003) and Ferns et al.) [2004] and the possibility of using it as a basis of practical methods for representation learning has been explored in long line of subsequent works [Castro| 2020 2020/Zhang et al.) 2021 Chen and Pan 2022
Kemertas and Jepson 2022]. Another popular framework for studying similarities between structured probability distributions that of optimal transport (OT, cf. Villani 2009), which has received serious attention within machine learning in the last decade, largely owing to the work of Cuturi [203)] Very recently, Calo al. [2024] pointed out that bisimulation metrics also fall within the family of OT distances, which not only allowed them to connect two distinct areas of mathematics but also import tools from the literature of computational optimal transport (Peyré and Cuturi 2019
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

and develop more efficient methods for computing bisimulation metrics. We refer to Appendix A of Calo al. (2024) for more historical details on the two extensive lines of literature on bisimulation metrics and optimal transport for stochastic processes.
In this paper, we extend this line of work and show that recasting bisimulation metrics OT distances allows not only computational advances, but the development of rigorous theory for statistical estimation of similarity metrics between stochastic processes. In particular, we build on the foundations laid down by Calo et al. [2024) and derive a new stochastic optimization algorithm for estimating bisimulation metrics based on sample observations only, and provide its complete computational and a new linear-progee program formulation of bisimulation metrics, which we solve via a stochastic saddle- point optimization method. For two Markov chains with state spaces X' and V the algorithm is guaranteed to return e-accurate estimate of the true similarity metric after ( |>|| (I-2) + [)1)/=") iterations, with each iteration making use of a single sample transition from each of the two chains, and costing ([x2) D) computation. This is the first result of its kind: no previous methods have successfully addressed this problem either in practice or in theory.
analysis for finite Markov chains. A core technical contribution
As mentioned above, the problem we study in this paper has been extensively studied in (at least) two major research communities. Within the optimal-transpon distances between stochastic processes has been studied under the names "adapted", "causal" or "bicausal" optimal transport (Pilug and Pichler 2012 Lassalle| 2018 Backhoff-
community, the problem of computing
et al.] [017] Eckstein and Pammer 2024). Considering the special case of Markov chains (as we do in the present paper). Moulos| [2021], O'Connor et al. [2022) and Brugère et al. (2024) proposed approximate
dynamic programming algorithms based on the observation that computing OT distances between Markov chains can be reduced to a problem of optimal control in Markov decision processes. Calo itt al. [2024] developed a novel linear programming framework for computing OT distances between Markov chains, and showed that such distances are equivalent to bisimulation metrics. However, previous approaches in this line of work all assume known transition dynamics.
Within the theoretical computer science community, the study of bisimulation metries progressed quite differently: after an initial flurry of foundational works of [Dehhannais et al.] [1999], van Breugel and Worrell [200T] and Desharnais et al. 22002. surprisingly few studies have addressed computational matters (one rare example being the work of Chen et al. (2012). Several recent works in reinforcement learning aim to learn approximate bisimulation metrics from sample transitions
using deep learning [Castro 2020 Zhang et al.) 2021 Chen and Pan 2022 Kemertas and Jepson 2022]. However, these approximate bisimulation metries are not well-founded in theory and as a consequence, do not enjoy the theoretical guarantees of the original metrics. In contrast, the stochastic- optimization method we develop in this paper is firmly rooted in a theoretical understanding of the problem and comes with provable computational and sample-compe complee-y guarantees.
The use of stochastic solvers compute OT distances has been explored in several past works, mostly in the context of static optimal transport between probability distributions. A good part of these methods are based on the observation that the static OT problem is formulated
linear program, and
the associated unconstrained dual optimization problem directly lends itself to numerical optimization. This view has been exploited by works like Genevay al. (0066) AArovsky et al. [2017), and (Seguy let al. 2018). with some rigorous performance guarantees provided by Ballu et al. [2020] Another line of work makes use of Monte Carlo estimates of OT distances [Genevay et al. 2018 Fatras et al.)
2019 |021 Mensch and Peyré 2020). To our knowledge, the idea of computing OT distances via stochastic primal-dus methods as we do in the present work has not been explored in this literature, and thus our contribution may be of independent interest within this context as well.
The rest of the paper is organized as follows. After formally defining our problem in Section 2 describe the foundations of our new algorithm and describe it in detail in Section 3 We state our main theoretical results in Section 4 where we also outline the main ideas of the analysis. We complement these with some empirical studies of the newly proposed method in Section 5 and conclude with a discussion of the results and open problems in Section 6
Notation. For a finite set S, we use As to denote the set of all probability distributions over S. For two sets X' and Y/ we will often write XY = X' x y to abbreviate their direct product. We will denote infinite sequences by z = (Ir, ..... and for any n the corresponding subsequences as Tn (zo..... I.).
2

---

2 Preliminaries
We study the problem of measuring distances between pairs of finite Markov chains. Specifically, we consider two stationary Markov processes Mx = (X, Px, v0,x) and My = (o,y), where
X' and y are the state spaces with finite cardinality,
the fransition kernels that specify the evolution of states as Px(2'1z> [x) = P [ = x' xx = xr] and Py(y' (y) = P((+1 = y Y y] (for all time indices t and state pairs I, x' and y, v,)
Px X Ax x and Py : + Ay
V0,x E A(X) and Vo.y E ())) are the initial-state distributions which specify the states at time t= 0 as Xo ~ vo.x and Yo Vo.9-
Without loss of generality,
will assume that Vo,x and Mo.y both Dirac measures respectively
supported on some fixed Zo and yo. and use 1/ = Vo,x @ Vo,y to denote the joint distribution of the pair of initial states (Xu, Yo) (which is of course Dirac measure Zo. Bo). For each 2 0, the objects above define a sequence of joint distributions IP [(,,, X1, XI.... X.)) = (Ip, .I,...,,,..)) and TP (Y,, Y1, ,,...., Yn)= sequences X (Xa, .... and Y (Y,, ......) )) via Kolmogorov's extension theorem. With a slight abuse of notation we use Mx and My to denote the corresponding measures satisfying Mx (T,) = PP [x, = #n] and My(U.)
(yo, y1,.... Bn)] These distributions in turn define the laws of the infinite
[Y, - Dn. for any X y E y and n 2
Our goal is to compute optimal transport distances between anfinite-horizon Markov chains. To this end, we will suppose access ground cost (or ground metric) с XY + E* that quantifies the (dis-)similarity between each state x E X and E y as c(r, For any two sequences I = (zo, 21, E x and y = (yo,31,...) 9l, 3I....) y we define the discounted total cost
(0, Σ 7c((z1,
1=
where 7 E (0, 1) is the discount factor (which emphasizes earlier differences between the two sequences, and serves to make sure that the distance is well-defined). As is usual in the optimal- transport literature, we will define the distance between the two stochastic processes Mxx and My by minimizing the expected cost over suitable class of couplings of the two joint distributions.
Formally, coupling of Mx and My is defined as a stochastic process on the joint space X' y whose law is defined for all n Mx.y(Tn Fn) P X Fn, Y yv] and satisfies E Ey.ey" Mx (Tn Mx (T,. and Σz,ex" Mx.y(In. Fn.) My(7,)) We denote the set of all couplings by II, and call coupling Mx,y E IT bicausal if and only if satisfies
Σ Mxxy(rylEnn-1b-11 = Mx(x((a-1) vEy
and Σ (x.y(ry/EnnIn...... = My(y)
TEX
respectively for all and y, and for all n. The set of all bicausal couplings will be denoted by The-- Intuitively, this is the class of couplings that respect the temporal structure of the Markov chains and only allow the distribution of each state X (resp. Y2+1) to be influenced by the past state pairs (Xe, Y.). The optimal transport distance between the two Markov chains Mx and My is then defined as
d,M,,, My) = inf
(1)
with the dependence on the cost function e suppressed for simplicity of notation. Following the observation made by [Calo et al. (2024), we will frequently refer to this distance as the bisimulation metric between Mxx and My.
3 Bisimulation metrics from sample streams
As observed by Calo et al. [2024), the bisimulation metric in 11 can be rewritten in terms of occupancy couplings. The occupancy coupling associated with the bicausal coupling E Ihe is distribution nº E Axyxy with entries
1d (x, s/) = (1 (1-))
Σ 1=
r Pr XX, = x, Y. v, v, = Y',Y,11 = =//,
3

---

where P. [-] denotes the probability law induced by the coupling л. Introducing the notation (u,c) = Enya.' 'v'' flz, y, 1/ )((z y), this means that the original optimization problem defining the distance can be obviously rewritten as linear function of n² as
(Mx, My)
inf (u" .c). mellbe
(2)
Calo et al. 2024] identified a set of linear constraints on 1r that are satisfied if and only if I E Hh which has effecctively reduced the problem of computing the distance to a linear program (LP). This formulation is closely related to the standard LP formulation of optimal control in Markov decision processes, where the primal variables 6.9 in Puterman] [194). As shown by Calo et al.] (2024), one of the linear constraints satisfied by any valid occupancy j is the following flow condition:
commonly called occupancy measures (see, e.g., Chapter
E(zz y. x', z'al"
Σ 2( + (1-7))o02, v) (VH.y E X))).
(3)
2.
Unfortunately, their other constraints explicitly feature the transition kernels Px and Py, which ultimately makes their LP unsuitable as a basis for stochastic optimization. Indeed, optimizing their LP via primal, dual, or primal-dual methods would require having least generative model of Px and Py that allows sampling from Px(-(zr and Py(-ly) at arbitrary states and y. In practice however, such models are rarely available and one has to make do with samples drawn directly from a stream of states generated by the two chains. We address this limitation by reformulating their linear constraints in a form that eliminates the transition kernels Px and Py from the constraints, and replaces them with joint state-next-state distribution from each chain that can be sampled from efficiently. In what follows, we first introduce our new LP formulation, and then provide a primal-dual stochastic optimization algorithm to approximately solve the resulting optimization problem along with its performance guarantees.
3.1 A new LP formulation of bisimulation metries
Our reformulation is based on the following observations. First, notice that any valid occupancy coupling has to arise as coupling of the marginal occupancy measures of the two chains Mx and My, defined respectively for each x, x' and y as
vx(x, x') = (1 - - IP [X, = xx, Xx+1= x'), =x1,
my(y, = (1- E. y YY y, Y++ v'].
Indeed, valid occupancy couplings respectively satisfy the coupling condition Ey. 1¹ (Ir, u',)) = (2)
(4) for all xr xx and y, y'. Second, the conditional occupancies induced by a bicausal coupling ग satisfy nº (z', y/x) (x(x'(=) (1)(=")"" (w/x) and
and
Ex.z' 1ºt (x. y, v.v = uy
1. (I, v/ly) = P>(""))) (zly),
due to the requirement of causality that the conditional law of YY, given X, (resp. X₂ given Yi) should be independent of the future state Xx+1 (resp. Y(+1). By multiplying both sides of these equations by Ex wx(2 x²) and Ey uy(u, v'), we obtain
Ey 10 (x, x', M. v/) vx (2, x' ) (ulx) and Ene 1et (z, x',,,'/))) my D' (xlp) (5)
Summing both sides for all y and x respectively recovers the coupling conditions of Equation
In this sense, both the causality and coupling conditions can be recovered by the single set of equations (5) The following key result shows that, together with the flow constraints of Equation this system of equations provides a complete characterization of occupancy couplings.
Proposition 1. The distribution μ the induced occupancy coupling of a bicausal coupling æ E II if and only if there exist Ax y and xy RR-Y such that the following equations hold:
Σ E1 v. x //)
U, x. v) + (1 7))o(z, v)
(Vx,V (x,,EED)
(6)
Σ
u(x, y. x. ..'.)/)) = vx (z, x (
(Va,z' . E XXY)
(7)
Σ u(x, v/) vylu. )
(Vn, E xyy)).
(8)

---

Furthermore, if the equations are satisfied for some μ., Ax and xy, we also have Eg (x(()x) = and En Ay(aly) for all x and y.
Thus, the set of equations ( uniquely identifies the complete set of occupancy couplings. In particular, whenever the constraints are satisfied for some there exists a bicausal coupling и inducing j as its occupancy coupling, and conversely all occupancy couplings satisfy the above equations. Further important side results can be read out from the proof, provided in Appendix A
3.2 A stochastic primal-dual method
An immediate consequence of Proposition is that the OT distance between Mx and My can be written as the solution of the minimization problem of Equation 2 with respect to je as the optimization variable, subject the constraints (6) (8) Equivalently, it can be written saddle point of the associated Lagrangian defined as
C(p, a V) = Σ
u(a, v, .'') (clx, v) + ax + tyy(.,))) +V V(x,,))
zyz's y
Σ
Ux (x (x --w.://wwww../-w/-//o/>//......
E ayy
wwryy
(1 7)
Σ vo( (0(
where ax RXXY and ay E XXY are the Lagrange multipliers associated with constraints 7 and (8) and V E xy are the multipliers for the flow constraint (6). Indeed, by the Lagrange multiplier theorem, the optimal value of the original LP can be written as dy (Mx. My) ming.A m v L(,, X; a, V). Importantly, the gradients of the Lagrangian with respect to dual variables ax and ay can be written as expectations with respect to the occupancy measures Vx and by: which suggests that the objective may be amenable to stochastic optimization given sample access to these distributions.
Inspired by this observation, we propose primal-dual stochastic optimization algorithm that aims to approximate the saddle point of the Lagrangian. In particular, we will suppose that we have sampling access to the occupancy measures V/ and Vy and use these samples to construct stochastic gradient estimators for in- crementally updating the primal and dual vari-
Algorithm 1 SOMCOT
Input: c, n, B, 7, K
Initialize: 11 U(XXX)), Ax(1 = (())
for all x, Ay = (x) for all y, a = 0, = 0.
For k: = 1,2...,
K:
Sample Xx, X ~ vx and YE, Yg ~ vy.
compute gradient estimators via Eqs. 110) 15
ables via variants of stochastic gradient descent-
update primal parameters via Eqs. (6 (18).
ascent. Concretely, the algorithm proceeds sequence of iterations & = 1, 1.2..... K, updat- ing the primal variables μk. Ax and Λy. via stochastic mirror descent (SMD) with entropic regularization and the dual variables VR, ax,, and ay.& via stochastic gradient ascent (SGA). We describe the paudient-estimatinn procedures and the update rules below, and provide high-level pseudocode as Algorithm I For brevity, we will refer to the algorithm as SOMCOT, for Stochastic Optimization for Markov Chain Optimal Transport. Further details about the derivation of SOMCOT and more detailed pseudocode can be found in Appendix B
update dual parameters via Eqs. (19
Output: T
>x= Pkk
K
The gradient estimators. For constructing the gradient estimators needed for the updates, we first sample transitions (X,, XX,) ~ vx and (Y,, YY) ~ vy from the marginal occupancy measures of Mx and My, and let Fs denote the record of all transition data drawn until the end round k. The primal updates are defined in terms of the following gradient estimates:
(x, x', v²) c(z,8) v) (x, v) a.,( v. + VVe( .') Va(x, v)
(10)
Te. (v/z) [xxxxe x a (x,
(11)
TK (xly) - 1 ay (2,
(12)
Clearly, we have gsS Vpl(in, Ax: 0n V&). Furthermore, is easy to check that E [rr,Ax FR-1] Dime, ARi ab. Vi) and E (r Fx-1] V Clue, dni on, Vo). Simi-
5

---

larly, we can define the gradient estimates for the dual variables as
ox (x,
Σ Mr v. ,', v/) (Xx,x; a
(13)
Dk.ay (2, V/ v/.5/)
Σ
Ma |, v', v/) 1 (n, )2,
(14)
(,, Σ x'y''
Enl x' 3²) (1 ())0(2, v) ? Σ 3.
Mu (2, 1)
(15)
which are again easily seen to satisfy E ox FR-1l Vox L(ues Ax: aep. Vx),, E gy F Vay C(us, Au; Ork, V&.) and gkv L(ug ds; Oie, Vi.)-
The update rules. The primal variables are updated via stochastic mirror descent with appropriately chosen entropic regularization functions. For μ. the updates are given as
Pa (2,, y. x', y²) exp( n9k,, (2, y,
(1 y,
(16)
Enger Engery PA (Z, D. F., 37) exp( -79k v (i, D. 5. 7))
with > 0 being a stepsize parameter, and the Ax variables are updated as
Xx )x exp(-1x...., (2)(z))
\x,0+1
(17)
Eg Ax, (1)
())
(x exp( ( ww)
Ay.x+10rly) = Σ> A ( (y) exp(- (zly))
(18)
with respective stepsize parameters 7x By O. Note that due to the design of the gradient estimators,
each iteration only needs to update these variables locally at xx (+(-){x)) and Ay(-(Y)) at the sampled
states X₁ and Yg. For the dual variables, we define IIp the orthogonal projection operator onto a
convex set D, and implement the following projected stochastic gradient ascent updates:
ax,++1 llp. [ax, A 3xāraxl.
(19)
ky,4+1 = ll., Jay, Byğk,ayl,
(20)
k+1 11 [Vx- 3gk.vl,
(21)
where Dy 8 (0, 1 and D. 8 (0, ) are the projection domains for each variable, and B, Sx By are the stepsize parameters.
The output. The algorithm terminates after K rounds, and produces the average of the primal iterates Fx R Σx= Me output. From this, estimate of the distance can be computed as (MM,, My) (ux; c). Averaging the output variables is motivated by the design of SOMCOT as a primal-dual method and its theoretical analysis, and it also helps stabilize the quality of the solution. Indeed, primal-dus methods are prone to instability and oscillations, which are smoothed out very effectively by averaging. We discuss the role of this step and other practical improvements the algorithm in Section bellow, and provide further comments on the potential usefulness of other side products computed by SOMCOT for downstream tasks.
4 Analysis
The following theorem is our main theoretical result about the performance of our algorithm.
Theorem 1. Suppose that ||c|loo < 1. Let F 1/2 Σx= Pk be the output of SOMCOT, let μ² be the optimal occupancy coupling achieving the minimum in Equation (2). and set the learning rates as
1 11 7)2
| log D (1 - 7)2
ny Do log | (1 7)2
17
11
K
K
K
| 1P
| |
12'
8x =
V (1-7)2K'
By- = V (1 >)2K*
B =
(1 (1-7)2K'

---

Then, the following bound is satisfied with probability at least - 6:
( -o-==OVXXX 1 (VK(1>)) (VI2) 21 (121+ (2]) + /IX + [V]) log (1/5))
Equivalently. for any E O, the output satisfies (7 - -,, etl - with probability at least - 8 if the number of iterations is at least K Ko = O (XIDD/1 log(1/
(1-7)-
guarantee is that scales with the state spaces as | [ (|x)) + (/]) instead of the full dimensionality of the decision variables, |x12 | Note however that each iteration has computational cost scaling with this full dimensionality, The scaling in terms of e optimal up to logarithmic factors, as can be deduced from well- known lower bounds for the static OT problem (see, e.g., Klatt et al. 2020). Finally, note that the big-O notation only hides numerical constants, and the bound features no problem-dependent factors whatsoever.
perhaps surprising feature of the
We provide the main idea of the proof below, and relegate the full analysis to Appendix C The main technical idea is relate the estimated transport cost (uK, c) to the true optimal transport cost via the analysis of the duality gap associated with the sequence of iterates computed by the algorithm. The duality gap 9x (j",, X";a", .V) against set of comparator points (n", >";a", V*) satisfies
K
GK(M A" Vo)
Σ (E(U), Au; V*) - L(1 a" Vi)) A=1
(22)
K
As is standard for analysis of primal-dual methods, the duality gap can be decomposed into the sum of the regrets of the minimizing player controlling l and X, and the maximizing player controlling o and V. which can be controlled using the well-established of online learning (Cesa-Bianchi and Lugosi 2006 Orabona 2019]. For the analysis, we will pick the comparator points follows. For
the primal variables, we let μ be the occupancy coupling achieving the minimum in Equation (2) and let the X variables be the conditional distributions of Y Y|X and X ( under the joint distribution 4². For the dual variables, we choose
(a", V*)=
arg max VEDY K
Σ
L(NR, A V).
1=1
Under these choices, the error can be upper bounded as follows. Lemma 1. KT - -=,".", e) S GK ("",'",a"" V**.
The proof of this lemma makes up the bulk of the analysis, and is thus relegated to Appendix C.I It then remains to upper-bound the regrets of the two sets of players, which is routine work that we execute in Appendix C.3
Experiments
We performed a suite of numerical experiments to study the empirical behavior of our newly proposed algorithm, as well as illustrate some potential applications that are enabled by our method. Due to space restrictions, we only show a small portion of the results here, and refer the reader to Appendix | for additional results and implementation details (most notably detailed discussion on byperparameetr-ttning)).
Several of the experiments are conducted with family of processes call block Markov chains, motivated by the framework of block Markov decision processes (or block MDPs, Du et al. 2019).
This framework commonly studied in the context of representation learning for reinforcement learning, where a standard postulate is that the dynamics of the environment are governed by simple latent structure. Block Markov chains formalize this setting by assuming the existence of a latent Markov chain with small discrete state space, with each latent state generating a unique set of observations. Formally, we emulate the block structure by fixing low diwmensiiiinal chain Mxx and another chain My that is copy of Mx up to an additional irrelevant noise variable. In our experiments, we Mxx be a uniform random walk on the state space X' = {1, ......................... ., n) and My is a Markov chain on the state space X' x 21,2..... ... B}. with the value in (1, 2. B) generated uniformly at random. In all experiments, we use a sparse cost function that only allows to clearly distinguish between states 1 and n, and treats all other states identical.
7

---

20 20
30
40
50
10
20
30
50
20
40
50
Conbbee
30
50
20
40
Figure 1: Encoder-decoder maps learned by the algorithm block Markov chain example (n 10,
B 5) for sample sizes 1000, 10000 and 100000.
Model selection: random walks
Model selection: pendulum
K-5000
K-1000
0.20
K-500
K=100
0.18
Distance
0.6
Ground truth
Distance
0.16
0.4
0.14
0.12
0.0
0.10
7.5
10.0
17.5
(a) Distances between true dynamics (0 = 0.5) and the different models in the model class, for dif- ferent values of the parameter 0 and the number of iterations K.
(b) Distances between the true dynamics (g 9.8) and the different models in the model class, for differ-
ent values of the gravity parameter g.
Figure 2: Model selection results for random walks and the pendulum environment
Representation learning. Within the family of block Markov processes, the task of representation learning is equivalent to finding the mapping between the latent states and the observations and vice versa. Our method is very well suited for this task, thanks to the following curious observation. Besides the estimated coupling Fg and the associated cost, the algorithm outputs other values that are potentially useful. Among these, the variables X = & > Ax, and λy = * Σk= Ay.x particularly interesting for purposes of representation,
these conditional distributions interpreted as an encoder-decoder pair, with (-Ix) and (-yy)ly giving the respective conditional and X|Y y under the estimate of the optimal coupling. To illustrate
be
distributions of Y|X the potential usefulness of these maps, we conducted a set of experiments on block Markov chains 10 and B 5, and show the encoder-decoder pairs computed by SOMCOT in Figure D Notably, the algorithm does not make use of any prior structural knowledge of the environment: each individual state y treated separate state. Despite this and the very limited information revealed by the cost function, block structure is clearly identified by SOMCOT after sufficiently many samples.
with parameters
Model selection. Another important use case identifying the hidden dynamics underlying re- alizations of stochastic processes. We model this scenario in two experiments. In the first one, we generate a struccst Markov chain Mj, and set low-dimensional Markov chains M,x with different transition kernels parameterized by 0 E [0, 1). This set contains the true model M; x underlying Mý. corresponding to 0.5. We compute estimates of the distances between My and all the candidates of the model class by running SOMCOT for various sample sizes, and show the results on Figure 2a Notably, the distance achieves its minimum for the true model, and increases as 0 is further separated from its true value.
8

---

We also conduct a second experiment on model selection in continuous state spaces. To this end, we consider the classic control environment Pendulum-vl from Gymnasium [Towers et al., 2024). We begin by training near-optimal policy using the DDPG algorithm [Lillicrap et al. 2015 and fix this policy to induce Markov chain over the environment. The continuous state variables of the pendulum are then discretized to bins each. We instantiate several copies of the environment by varying a hyperparameter: the acceleration constant 9, which by default is set to g = 9.8. The learning task we consider is to identify which of the class of candidate models best explains the unknown true dynamics corresponding to the default choice of g. Figure 26 shows the results obtained. Again, we observe that the distance achieves its minimum for the true model, and increases g departs from its true value. Notably, due to discretization of the state space, the observations are not Markovian, yet the results clearly indicate that SOMCOT is still able to produce meaningful distance estimates, thus illustrating the potential of this methodology for general representation learning tasks.
Discussion
In this work, we have explored the use of stochastic methods to compute distances between Markov chains. This is still a largely unexplored field, and we believe the results presented here open the door to many interesting advances. We outline some of these future research directions we consider be the most promising.
Most importantly, it remains unclear how to properly scale our algorithm to larger problems with potentially infinite state spaces. While we believe that our bounds cannot be improved significantly in the case of finite state spaces, addressing infinite state spaces should be possible under appropriate structural assumptions, One may take direct inspiration from the OT literature to extend our approach to these settings. For instance, parametrizing the dual variables via kernels or neural networks has been shown to be an effective approach to solve static OT problems (cf. Genevay et al. 2016 Seguy al. 22188, and extending this idea our setting is straightforward. The real challenge seems to be approximating the primal variables, which correspond to (conditional) probability distributions, which are not straightforward to parametrize via modern architectures (at least as long as one is interested in theoretically sound methods). We leave the investigation of this very interesting question open for future work.
Among all applications of optimal transport for Markov chains, its use in representation learning for RL is particularly interesting us. Many previous works this domain have highlighted the potential of bisimulation metrics for learning state abstractions, but all theoretically sound previous methods for computing such distances required full knowledge of the transition kernels. Removing this need brings us closer to realizing this potential. The experiments presented here demonstrate the effectiveness of bisimulation metrics in capturing symmetries and latent dynamics of Markov chains directly from sampled trajectories, both in random walks and discretized classical control environments. Incorporating function approximation along the lines mentioned above could significantly enhance these applications.
Besides the already-mentione there are other side products of SOMCOT that can prove useful for representation learning. Most notably, the optimal dual variables oxx and ay correspond to the derivatives of the distance with respect to the state-transition distributions ux and By, which is a fact that can prove extremely useful for the development of practical methods. Indeed, notice that these distributions themselves are differentiable with respect to the transition kernels, which altogether allows one to backpropagate loss function in representation learning tasks, Successful implementation
interpretation of the variables X and Xy as encoder-decoder maps,
through the OT distance of this idea may lead to strong theoretically sound alternatives to empirically successful methods such as MuZero [Schrittwieser et al., (202). This latter method uses a loss function remarkably similar to our OT distance, albeit with some limitations that disallow its application to stochastic environments (cf. Diang 2024). Once again, we leave this direction for future work.
The authors wish to thank Csaba Szepesvári for sions during the preparation of this manuscript. This project has received funding from the European Research Council (ERC), under the European Union's 's Horizon 2020 research and innovation pro- gramme (Grant agreement No. 950180). This work has been co- oo-uunee by CECUUU/EE/UE-PPD2022 447145NB-100 AGAUR SGR and MCIN/AEL/1O, WWW.///E111..133/50000 100010011000 under the Maria de Maeztu Units of Excellence Programme (CEX2021-000 195-M).
Acknovledgements.
discus-

---

References
Jason Altschuler, Jonathan Niles Weed, and Philippe Rigollet. Near-lineat time approximation algorithms for optimal transport via Sinkhorn iteration. In Neural Information Processing Systems, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML), pages 214-223, 2017.
Julio Backhoff-Veraguus, Veraguas, Mathias Beiglbock, Yiqing Lin, and Anastasiaa Zalashko. Causal transport in discrete time and applications. SIAM Journal on Optimization, 27(4):2528 25622, 2017.
Marin Ballu, Quentin Berthet, and Francis Bach. Stochastic optimization for regularized wasserstein estimators. In International Conference on Machine Learning, pages
2020.
Tristan Brugère, Zhengchao Wan, and Yusu Wang. Distances for Markov chains, and their differ- entiation. In International Conference on Algorithmic Learning Theory (ALT), pages 282-336, 2024.
Sergio Calo, Anders Jonsson, Gergely Neu, Ludovie Schwartz, and Javier Segovia-Aguas. Bisimula- tion metrics are optimal transport distances, and can be computed efficiently. In Neural Information Processing Systems, 2024.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In AAAI Conference on Artificial Intelligence (AAAI), pages 10069-10076, 2020.
Nicolò Desa-Bianchi and Gábor Lugosi. Prediction, Learning, and Games. Cambridge University Press, 2006.
Di Chen, Franck van Breugel, and James Worrell. On the complexity of computing probabilistic bisimilarity. In Foundations of Software Science and Computational Structures (FOSSACS), pages 437-451, 2012.
Jianda Chen and Sinno Jialin Pan. Learning Representations via Robust Behavioral Metric for Deep Reinforcement Learning. Neural Information Processing Systems, 2022.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Neural Informa- fion Processing Systems, 2013.
Josée Desharnais, Vineet Gupta, Radha Jagadeesan, and Prakash Panangaden. Metrics for labeled Markov systems. In International Conference on Concurrency Theory (CONCUR), pages 258-273, 1999.
Josée Desharnais, Radha Jagadeesan, Vineet Gupta, and Prakash Panangaden. The metric analogue of weak bisimulation for probabilistic processes. IEEE Symposium on Logic in Computer Science (LICS), pages 423-422,2 2002.
Simon Du Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Lang- ford. Provably Efficient RL with Rich Observations via Latent State Decoding, In International Conference on Machine Learning (ICML), pages 1665-1674, 2019.
Stephan Eckstein and Gudmund Pammer. Computational methods for adapted optimal transport. The Annals of Applied Probability, 34(15):675-713, 2024.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch Wasserstein: asymptotic and gradient properties. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792. 2021.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite Markov decision processes. In Uncertainty in Artificial Intelligence (UAI), pages 162-169, 2004.
10

---

Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic Optimization for Large- scale Optimal Transport. In Neural Information Processing Systems, 2016.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference Artificial Intelligence and Statistics (AISTATS), pages 1608-1617, 2018.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence, 147(1-2):163-22) 2003.
Nan Jiang. note loss functions and error compounding in model-based reinforcement learning. arXiv preprint arXiv:2404.09944, 2024.
Mete Kemertas and Allan Jepson. Approximate policy iteration with bisimulation metrics. Transac- tions on Machine Learning Research, 2022.
Marcel Klatt, Carla Tameling, and Axel Munk. Empirical regularized optimal transport: Statistical theory and applications. SIAM Journal on Mathematics of Data Science, 2(2):419 443, 2020.
Rémi Lassalle. Causal Transport Plans and Their Monge- Kantor-vnnt Problems. Taylor & Francis, 2018.
Timothy Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Manfred Otto Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1rX.vv:500002977 2015.
Arthur Mensch and Gabriel Peyré. Online sinkhorn: Optimal transport distances from sample streams. Advances in Neural Information Processing Systems, 33:1657-1667, 2020.
Robin Milner. Communication and Concurrency. Prentice Hall, 1989.
Vrettos Moulos. Bicausal optimal transport for Markov chains via dynamic programming. In IEEE International Symposium on Information Theory (ISIT), pages 1688-1693, 2021.
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 19:1574-1609, 2009.
Gergely Neu and Nneka Okolo. Dealing with unbounded gradients in stochastic saddle-pon point optimization. In International Conference on Machine Learning (ICML), 2024.
Kevin O'Connor, Kevin McGoff, and Andrew B Nobel. Optimal transport for stationary Markov chains via policy iteration. Journal of Machine Learning Research, 23(1):2175-2226, 2022.
Francesco Orabona. A modern introduction to online learning, arXiv preprint arXiv:1912.13213. 2019.
David M. Park. Concurrency and automata on infinite sequences. Symposium on Theoretical Computer Science, volume 104 of Lecture Notes in Computer Science, pages 167-183. Springer, 1981.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-6):355-600 2019.
Georg C. Pflug and Alois Pichler. A distance for multistage stochastic optimization models. SIAM Journal on Optimization, 22(1):1-23, 2012.
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wiley-Intersciencee, April 1994.
Alexander Rakhlin and Karthik Sridharan. On equivalence of martingale tail bounds and deterministic regret inequalities. In Conference on Learning Theory, pages 1704-1723 PMLR, 2017.
Julian Schritwieserr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-699, 2020.
11

---

Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Math- ieu Blondel. Large-scale optimal transport and mapping estimation. International Conference on Learning Representations (ICIR). 2018.
Jiajie Tao, Hao Ni, and Chong Liu. High rank path development: an approach to learning the filtration of stochastic processes. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Vayer Titouan, Nicolas Courty, Romain Tavenard, Chapel Laetitia, and Rémi Flamary. Optimal Transport for structured data with application on graphs. In International Conference on Machine Learning (ICML), pages 6275-6284, 2019.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032, 2024.
Franck van Breugel and James Worrell. An algorithm for quantitative verification of probabilistic transition systems. In International Conference Concurrency Theory (CONCUR), pages 336- 350, 2001.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Tianlin Xu, Li Kevin Wenliang, Michael Munn, and Beatrice Acciaio. COT-GAN: Generating sequential data via causal optimal transport. Advances in neural information processing systems, 33:8798-8809, 2020.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning Invariant Representations for Reinforeement Learning without Reconstruction. In International Conference on Learning Representations (ICLR), 2021.
12

---

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 long 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: The limitations of the work and the future research directions are included in the discussion.
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 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
13

---

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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch 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 ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless 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 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 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 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 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset 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 reprodueibility. the case of closed-source models, 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 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://nips... cc/
publ / guides Codeeuumie ss ssionn cisionnolicc
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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 (https //nips. cc/ wwwpiiclguiiioooooooooo / //odiddee/Cdeeeeeeee CooeeeeeeeSeeeeeee SoddSuumieeee Submiss wwwssssoiiccgg.e.. for more details.
The authors should provide instructions on data access and preparation, including how 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 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 possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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, 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: [Yes]
Justification:
We report 95% confidence intervals when applicable.
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, confi- dence intervals, statistical significance tests, at least for the experiments that support 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
15

---

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 96% CI, if the hypothesis of Normality of enrors 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 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 compute required for each of the individual experimental runs well 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 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 https: /ttpa:////urip..../m http://pp.....o.........ttaaaaaa
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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 direct 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
16

---

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 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.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperssi hopersswtccoee com/ dataset has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
17

---

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 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
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 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.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
18

---

Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: /tmpps://neee http://uuuriis..cooooeeeeeee LLM) for what should or should not be described.
2025/
19

---

A Equivalence of the LP formulation and the bisimulation metric
In this section we prove Proposition [ which together with Equation 2 implies that the novel LP formulation is equivalent to Equation (I) for computing the bisimulation metric. Our proof uses the linear programming formulation of Calo et al. [0224] starting point. Concretely, Calo al. (2024) prove that is the induced occupancy coupling of bicausal coupling E The if and only if ji satisfies the following set of constraints:
Eulo,
u(x, y. v,"' =7 >
E.ld....a, d'
Endd v,,,,) + (1-7))0(2,))) (V E xx0),
(23)
2 a/al
U1I,Y,x', v,,'8) = Σ
u(a,y, //)")) (VH,y, e xx)xx),
(24)
x" m/
Σ M, Σ "a..,,"'''""
(Vz,y.V' (xyx)).
(25)
Importantly, the above constraints provide complete characterization of occupancy couplings: not only do occupancy couplings satisfy all equations, but any satisfying the three linear systems of equations above is valid occupancy coupling (ef. Lemma in Calo al. 2024).
To prove the proposition it sufficient to show that ju satisfies Equations (6)
satisfies Equations 23 25 Equation (6) is identical to Equation 23 and thus if and only if left with showing that Equations and (8) are equivalent to Equations (24) and (25). respectively, We will
show the first of these claims, and note that the second claim will follow by symmetry. Within
the proof, we will repeatedly make use of the shorthand notation vx(x)) = Ex' Vxx (2, x') and the easy-to-see fact that Px (x' (x) Vx(x 2²) (z)
First, let us assume that f X-XXXXXX'X satisfies Equation 24 Then, it is easy check that Equation (7) is satisfied with the choice (x(v]=) = Ex , p(T, //////////((((( (()), making use of the relation between Vxx and Px stated above. Conversely, assume that μs, Ax satisfy Equation (7) Summing both sides over x' gives Ex',y' p(z, v,,' v/) = xxxx)xxxxla which can be plugged back into Equation (7) obtain
Σ (u,,,,, v/) ux(x, x' (x(()=)) = Px ((ld))))x(())))))))) Pa(dla))ex()))eell)) Px(x
Σ u(x, y, x", ./),, a'/a/
thus confirming that Equation 24 is indeed satisfied.
For the last part, let us define μx as px x²) ) = En, p(z, y, ,') for each (x, x²) whenever Equations (6) are satisfied. As per the above argument, Equations 3 and (24) are also satisfied, and thus summing both equations over y yields
Σ Eux(a,z2) = Eux(x',)) (1 7)),.x
(x(xxx)) Px(x'(<) ( 2xx(x,"").
A standard argument (provided as Lemma 12 in Appendix E} shows that the unique solution to this system of equations is equal to the marginal occupancy measure vxx This implies
Ex.. u(x, "(,,y,"'''' "x(xxx)) Ex vx(xx
Ax(()lz> Σ
vx(x)
vx(x)
vx(x))
y
which concludes the proof.
B Further details about the algorithm
In this section we describe some further details about the derivation of our algorithm (SOMCOT) that omitted from the main text. Algorithm provides full pseudocode for SOMCOT.
At high level, the algorithm aims to find the saddle point of the Lagrangian (9) by performing primal-dual updates for the two sets of variables (μ, A) and (a, V), referred to as minimizing and
20

---

maximizing players, respectively (or often simply call them min and max players). Both sets of players maintain sequence of iterates (ux, >y) and (as, Vx), which are updated using versions of online stochastic mirror descent, described below in detail. For the updates, the f and À players move in the direction of the negative gradient of the I Lagrangian evaluated (1k, Aes ock, Vi:,, and the a and V players move in the direction of the positive gradient.
Since some of these gradients involve the occupancy measures vx and My, they cannot be computed exactly without perfect knowledge of these distributions. However, since the dependence on vx and My is always linear, it straightforward to obtain unbiased gradient estimators given only sample access to the chains Mx and My. We provide a detailed guide for sampling from these distributions in Appendix [B.4]
The remainder of the section provides a detailed derivation of the gradients and update rules used in each iteration. We begin by introducing the Mirror Descent algorithm that forms the basis of the update rules for each variable.
Algorithm 2 Stochastic Optimization for Markov Chain Optimal Transport (SOMCOT)
Require: Convex sets Dax C XxXXX Day CRXXXX Dv R'x Initial values (11, Ax,1, Ay 1. 0y.1- V1, Learning rates 7): 1122, 1,,, 3x, By, 0.
1: for k k= 1,... K - 1 do
2: Step 1: Draw samples from the Markov chains
3: Receive Xx.XX ~ Vx,, (Yx, Yg Vy
4: Step 2: Compute gradients or stochastic gradients
5:
(2, 5') c(2, v) - ax (x,,,( y) ay,((z,.v))) + VVa('' 8') Vx(xx 8)
6:
Dk.Ax (+) <- (xx=o)axxal x)
7:
gR.Ay ( (xly) t
8:
Ik,ax (I, x,, .))
of (x y, x1,)) M
9:
Th.y (2, 3') + Ene (1, y, x' g') 1
gk.v v) t -nrr (x y. v/) (1 y)y((1y) - Ma
Mk (2.
11:
Step 3: Update primal variables
12:
Ma+1 y. x', y') x p4 (x, V. x',y/ exp(((( 199, (2,8,22, 222)
13:
(x.e+1(ylx) x Axx(()/))) exp(-7x"k.>x (
14:
Ay x Ay,A(=) (x(v) exp(-1)y (zlv))
Step 4: Update dual variables
16:
0x ll (ax (x
17:
ayA+1 + IID., (ay,x Bygk,ay)
18:
VK+1 II (Vk- Bgn.v)
19:
end for
Output F 1 Ek E Hk.
20:
B.1 Online Stochastic Mirror Descent
Online Stochastic Mirror Descent (OSMD) is an algorithm for the problem of online linear optimiza- tion, where in sequence of rounds A = 1, K, the following steps are repeated:
1. The online learner picks a decision zy taking values in the vector space Z,
2. the environment picks a linear function Ok 2 + R.
3. the online learner incurs loss 2
4. the online learner observes an unbiased estimate Đk E Z" of the loss function.
The sequence of steps above defines a filtration (Fe) and the loss estimate ãk is assumed to satisfy [ Fx-1] = fk- Typically, the vectors gk are subgradients of sequence of convex loss functions, and thus we will often refer to them with this term, and also call the vectors GJ stochastic subgradients (or simply stochastic gradients). OSMD computes sequence of updates based on these noisy gradient estimates and a convex and differentiable distance distance-generatiin generating function & Z R. Concretely,
21

---

OSMD operates with the Bregman divergence By of V defined for each pair 2, E 2 as
By (= = (z) - ((2)) = =')-
OSMD starts with an initial point ž E Z, and computes each subsequent iterate using the recursive update rule
-k arg min (Tk, z) + - (= 1)
(26)
.E2
where 1) > O is called the learning rate.
Each of the update rules used by SOMCOT follows from instantiating OSMD with a specific decision function vi and a noisy subgradient estimator. Concretely, we will make use of the following instances and corresponding update rules of MD, whose derivations are available in standard textbooks (e.g. Orabona 2019).
space Z, iisttace--gnerallee
Proposition 2. When 2 Rd and v is the squared Euclidean norm defined as (()) 1/2 || for each z E Z, the OSMD update reduces to the projected stochastic gradient descent update rule
2k+1
1lz (Zx: - n/gg))
where 11z is the orthogonal projection onto the set Z defined as Ilz(x)) arg min | vll2. Proposition 3. When Z = Ax is the probability simplex on finite set X' and V is the negative entropy defined (p) Eg (z) log)(z), p E Ax, the OSMD update reduces to
(x)) P=(z)e- (m)(
(Vx E 2)).
E,,ro(y) mo.tv) (
Proposition 4. When Z Aypx is the conditional simplex on finite sets X and y and Wo is the total negative entropy (p)) 2n, p(ul))) log p(yla)) p Ayyx, the OSMD update reduces to
Pk p (v)
Ph Px+1(
E, (o (5. [91m (Vx,YE E XY).
B.2 Primal updates
In this section we derive the gradients and update rules of the primal variables and Ax. The gradient and update rule of Ay follow by symmetry.
For μ4, first notice that any valid occupancy coupling ys is an element of the simplex A Axyxy: summing the flow constraint in (6 over x and y immediately yields En,,, yy p(x,, y,, yy) = 1. Thus, is natural enforce this constraint throughout the execution of the algorithm and use OSMD with the entropy regularizer given in Proposition 3 In order to derive the update rule, it remains to compute the gradients of the Lagrangian with respect μr, which is given
Oc Øp [n, X; a, VJ(I,, y, a', s/) = c(x, v) - (x x', ,) ay +VV.)'+>VVVVVV) Vxx v). (27)
In each iteration k. the algorithm computes the gradient 9k.,- ac Juae Ax Ay,x axx, ay.k V/2. which can be used as the unbiased estimator ğk- Altogether, this yields the update rule on line 12 of
Algorithm 2
As for the xx variables, notice that Proposition] | implies that Xx belongs to the conditional simplex
AyΙx = {A e R2x Vz E X. E, A(y|z) = 1}. Thus, it is natural to use the total negative entropy as regularization function (as suggested in Proposition 4.. For selecting the update direction, we note that the gradient of the Lagrangian with respect xx
ac 0xx jr, X: V((V = Σ
Ex (x X' .)]. (28)
Thus, we can obtain stochastic gradient estimate TR.Ax of ac by sampling a transition (X,, X²) from vx and setting Ax things together, this yields the update rules on lines B 14 of Algorithm
[uk, Xx Ay,ia aya Vil
X1,8). Putting
22

---

B.3 Dual updates
We now move our attention to the dual variables. Again, we will derive the gradients and update rules for ax and V, and the gradient and update rule for ay follow by symmetry. Since we are maximizing over the dual variables which are not restricted to any simplex, we update them using projected (stochastic) gradient ascent, which is why the gradients are negated below. The feasible sets for the dual variables are chosen to enable using Lemma 2 for bounding the estimation error- see Appendix C.1 for details.
The gradient of the Lagrangian with respect to ax is defined
dax [n, a, (((,z',)) =
ac
Σ
- vxel 2x(ex
- mix
Eu(o.8,''.1)) 1 [xx,, mn {x(1)z))
Our algorithm will use the update direction Tk.ax (x, x'.v)
Ey ux(2 v. x., [xxe.X ( (w)a). We will apply OSMD to update ox using learning rate Bx and the
regularizer in Proposition 2
which yields the update rule on lines 16 of Algorithm
The gradient of the Lagrangian with respect to V is given by
dc avlu (V, ax, VI(y)) =
Σ p(u y, x', v/) (1- - 7) 1) - Σ
p(G C( xr
R.
We use the update direction Pk.V = Exr.y (1x y, (1)-(1-7)u022 v)
Σ 2. MA (T. T. v)
and apply OSMD to update V using a learning rate B and the regularizer in Proposition 2 which yields the update rule on line of of Algorithm
B.4 Sampling from Vx and by
key step in constructing our gradient estimators (and thus running our algorithm) is drawing samples from the occupancy measures vx and vy. Here we provide further details about how to perform this operation in practice.
In order to generate a sample from the occupancy measure, we let G be geometric random variable with mean I and recall the definition of the marginal occupancy measure vx to write
20
(x(xxx')== = (1 7) Σ -P|x== - X
Σ
PIG= x, X = =x'
= PP X =Il
Thus, one can obtain independent samples from vx by first sampling geometric stopping time G, sample sequence
X6,
and keep the last pair of states Xo, XG+1-
We remark that the task of sampling from an occupancy measure is common reinforcement learning, and in particular it is necessary for correctly implementing policy gradient methods. To avoid sampling an entire sequence each iteration, it is standard practice to replace samples from the occupancy measure with arbitrary sample trajectories generated by the Markov chain. Specifically, it common to ignore discounting and draw samples directly from trajectories which consecutive state pairs are no longer independent. We expect that, like most other RI algorithms, our method is also resilient to such abuse, and can be fed with sample pairs drawn from longer trajectories without resets or throwing away samples to ensure independence.
C Analysis
This section provides the complete details for the proof of our main result, Theorem T the analysis, we will assume ||e|| 1. Completing the outline provided in Section requires filling
Throughout
23

---

two gaps: proving Lemmal 1 and bounding the duality gap in terms of the regrets of the two players. These are respectively done in Sections C.I and C.3 below (with Section S2ction[[. additional technical tools for the proof of Lemma Putting the two parts together complete the proof.
C.1 Proof of Lemma
The majority of our theoretical analysis dedicated to proving the error bound stated as Lemmal recalled here for convenience
-", S ""'""""","""""""
(29)
As first step towards this proof, we first need to define technical tool that will allow to quantify the constraint violations associated with the output pg. Indeed, one challenge in the analysis is that
Fk does not necessarily satisfy the constraints (6) (8) exactly. We quantify this effect by defining
total absolute constraint violations associated with the primal variables Js, Ax and xy respectively by
0F(μ)= = Σ
Σ
u(a, v, x'v/ - Σ (,, - (1 7)0 v)
2.31
0Cx Ax)= = Σ
1xx n. ./)//)/))010 ux(x, x' x"x)) (x(()) Σ
2,
aCy(H, = Σ
ula, y, x', .2) 2/ Σ
Z.D-y'
For the sake of analysis, we will make of rounding procedure that will convert μæ into valid occupancy coupling r(/kg) that satisfies all constraints. Importantly, this rounding procedure never has to be executed in reality: it is only used as device within the analysis. The details of this rounding process (which is an adaptation of a method developed by Calo et al. (024) are provided in Appendix C.2. The following lemma provides an upper bound on the rounding error in terms of the total absolute constraint violations.
Lemma 2. Let E Axyxy and r(u) be its rounding (as defined in Appendix C.2 and xx and Xy be arbitrary. Then, have
(p-r(n)]l S 3x((,, Ax) + 3Cy +
(30)
1-7
The proof is provided along with all relevant definitions in Appendix C.2
With this rounding process
and its guarantees hand, we can rewrite the absolute error between the cost estimate (7x c) and the true cost di (Mxx My) (1² c) as follows:
Ki 1², c) S (r(ux) -=,e) + [(((()) - |x llclooo T - -c + 1| - r(x))l, |clloo
(31)
Here, the first step follows from the triangle inequality and the crucially important fact that (r(Pxc 1=" c) 2 0 thanks to the feasibility of r((xx) and the optimality of μ²"
now only remains to relate the quantity appearing on the right-hand side the above bound with the duality gap. To this end, we define the shorthand Xx = K Ek Ax, and Ny. > xy, and recall the choice
K
(a", V*)
arg max VEDY K
L(ux, Aa; V).
k=
Then, by plugging these variables into the Lagrangian, is easy to check that
K
6Cx(TR Xx) + 60y Jy) + 20F (0x(
A(pi,de;a', VV) (xx +
1-7
24

---

which, by using Lemma 2 implies the following bound:
T + + - r( - K Clpx, "sia",
On the other hand, it is easily verified that (1²,,) L(u", a; V) holds for any choice of and V. thanks to the fact that μ and A verify all the constraints of the LP. Putting this together with Equation (3Ib. we obtain that the error can be bounded in terms of the duality gap at the above-defined comparator (f* X", o², V*) as
KīT ,,,c)) (,,'",
This concludes the proof of Lemma I1
C.2 Rounded coupling and rounding error
We describe the process and guarantees of rounding an (approximate) occupancy coupling μ E RXX' We note that computing this rounding requires knowledge of ux, but this does not cause any practical problems since the rounding is only ever executed in the analysis. For the rounding process itself, we first introduce the state-occupancy measure (z, y) = Env μ(zz, y, x'.8'), and we define the associated transition coupling Ja as the kernel Tp XY - A Xx with entries
playaa af2 if v v) # 0, (,(2))) Px Py(u' otherwise.
Th (1', s/ 12,
As shown by Calo et al. 22024], each transition coupling क XY + Axy induces unique occupancy coupling j²,, and that the occupancy induced by The is valid if and only if it equals μ (i.e., if μ²= holds). For more details, we refer to Appendix B.2 in Calo et al. [2024].
Following Calo et al. 2024). we will apply the rounding procedure of Altschuler et al. 2017
Algorithm to T to obtain valid transition coupling r(=x), and then extract the occupancy
coupling induced by r(Tp). More precisely, for two probability distributions E A (XX), E A()),
the set of valid couplings is defined as Up. {P E R P p; p 1 q). For
nonnegative matrix F' e IR xx the rounding procedure outputs valid coupling r(F, p. q) Up By Lemma of Altschuler et al. [2017). the rounded coupling satisfies
||r( p. Fll, S 2 1
||F" 1|]).
For completeness the procedure is detailed in Algorithm 3
Algorithm 3 Rounding procedure for couplings
Input:
approximate coupling F, marginals p, Q X f (iag(min(p//FF 1), (F.)).1)) 1))
F' + XF
Y diag(min(g///) (g 1))
F" + F'Y
err, pp- F" 1, err, - 9 F"T Output: G F" + enr, ||errp
This procedure is not symmetric, and thus we consider the following symmetrized procedure defined
sym (F.P r(F, + r(FT -q,p)²
2
To obtain the rounded transition coupling, we apply the rounding procedure individually for each pair of states x, y. In particular, for transition kernel n. XY Axy, we define its rounded counterpart ग = r(T) with entries
7 7||,,, v) sm (=([2r y), Px(-)x), (-1)))
With some abuse of notation, we will now denote as r(u) the occupancy coupling induced by r(ap)- Because r r(T,)) is valid transition coupling, r(u) is valid occupancy coupling. The following
25

---

derivations will relate the distance between r(u) and μ to the total absolute constraint violations of thus providing a proof for Lemma 2
To make the subsequent derivations easier, we will define some handy notation. We first define the operator E: Axyxy - Axy via its action any μ as (En)) y) Ex.5 (1 D. x, y), and
note that this allows us to rewrite the flow condition (6) in the form Vμ = E + (1-11)11) y)-- For state- doste-distribution V E Axy and kernel x X' 7 Axy, we define the composition 1/ o π as the distribution p with entries p(æ, yy) v(x y) y/ |x; y). We will specifically use the notation (μ²) = 44 o (r(TT)) - Tx). Armed with all this notation, we bound the & distance between
r(u) and j as
[1(()) |11
|Pr(N) (TTT) Uu Txll
|PA(N) r(Tx)) U/ r((T)) + up or(T)) - Uu Tlll
|PFCH) "/lll + +11, (r(TT) - mn]ll1
[P((N)) (1 -7) + (1 - y) - -/ull + ||ا
= 1|TEr(1)) - -7Ep + hT + (1 - -7)00- - + [|()))
S ()((()) - -111 + ||
where the second-to-last line uses the fact that r(u) is a valid occupancy coupling and as such satisfy the flow condition (6). and we have recalled the definition of aF(u) stated in the main text. After reordering, we obtain
aF + 1|
[[(()) mll1 S
1-y
and thus it remains to upper bound ||]())] This is done in the following lemma, using which concludes the proof of Lemma
Lemma 3. For any μ Axyxy, and any xx x + Dy and xy y Ax we have
||(()]] S 30Cy(u, Xy) + (u,
Proof. By Lemma 7 of Altschuler et al. [017], we have that for arbitrary state pairs x, y. the
following is satisfied:
[I((T)()(<r,))) v) T T,ll w)ll
S 221
Σ Px(('1x) - Σ 11,,, v)
+
Σ
Po(e'ly)
x
x²
By symmetry, this directly gives
||r"" (TT)())))))))))))))) v) Ta (. (z, 2)
3 Σ
Px(x1<2) - Ialle' 1/
+ +Σ Py(v'lv) -
- Eall'ss/la
2
Now, multiplying both sides by v (1 8). we get
(p) y) = mx(2z)) v) (-12 v) - T v)ll
3
Σ
Px(('(x)) - E Eula, x' v/) + Σ (oll))))
2
The first term on the right-hand side of the above expression can be bounded as follows:
Σ
Px(e'la)u,
Σ Iu(a,,,'',
Σ (x(2)) - 2x (xx((lx)) vx(x,
(oce)l/)==xxxxex) (xx(( - Σ p(a Em(Z,,,2',1) v,2'
26

---

(x(())))) - (ot - Σ u( yf y
Pacctpllaade-----00000000-11-111 2x)
Mula, y) - (xx=))Px(=1=))xx(00 ( +acx(u, Ax)
(1))
Σ
p(a.v.='.')
-x
+xx((1xxx)
20Cx ( Ax).
Here, we used the triangle inequality for (i) and the fact that Ene (x(x'1x)= for (i) and (ii). The proof is concluded by repeating the same argument for the constraint violations (py(u, Xy). and plugging the results back into the previous inequalities.
C.3 Regret bounds of the primal and dual sequence
This section provides an upper bound on the duality gap defined in Equation 22) in terms of the regrets of the two set of algorithms controlling the primal and dual variables. Recalling the convention established in Appendix B we will refer to the algorithms the min- and max- max-players, with their regrets respectively defined as
K [Cu,, Aus",, L(pu: Az Vi))
regret max
(a'vV*)
k=1
regret
min ,²)
Σ
(C(U,, Vi) C(u,, a" Vk)),
k=1
where our notation emphasizes that each regret is measured against the comparators a", V and 1 X*, With this notation, the duality gap can be rewritten as
GK(u", at;"", V*) regret max (a", V*) regret tin (u", X²)
(32)
K
With a mild abuse our earlier notation, we write the full expression of the Lagrangian in terms of the x,,,, ary and Ax ny variables as C(P, Ax xy; ax, ay, V). The regret terms that need to be bounded can be further decomposed in terms of the following individual terms defined for each set of primal and dual variables:
K
maa regret K
(axx)=
Σ Ax.. V*) - L(u2, Aysaara V*)
regret (ay)
max
V*)
K
max regret K
= (V**)))
Aln V*) L(uu, Xx, A 0y.k, Ve)
regret (μ²)
min
Vx) L(p", Ayx Vi)
K
regret min (Xx))
C(p", Ax.k. dy. ax, ay.k V4) C("", Nix Ny. ay.k Vi.
27

---

K Σ 1(0 Ap.iax,a V4) - L("", Xxx, Nzz ax, ay,ss V4)
min regret (15))
Thanks to the bilinearity of the Lagrangian, each of these terms be learning algorithm with linear loss / gain functions and decision variables taking values in a convex decision space Z (embedded within some Euclidean space). In particular, each of these regrets can be written in the following form for some sequences (gk)s E R2", (Zk) E 2 and & E Z:
the regret of online
K
regret K (=*)
Ok: 2 - =").
As noted in Section our algorithm can be understood as running an instance of Online Stochastic Mirror Descent (OSMD) for each set of variables, and thus each regret term can be bounded using standard results. One challenge for the analysis is that the comparator points a and V are chosen in data-dependent manner. This is not easily handled by standard tools in online learning, but can still be treated with some relatively more advanced tools that are common in the context of saddle-point optimization (most notably, using techniques of Nemirovkki et al. 2009 Rakhlin and Sridharan 2017). In particular, we will use the following general result to bound the regrets of each player in the analysis below.
Lemma 4. Let 2"" E Z be potentially data- data-dependent comparator and assume that V is X-strongly convex with respect some norm II-II whose dual is denoted by ||-Ila. Furthermore, suppose that sup,,z'z || z/11 & C holds for some constant >> O. Then, for any îj > 0, the sequence (z*) produced by OSMD satisfies the following bound with probability at least 1 5:
K
By (=11=1)
(gs, 2k =²) I
19
7
2A
k=1
By (= (=)l=1)
1ggh - -1 + C, 2 19 Tx 12 log g.
7
2X
While composed of standard elements, we provide the proof for the sake of completeness in Ap- pendix [D] The regret bound itself can be simplified in two different ways, depending on whether or not the algorithm in question uses deterministic or stochastic gradients: for deterministic updates, we have Oh 9t and we can choose 1/5j- = O, whereas for stochastic updates the choice îj = v/ is more natural. This is how we will apply the lemma to each regret term below. In what follows, we will instantiate this bound to bound the regrets of all players listed above, which will require establishing i) the
properties the regularization functions, ii) bounds on the Bregman divergences between the initial points and the comparators and iii bounds on the dual norms of the gradients and the gradient noise. This is done case by case in the following subsections.
C.3.1 Regret of the a-players
The policy of the a-players is to run projected online stochastic gradient ascent on the feasible set Z Boo 1-1g) ) and unbiased gradient estimators with elements defined respectively as
Dx.x (x, 1.. v)
Hu (x y, x' .') Axx..x/==xx)xxxxx
V
and
Jk,oy (I, y. 1') Σ
Me(x, y, x' ./) I N (y,
The following lemma provides an upper bound on each of the two a-players. Lemma 5. With probability at least 8, the regret of the x x-lyyrr is bounded as
18 18 |J|
72K
2 log
regret
max (axx)) K
S
43xK + 1 (1 7)2
(33)
(1 33
and the regret of the ay-player is bounded as
18 1 X
72K
regret K
max (a²) < (1-7)283
+ 4,3 K +
V
(1 71-
log 5
(34)
28

---

Proof. We prove the claim for axx, and the result for ay will follow by symmetry. We start by noting that the gradient estimators and the gradients satisfy |1 S 2 and I1gk,ox Isoall11 S 2. Indeed, this can be verified easily
0x 11a S Σ pu(a, 2') v/) + Σ 1xx.va//Axxa = 2,
K.N,Z' wy
2,2'M
because of the normalization of both μa and Ax, Similarly, we have
llgkaxx - 9k,axll1 1
Σ
(1(x,x + =x(xxx²)) (x,((1lz) = 2.
3n,t :9
Furthermore, |1o ax21 S I- trivially holds thanks to the definition of the domain of oxx and the choice 0x211 = 0. Finally, notice that v is I-strongly convex with respect to I1-I12- and thus Lemma 4 (with the choice î 1 3x) immediately implies the claim after using the relations 0 12 S 0 | - 2 and ||axx ax,1 21 z l1xx - axx,1l I 36)x1)211
(1-7)*
C.3.2 Regret of the V-player
Similarly to the a-players, the V-player employs online gradient ascent the feasible set 2 Bxo 2 with entries of the gradients given in each round as
(2,
Σ 1, 3²) (1 7)) v)
?
Σ 3.
Pa (2, x.
xt.y
The following lemma gives a bound on its regret.
Lemma 6. The regret of the V V-player is bounded as
4 4| |D] + 23K.
max (V*)= regret K
(35)
3(1 7)
Proof. Since the V-player employs deterministic gradients, we will apply Lemma 4 with 1/
and bound the Euclidean norms of the comparator V and the gradients. By the choice of the feasible
set for V and the choice Vi = O, we immediately have I|V* - Vill2 I | || - V1 S xx1y] Furhermore, evaluating the gradient of the Lagrangian with respect to V, we get
(1 (1-y)2
11 s (1 7) Σ vola, 8) + (1+7)
Σ
u(I,8, x.. v/) 22 2,
x.y
Z,y.0' (2/
which turn implies 11gh.vl1 S 11g S 2. Plugging these results in the bound of Lemma concludes the proof.
C.3.3 Regret of the j-player
The p-player plays OSMD with entropy regularization, and gradients with elements defined as Ik.p (z. y, x', y c(I, v) a xx (I, x' y) ary,k v.v') + VVa(=' v/) Ve(x,, The following bound gives bound on the regret of this player.
Lemma 7. The regret of the μ-player is bounded as
min regret K (a*) log log( 1)P)
2007
7)
Proof. The proof follows from noticing that the regularization function v is 1-strongly convex with respect to the norm II and that the dual norm of the gradients is bounded as Il S I-7 Indeed, this follows by upper-bounding each entry of the gradient as
20
19 (z, yr x', y/) < c(x, 1/ + |ax, (x, x', 3/)) I + (ay, y, v')l + (Va (V(x,,))
< + 12
4(7 +1 17 + 37
20
1 1-?
1 1- 1 1-
Finally, we recall the choice of μ1 being uniform over XYXX,, and the standard result that the relative entropy between any distribution and μ1 is equal to ? ?))
D
29

---

C.3.4 Regret of the X-players
The regret analysis of the X-players is slightly nonstandard. Focusing on the Ax-player here, we note that the updates correspond to using OSMD on the decision space = {> Vx, y (ylzz) Z
0; Ve N(yax) 1] with the following choice of regularization function:
(A) = EΣ ( log (og(A()(2z)).
As we show in Lemma 10 this regularization function is 1-strongly convex with respect to the 2- group norm defined for each À Z
1212,11
Σ
Σ
is easy to verify that the corresponding dual norm is the 2 oc group norm defined as || En (max, (g(2, v)l)a. We also recall that the updates make use of the following unbiased estimate of the gradient:
Ax (x, u) = (x, X',
(36)
With these facts at hand, we prove the following bound on the regret the X-players. Lemma 8. With probability at least 5, the regret of the Ax player and is bounded
log [V]
901x
288 X K log (3)
regret max (Axe)
(37)
Ux
(1 7)3
(1 7)-
and the regret of the xy-player is bounded
1og |x(|
907y (1 - 7)2 +
288
| K log (3)
regret (15)
max
(38)
Ty
7)2
Proof. We provide a complete proof for Xx, and note that the result for Ay is analogous. For this case, notice that Lemmas and 10 suggest that we should first obtain upper- bounds on the magnitude of the gradients terms of their 2, so-group norms, and thus we first establish that
| 00,000
Σ max ((x ax w1l
< 1x( |xxx ||axx
6 Note that the latter is upper- roper-bounded as ||ax,, S 1-y by construction. Further observing that the true gradient norm can be bounded via the same argument as Ilā 112,0c I I- we also have
12
1 - 0k.xxll2,0 S 11g,,xxl12.0 + 1 11,.00 S
1 1-
Finally, since (- is chosen as the uniform distribution over y for all x. we have By (Ax.)) < |x( log [3), and the primal-norm distance satisfies ||*- All < 2//2% Now, the claim follows from using Lemma with 7j = 7x.
C.4 Proof of Theorem
The proof of the theorem now follows from putting together Lemma | with the regret decomposition Taking a union bound over the two probabilistic claims of Lemma and 5 this gives that the following bound holds with probability at least 1 25:
in Equation (32 and combining Lemmas 5 g
18/
72
2 j
[px - 100-11,,00 ,,,,,) (x -2
+ 4,3x +
log
K(1 - 7)
18/2111112
72
+ 4,3y +
log
3y (1 7)²
V K(1 - 72
30

---

4/201107
+
3K(1 -7)2
2 2log( (111(1)1131)
2007
nk
2(1 7)²
X (og)
90nx
288 X' log
+
1xK
71
K 7)2
log |28|
90nly
288 y log 2
+
+
7yK
(1-7)² 7)2
K 7)2
Setting 8x
9 31 2(1-7)"K*
By
2(1-2) -I
91/11 DIF
2 (1-ナ)K*
7x
2) log |
90K
7x
11 2) log [2(|
90K
17
(1-7)²
100K
log(( [2I)
the bound becomes
12 12//2/2112/1 (V (VIX++ VIDI)
4 4
K7 ,,,,))-
(1 ->)V
+ 3v 20 log [39
3 3/10 log [
40 Tog( 111)))
(1 W)
(1 (1- 1->V
7)K K
12 log
12 2/20 Y/
12 / 2/
+
+
(1 WV
(1-7) (1>
- 7)>k
(1211 + DI) + log
1 8 + 13'| 12y
(1 - 2)
This concludes the proof.
D Online learning: The proof of Lemma
This section is dedicated to proving the general regret bound we use throughout the analysis for upper-bounding the regret of each player, Lemma 4 As mentioned in Appendix [C.3] the main challenge that we need to deal with is that the comparators for some of the regret terms are data dependent, which requires some additional steps that are typically not necessary in regret analyses. For concreteness, we adapt the notation of Lemma 4 and write the regret against comparator =" as
regretx((")
(gs Zh 2²)
Thy 2k 2
2²
Rx
Mr
where in the second equality we also added some terms corresponding to the stochastic gradient gi- Here, the first term Rg corresponds to the regret of the online learning algorithm on the sequence of stochastic gradients ÿk. which can be upper-bounded using standard tools of online learning. For the second term, notice that the stochastic gradient satisfies E [ FR-1] gk. and thus if =* is independent of the sequence of stochastic gradients, the second term Mr in the above decomposition is martingale. However, this is no longer true if 2" is statistically dependent the sequence. In order to account for this, we adopt an elegant technique by Rakhlin and Sridharan [017)] to control the resulting sequence of dependent random variables In particular, learning algorithm for the sake of analysis, and use its regret bound to account for the additional error terms in the above decomposition. For sake of concreteness, we define the sequence of decisions made by this algorithm by setting 2 = 21 and updating the parameters recursively via a mirror
introduce a second online
'This technique is commonly attributed to Nemirovski et al. 2009 but we find the connection with Rakhlin and Sridharan (0177) more illuminating. Otherwise, we learned this proof technique from Neu and Okolo 2024
31

---

descent scheme analogous to the one underlying the sequence zy::
1 (9x. Th. 2) By By (2llzz)
1++1 arg min
1)
Using this notation, the regret of the original algorithm can be rewritten as follows:
Tle 2k 2²)
gk. 2k 2
+
(gk
gk.
2ke
k +
gk
9k 2k
2²)
Rx
M
Rx
Thanks this construction, the term MR is martingale and Rx is the regret of the auxiliarly online learning algorithm in the newly defined online learning game.
For the concrete proof of Lemma 4 we will make of the following classic result regarding the regret of mirror descent.
Lemma 9. Lei 2 and assume that V is X-strongly convex with respect some norm ||-I1 whose dual is denoted by I|-Il,,- Consider the sequence with an arbitrary u1 E and all subsequent iterates defined as
(uk 2) + 1 Bu (elleu))))
uk+1
arg min 262
where aş is an arbitrary sequence in Z and w > 0. Then, for any u E 3, the sequence (ux) produced by OSMD satisfies the following bound:
(OR, -0-- u) S Bu [u]) [u])
W
lax12
(39)
2X
The proof is standard and can be found in many textbooks for concreteness. we refer to Theorem 6.10 of Orabona [2019]. To proceed, we apply this lemma to the standard sequence of iterates in our setting with dg Ok and w = 7 bound R and once again with Oş 9k Ik and 1j bound Rg. Finally, we use the Hoeffding- Laafidinn-Azuaaa inequality (Lemma
control the remaining term as
K Σ (gk gk. 2x -xx) S.C. C 2 Σ 19 1 log
1
M
8
with probability least 8. Indeed, notice that under the condition max 2- 2'|l C
each term satisfies Kge Gk. 24 - mx)1 S C l|gh - Txll.-- which allows using Lemma with Ci: 2C' [[gh
5x ll... Putting these results together concludes the proof of Lemma
E Technical Lemmas
Lemma 10. The function (p)) = Ej Et= p(ili) log p(ls)) is 1-strongly convex with respect
to the 2-1 group norm lipll2,1 = V E/ (E 1 on the set 2 = E RRIJ. p((]k) 2
(( El=1p(1l p(il)) = (Vj).
Proof. We first note that, by standard calculations, the Bregman divergence induced by
p(il)
(xll)))) =
p(15) log
j=1==1 q(ilj)
Now, by Pinsker's 's inequality, we have that
(xll)) (pp(-1)) a(-11)1
J=1
which is equivalent to the statement of the lemma.
32

---

Lemma 11. (Hoeffding Azuma inequality, see, e.g., Lemma A.7 in Cesa- Bianchi and Lugosi 2006) Let ( be martingale with respect to a filtration (FR) Assume that there are predictable processes (At), k and (Bx), and positive constant (cx) such that for all k 2 1, almost surely, As Zz
Zk-1 Bu and BR Au S ct.
O, Then, for all
2Ea
P(Z: Zo 2 c) exp
(40)
Σ c
or equivalently for all 6 E (0,
E & log(
P
Zi Zo 2
8..
(41)
Lemma 12. The occupancy measure VX R of the Markov chain Mx is uniquely defined by the two sets of equations
Σ xxx(x,x')= = >
-ux("", (x(x",,)) +1 x2)) (Vx),
(42)
vx(x,x0) = Px(x1(z) Σ vx(xxx²) (Vx,
(43)
Proof.
Using the definition of the occupancy measure Vx we obtain
px(x, x = (1-n)
7 [X, = x, X =x]
= =(1- Σ
7 >xx('' 2 [X+ = =x]
00
(x'(()
(1 (1-7)/0,x<)) + +(1
Σ
y 'P / [P,xx=0) =a]
1=1
2
(x((()
(1 (1-7)v0.x()) + E
=
= Px(el(>)
(1 (1-
+ Σ Exx(""
where we used the stationarity of the transition kernel Px, the definition of vo,x, the law of total probability, and the stationarity of the Markov chain to recognize vx(x"" (x",2) x)) in the last step. Summing the previous equation over x' yields (42). and substituting (42) into the previous equation yields (43).
In order to show that the solution V/x to (42 and (43 is unique, we introduce the notation &x as Ex(x) En. Vx (x, x²) for each x. Substituting
43 into 42
yields
Ex(xx) (z) = Σ Px(xla")
(x" + (1 (z) (Vx).
2
By defining {x and Vo.x as vectors and Px as a matrix, we can write this system of equations in matrix form εx =P (1 2, equivalently, (
(1 7)),,x Since Px is a positive matrix with spectral radius 1, the Perron-Froobenus theorem applies and the matrix (T->Px) is invertible. Hence there exists unique solution &x (1 (T PI) -PP)-lvoo which
together with (43) implies that vx is uniquely defined as vxx (z( x')= (x (x)()r).
Additional details on experiments
In this appendix we present further details about the experiments included in the main text, along with some additional empirical results. Along the way, we will also provide some further comments on best practices when implementing SOMCOT, including recommended hyperparameter settings.
33

---

Figure 3: Distance matrices between instances after running SOMCOT for 1000 and 10000 steps, and the ground truth obtained via Sinkhorn Value Iteration.
F1 Similarity metrics between parametric Markov chains
This experiment serves to illustrate the ability of bisimulation metrics to capture intuitive similarities between stochastic processes, well as the empirical behavior of SOMCOT when used to approximate such similarity metrics based on data. To this end, we generated several random walk instances from the same family used in the experiments in the main body (described in detail in Appendix FF For this experiment, we let X' y = {1. 2....... n. with = 1 and set the block size as B Deviating from the setup described in Appendix F3 we set the reward function as r(1) r(n)
for both extremes of the state space, which induces a symmetry on the state space. For generating the set of environments, we varied the initial states zo and yo between 12.3.....
1} and the bias parameter & in the set {0.05, 0.1, 0.15, ..... 0.95], thus resulting in 72 different instances. We then computed pairwise distances between these instances using SOMCOT with various sample sizes, and compared the results with the ground truth (computed by the Sinkhorn Value Iteration method of Calo al. 2024). Figure shows the similarity matrices obtained by these methods, showing that the distances computed by SOMCOT successfully capture the structure of the problem: even though the exact numerical values of the true distances are not approximated very accurately, the qualitative picture obtained by SOMCOT is very similar to the ground truth. In particular, the symmetry induced our choice reward function is clearly visible with the matrix being symmetric along the counter-diagonal as well as the main diagonal.
F.2 Practical implementation details
Being a primal-dual method, SOMCOT is not as easy to tune as a common stochastic optimization algorithm. There are several implementation details that one needs design carefully in order to make sure that the algorithm behaves in a stable way and outputs good estimates. This section describes our experience working with SOMCOT, and provides practical guidance for implementation.
SOMCOT has one tunable parameter per optimization variable: a positive learning rate that controls the magnitude of the updates during optimization. While our theoretical analysis suggests some specific values for these learning rates to guarantee convergence, such values are typically too conservative (as practice, using larger learning rates can significantly reduce
common in stochastic optimization). the number of iterations needed to reach good solutions. Since our problem involves optimizing six variables, this leads to six separate hyperparameters, which makes tuning a grueling task. To address this, we tie some of the learning rates together: all primal variables share single learning rate denoted by 7), and all dual variables share another one denoted by B.
we observed that in practice using a fixed value for the dual learning rate r often made difficult to achieve stable convergence across different problem instances. To address this, we introduced a decaying learning-rate scheme of the form 77 = 7p
Moreover,
VI+ak where & is the index of the
current iteration and > 0 is tunable parameter. This decay helps balance the need for large updates in early iterations with the stability required for convergence in later stages.
Figure T Lillustrates the performance of the algorithm under different learning rate settings. This shows that, even given the above choices, is not easy to pick hyperparameters that work uniformly well across problem instances. Even for single instance, the combination of 7/o and 3 that leads to the best performance requires careful hyperparameter search. order to understand the behavior of SOMCOT under different parameter choice, is helpful to remember the roles of the primal and dual variables,
34

---

vsi
states
0.0373
0.1
0.0002
error
0.0007
0.01
0.0001
0.0e11
0.033
0.807
0.019
0.0088
00139
0.araa
0.0001
0
0.2099
0
0
0.2001
0.007
0.009
0.2052
0./299
0.0053
6.
deta
Figure 4: The influence of the ratio between 17 and on the convergence of SOMCOT for different chain sizes. Error and learning rates shown a logarithmic scale. To produce this plot, decay rate 0.001 was used for 1). No decay was applied on B.
and in particular that the dual variables serve penalize the primal variables for violating the primal constraints. Thus, a value of 7 that is too high relative to / leads to large constraint violations, resulting ja values that yield very small distances but fall outside the feasible set. Ultimately, setting 33 too small results in gross underestimation estimates that are suspiciously close to zero, the value of should be increased the value of rj be decreased. The opposite scenario produces the inverse effect: B that is too large relative to causes the dual variables to update too quickly, leading to a resulting distance that overestimates the actual value. This issue can largely be mitigated by decaying 1) while keeping 3 constant (as described above). In our experience, it is often better to pick a large initial value for ny while this typically leads to rapid drop of the distance estimate to zero, the estimates eventually start increasing and converge toward the true cost.
the true distance. Thus, whenever one sees distance
Theorem provides guarantees for the averaged output |* = Ex= lkk, where fa the value of ju obtained at iteration k. This is commonly required for algorithms based regret analysis, least for the theoretical guarantees to go through. In typical applications of stochastic optimization, this strictly necessary and the final iteration can perform well enough. However, this is typically not the case for primal-dual algorithms like SOMCOT, where iterate averaging often makes a big difference the stability of algorithms. This is true in our case too: without averaging. the iterates typically fluctuate quite wildly around the optimum. Averaging makes the estimates much more stable, and thus strongly recommended (even if only for the last half of the iterates or less).
averaging step
Finally, we note that all our experiments have made use of i.i.d. transitions sampled from the occupancy measures of the chains. This falls in line perfectly with the theory, but may be impractical in applications where transitions may be dependent or be sampled from undiscounted trajectory distributions. While we have experimented with such data, we believe that SOMCOT should be able to deal with it as long as efforts are made to break the correlations between the consecutive samples, for instance by sampling the transitions randomly from a buffer (instead of processing them in their original order). In our experiments. we have sometimes made use of minibatch updates, which can affect computational efficiency and stability, but no major impact on the overall convergence properties has been observed. We display all hyperparameter choices we have made in the experiments in Table
F33 Details about the environments
Our experiments made use of two families of Markov chains: a collection of parametrized random walks, and several instances of the classic "inverted pendulum" environment. We describe the details of these settings below.
Parametrized random walks. We consider a one-dimensional random walk over a finite state space X' = 21, 2......) n) with biased transitions. A transition from state x moves to x + with probability 0 E 1) and æ I with probability 1 States and n are "sticky walls": the process remains there with probability 0.9 moves to the neighboring state with probability 0.1.
35

---

Experiment
%
8
b
Figure
40
0
0.2
0.99
Figure
20
0
0.5
0.99
Figure
2a
0.1
0.001
0.5
8
0.95
Figure
0.1
0.05
0.2
0.95
Table Table summarizing our hyperparameter choices for each experiment. Recall that the learning and the minibatch size is denoted by b.
rates follow the decaying scheme
VI+ak"
Additionally, we define a reward function on the state space, with values r(1) = r(n) = -1, and The initial state distribution is Dirac measure produce the plot shown in Figure 2a we generate a low-dimensional chain Mx with bias 0 = 0.5 following this setting. Then, produce set of chains My B, each of them with different bias parameter. In addition, all My are augmented with an additional irrelevant noise variable, producing B observations per each latent state in X' Formally, My is a Markov chain the state space y equal to X' x {1, , B}. reflecting the absolute difference in rewards between the states.
r(x) 0 for all 2, ......
1. To
The cost between x E X' and E ) is given by c( y) = Ir(*) r(y),,
Inverted pendulum. We begin by training a npar-optimal policy using DDPG, a widely known Deep RL algorithm, in the standard Pendulun-vV environment, which is then used to induce Markov chain. One could use any policy, but using a near-optimal policy produces richer dynamics (e.g., using random policy in the Pendulun-v/ environment reduces the effective state space the surroundings of the initial state). Once a policy is fixed, we discretized each state variable of the environment into bins. Since the Pendulum-vl environment has variables (angle 0 and angular velocity w), the resulting state space has n² states. In our experiments, we have chosen n = 7, which resulted in total of 49 states. Note that due the discretization, the resulting stochastic process is longer Markov chain, as the states are no longer sufficient to predict the distribution of the next state. Nevertheless, the conducted experiments follow the same principle as the aforementioned random walks: We will compare the (approximate) Markov chain Mx of discretized observations with set of parametrized models My. The parameter governing the dynamics the acceleration constant g. capturing the effect of gravity. We set the default value g = 9.8 in Mx, and choose values in [0.4, 20] in the model set. The cost function is given by c(x, y) |r(x) r(y)), where r(x:) the average reward in bin x (computed from all samples that fell into bin x along long simulated trajectory).
36


Reviews:
[{"review_id": "qcLd9Ec5Lu", "reviewer": "Reviewer_gruX", "summary": "Calo et al. (2024) showed that bisimulation metrics are a kind of optimal transport (OT) distance.\n\nThe present paper extends this by deriving a new stochastic optimization algorithm (SOMCOT) for estimating bisimulation metrics between Markov chains based on sample observations, without knowledge of the transition law. Key to this is proposition 1, which offers a new characterization of occupancy couplings, opening up a novel path to optimization through stochastic primal-dual methods. \n\nTheoretical guarantees are provided and the method is demonstrated on several toy examples", "questions": "I would certainly be willing to increase my evaluation if some or all of the following could be addressed. They are all essentially related to practical utility. My feeling is that this work may be useful to the community, and so worth publishing, but that it doesn't make this case clearly enough at the moment. \n\n1. Better comparison with existing methods\n\nI wouldn't mind if the example is very small, but I do think there should be at least some way of comparing the current approach to previous methods, even just to see what is lost in terms of accuracy moving to a sample-based approach from explicit transition kernel knowledge. Would it be possible to use the approach of Calo et al. (2024) or maybe Chen et al. (2012)? I'm not particularly familiar with those works, but there must be some other method that could be compared to. I'd be very happy to see competitive accuracy from SOMCOT and perhaps demonstrating a regime where SOMCOT is superior. \n\n2. Empirical scalability\n\nHow does performance degrade on practical problems when we are limited in computation?\n\n3. Could the authors elaborate on a plan towards extending the approach to infinite state spaces using function approximation etc.? This isn't vital, but if they can show anything relating to it then that does strengthen the claim that the paper is a foundational step into a promising direction, and not a dead end.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- This is a sophisticated contribution with extensive theoretical guarantees. It seems to me a rigorous and interesting piece of work. It's also well written.\n- The approach relies only on sampling and not on knowledge of the transition law. This relaxation is a meaningful step forward and non-trivial.\n- A key contribution is the new LP formulation in proposition 1 which, if I’ve understood it correctly, I find to be a clever result allowing us to describe occupancy couplings without using the transition kernel explicitly. It does so by introducing auxiliary variables representing conditional distributions under the optimal coupling.\n- The paper offers an analysis of sample complexity with explicit constants.\n\nWeaknesses\n\n- There are at least two clear limitations to the work: it relies on finite state spaces and is extremely computationally intensive.\n- The finite state space assumption. While the authors acknowledge this limitation, it really does limit the applicability of the methods and there is little in the way of evidence to suggest an extension is possible.\n- The computational complexity. The sample complexity [O(|X||Y|(|X| + |Y|)/eps^2)] is clearly highly undesirable, as is the cost per iteration [Theta(|X|^2 |Y|^2)]. This implies for example that in a fairly modest system of 100 states, you would need ~2M/eps^2 samples, which for eps=0.1 (a low bar in some problems) implies 200M samples. I would assume, though I’m no expert, in realistic reinforcement learning scenarios, this will balloon to trillions of samples.\n- As far as I can see, the paper does not compare the approach empirically to any other existing method. I take the authors’ point that this method is the first of its kind (p.2) in some respects, but without any kind of comparison it’s hard to say whether the algorithm offers any advantage over simpler baselines, or how its complexity compares in practice to alternatives. They could have restricted their attention, at least in one example, to a case where this comparison was possible, e.g. some tractable small example where many methods could be used. The lack of empirical comparison is an issue: (i) exact methods exist in some small state space examples, (ii) I cannot tell from this work when to prefer this method over alternatives that may be available.\n- Nit: is it standard in this field to use x bar to mean an infinite-length vector x? In statistics this commonly refers to a sample mean, so this seems a little confusing.", "was_revised": true, "final_justification": "On reflection, I take the authors' point that the work is quite original and I was unfair in my original assessment that they ought to compare it more thoroughly to other approaches. I would appreciate if they could spare a few more words to this effect in the final paper, as it was unclear to me. The work is overall very strong and I have a more positive outlook on its merits having read their response. Accordingly, I've upgraded my score.", "rebuttal": "Thank you for your careful reading and thoughtful comments on our work! We are glad to see that you have appreciated the originality of our contribution and the rigor of our theoretical analysis. We also acknowledge the criticism that our approach is currently limited to finite state spaces, and that it is not entirely obvious at this time how to scale the approach to problems with high-dimensional state spaces. We would nevertheless like to push back on some of the concrete points made in your review, hoping that we can reach agreement about the value and potential of our work.\n\nFirst of all, we would like to double down on our point that our method is “the first of its kind”: we truly are the first to rigorously consider the problem of estimating bisimulation metrics / OT distances between Markov chains, and therefore there are literally no other known methods that aim to solve the same problem. Even for small problems, there are no exact methods or known alternatives. The only possible alternative we can think of would be a simple “plug-in” approach: build an empirical estimate of the two transition kernels based on all observations, and compute the distance between the resulting estimates using a state-of-the-art solver (e.g., Sinkhorn value iteration by Calo et al., 2024). This approach has no guarantees on the estimation error, and has the same runtime as our newly proposed algorithm. We are happy to make this point more explicit in the final version.\n\nSecond, regarding the scaling with the sizes of the state spaces: we of course agree that a quadratic scaling does not look appealing at first sight. Nevertheless, occupancy couplings (the key component in the *definition* of our distance metric) are of size $|\\mathcal{X}|^2|\\mathcal{Y}|^2$, and thus it is clearly unrealistic to expect that such an object can be computed in sub-quadratic time. Just constructing the required data structure costs $|\\mathcal{X}|^2|\\mathcal{Y}|^2$ time and space. In light of this, we actually found it quite surprising that we could achieve a sample complexity guarantee of the improved order $|\\mathcal{X}||\\mathcal{Y}| (|\\mathcal{X}|+|\\mathcal{Y}|)$, and we believe that this guarantee is not improvable in general. Of course, improvements might be possible under further structural assumptions (e.g., if the transition kernels are low rank as another reviewer has suggested), but we feel that such extensions are well beyond the scope of this paper.\n\nFinally, about the assumptions about the finiteness of the state spaces and the potential to scale up our algorithmic ideas: these are indeed the questions we are planning to address next (as mentioned in Section 6). While you are correct to point out that the method is unsuitable for large state spaces in its current form, we have serious reasons to believe that it can serve as an important foundation for a more practical method that can deal with larger problems. In particular, since the algorithm is based on gradient descent (and not dynamic programming like previous methods for computing OT distances), there is a natural way to combine it with “function approximation”:  just parametrize the primal and dual variables and perform stochastic gradient descent/ascent on the resulting parametric approximation of the Lagrangian. Doing this extension well enough will require some further careful work that is beyond the scope of the present work.\n\nOnce again, we would like to recall that our work is the very first to provide theoretical guarantees for this problem setting. Our objective was to develop a robust set of foundational ideas and to perform a careful and complete mathematical analysis in the setting we have considered. We believe that these foundations will be very useful for further developments, and that many colleagues in the NeurIPS community will find them inspiring."}, {"review_id": "Up0SFByFTR", "reviewer": "Reviewer_52pg", "summary": "This paper studies the problem of estimating the adapted (bicausal) Wasserstein distance between two finite-state Markov chains from sample trajectories. Building on the observation that the bisimulation metric used in reinforcement learning coincides with the adapted OT distance under discounted cost, the authors propose a stochastic primal-dual algorithm (SOMCOT) that operates using only sample access to the underlying chains.\n\nA key contribution is a new linear programming formulation over occupancy couplings that eliminates dependence on the transition kernels, making the problem amenable to stochastic approximation. The algorithm performs mirror descent on the primal variables and projected gradient ascent on the duals. A non-asymptotic error bound of order O\\left( \\frac{1}{\\sqrt{K}(1 - \\gamma)} \\right) is established, with dimension dependence matching that of standard OT estimators.\n\nEmpirical results illustrate the method’s capacity to recover latent state structure and perform model selection from sample data. The formulation is theoretically sound and avoids the heuristics common in earlier deep RL applications of bisimulation. While the algorithm is framed in the language of bisimulation, the underlying quantity is precisely the discounted adapted Wasserstein distance, and the contribution is best viewed as a sample-based estimation method for this object.", "questions": "1.\tComparison with Calo et al. (2024):\nThe proposed LP formulation appears closely related to the one in Calo et al. (2024), replacing conditional transition kernels P(x{\\prime}|x) with the marginal occupancy distribution \\nu_X(x, x{\\prime}). While the authors argue that their formulation is more suitable for the sample-access setting, conditional distributions can also be estimated from trajectories. Could the authors clarify precisely why their formulation offers an advantage over that of Calo et al., especially since this work builds directly upon it? A more careful comparison—both conceptual and empirical—seems warranted.\n\t2.\tSampling Cost under Discounting:\nThe algorithm requires samples from the distribution\n(1 - \\gamma) \\sum_{t=0}^\\infty \\gamma^t \\mathbb{P}(X_t, X_{t+1}),\nwhich may involve simulating long trajectories when \\gamma is close to 1. How is this distribution sampled in practice? Does the sample complexity bound in Theorem 1 account for the cost of generating such samples, particularly in terms of the expected number of steps required per sample?\n\t3.\tNormalization and Adapted Wasserstein:\nIn Equation (2), the objective \\langle \\mu, c \\rangle lacks an explicit (1 - \\gamma) factor. This raises concerns about alignment with the standard definition of the adapted (bicausal) Wasserstein distance, which includes discounting in the cost. Are the authors implicitly including the (1 - \\gamma) normalization in the definition of \\mu? If so, could this be clarified more explicitly to avoid confusion?\n\t4.\tConnection to Optimal Control and Dynamic Programming:\nThe adapted Wasserstein distance between Markov chains can be interpreted as the value of an optimal control problem over causal couplings. The variable \\lambda_X(y|x) in the paper resembles a policy mapping in this context. Is there a direct connection between the proposed algorithm and dynamic programming or optimal control formulations of adapted OT, such as those explored in Moulos (2021) or O’Connor et al. (2022)? Clarifying this connection could help contextualize the algorithm within the broader literature.", "limitations": "None.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\t1.\tWell-Formulated Estimation Problem:\nThe paper develops a stochastic algorithm for estimating the adapted (bicausal) Wasserstein distance between two Markov chains from sample streams. The problem is well-posed and connected to classical constructions via occupancy measures and causal couplings.\n\t2.\tSample-Based Optimization Framework:\nBy introducing a linear program over normalized occupancy couplings, the authors construct a saddle-point formulation that can be optimized without access to transition kernels. The method—SOMCOT—uses mirror descent and projected gradient ascent, and is supported by finite-time convergence guarantees.\n\t3.\tSample Complexity Analysis:\nThe convergence rate in \\varepsilon is \\widetilde{O}(1/\\varepsilon^2), which is consistent with minimax rates for empirical optimal transport. The dependence on 1/(1 - \\gamma)^2 is correct under the authors’ normalization, and avoids the more pessimistic 1/(1 - \\gamma)^3 rate that arises in standard Monte Carlo estimation of discounted costs.\n\t4.\tEmpirical Utility and Side Products:\nThe dual variables and conditional marginals (e.g., \\lambda_X(y|x), \\lambda_Y(x|y)) admit interpretations as encoder-decoder maps, and are shown to capture latent structure in block Markov chains. The method also shows promise for model selection and abstraction tasks.\n\n⸻\n\nWeaknesses\n\t1.\tFraming Ambiguity:\nAlthough the paper cites the bicausal OT literature and works with bicausal couplings, it frames the contribution primarily in terms of bisimulation metrics. This obscures the broader significance of the estimator as a sample-based method for adapted Wasserstein distances.\n\t2.\tPer-Iteration Computational Burden:\nThe per-iteration cost is O(|X|^2 |Y|^2), due to updates over full occupancy tensors. This limits scalability and raises the question of whether dynamic-programming-based or low-rank approximations could yield more efficient alternatives.\n\t3.\tMissing Normalization Clarification:\nThe paper does not explain that the discounted cost is normalized by (1 - \\gamma) in the definition of occupancy measures, and that this normalization is what permits the 1/(1 - \\gamma)^2 rate. Without this, readers familiar with standard cost accumulation may incorrectly expect a 1/(1 - \\gamma)^3 dependence.\n\t4.\tNo Discussion of Known Low-Rank Cases:\nIn the special case where both chains have rank-1 transition matrices, the problem reduces to static OT between stationary distributions. This is a classical LP (packing type), and the paper does not compare its approach or complexity with recent optimal running times for optimal transport. \n\t5.\tAbsence of Lower Bounds:\nThe paper asserts that its rate is optimal (up to logarithmic factors), but does not prove a matching lower bound or explore dimension-dependent regimes where structural assumptions (e.g., rank, sparsity) could yield improved complexity.", "was_revised": true, "final_justification": "I've read the discussion and responses, I'm satisfied and will keep my positive score.", "rebuttal": "Thank you for your careful reading of our work, and a detailed appraisal of the strengths of the contribution! Your comments listed under “weaknesses” are also well-taken; we respond to these below.\n\n1. **Framing ambiguity:** We opted to mainly frame our contribution in terms of bisimulation metrics as we felt that repeatedly writing out all potential names of the distance (“bicausal OT distance”, “adapted OT distance”, …) would be cumbersome for the reader. We will revise the text to emphasize the connections to (computational) optimal transport more explicitly in the final version.\n2. **Per-Iteration Computational Burden:** We agree that the quadratic scaling of the runtime with the domain sizes does not look appealing at first sight. That said, we are not aware of any alternatives with better computational complexity, even for the case of known transition kernels. For instance, the Sinkhorn value iteration method of Calo et al. (2024) has the same runtime complexity. We are not aware of any other method with an explicit complexity guarantee. We believe that speedups should be possible under some more concrete assumptions on the kernels (e.g., low-rank), but we did not explore this question in depth so far.\n3. **Missing Normalization Clarification:** You are absolutely right, our definitions involve a normalization constant of $(1-\\gamma)$, and we probably didn’t do a good job at emphasizing the impact of this scaling factor on the rates. We will provide additional discussion in the final version.\n4. **No Discussion of Known Low-Rank Cases:** We did not provide a detailed review of computational OT to keep the paper more focused (and fit the page limit), but we otherwise agree that a discussion should be added. We will do so for the final version. For now, we just note that one can specialize our method to the “flat” OT problem without Markovian structure, and achieve a rate of order $n^2/\\varepsilon^2$ which matches the best rate that we are aware of for Sinkhorn’s method (Dvurechensky et al., 2018) — while requiring only sample access to the distributions! We are not aware of other methods achieving a faster rate. (E.g., naively combining the state of the art method of Blanchet et al. (2018) with a uniform-convergence argument to account for sampling errors would still result in a $n^2/\\varepsilon^2$ rate according to a quick back-of-the-envelope calculation.) We felt that this question was a bit niche to explore in depth in the present paper, but we are open to including some additional content in the final version.\n5. **Absence of Lower Bounds:** We agree that it would be interesting to pin down the exact minimax rates for the unstructured case we consider in this paper, or explore additional assumptions under which faster rates could be possible.\n\nWhile we agree that all these questions are very interesting (especially the last two), we also believe that our contribution is strong enough on its own right and the paper is quite well-rounded as it is to be worth publishing.\n\nRegarding your additional questions:\n1. **Comparison with Calo et al. (2024):** Indeed our LP is very similar to the one proposed by Calo et al. (2024), the main difference being that our constraints directly feature the marginal occupancy measures of the two chains. This allows us to develop a primal-dual stochastic optimization method that only requires samples from the occupancies. In contrast, primal-dual methods derived from the previous LP formulation would require sampling from the current estimate $\\mu_k$ of the joint occupancy coupling, and querying the transition kernels at these sampled states. This is entirely impractical when having access to trajectory data only. The only possible alternative we can think of would be a simple “plug-in” approach: build an empirical estimate of the two transition kernels based on all observations, and compute the distance between the resulting estimates using a state-of-the-art solver (e.g., Sinkhorn value iteration by Calo et al., 2024). It is not clear at all how the estimation errors would propagate through the analysis; the only analysis we can think of would require the sampling probabilities to be bounded away from zero in every state, and the inverse of the lower bound on the probabilities would appear in the sample complexity bounds. In any case, the resulting method would have the same runtime as our newly proposed algorithm. We are happy to develop this point further in the final version.\n2. **Sampling Cost under Discounting:** Theorem 1 states the sample complexity in terms of i.i.d. samples taken from the occupancy measures, and does not account for the cost of generating such samples. This is done in Appendix B.4. We believe that it should be possible to use all samples between resets and obtain an improved complexity bound, but we did not work out the details (which might be challenging due to the dependence between the samples).\n3. **Normalization and Adapted Wasserstein:** Thank you for catching this! We will make sure to include the normalization constants consistently everywhere in the definitions and the derivations (and in particular add the missing normalization constant to the first definition of Eq. 1).\n4. **Connection to Optimal Control and Dynamic Programming:** Our LP formulation is directly tailored to optimization via stochastic primal-dual methods, which obscures the connection to optimal control a little bit. The LP (and the entire approach) of Calo et al. (2024) makes the connection between the OT problem and the optimal control problem more explicit; we will point this out to the reader in the final version to facilitate navigating the literature."}, {"review_id": "u1ar8QHNwA", "reviewer": "Reviewer_bjRP", "summary": "The goal of this work is to approximate the Wasserstein distance between two distributions over an infinite-dimensional data space, from samples. Specifically, these distributions are discrete-time Markov processes. \n\nThe idea is as follows:\n- Classicaly write the Wasserstein divergence as minimizing a cost over couplings between the two distributions (Eq 1)\n- Decompose the distributions autoregressively using the Markov transition probabilities or occupancy measures (Eq 2)\n- Characterize the occupancy measures by three constraints (Eqs 6-7-8)\n- Dualize these constraints (Eq 9)\n- Solve the resulting min-max problem using a stochastic primal-dual method (Algorithm 1). \n\nThe authors provide computational and sample complexity guarantees (Section 4) and experimental validation (Section 5).", "questions": "Q1. Is this analysis amenable to continuous-time Markov chains, so as to compute the Wasserstein distance between two path measures?\n\nQ2. This work shows how to *estimate* the Wasserstein divergence between two Markov chains from samples, using a primal-dual optimization algorithm to solve a min-max formulation of the divergence. Suppose we additionally wanted *optimize* through the divergence, for example if one Markov chain is the \"ground truth\" and the other Markov chain is a \"parametric model\". Do the authors think  that the extra minimization would be numerically stable?", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths**. The writing is exceptionally clear. The main steps are clearly conveyed (see summary). Furthermore, the writing and notations make the parallels clear between this setup and the finite-dimensional case. It is appreciable to have a a theoretical analysis of the computational and sample complexity and Figure 2.b. is a particularly convincing use-case. \n\n**Weaknesses**. None in particular.", "was_revised": true, "final_justification": "I would like to maintain my score of \"strong accept\". \n\nI believe the authors' contribution is interesting, important and well-explained. I have read through their answers to other reviewers and am satisfied that they are addressing in good faith the issues that were brought up. \n\nIn particular, I saw some concerns on computational complexity and more extensive experiments. I do not believe they are central to this work, whose main contribution is proposing a relatively novel theoretical framework for computing optimal transport distances between Markov Chains and illustrating that it is applicable on reasonable experiments. How to further scale this algorithm would be the topic of another paper in my opinion. Also, in my understanding, there are not many competing algorithms to benchmark against nor is the point to claim a SOTA result. Rather, the experiments serve as a proof-of-concept that this quite theoretical framework can be used in practice.", "rebuttal": "Thank you for your positive evaluation of our work and thoughtful comments!\n\nRegarding your questions: \n\n1. In principle, many of the objects in our formulation (and particularly occupancy couplings) do generalize to continuous time and state spaces, but we believe that adapting our algorithms to such settings would be extremely challenging (at least for us, who are mostly computer scientists). For instance, the constraints in the LP would need to be replaced by differential equations, and we dread to think about the conditions required for Lagrangian duality to hold in this scenario. We leave the investigation of this setting to more qualified colleagues.\n2. Optimizing our distance estimates with respect to parametric models is one of the most exciting questions we would like to explore in the future! Note that there is a clear path towards doing this, at least conceptually: the dual variables at the optimum give the derivatives with respect to the occupancy measures, which themselves are differentiable with respect to the kernels. We have not considered proving guarantees about the accuracy of our estimates of the dual variables so far (which would give an idea about the accuracy of a gradient-based optimization method for optimizing the representations), but this is definitely on our long list of topics to explore! Thank you for the suggestion!"}, {"review_id": "2oKS72iNBQ", "reviewer": "Reviewer_KfLQ", "summary": "This paper proposes a new stochastic optimization method for computing bisimulation distances between finite Markov chains using only sample trajectories, without requiring knowledge of the transition kernels. They extend results from   Calo et al. [2024], where computation of this metric required explicit transition kernels. The method is grounded in a novel linear programming (LP) formulation that characterizes the so-called bicausal couplings, enabling an estimation of the optimal transport distance between Markov processes. The authors introduce a saddle-point stochastic optimization algorithm (SOMCOT), provide rigorous sample complexity guarantees, and validate the approach with experiments. The paper builds upon recent work connecting bisimulation metrics with optimal transport distances and aims to generalize this theory to settings with limited observability.", "questions": "* Q1:  page 3, line 3: Why do the authors refer to the processes as “stationary,” despite having deterministic initial states? Was it meant “ergodic” or “time-homogeneous”?\n\n* Q2: in “Preliminaries” there are notations $\\mathcal{X}^\\infty$ and $\\mathcal{X}^\\mathbb{N}$? Are they the same?\n\n\n* Q3: How does the proposed bicausal coupling compare to classical Markovian couplings? Could this be clarified more explicitly (is one a subset of the other?) (for Markovian vs non-Markovian couplings see e.g. [R1])\n\n\n* Q4: Could the authors provide a concrete example (e.g., from RL or model selection) where computing this distance would lead to a practically useful decision?\n\n\n* Q5: Why are the experiments focused solely on synthetic examples? Have any real-world applications been attempted or considered?\n\n \n* Q6:  The proposed formulation removes the need to sample from transition kernels \\( P_X(\\cdot|x) \\), \\( P_Y(\\cdot|y) \\), and instead relies on samples from the joint distributions over consecutive states. Could the authors clarify what assumptions are made about the availability of such samples? For example, if only a single long trajectory per process is available, under what mixing or subsampling conditions can these empirical joint distributions be used reliably in the stochastic optimization algorithm?\n\n* Q7: Each SOMCOT iteration costs $\\Theta(|X|^2|Y|^2)$ time and $O(|X|^2|Y|^2)$  memory, while the sample-complexity bound scales only with $|X||Y|(|X|+|Y|)$.  For state spaces beyond a few hundred states this becomes prohibitive. Do the authors foresee low-rank, sparse, or neural parameterizations that could make the algorithm practical? A brief discussion would strengthen the experimental narrative.\n\n* Q8: The distance is defined with a discount factor $\\gamma\\in(0,1)$. How sensitive is the estimator to the choice of $\\gamma$in practice? Are there guidelines (or theory) for selecting $\\gamma$ when the horizon of interest is finite or unknown?\n\n[R1] Czumaj, A., Vöcking, B. (2014). Thorp Shuffling, Butterflies, and Non-Markovian Couplings. In: Esparza, J., Fraigniaud, P., Husfeldt, T., Koutsoupias, E. (eds) Automata, Languages, and Programming. ICALP 2014. Lecture Notes in Computer Science, vol 8572. Springer, Berlin, Heidelberg.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n* **Theoretical rigor**: The proposed method is based on a solid and novel LP formulation of bisimulation metrics, with full theoretical guarantees and a clean stochastic primal-dual algorithm.\n\n* **Nontrivial contribution**: Reformulating the problem to eliminate explicit dependence on transition kernels and make it sample-based is meaningful and potentially useful.\n\n* **Relation to OT and RL**: The work connects two active lines of research—optimal transport and bisimulation in RL—and contributes a new theoretical perspective.\n\n\n**Weaknesses:**\n\n* **Unclear practical motivation**: Although the paper highlights real-world relevance (e.g., only having access to sample trajectories), it lacks concrete scenarios where this distance is actually needed. No concrete application or decision-making scenario is provided where this distance is clearly useful.\n\n* **Limited experimental insight**: The experimental section includes illustrative examples (e.g., block Markov chains, parameterized random walks), but it is unclear what questions these experiments are intended to answer or what insight is gained. No direct comparison with competing estimators or baselines is provided.\n\n* **No clear real-world application**: It remains open how this framework would be applied in, for instance, reinforcement learning, model comparison, or structure discovery—despite several such references in the introduction.\n\n\n*Comments*\n* While the paper focuses on optimal transport-based distances between Markov processes, it might benefit from briefly situating itself within the broader literature on estimating or testing properties of Markov chains from a single trajectory. You should e.g. consider mentioning:\n    * [a] Hsu, D., Kontorovich, A., and Szepesvári, C. (2019). Mixing Time Estimation in Reversible Markov Chains from a Single Sample Path. The Annals of Applied Probability, 29(4), 2439–2480.\n    * [b] Daskalakis, C., Dikkala, N., Kamath, G., and Tixier, R. (2022). Identity Testing of Reversible Markov Chains. AISTATS 2022,  151, 1516–1540.\n    * [c] Wolfer, G., and Kontorovich, A. (2019). Estimating the Mixing Time of Ergodic Markov Chains. Proceedings of the 32nd Annual Conference on Learning Theory (COLT), PMLR 99, 1–40.  \n    * [d]  Wolfer, G. (2020). Mixing Time Estimation in Ergodic Markov Chains from a Single Trajectory with Contraction Methods. Proceedings of the 31st International Conference on Algorithmic Learning Theory (ALT), PMLR 117, 1–15.  \n\n* **To clarify**: my main concern relates to the lack of a clearly articulated practical use case for the proposed method. However, since optimal transport and RL  are not my  area of expertise, it's possible that I’m overlooking applications that others may appreciate. If this point is addressed or judged less important by other reviewers, I would be glad to see the paper accepted.", "was_revised": true, "final_justification": "I maintain my borderline accept rating, but acknowledge that other reviewers - - likely more familiar with the application domain -- see stronger potential in this work. The rebuttal’s clarification of the RL motivation, the role of bisimulation metrics in representation learning, and the relationship between bicausal and Markovian couplings strengthen the case, and I believe the paper should be accepted.", "rebuttal": "Thank you for your critical reading of our work and your thoughtful comments!\n\nRegarding your remarks listed under “weaknesses”: we concede that we may have not spent too much effort on explaining the practical relevance of learning bisimulation metrics (a few references aside). This might have been because bisimulation metrics are very popular for the purpose of representation learning in our “home community” of RL, and their usefulness is widely acknowledged in this area—see, e.g., the well-known work of Zhang et al. (2021) cited in our submission. That work clearly demonstrates the usefulness of bisimulation metrics, and in fact has motivated some of us to explore this topic in the first place. Since then, our main focus has been developing theoretically grounded methods for estimating bisimulation metrics, complementing the existing range of works that almost exclusively use heuristics for this purpose. We will highlight this motivation more clearly in the final version of the paper.\n\nAs for the experiments: our purpose was to demonstrate some original ways in which bisimulation metrics could be used for representation learning. We (and other readers) have found these to be interesting. For instance, we would not be able to think of another tractable way of learning an “encoder-decoder” pair illustrated on Figure 1, and we believe that this could be turned into a useful method for representation learning. We do agree though that there is no concrete hypothesis that is being tested in these experiments, and we don’t make claims about beating any other baselines (mostly because there are no existing baselines to compare with). We simply believed that the theoretical claims made in the paper are strong enough to grant us a bit of a creative license with the experiments. We hope that you don’t disagree completely with our stylistic choices.\nIn the final version, we will also make sure to mention other works on testing properties of Markov chains — this literature is indeed relevant! We would concretely highlight the potential of our method for estimating mixing times of Markov chains by exploiting the link between Markovian couplings and mixing times.\n\nFinally, thank you for your detailed comments! We will update the paper to address all of them; we respond briefly to each point below.\n\n- **Q1:** Indeed, this should have been “time-homogenous”.\n- **Q2:** Thanks for catching this, these two notations indeed refer to the same idea.\n- **Q3:** Markovian couplings form a subset of bicausal couplings, see, e.g. the work of Moulos (2021) on the same problem as we consider. Thank you for the pointer to the “Thorp shuffling, etc” paper, it looks very interesting!\n- **Q4:** See our response above.\n- **Q5:** For now, our method is limited to small finite state spaces; we hope to extend it towards larger problems in future work. As noted above, we hoped that the experiments would pique the interest of the reader despite being limited to relatively small problems.\n- **Q6:** As explained in Appendix B.4, obtaining samples from the occupancy measures requires the ability to reset to the initial state. If such resets are not available, optimizing the discounted objective we consider is probably not a good idea in the first place, and it might be better to consider the undiscounted case (roughly corresponding to sending $\\gamma$ to 1). We believe that our analysis can be adapted to this case under commonly considered mixing conditions on the process, but we opted to focus on the discounted case for simplicity here.\n- **Q7:** This is one of the most interesting questions we aim to address in future work. We note that we find it unreasonable to expect that the sample and computational complexity can be improved without further assumptions on the processes. Identifying the correct assumptions that enable such improvements is a very important question, and we indeed have our eyes on structures such as low-rank or sparse transition matrices.\n- **Q8:** The discount factor $\\gamma$ is part of the problem definition, and its “correct” choice always depends on the specifics of the application one has in mind. The same considerations apply as in the broader field of dynamic programming and reinforcement learning: small discount factors make the objective easier to optimize but the solution to be more myopic, whereas large discounts can capture the dynamic more meaningfully at the expense of making the optimization process more challenging. The factor $1/(1-\\gamma)$ can be thought of as the time scale at which the differences between the two chains are effectively taken into account (the “effective horizon”). For time scales that are not known a priori, we believe that it may be more meaningful to consider the undiscounted case mentioned above."}]

Author rebuttal:

Meta-review: This paper discusses the problem of estimating bisimulation metrics (equivalently, adapted/bicausal Wasserstein distances) between Markov chains using only sample trajectories rather than full transition dynamics. The authors develop a stochastic primal-dual algorithm that they term SOMCOT, a linear programming (LP) reformulation that sidesteps the need for explicit transition kernels. The work is technically strong with finite-sample guarantees, though currently restricted to finite state spaces with quadratic computational complexity.

__Strengths__:
- First method to estimate these distances from samples alone - previous work required full transition dynamics
- Novel LP formulation over occupancy couplings enables stochastic optimization
- Tight sample complexity bounds with rigorous theoretical analysis
- Clear writing and mathematical presentation
- Connects optimal transport with RL bisimulation metrics

__Weaknesses__:
- $O(|X|²|Y|²)$ per-iteration cost limits scalability
- Restricted to finite state spaces with no clear extension path
- All experiments synthetic, no real-world demonstrations
- Millions of samples needed even for modest-sized problems

The consensus leans toward acceptance as a solid theoretical contribution that opens a new research direction, despite practical limitations that future work must address.
Decision: Accept (poster)