Exploratory Memory-Augmented LLM Agent via Hybrid On- and Off-Policy Optimization

Ce papier présente EMPO², un cadre d'apprentissage par renforcement hybride intégrant une mémoire pour l'exploration et combinant des mises à jour on- et off-policy, qui améliore significativement les performances et l'adaptabilité des agents LLM sur des tâches complexes et hors distribution.

L'essentiel

🧠 Le Problème : L'IA qui a peur de sortir de sa zone de confort

Imaginez que vous apprenez à cuisiner un plat complexe. Si vous ne faites que relire votre livre de recettes (les connaissances pré-entraînées), vous risquez de rater le plat si vous n'avez pas exactement les mêmes ingrédients que ceux décrits dans le livre.

C'est le problème actuel des "Agents IA" (des intelligences artificielles capables d'agir dans le monde réel ou virtuel). Ils sont très forts pour utiliser ce qu'ils savent déjà, mais ils sont très mauvais pour explorer. S'ils se retrouvent dans une situation nouvelle où ils ne connaissent pas la solution, ils paniquent, répètent les mêmes erreurs et n'arrivent pas à apprendre par essai-erreur. Ils sont comme un touriste qui refuse de sortir de l'hôtel par peur de se perdre, alors que l'aventure se trouve juste à l'extérieur.

💡 La Solution : EMPO2, le "Carnet de Voyage" Intelligent

Les chercheurs de Microsoft et de l'Université KAIST ont créé une nouvelle méthode appelée EMPO2. Pour comprendre comment ça marche, imaginons un agent IA qui apprend à résoudre des énigmes dans un jeu vidéo (comme ScienceWorld ou WebShop).

EMPO2 utilise deux outils magiques en même temps :

1. Le Cerveau (La Mémoire Paramétrique) : C'est le modèle lui-même, ses poids et ses connexions neuronales. C'est comme la mémoire à long terme de votre cerveau.
2. Le Carnet de Notes (La Mémoire Non-Paramétrique) : C'est un petit carnet externe où l'agent écrit ses réflexions après chaque essai. "J'ai essayé de tourner la clé à droite, ça n'a pas marché. La prochaine fois, je vais essayer à gauche."

🎭 Comment EMPO2 apprend : Le Duo Dynamique

La grande innovation, c'est que EMPO2 ne se contente pas d'écrire dans son carnet. Il joue à un jeu de rôle en deux temps, un peu comme un étudiant qui révise pour un examen :

• Phase 1 : L'Exploration avec le Carnet (Hors-Politique)
  L'agent essaie de résoudre le problème en consultant son carnet de notes. Il lit : "Ah oui, la dernière fois, j'ai oublié de chercher la clé sous le tapis." Grâce à ce conseil, il essaie une nouvelle stratégie. S'il réussit, il a gagné !
  L'analogie : C'est comme un élève qui regarde ses anciens devoirs corrigés pour comprendre ses erreurs avant de repasser le test.

• Phase 2 : L'Intériorisation (Sur-Politique)
  C'est là que la magie opère. Au lieu de simplement copier le conseil du carnet, l'agent va apprendre à faire cela tout seul. Il prend les bonnes actions qu'il a trouvées grâce au carnet et les grave directement dans son "cerveau" (ses paramètres).
  L'analogie : C'est comme si l'élève, après avoir lu ses notes, fermait son carnet et s'entraînait à faire le mouvement de mémoire, jusqu'à ce que cela devienne un réflexe naturel. Il n'a plus besoin du carnet pour réussir la prochaine fois.

🚀 Pourquoi c'est génial ? (Les Résultats)

Les chercheurs ont testé cette méthode sur deux terrains d'entraînement :

1. ScienceWorld : Un jeu où l'agent doit faire des expériences scientifiques (ex: allumer une ampoule rouge).
2. WebShop : Un jeu où l'agent doit acheter des objets en ligne en suivant des instructions précises.

Les résultats sont impressionnants :

• Sur ScienceWorld : EMPO2 a été 128 % plus performant que les méthodes précédentes. Là où les autres agents s'arrêtaient, bloqués par une erreur, EMPO2 a continué à explorer, a consulté son carnet, a appris, et a fini par réussir.
• Sur WebShop : Il a aussi gagné 11 % de plus, prouvant qu'il est plus efficace pour naviguer dans des environnements complexes.

🌍 Le Super-Pouvoir : L'Adaptation Instantanée

Le plus beau dans l'histoire, c'est la capacité d'adaptation.
Imaginez que vous apprenez à conduire une voiture, puis on vous donne un camion.

• Les anciennes IA devaient tout réapprendre de zéro.
• EMPO2, lui, utilise son carnet de notes. Il se souvient : "Ah, j'ai appris à freiner avec la voiture, mais pour le camion, il faut plus de distance." Il s'adapte en quelques essais, sans avoir besoin de réapprendre tout son cerveau.

🏁 En Résumé

EMPO2 est comme un explorateur très malin qui ne se contente pas de marcher au hasard.

1. Il note ses erreurs et ses découvertes dans un carnet (mémoire externe).
2. Il utilise ce carnet pour tester de nouvelles idées.
3. Il intègre ce qu'il a appris dans son cerveau pour devenir plus intelligent, même quand il n'a plus son carnet.

C'est une étape majeure pour créer des IA qui ne sont pas seulement de bons élèves qui révisent leurs leçons, mais de vrais aventuriers capables de découvrir de nouveaux mondes et de s'y adapter rapidement.

Résumé technique

1. Problématique : Le goulot d'étranglement de l'exploration

Les agents basés sur les grands modèles de langage (LLM) combinés à l'apprentissage par renforcement (RL) excellent souvent dans l'exploitation de connaissances pré-entraînées, mais ils échouent fréquemment dans des environnements nécessitant la découverte de nouveaux états ou l'acquisition active d'informations inconnues.

• Limites actuelles : Les méthodes existantes (comme GRPO) convergent prématurément vers des solutions sous-optimales car elles manquent de mécanismes d'exploration systématique. Elles peinent à sortir de la distribution des comportements familiers.
• Limites de la mémoire externe : Des approches antérieures comme Reflexion utilisent une mémoire externe non paramétrique pour stocker des réflexions verbales. Cependant, ces méthodes saturent rapidement car elles ne mettent pas à jour les paramètres du modèle. L'agent devient dépendant de la mémoire sans internaliser les connaissances, limitant sa généralisation à long terme.

2. Méthodologie : EMPO2

Les auteurs proposent EMPO2 (Exploratory Memory-Augmented On- and Off-Policy Optimization), un cadre d'apprentissage hybride qui combine la mise à jour des paramètres du modèle (paramétrique) et l'utilisation d'une mémoire externe (non paramétrique).

A. Double mise à jour (Paramétrique et Non-paramétrique)

Le cœur de la méthode réside dans l'interaction entre deux modes de mise à jour :

1. Mise à jour non-paramétrique (Mémoire) : L'agent génère des "conseils" (tips) réflexifs à partir de ses trajectoires passées (succès ou échecs) et les stocke dans un tampon de mémoire. Ces conseils servent de guidage pour les étapes suivantes.
2. Mise à jour paramétrique (RL) : Le modèle apprend à internaliser ces conseils pour améliorer sa politique intrinsèque, réduisant ainsi sa dépendance future à la mémoire externe.

B. Phases Hybrides

EMPO2 opère selon deux phases distinctes, chacune ayant deux modes possibles :

• Phase de Déploiement (Rollout) :

  • Sans mémoire : L'agent agit uniquement sur la base de l'état actuel et de la tâche.
  • Avec mémoire : L'agent récupère des conseils pertinents depuis la mémoire et les utilise comme contexte supplémentaire pour générer ses actions.
  • Sélection : Un échantillonnage probabiliste ($p$) détermine l'usage de la mémoire lors du déploiement.

• Phase de Mise à jour (Update) :

  • Apprentissage On-Policy : Les trajectoires générées avec mémoire sont utilisées pour mettre à jour le modèle en conservant le contexte des conseils (conditionnement identique).
  • Apprentissage Off-Policy (Innovation clé) : Les trajectoires générées avec mémoire sont utilisées pour mettre à jour le modèle, mais sans les conseils lors du calcul de la probabilité de l'action.
    • Mécanisme : Cela fonctionne comme une distillation de connaissances guidée par la récompense. Le modèle "élève" (teacher) a agi avec des conseils, et le modèle "élève" (student) apprend à reproduire ces actions réussies sans les conseils. Cela force le modèle à internaliser la logique de l'exploration directement dans ses poids.

C. Stabilisation et Récompenses Intrinsèques

• Masquage des tokens : Pour éviter l'instabilité lors de l'apprentissage off-policy (dû aux rapports de vraisemblance non bornés), un mécanisme de masquage supprime les termes d'avantage pour les tokens ayant une probabilité trop faible.
• Récompenses Intrinsèques : Une récompense basée sur la nouveauté de l'état (similarité cosinus avec les états visités) est ajoutée pour encourager l'exploration de zones inconnues, même en l'absence de récompense extrinsèque.

3. Contributions Clés

1. Cadre Unifié : EMPO2 est le premier cadre à intégrer simultanément la mémoire externe pour l'exploration immédiate et la mise à jour des paramètres pour l'apprentissage à long terme, créant un pont entre l'optimisation paramétrique et le raisonnement augmenté par la mémoire.
2. Apprentissage Off-Policy avec Mémoire : L'introduction d'une mise à jour off-policy où les conseils sont retirés lors de l'entraînement permet d'internaliser l'efficacité de l'exploration guidée par la mémoire.
3. Robustesse et Généralisation : La méthode permet aux agents de s'adapter à de nouvelles tâches (hors distribution) avec peu d'essais et sans mise à jour des poids, en utilisant simplement la mémoire accumulée.

4. Résultats Expérimentaux

Les expériences ont été menées sur deux environnements complexes nécessitant une exploration : ScienceWorld (expériences scientifiques textuelles) et WebShop (achat en ligne). Le modèle de base utilisé est Qwen2.5-7B-Instruct.

• Performance In-Distribution (ID) :
  • ScienceWorld : EMPO2 dépasse GRPO de 128,6 % en moyenne. Il résout des tâches où GRPO converge vers des solutions sous-optimales (ex: allumer une ampoule rouge en trouvant d'abord l'ampoule).
  • WebShop : Amélioration de 11,3 % par rapport à GRPO, surpassant également les méthodes hors ligne (Retrospex) et d'autres approches en ligne (GiGPO).
• Performance Out-of-Distribution (OOD) :
  • Dans des tests sur de nouvelles tâches, EMPO2 s'adapte rapidement (quelques essais) en utilisant la mémoire, sans mise à jour des paramètres. GRPO, en revanche, montre une grande variabilité et échoue souvent sur des tâches non vues.
• Études d'ablation :
  • La suppression de l'apprentissage off-policy ou de l'apprentissage on-policy avec mémoire dégrade les performances, confirmant la nécessité de l'hybridation.
  • La probabilité d'utilisation de la mémoire ($p$) et de mise à jour off-policy ($q$) a été optimisée pour assurer une convergence stable.

5. Signification et Impact

EMPO2 représente une avancée significative pour le développement d'agents LLM autonomes et généralisables.

• Efficacité de l'échantillonnage : En combinant l'exploration guidée par la mémoire et l'apprentissage par renforcement, EMPO2 apprend beaucoup plus efficacement que les méthodes purement paramétriques.
• Vers une intelligence généralisable : La méthode résout le dilemme entre l'utilisation de la mémoire (pour l'adaptation immédiate) et l'internalisation des connaissances (pour la robustesse à long terme). Elle permet aux agents de "grandir" en intégrant les bénéfices de l'exploration directement dans leur architecture, réduisant ainsi la dépendance à des modules externes.
• Futur : Ce travail ouvre la voie à des agents capables de maîtriser des tâches complexes (mathématiques, codage, raisonnement multi-hop) en apprenant par l'expérience sans nécessiter de données étiquetées massives ou de règles humaines explicites.

En résumé, EMPO2 transforme l'exploration d'un défi de RL en un processus structuré où la mémoire sert de tremplin pour améliorer durablement les capacités de raisonnement du modèle.