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

Summary

The goal of this work is to study optimal transport (OT) distances between pairs of finite Markov chains, providing a novel relation between OT distances and probabilistic bisimulation metrics. The proposed linear program builds on ideas from optimal control in Markov decision processes, and the designed algorithm for solving the proposed linear program combines Sinkhorn's algorithm with an entropy-regularized version of the classic Value Iteration algorithm. Convergence guarantees and computational complexity analysis of the method are provided in the main text and in the appendix.

Strengths

The paper is well-organized, the topic of study is very interesting, and the contribution is novel. Previous works on OT and bisimulation metrics are well-referenced and up-to-date.

Weaknesses

The main theoretical contributions are presented in Section 3, where the authors introduce the definition of occupancy couplings $\mu$. These are distributions over $\mathcal{XY}\times \mathcal{XY}$. It is not entirely clear why they need to duplicate the variables $(x,y)$ and consider quadrupes $(xy,x'y')$.

Questions

- In the abstract, the authors say "In this work, we develop an alternative perspective by considering couplings between a “flattened” version of the joint distributions that we call discounted occupancy couplings [...]" Could you expand on what do you mean by "flatten"? - In the introduction, the authors say that "The possibility that the objects in question may be random further complicates the picture, and in such cases it becomes more natural to measure distances between the underlying joint probability distributions." Could you explain/expand? - Section 2.1: Does $\gamma$ depend on $t$? Is $\gamma^t\in (0,1)$? - Line 99, if Markov chains are considered, isn't it that $M_{\mathcal{X}}$ $(x_n | \bar{x}_{n-1})$ $=$ $M_{\mathcal{X}}$ $(x_n|x_{n-1})$ ? - Line 148 and equation (4): Which is the dependence of $d_\gamma$ on the labeling function $r$? Line 149: Which is the dependence of the set $\mathcal{F}_\gamma$ on $\gamma$? It is clear from eq. (5) that $U^*$ does explicitly depend on $r$ and $\gamma$. Could you expand on the equality $U^*=d_\gamma$? - Explain the pevious to last equality in (6). How do you make the variables $x',y'$ to apprear? - Explain the last equality in (6). Why do you have $\langle \mu^\pi,c\rangle$ if c depends on $(x,y)$ and $\mu^\pi$ depends on $x,y,x',y'$? Comments: - The term "discounted occupancy couplings" mentioned in the abtract is not excatly used in the main text, the authors mainly use just "occupancy couplings". - $n$ is used for the size $|\mathcal{Y}|$ (line 87) and as an index. Isn't it $|\mathcal{X}|=m=|\mathcal{Y}|$? In the experiments, why the choice of $n$ is not analyzed too? Typos: - Abstract: line 17, "method that we call", instead of "method we call". - Notations: line 79, "the corresponding finite subsequences", instead of "the corresponding subsequences" - Line 39 (Introduction): "important" is repeated twice. - Notations: Add the definitions of $\mathcal{X}^\infty$, $\mathcal{X}^n$, $\Delta_{\mathcal{X}}$.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This submission Namely, the authors define a notion of optimal transport distance between Markov chains on state spaces with a ground metric. This notion of distance differs from standard optimal transport, as the set of couplings is restricted to the set of so-called bicausal couplings. Using the results of Moulos [2021], it is shown that this Markovian optimal transport can be characterized in terms of the solution of the Bellman optimality equations for a Markov decision process. Using this equivalence, it is demonstrated that, in certain cases, bisimulation metrics can be though of, equivalently, as Markovian optimal transport with a specific cost function by exploiting their connection to solutions of certain fixed-point equations. Next, it is shown that the optimal value for the derived Markovian optimal transport problem can be computed by solving a finite-dimensional linear program, where the constraint set consists of the intersection of three sets. Rather than directly solving this linear program, the authors propose to regularize the problem using a conditional entropy (by analogy with entropic regularization of optimal transport distances which can then be solved efficiently using Sinkhorn iterations). To solve the entropy regularized problem, the Sinkhorn Value Iteration algorithm is proposed and the number of steps required to obtain a desired accuracy in estimating the Markovian optimal transport distance is provided. The authors also propose an analyze an alternative algorithm which they dub Sinkhorn Policy Iteration. The paper concludes with some numerical experiments to illustrate the performance of the proposed algorithms.

Strengths

In my opinion, the submission is well-written and its contributions relative to the broader literature are clearly identified. The main contributions of this work, identifying bisimulation metrics as a type of optimal transport problem and providing some new algorithms for estimating bisimulation metrics by using this connection is of interest and is, to my knowledge, novel.

Weaknesses

1. It appears that the connection between the Markovian optimal transport problem and the finite-dimensional linear program provided in Theorem 1 only enables the computation of the optimal value for optimal transport problem, but does not allow the recovery of the optimal bicausal coupling. If this is the case, this should be further clarified in the text. 2. While the analysis for the Sinkhorn Value Iteration is a nice addition, the guarantees are a bit confusing. Notably, if we wish to attain a precision of $\epsilon$ in estimating the Markovian optimal transport one requires $K=O(1/\epsilon^2)$ iterations and setting the regularization parameter to $\eta =C/\sqrt K=O(\epsilon)$. However, if epsilon is sufficiently small, the objective in the update (10) will be dominated by the entropy which is minimized by $\mu=\mu_k$. Naively I would assume that $1/\eta$ should be small so that minimization of the linear term dominates.

Questions

Apart from the points mentioned above, it would be helpful for the authors to state the complexity of the Sinkhorn Value Iteration algorithm in the main text (rather than only in the appendix).

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors address the limitations of their work in the discussion.

Summary

This work integrates optimal transport with Markov chains by proposing an alternative joint distribution between Markov processes, namely "discounted occupancy couplings". They show that optimal transport distances can be computed as a a linear program (LP) in reduced space. This improves the computational efficiency of OT between markov decision process (MDP) over the previous methods. The previous methods, as they reviewed, often requires complex dynamic programming algorithms. They showed that the new formulation can be extended to the well-known entropy regularization problem, employing Sinkhorn-like iterations to solve, making the computation scalable with large problem. This paper provides both theoretical and experimental supports for the new formulation. In the end, they discuss the potential applications on RL models, limitations and challenges.

Strengths

The main contribution of this work is demonstration of optimal transport shares the same formulation as probabilistic bisimulation metrics, which is popular in the practice of reinforcement learning (RL). By establishing that solving for optimal transport is equivalent to computing bisimulation metrics, the authors creates a novel link between RL and optimal transport theory. - The formulation of OT on MDP is novel. This work provides an new prospective of Bisimulation metric between MDPs. - Theoretical analysis looks good. - Two computational solutions (Sinkhorn-like iterations) were proposed and tested in experiments sections.

Weaknesses

The theoretical guarantees are based on the assumption of perfect projection steps $m=\infty$, which is not practical. As stated in line 290, such exact computation is unnecessary in practice. This claim needs a further theoretical support.

Questions

- What is the best practice of bisimulation metric computation? Could your provide the comparison with your methods? - What does the entropy means in the sense of the coupling between MDPs? Some explanations on this will be appreciated - How does this metric works with data contains noise as we know OT is sensitive to outliers?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

None