Simulus: Combining Improvements in Sample-Efficient World Model Agents

Technische Zusammenfassung: Simulus – Kombination von Verbesserungen bei sample-effizienten Weltmodell-Agenten

Problemstellung
Weltmodelle (WMs) stellen eine Grenze im Bereich des sample-effizienten Reinforcement Learning (RL) dar, doch ihr Potenzial bleibt häufig ungenutzt aufgrund der Komplexität, vielversprechende Verbesserungen zu integrieren. Während einzelne Fortschritte existieren, erfordert ihre Kombination erhebliches Fachwissen und ingenieurtechnischen Aufwand. Darüber hinaus sind bestehende tokenbasierte Weltmodelle (TBWMs) weitgehend auf Bildbeobachtungen und diskrete Aktionen beschränkt, was ihre Anwendbarkeit auf kontinuierliche Steuerung und multimodale Szenarien einschränkt. Zusätzlich werden Komponenten wie intrinsische Motivation (zur Reduktion epistemischer Unsicherheit) und priorisierte Wiedergabe bei sample-effizienten Agenten häufig übersehen, teilweise aufgrund der hohen Integrationskosten und des Risikos, dass intrinsische Exploration knappe Interaktionsbudgets für aufgabenirrelevante Erfahrungen verschwendet. Dieses Landschaftsbild spiegelt die Ära vor Rainbow im Deep RL wider, in der einzeln validierte Verbesserungen nicht systematisch kombiniert worden waren.

Methodik
Die Autoren stellen Simulus vor, einen modularen, tokenbasierten Weltmodell-Agenten, der entwickelt wurde, um eine Reihe komplementärer Verbesserungen zu integrieren und zu validieren. Der Agent basiert auf der REM-Architektur und besteht aus vier distincten Modulen: einem Repräsentationsmodul ($V$), einem Weltmodell ($M$), einem Controller ($C$) und einem Replay-Puffer.

1. Flexibles Tokenisierungs-Framework ($V$): Simulus verwendet ein modulares Tokenisierungssystem, das beliebige Kombinationen von Beobachtungs- und Aktionsmodalitäten unterstützt. Es verarbeitet nativ Bilder, kontinuierliche Vektoren, kategorische Variablen und „2D-Kategorische" (bildähnliche Gitter aus kategorischen Variablen). Encoder erzeugen Token-Sequenzen fester Länge, die zu einer einheitlichen Repräsentation konkateniert werden. Dies ermöglicht dem Agenten den Betrieb in diversen Umgebungen, einschließlich kontinuierlicher Steuerung und symbolischer multimodaler Aufgaben.
2. Weltmodell ($M$): Das Dynamikmodell nutzt ein Retentive Network (RetNet), das um einen Parallel Observation Prediction (POP)-Mechanismus erweitert ist, um Umgebungs-dynamiken als Sequenzmodellierung über diskrete Token zu modellieren.
   • Intrinsische Motivation: Um epistemische Unsicherheit zu reduzieren, setzt Simulus ein Ensemble von Vorhersageköpfen für die nächste Beobachtung ein. Die Diskrepanz zwischen diesen Köpfen, gemessen mittels Jensen-Shannon-Divergenz (JSD), erzeugt intrinsische Belohnungen ($r_{int}$). Diese werden mit extrinsischen Belohnungen kombiniert, um den Controller während der „Imagination" (Planung innerhalb des Modells) in Regionen hoher Unsicherheit zu lenken.
   • Priorisierte Wiedergabe: Der Replay-Puffer verwahrt Weltmodell-Verlustwerte für gespeicherte Beispiele. Das Training nutzt eine Mischung aus uniformer und priorisierter Stichprobenziehung (gesteuert durch einen einzigen Parameter $\alpha$), wobei Proben proportional zu ihren softmax-transformierten Verlusten ausgewählt werden.
3. Controller ($C$): Der Controller ist eine erweiterte Actor-Critic-Architektur, die ausschließlich auf simulierten Erfahrungen arbeitet, die vom Weltmodell generiert werden.
   • Regression-as-Classification (RaC): Um unbeschränkte Belohnungs- und Return-Signale robust zu handhaben, adoptiert Simulus ein Klassifikationsnetzwerk mit exponentiellen Bins für die Wertvorhersage. Es nutzt die HL-Gauss-Verlustmethode (basiert auf Kreuzentropie) anstelle des traditionellen mittleren quadratischen Fehlers, was die Stabilität der Wertabschätzung verbessert.
4. Trainingszyklus: Der Agent folgt einem Zyklus aus Datensammlung, Repräsentationslernen, Weltmodelllernen und Steuerungslernen in der Imagination. Der Controller wird mit $\lambda$-Returns und einem REINFORCE-Ziel unter Verwendung einer Wert-Baseline trainiert.

Hauptbeiträge

• Integration übersehener Komponenten: Die Arbeit identifiziert und integriert intrinsische Motivation, priorisierte Weltmodell-Wiedergabe und einen hybriden RaC-Ansatz. Sie zeigt, dass diese Komponenten nicht nur einzeln effektiv sind, sondern synergistische Gewinne erzielen, wenn sie kombiniert werden.
• Vielseitige tokenbasierte Agenten: Durch die Anpassung der Multimodal-Tokenisierung an das sample-effiziente Regime etabliert Simulus TBWMs als vielseitige Agenten, die visuelle, kontinuierliche und multimodale Eingaben verarbeiten können, und geht damit über die Grenzen rein bildbasierter Modelle mit diskreten Aktionen hinaus.
• Validierung modellinduzierter Exploration: Die Arbeit zeigt, dass modellinduzierte Exploration die Leistung auch unter engen Interaktionsbudgets signifikant verbessert und damit die Annahme herausfordert, dass intrinsische Motivation für sample-effiziente Settings zu kostspielig sei.
• Open-Source-Fundament: Die Autoren bieten durch die Open-Source-Veröffentlichung von Code und Modellgewichten einen starken Ausgangspunkt für die TBWM-Forschung und senken die Einstiegshürden für die Entwicklung wettbewerbsfähiger Weltmodell-Agenten.

Ergebnisse
Simulus wurde über drei diverse Benchmarks mit strengen Interaktionsbudgets evaluiert:

• Visueller Atari 100K: Simulus erreichte die State-of-the-Art (SOTA) sample-Effizienz unter planungsfreien Methoden und übertraf Baselines wie DreamerV3, TWM, STORM und DIAMOND. Bemerkenswert ist, dass es das erste planungsfreie Weltmodell ist, das menschliche Interquantile-Mittel (IQM) und Medianwerte erreicht und auf 13 von 26 Spielen übermenschliche Leistung erzielt.
• Kontinuierliche Steuerung (DMC Proprioception 500K): Bei kontinuierlichen Steuerungsaufgaben entsprach Simulus konsistent DreamerV3 oder übertraf es leicht und demonstrierte die Effektivität tokenbasierter Architekturen in kontinuierlichen Domänen mit kompakten Vokabularen.
• Symbolisch Multimodal (Craftax-1M): In einer teilweise beobachtbaren 2D-Überlebensumgebung, die symbolische Gitter und kontinuierliche Merkmale kombiniert, übertraf Simulus gleichzeitige Weltmodellmethoden und alle modellfreien Baselines (einschließlich RND, PPO-RNN und E3B) und handhabte lange Token-Sequenzen effizient.

Ablationsstudien bestätigten, dass jede Komponente einzeln beiträgt, wobei sich intrinsische Belohnungen als besonders entscheidend erwiesen. Die Kombination aller Komponenten erzielte die stärkste Leistung, was auf eine signifikante Synergie hinweist.

Bedeutung
Die Arbeit argumentiert, dass das ungenutzte Potenzial von Weltmodell-Agenten aus der Schwierigkeit resultiert, modulare Verbesserungen zu kombinieren. Indem sie Inspiration aus dem Rainbow-Framework in DQN zieht, demonstriert Simulus, dass eine systematische Kombination von Tokenisierungsflexibilität, intrinsischer Motivation, priorisierter Wiedergabe und fortschrittlichen Wertvorhersagetechniken zu State-of-the-Art sample-Effizienz führt. Die Arbeit etabliert, dass tokenbasierte Weltmodelle ein vielseitiger und skalierbarer Ansatz für diverse RL-Domänen sind, sofern die Komponenten sorgfältig integriert werden, um die Interferenzen zu vermeiden, die oft in monolithischen Architekturen auftreten. Die Autoren betonen, dass tokenbasierte Agenten zwar derzeit im Vergleich zu einigen Baselines mit Trainingsgeschwindigkeitsbeschränkungen konfrontiert sind, ihr modulares Design jedoch deutliche Vorteile für die Policy-Inferenz und die Skalierbarkeit zukünftiger Forschung bietet.