Evolving Diffusion and Flow Matching Policies for Online Reinforcement Learning

Le papier présente GoRL, un cadre d'apprentissage par renforcement en ligne qui stabilise l'entraînement de politiques génératives expressives en découplant l'optimisation dans un espace latent tractable de la synthèse d'actions, surpassant ainsi les méthodes existantes sur des tâches de contrôle continu complexes.

L'essentiel

Le Problème : Le Dilemme du Chef Cuisinier

Imaginez que vous essayez d'apprendre à un robot à jouer à un jeu vidéo complexe ou à marcher comme un humain. Pour cela, le robot doit prendre des décisions (des actions) à chaque instant.

Il existe deux façons principales d'enseigner ces décisions, et elles ont chacune un gros défaut :

1. La méthode "Sûre mais Bête" (Les Gaussiennes) :
   C'est comme si le robot avait une seule recette de base. Il sait faire un mouvement, mais il ne peut imaginer que des variations très proches de cette recette. C'est stable et facile à apprendre, mais si le jeu demande de faire deux choses très différentes et opposées en même temps (comme sauter à gauche ou à droite selon la situation), le robot va essayer de faire un mouvement "moyen" qui ne fonctionne ni d'un côté ni de l'autre. Il est trop rigide.

2. La méthode "Créative mais Chaotique" (Les Modèles Génératifs) :
   C'est comme donner au robot un chef cuisinier génial capable de créer des millions de recettes complexes et originales. Il peut imaginer des mouvements très précis et variés. Le problème ? Ce chef est très difficile à entraîner en direct. Si vous lui donnez un conseil pendant qu'il cuisine, il se perd, il panique, et tout le plat est gâché. En langage technique, l'apprentissage devient instable et le robot "s'effondre".

Le défi des chercheurs : Comment avoir la créativité du chef génial sans la panique de l'entraînement ?

---

La Solution : GORL (Le Duo Parfait)

Les auteurs proposent une nouvelle méthode appelée GORL. Leur idée brillante est de séparer la direction de la cuisine.

Imaginez une grande entreprise avec deux rôles distincts :

1. Le Chef de Projet (Le "Latent Policy") :
   C'est un manager très simple, calme et rationnel. Il ne cuisine pas. Il ne fait que choisir une "intention" ou un "plan" simple (par exemple : "Aujourd'hui, on veut être agressif" ou "Aujourd'hui, on veut être prudent").

   • Pourquoi c'est bien ? Comme ses plans sont simples, on peut le corriger facilement et rapidement sans qu'il ne s'effondre. C'est la partie stable.

2. Le Chef Cuisinier (Le "Décodeur Génératif") :
   C'est le génie créatif. Il reçoit le plan simple du Chef de Projet et le transforme en une action complexe et précise (le mouvement réel du robot).

   • Pourquoi c'est bien ? Il peut créer des mouvements incroyablement complexes et variés. C'est la partie puissante.

Comment ça marche ? (La Danse en Deux Temps)

Au lieu de faire travailler les deux en même temps (ce qui crée le chaos), GORL utilise une méthode en deux étapes qui se répète :

• Étape 1 : Le Chef de Projet apprend.
  On fige le Chef Cuisinier. Le Chef de Projet essaie de trouver les meilleurs plans simples pour gagner le jeu. Comme le plan est simple, il apprend vite et sûrement.
• Étape 2 : Le Chef Cuisinier s'entraîne.
  On fige le Chef de Projet. On regarde les actions qu'il a produites (les bons plans) et on dit au Chef Cuisinier : "Regarde, c'est ce qu'il faut faire pour réussir". Le Chef Cuisinier s'entraîne à reproduire ces bons résultats, mais en partant d'une base neutre (comme s'il écoutait une radio fixe) pour ne pas se copier lui-même.

Le secret : En alternant ces deux phases, le Chef de Projet devient de plus en plus intelligent, et le Chef Cuisinier devient de plus en plus capable d'exécuter des mouvements complexes. Ils s'améliorent l'un l'autre sans se faire peur.

L'Analogie du "Brouillon" et de la "Peinture Finale"

Pensez à un artiste qui peint un tableau :

• Le Chef de Projet, c'est celui qui fait le croquis au crayon. Il peut effacer, recommencer, et corriger ses erreurs facilement. C'est facile à optimiser.
• Le Chef Cuisinier, c'est celui qui prend ce croquis et le transforme en une peinture à l'huile magnifique et détaillée.
• Avec GORL, on ne demande pas à l'artiste de corriger la peinture à l'huile directement (ce qui serait un désastre). On corrige d'abord le croquis, puis on laisse le peintre mettre à jour sa technique pour mieux suivre les nouveaux croquis.

Les Résultats : Pourquoi c'est impressionnant ?

Les chercheurs ont testé cette méthode sur des robots virtuels (comme un petit robot sauteur appelé "Hopper").

• Les méthodes anciennes (le "Chef unique") arrivaient à peine à faire le robot sauter un peu, puis ils se bloquaient.
• Avec GORL, le robot a appris à sauter de manière incroyable, atteignant des scores 3 fois supérieurs à la meilleure méthode précédente. Il a réussi à trouver des stratégies complexes que les autres méthodes n'osaient même pas imaginer.

En Résumé

GORL résout le vieux problème de l'intelligence artificielle : comment être à la fois stable (ne pas s'effondrer quand on apprend) et créatif (pouvoir faire des choses complexes) ?

La réponse est : Ne faites pas tout faire à la même personne. Séparez la prise de décision simple (stable) de l'exécution complexe (créative), et faites-les travailler en équipe par petits pas. C'est ainsi que l'on obtient des robots qui apprennent vite, sans paniquer, et qui deviennent des champions.

Résumé technique

1. Problématique : Le Dilemme Stabilité-Expressivité

L'apprentissage par renforcement (RL) en ligne pour le contrôle continu repose traditionnellement sur des politiques paramétrées par des distributions unimodales simples (comme des Gaussiennes diagonales). Ces distributions offrent une stabilité d'optimisation grâce à des vraisemblances (likelihoods) analytiques et des gradients lisses, essentiels pour des algorithmes comme PPO ou SAC.

Cependant, ces politiques souffrent d'un goulot d'étranglement en termes d'expressivité. Dans des environnements complexes, la distribution d'actions optimale est souvent multimodale (plusieurs stratégies distinctes peuvent mener à une haute récompense). Une politique unimodale est contrainte de placer une masse de probabilité significative dans les régions à faible récompense entre les modes (effet de "mode-covering"), ce qui entraîne des performances sous-optimales et une fragilité.

Les modèles génératifs récents (Diffusion et Flow Matching) offrent une expressivité supérieure pour modéliser ces distributions multimodales. Pourtant, leur application au RL en ligne se heurte à deux obstacles majeurs :

1. Vraisemblances intraitables : Le calcul de la vraisemblance exacte pour les modèles de diffusion ou de flux est coûteux ou impossible, rendant les mises à jour basées sur le rapport de vraisemblance (comme dans PPO) impraticables.
2. Instabilité des gradients : L'optimisation directe nécessite de rétropropager les gradients à travers de longues chaînes d'échantillonnage (dizaines ou centaines d'étapes de débruitage ou d'ODE), ce qui amplifie la variance et provoque souvent l'effondrement de l'apprentissage sous des distributions de données non stationnaires.

2. Méthodologie : GORL (Generative Online Reinforcement Learning)

Pour résoudre cette tension, les auteurs proposent GORL, un cadre algorithmique agnostique fondé sur un principe structurel clé : découpler l'optimisation de la génération.

A. Factorisation Latente-Générative

GORL décompose la politique $\pi(a|s)$ en deux composants distincts :

1. Un Encodeur (Politique Latente) $\pi_\theta(\varepsilon | s)$ : Une politique tractable (généralement une Gaussienne) qui mappe l'état $s$ vers un espace latent $\varepsilon$. C'est ici que se fait l'optimisation par RL.
2. Un Décodeur Génératif Conditionnel $g_\phi(s, \varepsilon)$ : Un modèle génératif puissant (Diffusion ou Flow Matching) qui mappe le latent $\varepsilon$ vers l'action $a$. Il est responsable de l'expressivité multimodale mais n'est pas directement optimisé par RL.

La politique globale est définie par :
$$\pi(a | s) = \int \pi_\theta(\varepsilon | s) \pi_\phi(a | s, \varepsilon) d\varepsilon$$

B. Algorithme d'Optimisation Alternée (Deux Échelles de Temps)

L'entraînement suit un calendrier alterné pour garantir la stabilité tout en augmentant l'expressivité :

• Phase 1 : Optimisation de l'Encodeur (Fixer $\phi$, Mettre à jour $\theta$)

  • Le décodeur $g_\phi$ est gelé.
  • L'encodeur $\pi_\theta$ est optimisé en utilisant un algorithme RL standard (ex: PPO) dans l'espace latent.
  • Les gradients sont calculés uniquement par rapport à $\theta$, évitant ainsi la rétropropagation à travers la chaîne générative complexe. La distribution d'actions induite est traitée comme une dynamique déterministe (via DDIM ou ODE) pendant la collecte de données.

• Phase 2 : Raffinement du Décodeur (Fixer $\theta$, Mettre à jour $\phi$)

  • L'encodeur est gelé.
  • Le décodeur est raffiné via un apprentissage supervisé (imitation) sur les données de rollouts récents collectés par l'encodeur optimisé.
  • Point Critique (Ancrage à Prior Fixe) : Pour éviter un bouclage de rétroaction où le décodeur apprendrait simplement à reproduire son propre comportement récent (auto-reconstruction), le décodeur est entraîné en utilisant des échantillons de bruit $\varepsilon$ tirés d'un prior fixe (Gaussien standard $N(0, I)$), et non de la politique latente courante. Cela force le décodeur à consolider les progrès d'exploration de l'encodeur en une capacité générative plus forte.

• Réinitialisation par Étapes : À la fin de chaque cycle, l'encodeur est réinitialisé au prior $N(0, I)$. Cela permet de réaligner l'encodeur avec la nouvelle carte de transport du décodeur, assurant une stabilité sur le long terme.

3. Contributions Clés

1. Analyse Théorique de l'Instabilité : Les auteurs identifient et formalisent pourquoi les politiques génératives échouent en RL en ligne (violation des conditions de vraisemblance tractable et de stabilité des gradients).
2. Cadre GORL : Introduction d'une architecture modulaire qui sépare l'optimisation (espace latent tractable) de la génération (espace d'action expressif).
3. Garanties Théoriques : Démonstration que les gradients de politique dans l'espace latent sont des estimateurs non biaisés pour la politique composite et que la divergence dans l'espace latent contrôle la différence de performance dans l'espace d'action.
4. Mécanisme d'Ancrage : Proposition de l'ancrage à un prior fixe pour le raffinement du décodeur, brisant le cycle de rétroaction néfaste et permettant une amélioration itérative de l'expressivité.

4. Résultats Expérimentaux

Les expériences ont été menées sur six tâches de contrôle continu du DMControl Suite (incluant HopperStand, CheetahRun, WalkerWalk, etc.) avec un budget d'interaction de 180M de pas.

• Performance Supérieure : GORL (avec décodeurs Diffusion ou Flow Matching) surpasse systématiquement les politiques Gaussiennes (PPO) et les méthodes génératives existantes (FPO, DPPO).
• Cas HopperStand : Sur cette tâche particulièrement difficile nécessitant un équilibre précis, GORL atteint un retour épisodique moyen de > 870, soit plus de 3 fois le résultat des meilleures baselines (qui plafonnent souvent autour de 286 pour PPO).
• Stabilité : Contrairement aux méthodes génératives directes (comme FPO) qui montrent des effondrements de performance au milieu ou en fin d'entraînement, GORL maintient une courbe d'apprentissage stable et croissante.
• Analyse Qualitative : La visualisation des distributions d'actions montre que GORL évolue d'une forme unimodale initiale vers une structure bimodale claire à la fin de l'entraînement, capturant plusieurs stratégies stables, là où la politique Gaussienne reste bloquée sur un seul mode.
• Compatibilité : Le cadre fonctionne aussi bien avec des algorithmes on-policy (PPO) que off-policy (SAC), prouvant son agnosticisme algorithmique.

5. Signification et Impact

GORL représente une avancée significative en RL en ligne en résolvant le compromis historique entre la stabilité de l'optimisation et la richesse de la modélisation des actions.

• Pratique : Il offre une voie concrète pour déployer des politiques génératives complexes (Diffusion/Flow) dans des scénarios d'apprentissage en ligne, là où elles étaient auparavant considérées comme trop instables.
• Théorique : Il démontre que la séparation structurelle entre l'agent de décision (optimisation) et le générateur de comportements (génération) est une clé pour stabiliser l'apprentissage profond dans des environnements non stationnaires.
• Futur : Ce travail ouvre la porte à l'application de ces techniques sur des tâches robotiques réelles et des environnements à haute dimensionnalité visuelle, où la multimodalité est cruciale pour la robustesse.

En résumé, GORL permet d'obtenir des politiques à la fois stables (grâce à l'optimisation dans un espace latent simple) et hautement expressives (grâce à un décodeur génératif évolutif), dépassant les limites des approches unimodales et des méthodes génératives couplées.