Résumé Technique : BayesFP : Estimation de l'a posteriori pour les politiques de flux via l'échantillonnage de Feynman-Kac

Énoncé du Problème

Les robots opérant dans des environnements réels doivent générer des trajectoires qui soient à la fois assez expressives pour capturer des comportements d'experts multimodaux et assez réactives pour satisfaire des contraintes de sécurité et des objectifs spécifiques à la tâche. Un défi critique surgit lorsque ces contraintes (par exemple, l'évitement d'obstacles, les limites articulaires, les coûts énergétiques) sont inconnues au moment de l'entraînement ou changent entre les déploiements.

Les approches existantes pour la génération de trajectoires contraintes tombent généralement dans deux catégories :

1. Filtrage/Optimisation post-hoc : Des méthodes qui rectifient les déroulements non contraints à l'aide de planificateurs basés sur le gradient ou de filtres de sécurité. Celles-ci découpent souvent la génération de la correction, ceما entraînant des trajectoires hors distribution, des minima locaux et une charge de calcul élevée.
2. Guidage génératif : Des méthodes qui modifient directement le processus d'échantillonnage (par exemple, guidage basé sur un classificateur, projection, dynamique réfléchie). Bien qu'efficaces, elles sont typiquement des modifications heuristiques des dynamiques d'échantillonnage. Elles manquent d'une relation fondée sur des principes avec une distribution a posteriori bien définie qui respecte conjointement la distribution a priori apprise et les contraintes au moment de l'inférence.

La question centrale adressée est la suivante : Étant donné une politique de diffusion ou de correspondance de flux (flow-matching) pré-entraînée, comment peut-on échantillonner des trajectoires qui respectent un problème d'optimisation sous contrainte spécifié uniquement au moment de l'inférence, tout en restant fidèle à la distribution de démonstration apprise ?

Méthodologie

Les auteurs formulent la génération de trajectoires contraintes comme un échantillonnage de l'a posteriori bayésien. La distribution de démonstration apprise $p_{data}(x)$ sert d'a priori, et une fonction de coût $J(x)$ encodant les contraintes au moment de l'inférence sert de vraisemblance. La cible est d'échantillonner à partir de l'a posteriori :
$$p(x | O = 1) \propto p_{data}(x) \exp(\beta J(x))$$
où $O=1$ indique l'optimalité et $\beta < 0$ contrôle la force de l'inclinaison induite par le coût.

Pour échantillonner à partir de cet a posteriori sans réentraîner la politique de base, l'article exploite le cadre du Correcteur de Feynman-Kac (FKC). Les principales contributions méthodologiques sont :

1. Dynamique pondérée de Feynman-Kac

Au lieu d'ajouter heuristiquement des gradients de coût au score ou au champ de vitesse de la politique, les auteurs dérivent une SDE (Équation Différentielle Stochastique) pondérée. Cette SDE régit l'évolution d'une population de particules où chaque particule porte un poids logarithmique $w_t$. Les dynamiques sont dérivées de l'EDP de Fokker-Planck de la distribution cible inclinée de Gibbs.

• Correction de dérive (Drift Correction) : Le terme de dérive inclut une correction sensible au coût proportionnelle à $\nabla J(x)$.
• Évolution des poids : Les poids accumulent un terme de repondération dérivé de l'écart entre les dynamiques simulées et l'évolution de la densité cible. Cela garantit que la population de particules est un estimateur asymptotiquement non biaisé de l'a posteriori, même si la correction de dérive est imparfaite.
• Rééchantillonnage : Pour atténuer la dégénérescence des poids sur de longs horizons, la méthode emploie des techniques de Monte Carlo séquentiel (SMC), spécifiquement le rééchantillonnage systématique sur un intervalle de temps actif.

2. Extension à la correspondance de flux (Flow-Matching) déterministe

Le cadre FKC a été initialement formulé pour les processus de diffusion stochastiques. Les auteurs l'étendent aux politiques de correspondance de flux déterministes (qui résolvent une ODE $dx = -v_t(x)dt$) en recadrant le flux comme une SDE préservant la marge.

• En suivant Singh et Fischer [16], ils introduisent un terme de diffusion auxiliaire $\tilde{\sigma}(t)$ et une correction de dérive basée sur le score correspondante.
• Crucialement, pour les flux rectifiés avec une source gaussienne, le score $\nabla \log q_t(x)$ peut être imputé sous forme fermée à partir du champ de vitesse appris $v_t(x)$, éliminant ainsi le besoin d'un réseau de score séparé.
• Cette construction permet d'appliquer le cadre FKC de manière uniforme aux politiques de diffusion et de correspondance de flux.

3. Garanties théoriques

L'article fournit une garantie théorique (Corollaire 4.2) que lorsque les poids de pénalité et l'amplitude de la température inverse $|\beta|$ augmentent, l'a posteriori de Gibbs se concentre de manière arbitrairement serrée sur des trajectoires qui sont réalisables et proches de l'optimalité. Cela implique que la méthode peut satisfaire les contraintes sans nécessiter de filtres de sécurité externes ou de planificateurs de secours.

Principales Contributions

1. Formulation Fondée sur des Principes : Ce travail reformule la génération de trajectoires contraintes comme une inférence bayésienne, offrant une alternative principled aux guidages heuristiques et aux projections post-hoc.
2. Échantillonneur Unifié : Il instancie l'a posteriori comme une SDE pondérée de Feynman-Kac, étendant le cadre des modèles de diffusion stochastiques aux politiques de correspondance de flux via un recadrage de SDE préservant la marge.
3. Convergence Théorique : Il prouve que l'a posteriori résultant se concentre sur des solutions réalisables et quasi-optimales sous une force de pénalité suffisante.
4. Validation Empirique : L'approche est validée sur plusieurs benchmarks utilisant des modèles pré-entraînés (Diffusion Policy, GR00T-N1.6, et $\pi0.5$), démontrant un succès dans la planification autour d'obstacles non convexes introduits lors de l'inférence.

Résultats Expérimentaux

Les auteurs ont évalué BayesFP sur plusieurs benchmarks, comparant l'approche à des bases telles que le Guidage par Combinaison Linéaire, CASF (Constraint-Aware Streaming Flow) et JM2D (Joint Model-based Model-free Diffusion).

• Tâches de Simulation :
  • Diffusion Policy : Sur les tâches RoboMimic (Can, Transport) avec des obstacles cylindriques ajoutés, BayesFP a atteint un taux de collision de 0 % contre 30 % (Sans Guidage) et 2 % (Combinaison Linéaire).
  • GR00T-N1.6 : Sur les tâches LIBERO-Object avec des obstacles cylindriques et des obstacles en forme de V non convexes, BayesFP a atteint des taux de succès de 84 % et 75 % respectivement, surpassant significativement les bases qui peinent avec les géométries non convexes.
  • $\pi0.5$ : Sur les tâches zero-shot RoboLab (ex: FoodPacking, BBQSauce), BayesFP a amélioré les taux de succès (ex: 36 % contre 28 % pour HammersInLeftBin) et réduit les taux de collision par rapport aux bases de combinaison linéaire.
• Expériences en Monde Réel :
  • En utilisant un bras SO101 affiné avec seulement 20 démonstrations, BayesFP a géré avec succès une tâche de PickAndPlace avec un obstacle (mug) introduit à l'inférence (réduisant la collision de 98 % à 6 %) et une tâche de BimodalMugSelection où une contrainte de langage est renforcée par une contrainte de coût pour supprimer un mode indésirable (atteignant 100 % de sélection du bon mug contre 72 % pour la base).
• Ablations :
  • Augmenter le nombre de particules ($K$) a amélioré les performances, mais la saturation est atteinte autour de $K=4$ pour des obstacles simples.
  • Augmenter la force de guidage $|\beta|$ a conduit à une décroissance exponentielle des violations de contraintes, cohérent avec les prédictions théoriques.

Signification et Revendications

L'article affirme que BayesFP fournit un échantillonneur unifié, au moment de l'inférence et sans réentraînement, pour les politiques de diffusion et de flux. Sa principale signification réside dans :

• Guidage Fondé sur des Principes : Contrairement aux méthodes de combinaison linéaire qui échantillonnent à partir d'une somme de distributions, BayesFP échantillonne à partir du produit réel (a posteriori) de l'a priori et de la vraisemblance du coût, assurant un traitement mathématiquement cohérent des contraintes.
• Gestion de la Non-Convexité : En échantillonnant des trajectoires entières à partir de l'a posteriori plutôt qu'en appliquant des corrections de gradient locales, la méthode gère naturellement les obstacles non convexes et les espaces de solutions multimodaux.
• Généralité : Le cadre est agnostique à l'architecture de la politique de base (diffusion ou flux) et à la nature spécifique des contraintes, à condition qu'elles puissent être encodées dans une fonction de coût dérivable.

Les auteurs reconnaissent des limites, notant un ralentissement de calcul (1,25–2$\times$) dû à la nature séquentielle de la construction de la fonction de distance signée (SDF) et à la surcharge de la simulation de particules, bien que cela soit atténué par l'exécution parallèle sur GPU. Ils notent également la nécessité que la fonction de coût soit dérivable. Des travaux futurs sont suggérés pour étendre le cadre aux politiques de flux en streaming pour un contrôle à faible latence et pour incorporer des représentations de contraintes plus riches.