Differentiating Bisimulation Metrics: A Framework for Parametric Markov Chain Fitting via Bicausal Optimal Transport
On the map: World models → Latent predictive world models also fits: Latent dynamics models and model-based RL
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Latent predictive world models, one of 11 directions in World models. Builds on 14 of its 25 references that are in the map. First posted on arXiv 29 Sep 2026.
Builds on
14 of 25 references.
References · 14
- 2018Behavioral Cloning from Observationcited by 482
- 2020Learning Invariant Representations for Reinforcement Learning without Reconstructioncited by 424
- 2019Imitation Learning via Off-Policy Distribution Matchingcited by 182
- 2019Scalable methods for computing state similarity in deterministic Markov Decision Processescited by 133
- 2020Primal Wasserstein Imitation Learningcited by 120
- 2019Provably Efficient Imitation Learning from Observation Alonecited by 78
- 2023Optimal Transport for Offline Imitation Learningcited by 42
- 2022Approximate Policy Iteration with Bisimulation Metricscited by 11
- 2024A Dual Approach to Imitation Learning from Observations with Offline Datasetscited by 10
- 2020Stochastic Optimization for Regularized Wasserstein Estimatorscited by 9
- 2023Imitation Learning from Observation through Optimal Transportcited by 7
- 2024Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficientlycited by 6
2 more
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2020Provable Representation Learning for Imitation Learning via Bi-level Optimization
- 2021Provably Breaking the Quadratic Error Compounding Barrier in Imitation Learning, Optimally
- 2025Learning Causal States Under Partial Observability and Perturbation
- 2026Optimal Transport Meets Reinforcement Learning: A Survey
- 2025CCDP: Composition of Conditional Diffusion Policies with Guided Sampling
- 2025Inverse Q-Learning Done Right: Offline Imitation Learning in Qπ-Realizable MDPs
- 2026Conditional Flow Matching for Transport Between Markov Processes
- 2025Dynamic Rank Adjustment in Diffusion Policies for Efficient and Flexible Training