A McKean-Pontrygin maximum principle for entropic-regularized optimal transport

1. Problématique

L'article s'attaque au problème du transport optimal dynamique entre deux distributions de probabilité données, $\pi_0$ et $\pi_T$, définies sur un espace d'état $x \in \mathbb{R}^{d_x}$. L'objectif est de trouver un contrôle $U_t$ qui guide un système stochastique d'un état initial à un état final tout en minimisant un coût quadratique.

Le système est régi par l'équation différentielle stochastique (EDS) contrôlée suivante :
$$d\tilde{X}_t = U_t dt + \sqrt{2}\Sigma^{1/2} d\tilde{B}_t$$
où $\tilde{B}_t$ est un mouvement brownien et $\Sigma$ est une matrice de diffusion.

Le coût à minimiser est :
$$J(U) = \frac{1}{2} \int_0^T \mathbb{E}[\|U_t\|_R^2] dt$$
avec $\|u\|_R^2 = u^T R^{-1} u$, où $R$ est une matrice définie positive.

Ce cadre général englobe plusieurs problèmes classiques :

• Transport optimal classique : lorsque $\Sigma = 0$ et $R = I$.
• Pont de Schrödinger : lorsque $\Sigma = R = I$ (correspondant à une régularisation entropique).
• Dynamiques de Langevin sous-amorties : via des choix spécifiques de $b(x)$, $G$ et $\Sigma$.

L'objectif principal est de formuler une approche de champ moyen (mean-field) pour ce problème, basée sur une extension du principe du maximum de Pontryagin, connue sous le nom de principe de McKean–Pontryagin.

2. Méthodologie

L'auteur propose une formulation variationnelle qui évite l'échantillonnage direct sur les trajectoires stochastiques, se concentrant plutôt sur des équations déterministes pour les états et les co-états.

A. Formulation Variationnelle de Champ Moyen

Au lieu de traiter directement les trajectoires stochastiques $\tilde{X}_t$, l'auteur introduit des états $X_t(a)$ et des co-états (co-variables) $P_t(a)$, où $a \in \mathbb{R}^{d_x}$ est une étiquette (label) constante servant de variable indépendante. La loi de $X_t(a)$ est notée $\rho_t$.

Le problème est formulé via un fonctionnel d'action $S$ défini sur l'espace des états, co-états, contrôles et potentiels :
$$S = \int_0^T \int_{\mathbb{R}^{d_x}} \left( \langle P_t, \dot{X}_t - U_t \rangle - \Sigma : D_x^2 \psi_t(X_t) + \frac{1}{2}\|U_t\|_R^2 \right) \rho_0(a) da dt + \text{termes de bord}$$
Ce fonctionnel inclut des multiplicateurs de Lagrange ($\beta_t, \psi_t$) pour imposer les contraintes dynamiques et les conditions aux limites sur les densités marginales $\pi_0$ et $\pi_T$.

B. Principe de McKean–Pontryagin

En appliquant le calcul des variations au fonctionnel d'action, l'auteur dérive un système d'équations couplées (équations d'Euler-Lagrange) qui constituent le principe de McKean–Pontryagin :

1. Équations d'évolution (Hamiltoniennes) :
   • $\dot{X}_t = U_t + \beta_t$
   • $-\dot{P}_t = \nabla_x(\Sigma : D_x^2 \psi_t(X_t)) - D_x^2 \psi_t(X_t) \beta_t$
2. Contraintes et lois de contrôle :
   • $P_t = \nabla_x \psi_t(X_t)$ (Relation entre co-état et potentiel)
   • $U_t = R \nabla_x \psi_t(X_t)$ (Loi de contrôle optimale)
   • Une équation de continuité pour la densité $\rho_t$ impliquant $\beta_t$.

Un point crucial est la liberté de jauge dans le choix du terme $\beta_t$. L'auteur montre que la loi de contrôle optimale $u_t(x) = R \nabla_x \psi_t(x)$ est indépendante du choix spécifique de $\beta_t$, tant que les conditions aux limites sont respectées. Un choix naturel est $\beta_t = -\Sigma \nabla_x \log \rho_t(X_t)$.

C. Connexion aux Équations de Hamilton-Jacobi-Bellman (HJB)

Le système dérivé montre que le potentiel $\psi_t(x)$ satisfait l'équation de Hamilton-Jacobi-Bellman associée au contrôle stochastique :
$$-\partial_t \psi_t = \mathcal{L}_{u_t} \psi_t - \frac{1}{2}\|u_t\|_R^2$$
où $\mathcal{L}_{u_t}$ est le générateur du processus. De plus, la densité $\rho_t$ satisfait l'équation de Fokker-Planck. Cela établit un lien direct avec la formulation classique du pont de Schrödinger.

3. Résultats Clés

• Unification : La méthodologie offre un traitement unifié des problèmes de transport déterministes et stochastiques (régularisés par entropie) dans un cadre de champ moyen.
• Équivalence avec les formulations existantes :
  • Pour $\Sigma \to 0$, on retrouve la formulation de Benamou-Brenier.
  • Pour $\Sigma = R = I$, on obtient une formulation de champ moyen déterministe du pont de Schrödinger dynamique.
  • Le potentiel $\psi_t$ est relié aux potentiels de Schrödinger classiques $\phi_t$ et $\hat{\phi}_t$ par $\hat{\phi}_t = e^{\psi_t}$ et $\rho_t = \phi_t \hat{\phi}_t$.
• Approche Non-Variationnelle et Stochastique :
  • L'article discute une approche lagrangienne non variationnelle (inspirée de [12]) et une formulation stochastique utilisant des intégrales de Stratonovitch ou Itô.
  • Il est démontré que bien que la formulation stochastique (basée sur des EDS couplées forward-backward) soit valide, l'approche variationnelle déterministe proposée (équations (12)-(13)) est préférable car elle ne nécessite pas d'échantillonnage de trajectoires stochastiques et conduit à des approximations moins bruitées.
• Extension au contrôle de la diffusion : L'auteur propose une extension où la matrice de diffusion $\Sigma$ est elle-même un terme de contrôle, pénalisé par une norme de Frobenius. Cela conduit à une pénalité de régularité sur la convexité du potentiel $\psi_t$.

4. Contributions Principales

1. Extension du Principe de Pontryagin : Généralisation du principe de McKean-Pontryagin (introduit précédemment dans [5]) au problème de couplage général avec régularisation entropique.
2. Évitement de l'échantillonnage stochastique : La méthode repose sur des équations différentielles ordinaires (EDO) pour les états et co-états dans l'espace des labels, évitant ainsi la simulation coûteuse de nombreuses trajectoires stochastiques.
3. Cadre Variationnel Unifié : Fournit une base variationnelle rigoureuse pour les équations de mouvement contraintes, reliant directement les conditions aux limites sur les densités aux potentiels de contrôle.
4. Flexibilité de la Jauge : Démonstration que le terme de bruit artificiel $\beta_t$ dans les équations de mouvement peut être choisi arbitrairement sans affecter la solution optimale, offrant une flexibilité pour les implémentations numériques.

5. Signification et Perspectives

Ce travail est significatif car il propose une alternative puissante aux méthodes basées sur les équations différentielles stochastiques rétrogrades (FBSDE) pour le transport optimal et le pont de Schrödinger.

• Avantages Numériques : En se basant sur des EDO de champ moyen plutôt que sur des FBSDE, la méthode promet des approximations numériques moins bruitées et potentiellement plus stables.
• Applications Potentielles : Le cadre est suffisamment général pour inclure des dynamiques de Langevin sous-amorties, des flux de gradient de Wasserstein-Kalman, et des problèmes de dynamique des fluides barotropes (en posant $\Sigma=0$).
• Travaux Futurs : L'auteur suggère que des comparaisons numériques détaillées entre les méthodes de champ moyen (EDO) et les méthodes stochastiques (FBSDE) sont nécessaires pour valider l'efficacité computationnelle de cette approche dans des applications pratiques.

En résumé, Sebastian Reich présente une reformulation élégante et unifiée du transport optimal régularisé, transformant un problème stochastique complexe en un système déterministe de champ moyen via le principe de McKean-Pontryagin, ouvrant la voie à de nouvelles stratégies de résolution numérique.