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

Die Arbeit stellt EMPO$^2$ vor, ein hybrides Reinforcement-Learning-Framework, das durch die Kombination von On- und Off-Policy-Optimierung sowie einem Gedächtnismechanismus die Explorationsfähigkeit von LLM-Agenten signifikant verbessert und so zu überlegener Leistung und Anpassungsfähigkeit in neuen Umgebungen führt.

Das Wesentliche

Stell dir vor, du hast einen sehr klugen, aber etwas sturen Freund, der ein riesiges Wissen in seinem Kopf hat (das ist unser KI-Modell). Wenn du ihn in eine neue, unbekannte Stadt schickst, um etwas Bestimmtes zu finden (z. B. einen roten Lichtschalter zu finden), nutzt er sofort alles, was er aus Büchern weiß.

Das Problem: In der echten Welt (oder in Computerspielen wie ScienceWorld) funktionieren die Regeln oft nicht so, wie in den Büchern. Dein Freund versucht, den Schalter zu drücken, aber er ist gar nicht da. Er bleibt stecken, wiederholt denselben Fehler und lernt nichts dazu. Er ist zu sehr damit beschäftigt, sein altes Wissen anzuwenden, anstatt die Umgebung wirklich zu erkunden.

Die Forscher aus diesem Papier haben eine Lösung namens EMPO2 entwickelt. Hier ist eine einfache Erklärung, wie das funktioniert, mit ein paar kreativen Vergleichen:

1. Das Problem: Der "sture" Experte

Stell dir vor, du spielst ein Videospiel. Dein Charakter ist ein Experte, der alles über Physik weiß. Aber im Spiel musst du erst herausfinden, wo die Batterie liegt, bevor du sie anschließen kannst.

• Der alte Weg (GRPO): Der Charakter versucht immer wieder, die Batterie anzuschließen, obwohl er sie nicht findet. Er lernt nur langsam durch Zufall, wenn er mal einen anderen Weg probiert. Das dauert ewig.
• Das neue Problem: Wenn er scheitert, vergisst er sofort, warum er gescheitert ist, und macht beim nächsten Versuch genau denselben Fehler.

2. Die Lösung: EMPO2 – Der "Notizbuch"-Agent

EMPO2 gibt dem Agenten zwei superkräftige Werkzeuge: ein Gedächtnis (ein Notizbuch) und eine doppelte Lernstrategie.

Werkzeug A: Das Notizbuch (Das "Gedächtnis")

Stell dir vor, dein Freund hat ein Notizbuch bei sich. Jedes Mal, wenn er scheitert (z. B. "Ich habe den Schalter nicht gefunden"), schreibt er sich einen Tipp auf: "Hey, ich war im Flur, aber der Schalter war nicht da. Ich muss ins Labor gehen!"

• Der Clou: Er nutzt dieses Notizbuch nicht nur, um sich zu erinnern, sondern er nutzt die Tipps, um beim nächsten Versuch besser zu suchen. Er liest: "Aha, ich war schon im Flur, also gehe ich jetzt ins Labor."
• Der Trick: Diese Tipps werden nicht von einem anderen Menschen geschrieben, sondern vom Agenten selbst generiert. Er lernt also, sich selbst zu belehren.

Werkzeug B: Die zwei Lern-Modi (On-Policy & Off-Policy)

Hier wird es clever. Der Agent lernt auf zwei verschiedene Arten gleichzeitig, wie ein Schüler, der sowohl aus dem Unterricht als auch aus Hausaufgaben lernt:

1. Der "Mit-Notizbuch"-Modus (On-Policy):
   Der Agent spielt das Spiel und liest dabei sein Notizbuch. Er nutzt die Tipps, um bessere Entscheidungen zu treffen. Das hilft ihm, schnell neue Wege zu finden (Exploration).
2. Der "Ohne-Notizbuch"-Modus (Off-Policy):
   Das ist der geniale Teil. Der Agent spielt eine Runde mit Tipps, aber wenn er danach lernt, tut er so, als hätte er keine Tipps gehabt.
   • Die Analogie: Stell dir vor, ein Koch (der Agent) kocht ein Gericht mit einer geheimen Zutat (dem Tipp), die ihm der Chef gegeben hat. Das Gericht schmeckt toll. Jetzt will der Koch lernen, wie man es selbst kocht, ohne dass der Chef ihm die Zutat gibt. Er schaut sich an: "Okay, ich habe mit der Zutat gekocht und es war gut. Ich muss also herausfinden, wie ich das ohne die Zutat genauso gut hinbekomme."
   • Das Ergebnis: Der Agent "verinnerlicht" die Tipps. Er lernt, dass der Weg ins Labor der richtige ist, und merkt sich das in seinem Gehirn (den neuronalen Gewichten), statt nur auf das Notizbuch zu schauen.

3. Das Ergebnis: Ein besserer Entdecker

In den Tests (wie im Spiel ScienceWorld oder beim Online-Shopping in WebShop) hat sich gezeigt:

• Schnelleres Lernen: EMPO2 findet Lösungen viel schneller als andere Methoden. Es ist wie ein Detektiv, der sich Notizen macht und daraus lernt, statt immer wieder in dieselbe Sackgasse zu laufen.
• Anpassungsfähigkeit: Wenn man den Agenten in eine völlig neue Umgebung schickt (z. B. von "Biolabor" zu "Elektrik"), kann er sich sofort anpassen. Er nutzt sein Notizbuch, um die ersten paar Versuche zu überstehen, und lernt dann so schnell, dass er das Notizbuch bald gar nicht mehr braucht.
• Robustheit: Selbst wenn das Notizbuch weg ist (z. B. wenn der Agent im echten Leben ohne externe Hilfe agieren muss), ist er immer noch viel schlauer als vorher, weil er die Tipps in sein eigenes Gehirn integriert hat.

Zusammenfassung in einem Satz

EMPO2 ist wie ein Schüler, der nicht nur aus Fehlern lernt, sondern sich selbst Notizen macht, diese Notizen nutzt, um bessere Versuche zu machen, und dann das Gelernte so fest in seinem Kopf verankert, dass er am Ende auch ohne Notizen ein Meister seines Fachs ist.

Es ist der Schlüssel, um KI-Agenten nicht nur "klug" zu machen, sondern sie zu neugierigen Entdeckern zu verwandeln, die sich in unbekannten Welten zurechtfinden können.

Technische Zusammenfassung

1. Problemstellung

Große Sprachmodelle (LLMs), die als Agenten in Reinforcement Learning (RL)-Umgebungen eingesetzt werden, leiden oft unter einem fundamentalen Mangel: Exploration.

• Ausnutzung vs. Erkundung: Bestehende LLM-Agenten verlassen sich stark auf ihr vortrainiertes Wissen (Ausnutzung/Exploitation) und scheitern in Umgebungen, die das Entdecken neuer Zustände oder das aktive Sammeln unbekannter Informationen erfordern.
• Limitationen bestehender Ansätze:
  • Reine Online-RL-Methoden (z. B. GRPO): Convergenz oft in suboptimalen Lösungen, da sie ohne systematische Exploration stecken bleiben.
  • Externe Speicher (Non-Parametric Memory, z. B. Reflexion): Zwar können sie Fehler korrigieren, aber da die Modellparameter statisch bleiben, ist die Anpassung kurzlebig. Das Wissen wird nicht dauerhaft im Modell verankert, was die Generalisierung auf neue Szenarien einschränkt.
• Ziel: Ein Framework zu entwickeln, das die Vorteile von externem Gedächtnis für die Exploration nutzt, diese Erkenntnisse aber gleichzeitig in die Modellparameter integriert, um robuste, generalisierbare Agenten ohne ständige Abhängigkeit von externen Hinweisen zu schaffen.

2. Methodik: EMPO2

Das vorgeschlagene Framework EMPO2 (Exploratory Memory-Augmented On- and Off-Policy Optimization) ist ein hybrides RL-System, das parametrische Updates (Lernen der LLM-Parameter) und nicht-parametrische Updates (externer Speicher) kombiniert.

Kernkomponenten:

1. Selbstgenerierter Speicher (Self-Generated Memory):

   • Der Agent generiert nach jedem Trajektorien-Abschluss „Tips" (Hinweise/Reflexionen) basierend auf seinen Erfahrungen.
   • Diese Tips werden in einem externen Speicher (Memory Buffer) gespeichert und können bei zukünftigen Aktionen abgerufen werden.
   • Im Gegensatz zu reinen Reflexions-Methoden dienen diese Tips hier als Zwischenschritt, um Exploration zu fördern, bevor das Wissen in die Parameter überführt wird.

2. Hybride Rollout-Modi (Rollout Phase):
   Der Agent wählt während der Interaktion mit der Umgebung zufällig zwischen zwei Modi:

   • Ohne Speicher: Der Agent agiert basierend nur auf dem aktuellen Zustand und der Aufgabe (Standard-Prompting).
   • Mit Speicher: Der Agent ruft relevante Tips aus dem Speicher ab und nutzt diese als zusätzlichen Kontext für die Aktionsgenerierung. Dies fördert die Erkundung neuer Strategien.

3. Hybride Update-Modi (Update Phase):
   Die gesammelten Daten werden in zwei Modi aktualisiert, um sowohl Stabilität als auch Wissensinternalisierung zu gewährleisten:

   • On-Policy Update: Die Tips werden beibehalten. Das Modell lernt, wie man mit Tips agiert (stabilisiert das Lernen).
   • Off-Policy Update (Schlüsselinnovation): Hier werden die während des Rollouts mit Tips generierten Trajektorien genutzt, aber beim Update werden die Tips aus dem Prompt entfernt.
     • Das Modell lernt also, die Ergebnisse (die erfolgreichen Aktionen), die durch Tips ermöglicht wurden, auch ohne diese Tipps zu reproduzieren.
     • Dies funktioniert als wissensbasierte Distillation: Die Tipps dienen als „Lehrer", der dem „Schüler" (dem Basis-Modell) zeigt, welche Aktionen erfolgreich sind, damit das Modell diese Fähigkeiten internalisiert.

4. Stabilisierungstechniken:

   • Masking: Um Instabilitäten bei Off-Policy-Lernen zu vermeiden (durch extrem unwahrscheinliche Tokens), werden Advantages für Tokens unter einem Wahrscheinlichkeitsschwellenwert maskiert.
   • Intrinsische Belohnung: Ein intrinsischer Reward wird basierend auf der Neuheit des Zustands hinzugefügt, um die Exploration zu fördern, wenn externe Belohnungen fehlen.

3. Wichtige Beiträge

• Einheitliches Framework: EMPO2 verbindet erstmals parametrisches RL (On/Off-Policy) mit nicht-parametrischem Gedächtnis in einem einzigen Algorithmus.
• Internalisierung von Exploration: Der Mechanismus des Off-Policy-Updates ohne Tips ermöglicht es dem Agenten, die durch Gedächtnis erlernten Erkundungsstrategien dauerhaft in seine Parameter zu integrieren.
• Robustheit: Das System ist sowohl in Umgebungen mit als auch ohne Zugriff auf den externen Speicher leistungsfähig, da das Modell die Vorteile des Speichers „gelernt" hat.
• Effizienz: Durch die Kombination von On- und Off-Policy-Lernen wird die Sample-Effizienz im Vergleich zu reinen Online-RL-Methoden erheblich gesteigert.

4. Ergebnisse

Die Evaluation erfolgte auf zwei etablierten Benchmarks: ScienceWorld (wissenschaftliche Experimente) und WebShop (E-Commerce). Als Basis-Modell diente Qwen2.5-7B-Instruct.

• In-Distribution (ID) Leistung:
  • Auf ScienceWorld erreichte EMPO2 eine Verbesserung von 128,6 % gegenüber dem starken GRPO-Baseline.
  • Auf WebShop wurde eine Steigerung von 11,3 % gegenüber GRPO erzielt.
  • EMPO2 übertraf auch Offline-RL-Ansätze (Retrospex) und reine Speicher-Methoden (Reflexion) deutlich.
• Out-of-Distribution (OOD) Anpassung:
  • In Tests mit neuen, unbekannten Aufgaben zeigte EMPO2 eine überlegene Anpassungsfähigkeit.
  • Das Modell konnte sich mit nur wenigen Versuchen und ohne Parameter-Updates an neue Domänen anpassen, indem es den externen Speicher nutzte. Dies unterstreicht die Fähigkeit des Modells, Exploration effektiv zu nutzen.
• Ablationsstudien:
  • Die Kombination aus On- und Off-Policy-Lernen ist entscheidend; das Entfernen eines der beiden Modi führt zu suboptimalen Ergebnissen.
  • Intrinsische Belohnungen sind notwendig, um das Lernen vor dem Zusammenbruch in homogene Verhaltensweisen zu schützen.

5. Bedeutung und Ausblick

EMPO2 adressiert eine der größten Schwächen aktueller LLM-Agenten: die Unfähigkeit, systematisch in unbekannten Umgebungen zu explorieren und dieses Wissen langfristig zu speichern.

• Paradigmenwechsel: Statt nur externe Hinweise zu nutzen, lernt das Modell, diese Hinweise zu internalisieren. Dies führt zu Agenten, die nicht nur kurzfristig adaptiv sind, sondern langfristig „klüger" werden.
• Skalierbarkeit: Der Ansatz ist unabhängig von der Modellgröße und könnte durch Skalierung auf größere Modelle weiter verbessert werden.
• Zukünftige Anwendungen: Die Methode bietet einen vielversprechenden Weg für die Entwicklung robusterer Agenten in komplexen Domänen wie Mathematik, Coding und multimodalen Aufgaben, wo reine Prompting-Strategien oft an ihre Grenzen stoßen.

Zusammenfassend stellt EMPO2 einen signifikanten Fortschritt dar, der die Lücke zwischen kurzfristiger Gedächtnisnutzung und langfristiger parametrischer Lernfähigkeit schließt und damit die Generalisierungsfähigkeit von LLM-Agenten in dynamischen Umgebungen massiv erhöht.