Differentiable bisimulation metrics improve parametric markov chain fitting methods

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

Machine Learning

Summary

Many tasks like teaching robots by showing examples, simplifying complex systems, and matching models to real-world behaviors need a way to measure how close a model is to the reality it tries to mimic. The authors focus on a special kind of distance called the bisimulation metric, which can also be seen as a way to optimally compare processes step-by-step. They show this measurement can be differentiated mathematically with respect to the model's rules, enabling better model training. Their approach, called Differentiable Bicausal Optimal Transport (D-BOT), alternates between calculating this distance and updating the model to improve it. They demonstrate this method helps in simplifying models, learning model parameters, and imitation learning from observations.

What this means in practice

  • For robotics engineers: Improve imitation learning by better matching robot models to observed behaviors without direct action data.
  • For simulation developers: Enhance parametric model fitting for simulations to produce more accurate and differentiable distance measurements.

Authors

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

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.
View PDFOpen arXiv