Het was in 2018, toen het idee van versterking van het leren in de context van een neurale netwerkwereldmodel voor het eerst werd geïntroduceerd, en al snel werd dit fundamentele principe toegepast op wereldmodellen. Sommige van de prominente modellen die versterking van het leren implementeren, waren het Dreamer-framework, dat versterking van het leren introduceerde vanuit de latent ruimte van een recurrente staatruimtemodel. De DreamerV2 demonstreerde dat het gebruik van discrete latenten kan resulteren in een vermindering van de accumulatie van fouten, en het DreamerV3-framework was in staat om mensachtige prestaties te bereiken op een reeks taken over verschillende domeinen met vaste hyperparameters.

Bovendien kunnen parallellen worden getrokken tussen beeldgeneratiemodellen en wereldmodellen, waaruit blijkt dat de vooruitgang in generatieve visiemodellen kan worden gerepliceerd om wereldmodellen te bevoordelen. Sinds het gebruik van transformers in natuurlijke taalverwerking frameworks populair werd, zijn DALL-E en VQGAN-frameworks ontstaan. Deze frameworks implementeerden discrete auto-encoders om beelden om te zetten in discrete tokens en waren in staat om krachtige en efficiënte tekst-naar-beeld-generatiemodellen te bouwen door de sequentiële modelleringsmogelijkheden van de autoregressieve transformers te benutten. Tegelijkertijd wonnen diffusiemodellen aan populariteit, en vandaag de dag hebben diffusiemodellen zichzelf gevestigd als een dominante paradigma voor het genereren van hoge-resolutiebeelden. Vanwege de mogelijkheden die diffusiemodellen en versterking van het leren bieden, worden er pogingen gedaan om deze twee benaderingen te combineren, met als doel om te profiteren van de flexibiliteit van diffusiemodellen als trajectoriemodellen, beloningsmodellen, planners en als beleid voor gegevensverrijking in offline-versterking van het leren.

Wereldmodellen bieden een veelbelovende methode voor het trainen van versterking van het leren agenten op een veilige en efficiënte manier. Traditioneel gebruiken deze modellen sequenties van discrete latent variabelen om omgevingsdynamica te simuleren. Echter, deze compressie kan visuele details die cruciaal zijn voor versterking van het leren over het hoofd zien. Tegelijkertijd zijn diffusiemodellen in populariteit gestegen voor beeldgeneratie, waardoor traditionele methoden die discrete latenten gebruiken, worden uitgedaagd. Geïnspireerd door deze verschuiving, zullen we in dit artikel praten over DIAMOND (DIffusion As a Model Of eNvironment Dreams), een versterking van het leren agent getraind binnen een diffusiewereldmodel. We zullen de noodzakelijke ontwerpkansen onderzoeken om diffusie geschikt te maken voor wereldmodellering en laten zien dat verbeterde visuele details leiden tot betere agentprestaties. DIAMOND stelt een nieuwe benchmark vast op de concurrerende Atari 100k-test, met een gemiddelde menselijke genormaliseerde score van 1,46, het hoogste voor agenten getraind volledig binnen een wereldmodel.

DIAMOND: Diffusie als een Model van Omgevingsdromen

Wereldmodellen of generatieve modellen van omgevingen zijn een van de belangrijkste componenten voor generatieve agenten om te plannen en te redeneren over hun omgevingen. Hoewel het gebruik van versterking van het leren aanzienlijk succes heeft behaald in recente jaren, zijn modellen die versterking van het leren implementeren bekend om steekproefinefficiënt te zijn, wat hun reële wereldtoepassingen aanzienlijk beperkt. Aan de andere kant hebben wereldmodellen aangetoond dat ze in staat zijn om versterking van het leren agenten efficiënt te trainen over diverse omgevingen met een aanzienlijk verbeterde steekproefefficiëntie, waardoor het model kan leren van reële wereldervaringen. Recentere wereldmodelleringframeworks modelleren omgevingsdynamica meestal als een sequentie van discrete latent variabelen, met het model dat de latent ruimte discretiseert om accumulatie van fouten over meerdere tijdstappen te voorkomen. Hoewel deze benadering aanzienlijke resultaten kan opleveren, is het ook geassocieerd met een verlies van informatie, wat leidt tot een verlies van reconstructiekwaliteit en verlies van generaliteit. Het verlies van informatie kan een aanzienlijke belemmering vormen voor reële scenario’s die informatie nodig hebben om goed gedefinieerd te zijn, zoals het trainen van autonome voertuigen. In dergelijke taken kunnen kleine veranderingen of details in de visuele invoer, zoals de kleur van het verkeerslicht of de richtingsaanwijzer van het voertuig voor je, het beleid van een agent veranderen. Hoewel het verhogen van het aantal discrete latenten kan helpen om informatieverlies te voorkomen, verhoogt het de berekeningskosten aanzienlijk.

Bovendien zijn diffusiemodellen in de afgelopen jaren opgekomen als de dominante benadering voor hoge-kwaliteit beeldgeneratiemodellen, omdat frameworks die zijn gebouwd op diffusiemodellen leren om een verstorend proces om te keren en concurreren met enkele van de meer gevestigde benaderingen die discrete tokens modelleren, en bieden daarmee een veelbelovende alternatief om de noodzaak voor discretisatie in wereldmodellering te elimineren. Diffusiemodellen zijn bekend om hun vermogen om gemakkelijk te worden geconditioneerd en om complexe, multimodale distributies flexibel te modelleren zonder modusinstorting. Deze kenmerken zijn cruciaal voor wereldmodellering, omdat conditionering een wereldmodel in staat stelt om nauwkeurig de acties van een agent te weerspiegelen, wat leidt tot een betrouwbare toewijzing van credits. Bovendien biedt het modelleren van multimodale distributies een grotere diversiteit aan trainingsscenario’s voor de agent, waardoor de algehele prestatie van de agent wordt verbeterd.

Op basis van deze kenmerken is DIAMOND (DIffusion As a Model Of eNvironment Dreams) een versterking van het leren agent getraind binnen een diffusiewereldmodel. Het DIAMOND-framework maakt zorgvuldige ontwerpkansen om ervoor te zorgen dat het diffusiewereldmodel efficiënt en stabiel blijft over lange tijdsperioden. Het framework biedt een kwalitatieve analyse om de belangrijkheid van deze ontwerpkansen te demonstreren. DIAMOND stelt een nieuwe standaard vast met een gemiddelde menselijke genormaliseerde score van 1,46 op de gevestigde Atari 100k-benchmark, het hoogste voor agenten getraind volledig binnen een wereldmodel. Het opereren in beeldruimte stelt het diffusiewereldmodel van DIAMOND in staat om de omgeving naadloos te vervangen, waardoor meer inzicht wordt verkregen in wereldmodel- en agentgedrag. Opvallend is dat de verbeterde prestatie in bepaalde games wordt toegeschreven aan een betere modellering van kritische visuele details. Het DIAMOND-framework modelleert de omgeving als een standaard POMDP of gedeeltelijk waarneembaar Markov-beslissingsproces met een verzameling staten, een verzameling discrete acties en een verzameling beeldobservaties. De overgangsfuncties beschrijven de omgevingsdynamica, en de beloningsfunctie kaart de overgangen naar schaalbare beloningen.

DIAMOND: Methodologie en Architectuur

In zijn kern zijn diffusiemodellen een klasse van generatieve modellen die een steekproef genereren door het verstorende proces om te keren en putten zwaar uit niet-evenwicht thermodynamica. Het DIAMOND-framework beschouwt een diffusieproces dat wordt geïndexeerd door een continue tijdsvariabele met overeenkomstige marginaals en randvoorwaarden met een traceerbare ongestructureerde prior-verdeling. Bovendien moet het DIAMOND-framework, om een generatief model te verkrijgen dat kaart van ruis naar gegevens, het proces omkeren, waarbij het omkeren ook een diffusieproces is dat achteruit loopt in de tijd. Bovendien is het, op een gegeven moment in de tijd, niet triviaal om de score-functie te schatten, omdat het DIAMOND-framework geen toegang heeft tot de werkelijke score-functie, en het model overwint deze hindernis door een score-gebaseerd doel te implementeren, een benadering die een framework in staat stelt om een score-model te trainen zonder de onderliggende score-functie te kennen. Het score-gebaseerde diffusiemodel biedt een onvoorwaardelijk generatief model. Echter, een voorwaardelijk generatief model van omgevingsdynamica is vereist om als een wereldmodel te dienen, en om deze reden kijkt het DIAMOND-framework naar het algemene geval van de POMDP-benadering, waarin het framework gebruik kan maken van eerdere observaties en acties om de onbekende Markoviaanse staat te benaderen. Zoals wordt aangetoond in Figuur 1, maakt het DIAMOND-framework gebruik van deze geschiedenis om een diffusiemodel te conditioneren, om de volgende observatie rechtstreeks te schatten en te genereren. Hoewel het DIAMOND-framework in theorie gebruik kan maken van elke SDE- of ODE-oplosser, is er een compromis tussen NFE of het aantal functie-evaluaties en de kwaliteit van de steekproef, die de inferentiekosten van diffusiemodellen aanzienlijk beïnvloedt.

Op basis van deze inzichten, laten we nu kijken naar de praktische realisatie van het DIAMOND-framework van een diffusie-gebaseerd wereldmodel, inclusief de drift- en diffusiecoëfficiënten die overeenkomen met een bepaalde keuze van diffusiebenadering. In plaats van te kiezen voor DDPM, een natuurlijk geschikte kandidaat voor de taak, bouwt het DIAMOND-framework voort op de EDM-formulering en beschouwt een perturbatiekernel met een reëelwaardige functie van diffusietijd, genaamd de ruisplanning. Het framework selecteert de voorwaardelijke coëfficiënten om de invoer- en uitvoervariatie voor elke spraakniveau te behouden. De netwerktraining combineert signaal en ruis adaptief, afhankelijk van het degradatieniveau, en wanneer de ruis laag is, wordt het doel de verschil tussen het schone en het verstoord signaal, d.w.z. de toegevoegde Gaussische ruis. Intuïtief voorkomt dit dat het trainingsdoel triviaal wordt in de lage-ruisregime. In de praktijk is dit doel van hoge variantie aan de uiteinden van de ruisplanning, dus het model steekt de ruisniveau van een log-normale verdeling, die empirisch is gekozen, om de training rond de middelste ruisgebieden te concatenen. Het DIAMOND-framework maakt gebruik van een standaard U-Net 2D-component voor het vectorveld en houdt een buffer van eerdere observaties en acties die het framework gebruikt om zichzelf te conditioneren. Het DIAMOND-framework concateneert vervolgens deze eerdere observaties met de volgende luide observatie en invoeracties via adaptieve groepsnormalisatielagen in de resterende blokken van de U-Net.

DIAMOND: Experimenten en Resultaten

Voor een uitgebreide evaluatie kiest het DIAMOND-framework voor de Atari 100k-benchmark. De Atari 100k-benchmark bestaat uit 26 games die zijn ontworpen om een breed scala aan agentmogelijkheden te testen. In elke game is een agent beperkt tot 100k acties in de omgeving, wat ongeveer gelijk is aan 2 uur van menselijke gameplay, om het spel te leren voordat het wordt geëvalueerd. Voor vergelijking trainen onbeperkte Atari-agenten meestal 50 miljoen stappen, wat een 500-voudige toename in ervaring vertegenwoordigt. We hebben DIAMOND getraind vanaf scratch met 5 willekeurige zaden voor elke game. Elke trainingsrun vereiste ongeveer 12 GB VRAM en duurde ongeveer 2,9 dagen op een enkele Nvidia RTX 4090, wat neerkomt op 1,03 GPU-jaren in totaal. De volgende tabel biedt de score voor alle games, het gemiddelde en de IQM of interquartielgemiddelde van mens-genormaliseerde scores.

Volgend op de beperkingen van puntgeschattingen, biedt het DIAMOND-framework gestratificeerde bootstrap-vertrouwen in het gemiddelde en de IQM of interquartielgemiddelde van mens-genormaliseerde scores, evenals prestatieprofielen en aanvullende metrieken, zoals samengevat in de volgende figuur.

De resultaten laten zien dat DIAMOND uitzonderlijk goed presteert over de benchmark, waarbij het de menselijke spelers overtreft in 11 games en een bovengemiddelde HNS van 1,46 bereikt, waarmee het een nieuw record vestigt voor agenten getraind volledig binnen een wereldmodel. Bovendien is de IQM van DIAMOND vergelijkbaar met STORM en overtreft het alle andere baselines. DIAMOND blinkt uit in omgevingen waarbij het vastleggen van kleine details cruciaal is, zoals Asterix, Breakout en RoadRunner. Bovendien, zoals eerder besproken, heeft het DIAMOND-framework de flexibiliteit om elke diffusiemodel in zijn pijplijn te implementeren, hoewel het kiest voor de EDM-benadering, het zou een natuurlijke keuze zijn geweest om te kiezen voor het DDPM-model, aangezien het al wordt geïmplementeerd in tal van beeldgeneratie-toepassingen. Om de EDM-benadering te vergelijken met de DDPM-implementatie, traint het DIAMOND-framework beide varianten met dezelfde netwerkarchitectuur op dezelfde gedeelde statische dataset met meer dan 100k frames verzameld met een expertbeleid. Het aantal denoiseringsstappen is rechtstreeks gerelateerd aan de inferentiekosten van het wereldmodel, en dus zullen minder stappen de kosten van het trainen van een agent op verbeelde trajecten verlagen. Om ervoor te zorgen dat ons wereldmodel computationeel vergelijkbaar blijft met andere baselines, zoals IRIS, die 16 NFE per tijdstap vereist, streven we ernaar om niet meer dan tientallen denoiseringsstappen te gebruiken, het liefst minder. Echter, het instellen van het aantal denoiseringsstappen te laag kan de visuele kwaliteit verslechteren, waardoor accumulatie van fouten optreedt. Om de stabiliteit van verschillende diffusievarianten te beoordelen, tonen we verbeelde trajecten die autoregressief zijn gegenereerd tot t = 1000 tijdstappen in de volgende figuur, met behulp van verschillende aantallen denoiseringsstappen n ≤ 10.

We observeren dat het gebruik van DDPM (a) in deze regime resulteert in ernstige accumulatie van fouten, waardoor het wereldmodel snel uit de distributie gaat. In tegenstelling tot het EDM-gebaseerde diffusiewereldmodel (b) blijft het veel stabiler over lange tijdsperioden, zelfs met slechts één denoiseringsstap. Verbeelde trajecten met diffusiewereldmodellen op basis van DDPM (links) en EDM (rechts) worden getoond. De initiële observatie bij t = 0 is hetzelfde voor beide, en elke rij correspondeert met een afnemend aantal denoiseringsstappen n. We observeren dat DDPM-gebaseerde generatie lijdt aan accumulatie van fouten, waarbij kleinere aantallen denoiseringsstappen snellere accumulatie van fouten veroorzaken. In tegenstelling tot het EDM-gebaseerde wereldmodel van DIAMOND blijft het veel stabiler, zelfs voor n = 1. De optimale enkele-stapvoorspelling is de verwachting over mogelijke reconstructies voor een gegeven luid signaal, die buiten de distributie kan liggen als de posterieure verdeling multimodaal is. Terwijl sommige games, zoals Breakout, deterministische overgangen hebben die nauwkeurig kunnen worden gemodelleerd met één denoiseringsstap, vertonen andere games partiële observabiliteit, waardoor multimodale observatieverdelingen ontstaan. In deze gevallen is een iteratieve oplosser nodig om de steekproefprocedure naar een specifieke modus te leiden, zoals wordt geïllustreerd in de game Boxing in de volgende figuur. Derhalve stelt het DIAMOND-framework n = 3 in alle onze experimenten.

De bovenstaande figuur vergelijkt enkele-stap (bovenste rij) en multi-stap (onderste rij) sampling in Boxing. De bewegingen van de zwarte speler zijn onvoorspelbaar, waardoor enkele-stap denoiseringsinterpolatie tussen mogelijke resultaten leidt, resulterend in vage voorspellingen. In tegenstelling tot multi-stap sampling produceert een duidelijk beeld door de generatie naar een specifieke modus te leiden. Interessant is dat, aangezien het beleid de witte speler controleert, zijn acties bekend zijn bij het wereldmodel, waardoor twijfel wordt geëlimineerd. Derhalve voorspellen zowel enkele-stap als multi-stap sampling correct de positie van de witte speler.

In de bovenstaande figuur vertonen de verbeelde trajecten van DIAMOND over het algemeen een hogere visuele kwaliteit en zijn ze getrouwer aan de werkelijke omgeving in vergelijking met die van IRIS. De trajecten gegenereerd door IRIS bevatten visuele inconsistenties tussen frames (gemarkeerd door witte dozen), zoals vijanden die worden weergegeven als beloningen en vice versa. Hoewel deze inconsistenties mogelijk alleen een paar pixels beïnvloeden, kunnen ze de versterking van het leren aanzienlijk beïnvloeden. Een agent streeft meestal naar het richten van beloningen en het vermijden van vijanden, dus deze kleine visuele discrepanties kunnen het leren van een optimaal beleid moeilijker maken. De figuur toont opeenvolgende frames verbeeld met IRIS (links) en DIAMOND (rechts). De witte dozen benadrukken inconsistenties tussen frames, die alleen voorkomen in trajecten gegenereerd met IRIS. In Asterix (bovenste rij) wordt een vijand (oranje) een beloning (rood) in de tweede frame, vervolgens terugkeert naar een vijand in de derde en opnieuw naar een beloning in de vierde. In Breakout (middelste rij) zijn de stenen en de score inconsistent tussen frames. In Road Runner (onderste rij) worden de beloningen (kleine blauwe stippen op de weg) inconsistent weergegeven tussen frames. Deze inconsistenties komen niet voor bij DIAMOND. In Breakout wordt de score betrouwbaar bijgewerkt met +7 wanneer een rode steen wordt gebroken.

Conclusie

In dit artikel hebben we gesproken over DIAMOND, een versterking van het leren agent getraind binnen een diffusiewereldmodel. Het DIAMOND-framework maakt zorgvuldige ontwerpkansen om ervoor te zorgen dat het diffusiewereldmodel efficiënt en stabiel blijft over lange tijdsperioden. Het framework biedt een kwalitatieve analyse om de belangrijkheid van deze ontwerpkansen te demonstreren. DIAMOND stelt een nieuwe standaard vast met een gemiddelde menselijke genormaliseerde score van 1,46 op de gevestigde Atari 100k-benchmark, het hoogste voor agenten getraind volledig binnen een wereldmodel. Het opereren in beeldruimte stelt het diffusiewereldmodel van DIAMOND in staat om de omgeving naadloos te vervangen, waardoor meer inzicht wordt verkregen in wereldmodel- en agentgedrag. Opvallend is dat de verbeterde prestatie in bepaalde games wordt toegeschreven aan een betere modellering van kritische visuele details. Het DIAMOND-framework modelleert de omgeving als een standaard POMDP of gedeeltelijk waarneembaar Markov-beslissingsproces met een verzameling staten, een verzameling discrete acties en een verzameling beeldobservaties. De overgangsfuncties beschrijven de omgevingsdynamica, en de beloningsfunctie kaart de overgangen naar schaalbare beloningen.