Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficiently (2406.04056v2)

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

• The paper demonstrates the equivalence between bisimulation metrics and optimal transport distances, unifying methods in reinforcement learning and transport theory.
• The paper introduces a linear programming reformulation along with the Sinkhorn Value Iteration algorithm to efficiently compute these metrics in Markov chains.
• The paper provides both theoretical convergence guarantees and empirical validation, underlining its significant impact on future research in reinforcement learning and computational optimal transport.

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

The paper "Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficiently" by Sergio Calo, Anders Jonsson, Gergely Neu, Ludovic Schwartz, and Javier Segovia-Aguas presents a novel framework for formulating optimal transport distances between Markov chains (MCs). This work establishes that bisimulation metrics, traditionally studied in theoretical computer science and reinforcement learning (RL), are equivalent to optimal transport distances. Furthermore, the authors propose efficient algorithms for computing these distances, significantly improving computational feasibility.

Key Contributions

1. Equivalence of Bisimulation Metrics and Optimal Transport Distances:
   • The authors demonstrate that bisimulation metrics, which quantify the similarity between states in Markov chains, are intrinsically optimal transport distances. This observation bridges two previously distinct notions, offering a unified perspective with broad implications.
2. Efficient Linear Program (LP) Reformulation:
   • The paper reformulates the computation of optimal transport distances as an LP in the space of discounted occupancy couplings. This reformulation enables the application of several algorithmic techniques from optimal transport theory, making the problem more tractable.
3. Sinkhorn Value Iteration (SVI) Algorithm:
   • A novel algorithm, Sinkhorn Value Iteration, is introduced. Using an entropy-regularized optimization approach, SVI efficiently computes the optimal transport distances. This method combines Sinkhorn's algorithm with the classical Value Iteration, ensuring fast convergence with reduced computational overhead.
4. Theoretical and Empirical Validation:
   • Both theoretical guarantees and empirical studies validate the effectiveness of SVI. Theoretical results showcase the convergence speed and computational efficiency, while empirical experiments demonstrate the practical applicability across various scenarios.

Detailed Insights

Equivalence of Metrics

The equivalence established between bisimulation metrics and optimal transport distances facilitates a new understanding of both concepts. In particular, bisimulation metrics in RL, traditionally used for state aggregation and representation learning, can now be interpreted and computed through the lens of optimal transport theory. This equivalence is formalized by showing that both metrics minimize a common objective, albeit through different formulations.

Linear Program Formulation

The classic dynamic programming approach to computing optimal transport distances often suffers from inefficiency due to its high computational cost. By introducing the concept of discounted occupancy couplings, the authors reformulate the problem as an LP. This reduced-dimensional representation preserves the essential structure while significantly lowering computational complexity.

Sinkhorn Value Iteration Algorithm

The SVI algorithm is a pivotal contribution that leverages entropy regularization to solve the LP efficiently. Here's a breakdown of its components:

• Entropy Regularization: Adds a regularization term to the optimization problem, promoting faster convergence and numerical stability.
• BeLLMan–Sinkhorn Operators: Used in the iterative update process, these operators ensure that the estimates converge to the true optimal coupling.
• Mirror Descent Updates: The algorithm incorporates mirror descent-style updates, ensuring each iteration refines the policy towards the optimal solution.

The theoretical analysis provides convergence guarantees, showing that SVI reaches an $\varepsilon$-optimal solution in a number of iterations dependent on the problem parameters and desired accuracy.

Implications and Future Directions

The implications of this work are manifold:

• Practical Algorithms for RL: The equivalence and the proposed efficient algorithms can significantly impact RL, particularly in state representation and transfer learning, enabling better scaling and more reliable performance.
• Interdisciplinary Connections: By bridging concepts from theoretical computer science and optimal transport, this work fosters interdisciplinary research, encouraging the application of optimal transport methods in other fields where bisimulation metrics are useful.
• Algorithmic Design: The LP formulation and SVI provide a foundation for the design of new algorithms in computational optimal transport, potentially inspiring further innovations in this space.

Future research could explore:

• Extended Applications: Applying the proposed framework and algorithms to other domains, such as formal verification and concurrency theory.
• Stochastic Learning Methods: Developing online or reinforcement learning-based approaches that integrate the LP formulation for scenarios where transition kernels are unknown or partially observable.
• Scalable Implementations: Enhancing scalability and efficiency through distributed computing or sophisticated approximation methods, enabling the analysis of large-scale Markov processes.

Conclusion

The paper presents a significant advancement in understanding and computing similarity metrics between Markov chains. By establishing the equivalence between bisimulation metrics and optimal transport distances and introducing the efficient SVI algorithm, the authors address both theoretical and practical challenges. These contributions hold substantial potential for advancing research and applications across machine learning, formal verification, and beyond.