Differentiating Bisimulation Metrics: A Framework for Parametric Markov Chain Fitting via Bicausal Optimal Transport

Sergio Calo; Amy Zhang; Javier Segovia-Aguas; Anders Jonsson

cs.LGSep 29, 2026

Differentiating Bisimulation Metrics: A Framework for Parametric Markov Chain Fitting via Bicausal Optimal Transport

Authors: Sergio Calo, Amy Zhang, Javier Segovia-Aguas, Anders Jonsson

Organizations: Universitat Pompeu Fabra, Barcelona, Spain · University of Texas at Austin

Abstract

Many problems in sequential decision-making, such as imitation learning from observations, state-space compression, world-model learning, and sim-to-real transfer, can be reduced to learning a model such that a notion of distance with respect to the target process is minimized. We consider this general framework and consider the bisimulation metric, equivalently Bicausal Optimal Transport (BOT), as the notion of distance to minimize. We show that BOT, since it can be formulated as a linear program (LP), is differentiable with respect to the model dynamics. We then derive an exact closed-form gradient via the envelope theorem applied to the LP saddle point. The result is a general algorithm, Differentiable Bicausal Optimal Transport (D-BOT), that can be applied to each of the problems above. The proposed algorithm learns the best model by alternating between distance computation and gradient steps. We apply D-BOT for three different settings: state-space compression, parametric model learning, and imitation learning from observations (ILfO). We show empirical results that confirm the viability of all three instantiations.

Figures & tables

Figure 1: Model learning ( $n_{Y}=n_{X}=16$ , $f_{\mathrm{true}}=1.5$ , $\gamma=0.95$ ). Left : learned drift $f$ versus outer iteration $k$ ; dashed red line marks the true value. Right : SOMCOT estimate of $d_{\gamma}(\mathcal{M}_{X},\mathcal{M}_{\theta})$ .

Figure 2: State-space compression ( $n_{X}=25$ , $n_{Y}=5$ , $\gamma=0.95$ ): bisimulation distance $d_{\gamma}(\mathcal{M}_{X},\mathcal{M}_{\theta})$ versus outer iterations. The 5-state abstract chain faithfully represents the original 25-state chain.

Figure 3: ILfO results with a stochastic expert policy.

Figure 4: ILfO results with a deterministic expert policy.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

(\bar{\lambda}_{Y,k},\,\bar{\alpha}_{Y,k})\leftarrow\textsc{SOMCOT}\!\left(\mathcal{M}_{E},\,\mathcal{M}_{\pi_{\theta_{k-1}}},\,c\right)

Appendix

Algorithm 3 D-BOT-ILfO : ILfO via Bicausal OT

Figure 5: Recovered soft encoder $\lambda_{X}(y|x)$ at convergence on the block-chain environment ( $n_{X}=25$ , $n_{Y}=5$ , $\gamma=0.95$ ). Rows correspond to original states $x\in\mathcal{X}$ ; columns to abstract states $y\in\mathcal{Y}$ . The encoder assigns each group of five original states to a distinct abstract state, recovering the true block structure.

Figure 6: Effect of the outer learning rate $\eta_{\mathrm{out}}$ on the convergence of D-BOT-Compress (block-chain, $n_{Y}=5$ , $\gamma=0.95$ ). Moderate values converge reliably; an excessively large rate introduces oscillations.

Figure 7: Effect of the abstract space size $n_{Y}$ on compression quality (block-chain, $n_{X}=25=5\times 5$ , true latent dim $=5$ , $\gamma=0.95$ ). Left: Final bisimulation distance $d_{\gamma}(\mathcal{M}_{X},\mathcal{M}_{\theta^{*}})$ versus $n_{Y}$ . The distance drops sharply from the under-parametrized regime ( $n_{Y}<5$ ) to the true latent dimension ( $n_{Y}=5$ , orange marker), then remains flat as $n_{Y}$ increases, this confirms that over-parametrization does not hurt the output.

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 296 papers

Optimal Transport Meets Reinforcement Learning: A Survey
Oct 1, 2026Yujie Zhu, Charles A. Hepburn, Matthew Thorpe +1Reinforcement Learning Optimal Transport
University of Warwick
Scalable Bi-causal Optimal Transport via KL Relaxation and Policy Gradients
May 17, 2026Haoyang Cao, Jesse Hoekstra, Renyuan Xu +2Stochastic Optimization Optimal Transport
Department of Applied Mathematics and Statistics, Data Science and AI Institute, and Mathematical Institute for Data Science, Johns Hopkins University. · Department of Statistics, Oxford-Man Institute of Quantitative Finance, and Nuffield College, University of Oxford. · Management Science & Engineering Department, Stanford University. +2
Task Inference Beyond Least Squares in Behavioral Foundation Models
Oct 4, 2026Kuan-Hsun Tu, Chien-Sheng Chiang, Hsin-Wei Chen +2Reinforcement Learning Multi-Task RL
National Taiwan University

Oct 1, 2026J/K move · Enter open · S save

Optimal Transport Meets Reinforcement Learning: A Survey

Yujie Zhu, Charles A. Hepburn, Matthew Thorpe, Giovanni Montana

University of Warwick

Reinforcement Learning Optimal Transport

Open paper AskPDF arXiv

Reinforcement learning (RL) algorithms frequently compare probability distributions, such as state visitation distributions induced by policies and experts, action distributions from learned policies and offline datasets, or transition distributions from learned models and environments. However, commonly used divergences may become ineffective when these distributions overlap weakly, which is frequently encountered in imitation learning, offline RL, and deployment under distribution shift. Optimal transport (OT) offers an alternative by measuring the cost of moving probability mass from one distribution to another under a ground cost that encodes task geometry. This survey covers how OT is used inside RL objectives and algorithms. For each method, we identify: the role OT plays, the distributions compared, the OT formulation used, and the treatment of temporal structure. Beyond categorising existing methods, we discuss the motivations behind different OT choices, practical considerations such as cost design and computational challenges, and highlight open problems including scalable trajectory-level transport, principled handling of mass mismatch, and theoretical analysis for OT-regularised RL.

Abstract

Optimal Transport

Explore similar work

Optimal Transport Meets Reinforcement Learning: A Survey

Scalable Bi-causal Optimal Transport via KL Relaxation and Policy Gradients

Task Inference Beyond Least Squares in Behavioral Foundation Models