Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Ce papier propose SAGE, une méthode d'inférence qui améliore la robustesse des planificateurs de diffusion en rééchantillonnant les actions grâce à un signal de cohérence dynamique appris via une architecture JEPA, sans nécessiter de réentraînement ni de rollouts environnementaux.

L'essentiel

🚀 Le Problème : Le Chef Cuisinier qui rêve trop

Imaginez que vous avez un chef cuisinier robot (c'est le "planificateur par diffusion") qui doit préparer un grand repas complexe. Ce chef a lu des milliers de livres de recettes (les données d'entraînement) et est capable d'inventer des milliers de plats différents en une seconde.

Le problème, c'est que ce chef est un rêveur. Parfois, il imagine un plat magnifique qui a l'air délicieux sur le papier (il a un très bon "score de valeur"), mais qui est impossible à cuisiner dans la réalité.

• Exemple : Il imagine un gâteau qui flotte dans les airs. Le livre de recettes dit "c'est un gâteau de 10/10", mais si vous essayez de le faire, il s'effondre immédiatement parce que la physique ne le permet pas.

Dans le monde de la robotique, c'est pareil : le robot planifie un chemin qui semble parfait pour atteindre son but, mais dès qu'il commence à bouger, il se cogne contre un mur ou tombe, car le premier mouvement de son plan était physiquement impossible.

💡 La Solution : SAGE, le "Garde-Manger" Intelligents

Les auteurs proposent une nouvelle méthode appelée SAGE (Self-supervised Action Gating with Energies).

Pour continuer l'analogie, imaginez que SAGE est un assistant de cuisine très vigilant qui se tient juste à côté du chef. Son travail n'est pas de cuisiner, ni de dire si le plat est bon, mais de vérifier une seule chose : "Est-ce que ce plat est physiquement réalisable avec nos ingrédients et notre four ?"

Comment ça marche ? (L'Analogie de l'Énergie)

1. Le Chef propose (Génération) : Le robot génère 50 idées de mouvements (50 futurs possibles).
2. L'Assistant vérifie (SAGE) : Au lieu de tout cuisiner, l'assistant regarde seulement les 5 premières étapes de chaque idée.
   • Il utilise une sorte de "radar de réalité" (appelé Énergie).
   • Si l'idée est cohérente avec la réalité (le robot peut faire ce mouvement sans tomber), le radar dit : "Bonne énergie, c'est stable".
   • Si l'idée est folle (le robot doit traverser un mur), le radar s'emballe : "Mauvaise énergie, c'est instable !"
3. Le Choix Final (Gating) : L'assistant élimine les idées trop "instables" (celles qui ont une mauvaise énergie). Ensuite, parmi les idées qui restent, on choisit celle qui a le meilleur score de goût (le but à atteindre).

🔍 La Magie : Apprendre sans Maître

Ce qui est génial avec SAGE, c'est qu'il n'a pas besoin qu'on lui explique ce qui est "bien" ou "mal" (pas besoin de récompenses ou de punitions).

• L'Analogie du Miroir : Imaginez que SAGE est un élève qui regarde des milliers d'heures de vidéos de robots qui marchent bien. Il apprend à reconnaître le "rythme" normal de la marche.
• Si un robot essaie de marcher en faisant des sauts de kangourou alors qu'il a des jambes de robot, SAGE dit : "Attends, ce mouvement ne correspond pas au rythme que j'ai vu dans les vidéos. C'est bizarre, c'est de la 'mauvaise énergie'."

Il apprend tout seul, en observant les données, sans jamais avoir besoin de se cogner contre un mur pour comprendre.

🏆 Les Résultats : Plus Robuste, Plus Rapide

Les chercheurs ont testé SAGE sur des robots qui doivent :

• Marcher (comme un chien ou un humain).
• Naviguer dans des labyrinthes complexes.
• Manipuler des objets (ouvrir un four, tourner un robinet).

Résultat : En ajoutant SAGE, les robots deviennent beaucoup moins "fragiles". Ils ne s'arrêtent plus au milieu de leur chemin parce qu'ils ont essayé un mouvement impossible. Ils réussissent mieux leurs tâches, même dans des environnements difficiles, et ce, sans avoir besoin de réapprendre à marcher ou de changer le cerveau du robot principal.

En Résumé

SAGE, c'est comme ajouter un filtre de réalité à un rêveur.

• Le Planificateur (le rêveur) imagine des milliers de futurs brillants.
• SAGE (le garde-fou) vérifie que le premier pas de chaque futur est solide.
• Ensemble, ils permettent au robot de choisir le meilleur chemin qui est aussi possible à parcourir.

C'est une méthode simple, efficace et qui ne demande pas de réapprendre tout le système, ce qui la rend très utile pour faire avancer la robotique de demain ! 🤖✨

Résumé technique

Titre : Amélioration des Planificateurs Diffusion par un Filtrage d'Actions Auto-Supervisé avec Énergies (SAGE)

1. Problématique

Les planificateurs basés sur des modèles de diffusion (Diffusion Planners) sont devenus une approche puissante pour l'apprentissage par renforcement hors ligne (Offline RL). Ils génèrent des trajectoires d'actions futures en modélisant la distribution conditionnelle des états et des actions. Cependant, ces méthodes souffrent d'un mode de défaillance critique :

• Incohérence Dynamique Locale : Les planificateurs sélectionnent souvent des trajectoires qui obtiennent un score élevé selon un modèle de valeur (critique), mais qui sont localement incohérentes avec la dynamique de l'environnement. Autrement dit, le premier pas de la trajectoire planifiée peut être impossible à exécuter depuis l'état actuel, car il ne correspond à aucune transition réaliste observée dans les données d'entraînement.
• Fragilité de l'Exécution : Lorsque l'agent tente d'exécuter un tel plan, il échoue immédiatement, ce qui entraîne une cascade d'erreurs lors des replanifications.
• Limitation des Approches Actuelles : Les méthodes existantes tentent de résoudre ce problème soit en intégrant des contraintes pendant la génération (guidage), ce qui peut déformer la distribution apprise, soit en utilisant des vérificateurs complexes qui nécessitent un réentraînement ou des interactions avec l'environnement.

2. Méthodologie : SAGE (Self-supervised Action Gating with Energies)

SAGE est une méthode de réordonnancement au moment de l'inférence (inference-time re-ranking). Elle ne modifie ni le générateur de diffusion ni le modèle de valeur, mais ajoute une couche de filtrage basée sur la faisabilité locale.

Architecture et Entraînement (Purement Hors Ligne) :
SAGE se compose de deux modules entraînés uniquement sur des données offline :

1. Encodeur JEPA (Joint-Embedding Predictive Architecture) :
   • Un encodeur $e_\theta$ est entraîné sur des séquences d'états pour apprendre une représentation latente.
   • Il utilise une architecture prédictive où l'encodeur prédit l'embedding latent d'un état futur (cible) à partir d'une fenêtre de contexte masquée.
   • Un "enseignant" (Teacher) avec une moyenne mobile exponentielle (EMA) fournit les cibles stables.
   • L'objectif est d'apprendre une dynamique cohérente avec les données sans utiliser de récompenses.
2. Prédicteur Latent Conditionné par l'Action :
   • Un prédicteur $f_\eta$ est entraîné dans l'espace latent figé de l'encodeur JEPA.
   • Il prend un état latent $z_t$ et une action $a_t$ pour prédire l'état latent suivant $\hat{z}_{t+1}$.
   • Il est entraîné avec trois objectifs :
     • Perte Teacher-forced : Prédire précisément le prochain état latent.
     • Perte de déroulement (Rollout) : Assurer la cohérence sur un horizon court.
     • Perte de hinge (Usage de l'action) : Empêcher le modèle d'ignorer l'action en pénalisant les prédictions qui restent précises même si les actions sont mélangées (permutées) dans le batch.

Phase d'Inférence (Gating) :
À chaque étape de décision :

1. Le planificateur diffusion génère $C$ candidats de trajectoires.
2. Pour chaque candidat, SAGE calcule une énergie de cohérence latente sur un préfixe court (les $K$ premières transitions) :
   $$E(\hat{\tau}) = \frac{1}{K} \sum_{k=0}^{K-1} \| f_\eta(z_{t+k}, a_{t+k}) - z_{t+k+1} \|_1$$
   Une énergie faible indique que la transition est cohérente avec la dynamique du dataset.
3. Filtrage et Sélection :
   • On conserve le pourcentage $P$ des candidats ayant l'énergie la plus faible (filtrage de faisabilité).
   • Parmi ces candidats restants, on sélectionne celui qui maximise un score combiné :
     $$i^* = \arg\max_{i} (J(\hat{\tau}^{(i)}) - \lambda E(\hat{\tau}^{(i)}))$$
     Où $J$ est le score de valeur original et $\lambda$ un poids de pénalité.

3. Contributions Clés

• Séparation Faisabilité/Valeur : SAGE traite la faisabilité locale comme un signal distinct de la valeur à long terme, évitant ainsi les conflits d'optimisation où un critique essaie de faire les deux tâches simultanément.
• Apprentissage Auto-Supervisé Sans Récompense : La faisabilité est apprise uniquement via la cohérence prédictive dans l'espace latent, sans nécessiter d'interactions avec l'environnement ni de données de récompense négatives.
• Modularité et Efficacité : SAGE est un module plug-and-play qui s'intègre à n'importe quel planificateur diffusion existant (comme Diffuser ou DV) sans réentraînement du générateur ou du critique.
• Robustesse aux Artefacts : La méthode détecte et élimine les trajectoires "fantômes" qui semblent bonnes en valeur mais sont physiquement irréalisables.

4. Résultats Expérimentaux

Les expériences ont été menées sur la suite de benchmarks D4RL, couvrant la locomotion (MuJoCo), la manipulation (Kitchen) et la navigation (AntMaze, Maze2D).

• Performance Globale : SAGE améliore systématiquement les performances des planificateurs diffusion de pointe (notamment DV - Diffusion Value).
  • Locomotion (MuJoCo) : Amélioration moyenne de +1.5 points par rapport au meilleur planificateur de base (DV).
  • Manipulation (Kitchen) : Améliorations significatives, passant de 73.6 à 74.5 (Mixed) et de 90.0 à 96.6 (Partial).
  • Navigation (AntMaze/Maze2D) : SAGE surpasse toutes les méthodes de base, y compris celles spécifiquement conçues pour la faisabilité (RGG, LoMAP, LDCQ). Par exemple, sur AntMaze, la moyenne passe de 81.6 à 84.5.
• Diagnostic de Faisabilité : Des tests de corruption (brouillage d'actions dans une fenêtre temporelle) montrent que l'énergie de SAGE détecte localement ces incohérences avec une précision exceptionnelle (AUROC > 0.94 sur tous les domaines).
• Efficacité Computationnelle : L'ajout de SAGE n'entraîne qu'une surcharge d'inférence minime (~6.8%), car il ne nécessite que quelques évaluations d'encodeur/prédicteur sur un préfixe court.
• Significativité Statistique : Les gains sont statistiquement significatifs (p-value < 0.05) sur la majorité des tâches, confirmant que l'amélioration n'est pas due au hasard.

5. Signification et Impact

Cet article propose un changement de paradigme dans la conception des planificateurs hors ligne. Au lieu de tenter de forcer le générateur à produire uniquement des trajectoires faisables (ce qui est difficile et coûteux), SAGE propose de séparer la génération de la sélection.

• Il permet d'utiliser des générateurs de diffusion puissants et flexibles tout en garantissant que les actions exécutées sont ancrées dans la réalité des données observées.
• La méthode offre une voie pratique et évolutive pour rendre l'apprentissage par renforcement hors ligne plus robuste, en particulier pour des tâches à long horizon où les erreurs de dynamique s'accumulent rapidement.
• SAGE démontre que l'utilisation de l'apprentissage auto-supervisé pour la cohérence dynamique est une alternative supérieure aux méthodes de vérification explicites ou de guidage complexe pendant la génération.

En résumé, SAGE améliore la fiabilité des agents autonomes en agissant comme un "gardien" intelligent qui rejette les plans séduisants mais irréalistes, sans alourdir le processus d'entraînement ni nécessiter de nouvelles interactions avec l'environnement.