Distances for Markov chains from sample streams

Thank you for your critical reading of our work and your thoughtful comments!

Regarding 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.

As 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. In 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.

Finally, thank you for your detailed comments! We will update the paper to address all of them; we respond briefly to each point below.

• Q1: Indeed, this should have been “time-homogenous”.
• Q2: Thanks for catching this, these two notations indeed refer to the same idea.
• 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!
• Q4: See our response above.
• 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.
• 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.
• 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.
• 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.

Thank you for your positive evaluation of our work and thoughtful comments!

Regarding your questions:

1. 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.
2. 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!

Thank you for your answer which I agree with. I too remain positive about the paper overall.

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.

1. 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.
2. 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.
3. 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.
4. 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.
5. 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.

While 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.

Regarding your additional questions:

1. 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.
2. 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).
3. 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).
4. 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.

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.

First 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.

Second, 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.

Finally, 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.

Once 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.