Es war 2018, als die Idee des Reinforcement Learning im Kontext eines neuronalen Netzwerks-Weltmodells erstmalig vorgestellt wurde, und bald darauf wurde dieses grundlegende Prinzip auf Weltmodelle angewendet. Einige der prominenten Modelle, die Reinforcement Learning implementieren, waren das Dreamer-Framework, das Reinforcement Learning aus dem latenten Raum eines rekurrenten Zustandsraummodells einführte. Das DreamerV2-Framework zeigte, dass die Verwendung diskreter Latente zu reduzierten Fehlerverbindungen führen kann, und das DreamerV3-Framework konnte eine menschliche Leistung auf einer Reihe von Aufgaben in verschiedenen Domänen mit festen Hyperparametern erreichen.

Darüber hinaus können Parallelen zwischen Bildgenerierungsmodellen und Weltmodellen gezogen werden, was darauf hindeutet, dass der Fortschritt in generativen Bildmodellen auf Weltmodelle übertragen werden kann. Seitdem die Verwendung von Transformern in natürlichen Sprachverarbeitungsframeworks populär wurde, sind Frameworks wie DALL-E und VQGAN entstanden. Diese Frameworks implementierten diskrete Autoencoder, um Bilder in diskrete Token umzuwandeln, und konnten hochleistungsfähige und effiziente Text-Bild-Generierungsmodelle erstellen, indem sie die sequenziellen Modellierungsabilities der autoregressiven Transformer nutzten. Gleichzeitig gewannen Diffusionsmodelle an Popularität, und heute haben sich Diffusionsmodelle als dominierendes Paradigma für die Generierung von hochauflösenden Bildern etabliert. Aufgrund der Fähigkeiten, die Diffusionsmodelle und Reinforcement Learning bieten, werden Versuche unternommen, diese beiden Ansätze zu kombinieren, um die Flexibilität von Diffusionsmodellen als Trajektionsmodelle, Belohnungsmodelle, Planer und Richtlinien für Datenverstärkung in Offline-Reinforcement Learning zu nutzen.

Weltmodelle bieten eine vielversprechende Methode für die sichere und effiziente Ausbildung von Reinforcement Learning-Agenten. Traditionell verwenden diese Modelle Sequenzen von diskreten latenten Variablen, um Umgebungs dynamik zu simulieren. Allerdings kann diese Kompression visuelle Details übersehen, die für Reinforcement Learning wichtig sind. Gleichzeitig haben Diffusionsmodelle an Popularität gewonnen, um traditionelle Methoden, die diskrete Latente verwenden, herauszufordern. Angeregt durch diese Entwicklung, werden wir in diesem Artikel über DIAMOND (DIffusion As a Model Of eNvironment Dreams) sprechen, einem Reinforcement Learning-Agenten, der in einem Diffusions-Weltmodell ausgebildet wird. Wir werden die notwendigen Designentscheidungen untersuchen, um Diffusion für Weltmodellierung geeignet zu machen, und zeigen, dass verbesserte visuelle Details zu besserer Agentenleistung führen. DIAMOND setzt einen neuen Benchmark auf dem Atari 100k-Test, mit einem mittleren menschlichen Normalisierungsscore von 1,46, dem höchsten für Agenten, die ausschließlich in einem Weltmodell ausgebildet wurden.

DIAMOND: Diffusion als Modell von Umgebungs-Träumen

Weltmodelle oder generative Modelle von Umgebungen sind als eines der wichtigsten Komponenten für generative Agenten, um über ihre Umgebungen zu planen und zu reflektieren, im Entstehen begriffen. Obwohl die Verwendung von Reinforcement Learning in den letzten Jahren erheblichen Erfolg erzielt hat, sind Modelle, die Reinforcement Learning implementieren, bekannt für ihre geringe Stichproben-Effizienz, was ihre realen Anwendungen erheblich einschränkt. Andererseits haben Weltmodelle ihre Fähigkeit demonstriert, Reinforcement Learning-Agenten effizient in verschiedenen Umgebungen mit wesentlich verbesserter Stichproben-Effizienz auszubilden, was es dem Modell ermöglicht, aus realen Erfahrungen zu lernen. Aktuelle Weltmodell-Frameworks modellieren Umgebungs-Dynamiken normalerweise als Sequenz von diskreten latenten Variablen, wobei das Modell den latenten Raum diskretisiert, um Fehlerverbindungen über mehrstufige Zeithorizonte zu vermeiden. Obwohl dieser Ansatz substantielle Ergebnisse liefern kann, ist er auch mit Informationsverlust verbunden, was zu einem Verlust der Rekonstruktionsqualität und der Allgemeingültigkeit führt. Der Informationsverlust kann zu einem erheblichen Hindernis für reale Szenarien werden, die gut definierte Informationen erfordern, wie z.B. das Training von autonomen Fahrzeugen. In solchen Aufgaben können kleine Änderungen oder Details in der visuellen Eingabe, wie die Farbe der Ampel oder das Blinklicht des Fahrzeugs vorne, die Richtlinie eines Agenten ändern. Obwohl die Erhöhung der Anzahl diskreter Latente den Informationsverlust vermeiden kann, erhöht es die Rechenkosten erheblich.

Darüber hinaus haben Diffusionsmodelle in den letzten Jahren als dominierendes Paradigma für hochwertige Bildgenerierungsframeworks an Bedeutung gewonnen, da Frameworks, die auf Diffusionsmodellen basieren, lernen, einen Rauschprozess umzukehren, und direkt mit etablierten Ansätzen konkurrieren, die diskrete Token modellieren, und somit eine vielversprechende Alternative bieten, um die Notwendigkeit der Diskretisierung in der Weltmodellierung zu eliminieren. Diffusionsmodelle sind bekannt für ihre Fähigkeit, leicht zu konditionieren und komplexe, multimodale Verteilungen ohne Moden-Kollaps zu modellieren. Diese Attribute sind für die Weltmodellierung von entscheidender Bedeutung, da die Konditionierung es einem Weltmodell ermöglicht, die Aktionen eines Agenten genau widerzuspiegeln, was zu einer zuverlässigeren Zuweisung von Krediten führt. Darüber hinaus bietet die Modellierung multimodaler Verteilungen eine größere Vielfalt von Trainings-Szenarien für den Agenten, was seine Gesamtleistung verbessert.

Basierend auf diesen Merkmalen ist DIAMOND (DIffusion As a Model Of eNvironment Dreams) ein Reinforcement Learning-Agent, der in einem Diffusions-Weltmodell ausgebildet wird. Das DIAMOND-Framework trifft sorgfältige Designentscheidungen, um sicherzustellen, dass sein Diffusions-Weltmodell effizient und stabil über lange Zeit horizonte bleibt. Das Framework bietet eine qualitative Analyse, um die Bedeutung dieser Designentscheidungen zu demonstrieren. DIAMOND setzt einen neuen Benchmark mit einem mittleren menschlichen Normalisierungsscore von 1,46 auf dem etablierten Atari 100k-Benchmark, dem höchsten für Agenten, die ausschließlich in einem Weltmodell ausgebildet wurden. Durch die Arbeit im Bildraum ermöglicht es DIAMONDS Diffusions-Weltmodell, die Umgebung nahtlos zu ersetzen, was zu einem besseren Verständnis von Weltmodell- und Agenten-Verhalten führt. Bemerkenswerterweise wird die verbesserte Leistung in bestimmten Spielen der besseren Modellierung kritischer visueller Details zugeschrieben. Das DIAMOND-Framework modelliert die Umgebung als Standard-POMDP oder Partially Observable Markov Decision Process mit einer Menge von Zuständen, einer Menge diskreter Aktionen und einer Menge von Bildbeobachtungen. Die Übergangsfunktionen beschreiben die Umgebungs-Dynamiken, und die Belohnungsfunktion ordnet die Übergänge skalaren Belohnungen zu. Die Beobachtungsfunktion beschreibt die Beobachtungswahrscheinlichkeiten und emittiert Bildbeobachtungen, die von den Agenten verwendet werden, um die Umgebung zu sehen, da sie nicht direkt auf die Zustände zugreifen können. Das primäre Ziel des Ansatzes war es, eine Richtlinie zu erhalten, die Beobachtungen zu Aktionen abbildet, um den erwarteten diskontierten Rückgang mit einem Diskontfaktor zu maximieren. Weltmodelle sind generative Modelle der Umgebung, und Weltmodelle können verwendet werden, um simulierte Umgebungen zu erstellen, um Reinforcement Learning-Agenten in der realen Umgebung und im Weltmodell auszubilden. Abbildung 1 zeigt die unrollierte Vorstellung des DIAMOND-Frameworks über die Zeit.

DIAMOND: Methodik und Architektur

Im Kern sind Diffusionsmodelle eine Klasse von generativen Modellen, die eine Stichprobe durch Umkehren des Rauschprozesses erzeugen und starke Inspirationen aus der nicht-gleichgewichtigen Thermodynamik ziehen. Das DIAMOND-Framework betrachtet einen Diffusionsprozess, der durch eine kontinuierliche Zeitvariable indexiert ist, mit entsprechenden Randbedingungen und einer tragbaren unstrukturierten Verteilung. Darüber hinaus muss das DIAMOND-Framework, um ein generatives Modell zu erhalten, das von Rauschen zu Daten abbildet, den Prozess umkehren, wobei der Umkehrprozess auch ein Diffusionsprozess ist, der rückwärts in der Zeit läuft. Darüber hinaus ist es an jedem gegebenen Zeitpunkt nicht trivial, die Score-Funktion zu schätzen, da das DIAMOND-Framework keinen Zugriff auf die tatsächliche Score-Funktion hat, und das Modell überwindet diese Hürde, indem es ein Score-Matching-Objekt implementiert, ein Ansatz, der es einem Framework ermöglicht, ein Score-Modell ohne Kenntnis der zugrunde liegenden Score-Funktion auszubilden. Das score-basierte Diffusionsmodell bietet ein unbedingtes generatives Modell. Allerdings ist ein bedingtes generatives Modell der Umgebungs-Dynamik erforderlich, um als Weltmodell zu dienen, und um diesem Zweck zu dienen, betrachtet das DIAMOND-Framework den allgemeinen Fall des POMDP-Ansatzes, bei dem das Framework die Vergangenheit von Beobachtungen und Aktionen nutzen kann, um den unbekannten Markov-Zustand zu approximieren. Wie in Abbildung 1 gezeigt, nutzt das DIAMOND-Framework diese Geschichte, um ein Diffusionsmodell zu konditionieren, um die nächste Beobachtung direkt zu schätzen und zu erzeugen. Obwohl das DIAMOND-Framework theoretisch auf jedes SDE- oder ODE-Verfahren zurückgreifen kann, gibt es einen Kompromiss zwischen NFE oder Anzahl der Funktionsauswertungen und der Beispielfähigkeit, der die Inferenzkosten von Diffusionsmodellen erheblich beeinflusst.

Basierend auf diesen Erkenntnissen betrachten wir nun die praktische Realisierung des DIAMOND-Frameworks eines diffusionsbasierten Weltmodells, einschließlich der Drift- und Diffusionskoeffizienten, die einer bestimmten Diffusionsansatzentscheidung entsprechen. Anstelle der Wahl von DDPM, einem natürlichen Kandidaten für die Aufgabe, baut das DIAMOND-Framework auf der EDM-Formulierung auf und betrachtet einen Störkern mit einer reellen Funktion der Diffusionszeit, die als Rauschplan bezeichnet wird. Das Framework wählt die Vorkonditionierer, um die Eingabe- und Ausgabevarianz für jede Stimmebene beizubehalten. Die Netzwerktrainingsmischung signalisiert und Rauschen adaptiv in Abhängigkeit vom Verschlechterungsgrad, und wenn das Rauschen niedrig ist, wird das Ziel zur Differenz zwischen dem sauberen und dem gestörten Signal, d.h. dem hinzugefügten Gaußschen Rauschen. Intuitiv verhindert dies, dass das Trainingsziel im niedrigen Rauschregime trivial wird. In der Praxis ist dieses Ziel bei den Extremen des Rauschplans hochvarianz, sodass das Modell den Rauschpegel aus einer log-normalen Verteilung auswählt, die empirisch ausgewählt wird, um die Trainingsdaten um die mittleren Rauschregionen zu konzentrieren. Das DIAMOND-Framework nutzt ein Standard-U-Net-2D-Komponente für den Vektorfeld und behält einen Puffer von vergangenen Beobachtungen und Aktionen, den das Framework verwendet, um sich zu konditionieren. Das DIAMOND-Framework konkatiniert dann diese vergangenen Beobachtungen mit der nächsten geräuschten Beobachtung und den Eingabeaktionen durch adaptive Gruppennormierungsschichten in den Residualblöcken des U-Nets.

DIAMOND: Experimente und Ergebnisse

Für eine umfassende Bewertung wählt das DIAMOND-Framework den Atari 100k-Benchmark. Der Atari 100k-Benchmark besteht aus 26 Spielen, die darauf ausgelegt sind, eine breite Palette von Agentenfähigkeiten zu testen. In jedem Spiel ist ein Agent auf 100.000 Aktionen in der Umgebung beschränkt, was in etwa zwei Stunden menschlicher Spielzeit entspricht, um das Spiel zu lernen, bevor es bewertet wird. Zum Vergleich trainieren unbeschränkte Atari-Agenten normalerweise 50 Millionen Schritte, was einer 500-fachen Erhöhung der Erfahrung entspricht. Wir trainierten DIAMOND von Grund auf mit 5 zufälligen Samen für jedes Spiel. Jeder Trainingslauf erforderte etwa 12 GB VRAM und dauerte ungefähr 2,9 Tage auf einem einzelnen Nvidia RTX 4090, was insgesamt 1,03 GPU-Jahre entspricht. Die folgende Tabelle zeigt den Score für alle Spiele, den Mittelwert und den IQM oder interquartilen Mittelwert der humanisierten Scores.

Im Anschluss an die Einschränkungen von Punktschätzungen bietet das DIAMOND-Framework stratifizierte Bootstrap-Vertrauensintervalle für den Mittelwert und den IQM oder interquartilen Mittelwert der humanisierten Scores sowie Leistungsprofile und weitere Metriken, wie in der folgenden Abbildung zusammengefasst.

Die Ergebnisse zeigen, dass DIAMOND außergewöhnlich gut über den Benchmark hinweg abschneidet, menschliche Spieler in 11 Spielen übertrifft und einen übermenschlichen Mittel-HNS von 1,46 erreicht, was einen neuen Rekord für Agenten darstellt, die ausschließlich in einem Weltmodell ausgebildet wurden. Darüber hinaus ist DIAMONDS IQM mit STORM vergleichbar und übertrifft alle anderen Baselines. DIAMOND ragt in Umgebungen hervor, in denen die Erfassung kleiner Details von entscheidender Bedeutung ist, wie z.B. Asterix, Breakout und RoadRunner. Darüber hinaus, wie bereits diskutiert, hat das DIAMOND-Framework die Flexibilität, jedes Diffusionsmodell in seiner Pipeline zu implementieren, obwohl es sich für den EDM-Ansatz entscheidet, wäre es eine natürliche Wahl gewesen, den DDPM-Modell zu wählen, da es bereits in zahlreichen Bildgenerierungsanwendungen implementiert wird. Um den EDM-Ansatz mit der DDPM-Implementierung zu vergleichen, trainiert das DIAMOND-Framework beide Varianten mit der gleichen Netzwerkarchitektur auf dem gleichen statischen Datensatz mit über 100.000 Frames, die mit einer Expertenrichtlinie gesammelt wurden. Die Anzahl der Denoising-Schritte ist direkt mit den Inferenzkosten des Weltmodells verbunden, und somit reduzieren weniger Schritte die Kosten für die Ausbildung eines Agenten auf imaginierten Trajektionen. Um sicherzustellen, dass unser Weltmodell computativ mit anderen Baselines wie IRIS vergleichbar bleibt, die 16 NFE pro Timestep erfordern, zielen wir darauf ab, nicht mehr als einige Dutzend Denoising-Schritte zu verwenden, vorzugsweise weniger. Allerdings kann die Festlegung der Anzahl der Denoising-Schritte auf ein zu niedriges Niveau die visuelle Qualität verschlechtern und zu Fehlerverbindungen führen. Um die Stabilität verschiedener Diffusionsvarianten zu bewerten, zeigen wir imaginierte Trajektionen, die autoregressiv bis zu t = 1000 Timesteps in der folgenden Abbildung generiert werden, unter Verwendung verschiedener Anzahl von Denoising-Schritten n ≤ 10.

Wir beobachten, dass die Verwendung von DDPM (a) in diesem Regime zu schweren Fehlerverbindungen führt, was dazu führt, dass das Weltmodell schnell aus der Verteilung driftet. Im Gegensatz dazu bleibt das EDM-basierte Diffusions-Weltmodell (b) viel stabiler über lange Zeit horizonte, sogar mit einem einzigen Denoising-Schritt. Imaginierte Trajektionen mit Diffusions-Weltmodellen basierend auf DDPM (links) und EDM (rechts) sind gezeigt. Die anfängliche Beobachtung bei t = 0 ist für beide gleich, und jede Zeile entspricht einer abnehmenden Anzahl von Denoising-Schritten n. Wir beobachten, dass die DDPM-basierte Generierung unter Fehlerverbindungen leidet, wobei kleinere Zahlen von Denoising-Schritten zu einer schnelleren Fehleransammlung führen. Im Gegensatz dazu bleibt das EDM-basierte Weltmodell von DIAMOND viel stabiler, sogar für n = 1. Die optimale Einzelschritt-Vorhersage ist die Erwartung über mögliche Rekonstruktionen für eine gegebene geräuschte Eingabe, die außerhalb der Verteilung liegen kann, wenn die posteriore Verteilung multimodal ist. Während einige Spiele wie Breakout deterministische Übergänge haben, die mit einem einzigen Denoising-Schritt genau modelliert werden können, zeigen andere Spiele partielle Beobachtbarkeit, was zu multimodalen Beobachtungsverteilungen führt. In diesen Fällen ist ein iterativer Solver notwendig, um das Stichprobenverfahren in Richtung eines bestimmten Modus zu lenken, wie in dem Spiel Boxing in der folgenden Abbildung illustriert. Folglich legt das DIAMOND-Framework n = 3 in all unseren Experimenten fest.

Die obige Abbildung vergleicht Einzelschritt- (obere Zeile) und Mehrschritt- (untere Zeile) Stichproben in Boxing. Die Bewegungen des schwarzen Spielers sind unvorhersehbar, was dazu führt, dass die Einzelschritt-Denoising zwischen möglichen Ergebnissen interpoliert, was zu verschwommenen Vorhersagen führt. Im Gegensatz dazu produziert die Mehrschritt-Stichprobenbildung ein klares Bild, indem sie die Generierung in Richtung eines bestimmten Modus lenkt. Interessanterweise eliminiert die Tatsache, dass die Richtlinie den weißen Spieler kontrolliert, die Unbestimmtheit, da die Aktionen des weißen Spielers dem Weltmodell bekannt sind. Somit können sowohl Einzelschritt- als auch Mehrschritt-Stichprobenbildung die Position des weißen Spielers korrekt vorhersagen.

In der obigen Abbildung zeigen die von DIAMOND imaginierten Trajektionen im Allgemeinen eine höhere visuelle Qualität und sind treuer zur wahren Umgebung im Vergleich zu denen, die von IRIS imaginiert werden. Die von IRIS generierten Trajektionen enthalten visuelle Inkonsistenzen zwischen den Frames (hervorgehoben durch weiße Kästchen), wie z.B. Feinde, die als Belohnungen dargestellt werden und umgekehrt. Obwohl diese Inkonsistenzen nur wenige Pixel betreffen können, können sie die Ausbildung eines optimalen Agenten erheblich beeinträchtigen. Die Abbildung zeigt aufeinanderfolgende Frames, die mit IRIS (links) und DIAMOND (rechts) imaginiert werden. Die weißen Kästchen heben Inkonsistenzen zwischen den Frames hervor, die nur in Trajektionen auftreten, die mit IRIS generiert werden. In Asterix (obere Zeile) wird ein Feind (orange) zu einer Belohnung (rot), dann wieder zu einem Feind in der dritten und wieder zu einer Belohnung in der vierten. In Breakout (mittlere Zeile) sind die Steine und der Score inkonsistent zwischen den Frames. In Road Runner (untere Zeile) werden die Belohnungen (kleine blaue Punkte auf der Straße) inkonsistent zwischen den Frames gerendert. Diese Inkonsistenzen treten nicht mit DIAMOND auf. In Breakout wird der Score zuverlässig um +7 aktualisiert, wenn ein roter Stein gebrochen wird.

Schlussfolgerung

In diesem Artikel haben wir über DIAMOND gesprochen, einen Reinforcement Learning-Agenten, der in einem Diffusions-Weltmodell ausgebildet wird. Das DIAMOND-Framework trifft sorgfältige Designentscheidungen, um sicherzustellen, dass sein Diffusions-Weltmodell effizient und stabil über lange Zeit horizonte bleibt. Das Framework bietet eine qualitative Analyse, um die Bedeutung dieser Designentscheidungen zu demonstrieren. DIAMOND setzt einen neuen Benchmark mit einem mittleren menschlichen Normalisierungsscore von 1,46 auf dem etablierten Atari 100k-Benchmark, dem höchsten für Agenten, die ausschließlich in einem Weltmodell ausgebildet wurden. Durch die Arbeit im Bildraum ermöglicht es DIAMONDS Diffusions-Weltmodell, die Umgebung nahtlos zu ersetzen, was zu einem besseren Verständnis von Weltmodell- und Agenten-Verhalten führt. Bemerkenswerterweise wird die verbesserte Leistung in bestimmten Spielen der besseren Modellierung kritischer visueller Details zugeschrieben. Das DIAMOND-Framework modelliert die Umgebung als Standard-POMDP oder Partially Observable Markov Decision Process mit einer Menge von Zuständen, einer Menge diskreter Aktionen und einer Menge von Bildbeobachtungen. Die Übergangsfunktionen beschreiben die Umgebungs-Dynamiken, und die Belohnungsfunktion ordnet die Übergänge skalaren Belohnungen zu.