Contact Wasserstein Geodesic

Code Review Version

About The Project

Code repository for the paper Contact Wasserstein Geodesics for Non-Conservative Schrödinger Bridges.

Contact Wasserstein Geodesics (CWG) provides a flexible framework for computing the most likely stochastic trajectory between two datasets representing discretized probability distributions.

Beyond the classical Schrödinger Bridge problem with a Brownian prior, CWG supports:

• Generalized Schrödinger Bridges, incorporating conservative forces via a potential function.
• Multi-marginal Schrödinger Bridges, using intermediate distributions as additional constraints.
• Non-Conservative Schrödinger Bridges (introduced in the paper), which model non-conservative forces to better capture complex real-world dynamics.

Content

ContactWassersteinGeodesic/
├── cwg/
│ ├── main.py
│ ├── networks/
│ └── utils/
├── data/
├── dataload/
└── params/

The script main.py is used to train the CWG solver or to load a previously trained model.

• The network folder contains the neural architectures and training algorithms.
• The utils folder includes functions for visualization as well as loading potential and guidance functions.
• The dataloader folder provides functions for sampling from the specified datasets.
• The params folder provides the parameters for reproducing the paper results.

The Wasserstein distance between data samples is computed using the Sinkhorn algorithm, implemented via the external library geomloss, in external/geomloss.

Getting Started

This project uses uv to manage dependencies and the virtual environment.

Installation

git clone --recurse-submodules https://github.boschdevcloud.com/TEA1RNG/ContactWassersteinGeodesic.git
cd ContactWassersteinGeodesic
uv sync
uv pip install diffusers transformers

Run the experiments

Proof of Concept - Two Moons Dataset

uv run python -m cwg.main --params poc_moons

Proof of Concept - Checkerboard Dataset

uv run python -m cwg.main --params poc_checkerboard

LiDAR Manifold Navigation Experiment

uv run python -m cwg.main --params lidar

Cell Sequencing Experiment - Embryoid Body (EB) Stem Cell Differentiation

uv run python -m cwg.main --params eb

Image Generation - Sea Temperature Prediction Task

uv run python -m cwg.main --params temperature

Image Generation - Robot Task Reconstruction

uv run python -m cwg.main --params robot

Inference of a previously trained model

uv run python -m cwg.main --params poc_moons --model_name [MODEL_NAME]

Guidance

uv run python -m cwg.main --params poc_moons --guidance True --model_name [MODEL_TO_FINE_TUNE]

uv run python -m cwg.main --params poc_checkerboard --guidance True --model_name [MODEL_TO_FINE_TUNE]

uv run python -m cwg.main --params lidar --guidance True --model_name [MODEL_TO_FINE_TUNE]

Contact

Andrea Testa - email andrea3.testa@gmail.com

Citation

@inproceedings{
testa2026contact,
title={Contact Wasserstein Geodesics for Non-Conservative Schr\"odinger Bridges},
author={Andrea Testa and S{\o}ren Hauberg and Tamim Asfour and Leonel Rozo},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=IaEohEBUgi}
}