Recovering Hidden Reward in Diffusion-Based Policies

1. Énoncé du problème

Les politiques basées sur la diffusion sont devenues un paradigme de pointe pour l'apprentissage de compétences de manipulation robotique à partir de démonstrations d'experts, grâce à leur capacité à modéliser des distributions d'actions complexes et multi-modales. Cependant, ces politiques reposent généralement sur le Clonage Comportemental (BC), qui apprend à imiter des actions sans modéliser explicitement l'intention sous-jacente ou la structure de récompense.

• Limites de la diffusion standard : Elles manquent d'une représentation énergétique explicite, ce qui les rend sujettes à l'échec lorsque les situations de test s'écartent de la distribution d'entraînement (Hors Distribution ou OOD). Elles ne peuvent pas facilement fournir un signal de récompense pour l'apprentissage par renforcement (RL) en aval sans recourir à un entraînement adversarial instable ou à un Apprentissage par Renforcement Inverse (IRL) coûteux en calcul.
• Le vide : Il existe un besoin d'un cadre unifiant la modélisation générative des actions avec la récupération de récompense, permettant à la fois une imitation de haute qualité et l'extraction d'un signal de récompense valide pour l'affinement de la politique, tout en garantissant des garanties théoriques sur la généralisation.

2. Méthodologie : ENERGYFLOW

Les auteurs proposent ENERGYFLOW, un cadre qui paramétrise une fonction d'énergie scalaire $E_\theta(s, a)$ dont le gradient sert de champ de débruitage pour la génération d'actions.

Architecture de base

• Paramétrisation de l'énergie scalaire : Contrairement aux politiques de diffusion standard qui régressent directement une prédiction de bruit vectorielle $\epsilon_\theta(s, a)$, ENERGYFLOW apprend une fonction d'énergie scalaire $E_\theta(a, s, t)$.
• Champ vectoriel conservatif : La fonction de score (champ de débruitage) est dérivée par différenciation automatique :
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  Par construction, cela garantit que le champ vectoriel est conservatif (sans rotationnel, $\nabla \times S = 0$), ce qui signifie qu'il est le gradient d'un potentiel scalaire. Cela impose l'intégrabilité, assurant que les préférences d'actions sont transitives et physiquement réalisables.
• Entraînement : Le modèle est entraîné en utilisant l'Appariement de Score de Débruitage (DSM). L'objectif minimise la différence entre le gradient appris et la vraie direction du bruit :
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• Inférence : Les actions sont générées en intégrant l'EDO de flux de probabilité en utilisant le gradient de la fonction d'énergie apprise.

Extraction de récompense

Le cadre extrait un signal de récompense directement à partir du paysage d'énergie appris. Pour gérer la constante d'intégration dépendante de l'état $c(s)$ (qui rend les valeurs d'énergie brutes ambiguës entre différents états), les auteurs proposent un Façonnage Centré :
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
Cela soustrait une ligne de base dépendante de l'état (estimée par échantillonnage de Monte Carlo), normalisant la récompense pour refléter uniquement les préférences d'actions relatives au sein d'un état.

3. Contributions théoriques clés

Le papier fournit des preuves théoriques rigoureuses établissant le lien entre l'appariement de score de diffusion et la récupération de récompense :

1. Équivalence Score-Récompense : Sous l'hypothèse d'Optimalité à Entropie Maximale, le papier démontre que la fonction de score apprise par les modèles de diffusion est proportionnelle au gradient de la fonction Q douce de l'expert ($\nabla_a Q^* \propto S^*$). Ainsi, l'appariement de score récupère implicitement le gradient de la récompense sans optimisation min-max adversarial.
2. Réduction de complexité via des contraintes conservatrices : Les auteurs démontrent que contraindre l'espace d'hypothèses aux champs conservatifs (gradients de scalaires) réduit strictement la complexité de Rademacher par rapport aux champs vectoriels non contraints. Plus précisément, la borne de complexité pour les champs conservatifs évolue avec la régularité de la représentation ($L$), tandis que les champs non contraints évoluent avec la dimension de sortie ($\sqrt{d}$).
3. Bornes de généralisation OOD : En tirant parti de la complexité réduite, le papier dérive des bornes de généralisation plus serrées pour les scénarios Hors Distribution (OOD). La contrainte conservatrice agit comme un régularisateur géométrique, empêchant le modèle d'apprendre des directions vectorielles arbitraires et le forçant à extrapoler la forme du paysage d'énergie, conduisant à une meilleure robustesse face aux décalages de distribution.
4. Identifiabilité et propagation d'erreur : Le papier établit que si la valeur Q absolue ne peut pas être identifiée de manière unique (en raison de la constante inconnue dépendante de l'état), le classement des actions au sein d'un état est exact. De plus, il borne la propagation des erreurs d'estimation de score sur les préférences d'actions, montrant que les erreurs évoluent linéairement avec la distance entre les actions (continuité de Lipschitz).

4. Résultats expérimentaux

Les auteurs ont évalué ENERGYFLOW sur les benchmarks RoboMimic et Meta-World, ainsi que sur des tâches robotiques réelles.

• Performance d'imitation (RQ1) : ENERGYFLOW a atteint des performances State-of-the-Art (SOTA).
  • RoboMimic : 93,8 % de taux de réussite moyen (contre 91,2 % pour la politique de diffusion).
  • Meta-World : 92,5 % de taux de réussite moyen (contre 90,7 % pour la politique de diffusion).
  • Il surpasse significativement d'autres méthodes basées sur l'énergie (comme le BC implicite) et les bases de référence IRL.
• Déploiement réel (RQ2) : Déployé sur un robot AGIBOT G1 pour les tâches "Bouteille" et "Tiroir". ENERGYFLOW a atteint un taux de réussite de 100 % avec des trajectoires fluides et moins d'hésitations par rapport aux bases de référence.
• Qualité de la récompense pour le RL en aval (RQ3) : Utilisée comme signal de récompense pour entraîner un agent Soft Actor-Critic (SAC) :
  • La récompense "Énergie Centrée" a permis à SAC d'atteindre des performances proches de l'oracle, surpassant significativement les récompenses éparses et les récompenses d'énergie brutes.
  • Elle a fourni une guidance dense, étape par étape, empêchant les plateaux précoces dans l'entraînement.
• Généralisation OOD (RQ4) : Sous des perturbations croissantes de la position initiale, ENERGYFLOW a maintenu des taux de réussite robustes là où les politiques de diffusion et de flux non contraintes se dégradaient significativement. Cela confirme l'avantage théorique de la contrainte conservatrice.
• Efficacité (RQ6) : ENERGYFLOW a atteint des vitesses d'inférence comparables aux politiques de flux non basées sur l'énergie (~10 ms de latence) tout en offrant les avantages d'une fonction d'énergie explicite, évitant les temps d'échantillonnage lents des Modèles Basés sur l'Énergie (EBM) traditionnels.

5. Signification et impact

• Unification de l'imitation et du RL : ENERGYFLOW comble le fossé entre le Clonage Comportemental et l'Apprentissage par Renforcement. Il permet à une politique d'être entraînée par imitation tout en exposant simultanément un signal de récompense valide pour un affinement ultérieur, éliminant le besoin de conception manuelle de récompense ou d'IRL adversarial instable.
• Garanties théoriques pour la diffusion : Ce travail fournit la première justification théorique de l'intérêt à imposer un champ conservatif (intégrabilité) dans les politiques de diffusion, le reliant à une complexité d'hypothèse réduite et à une meilleure généralisation OOD.
• Utilité pratique : La méthode est efficace en calcul, évolutive vers des espaces d'actions de haute dimension et efficace dans des environnements robotiques réels, offrant une solution robuste pour l'apprentissage de compétences de manipulation complexes à partir de démonstrations limitées.

En résumé, ENERGYFLOW démontre que les contraintes structurelles requises pour une extraction valide de récompense (champs conservatifs) servent simultanément de biais inductifs puissants qui améliorent la généralisation et la robustesse des politiques basées sur la diffusion.