Det var i 2018, da ideen om forsterkingslæring i sammenheng med en neural nettverksverdensmodell først ble introdusert, og snart ble denne grunnleggende prinsippet anvendt på verdensmodeller. Noen av de fremtredende modellene som implementerer forsterkingslæring var Dreamer-rammen, som innførte forsterkingslæring fra den latente rommet til en rekurrerende tilstand rommodell. DreamerV2 viste at bruk av diskrete latenter kunne resultere i reduserte feil, og DreamerV3-rammen kunne oppnå menneske-lignende ytelse på en rekke oppgaver over forskjellige domener med faste hyperparametere.

Videre kan det trekkes paralleller mellom bilde-genereringsmodeller og verdensmodeller, som indikerer at fremgangen i generative visningsmodeller kunne replikeres for å惠 verdensmodellene. Siden bruk av transformatorer i naturlig språkbehandling rammer vant popularitet, oppstod DALL-E og VQGAN-rammene. Disse rammene implementerte diskrete autoencodere for å konvertere bilder til diskrete token, og kunne bygge svært kraftfulle og effektive tekst-til-bilde-genereringsmodeller ved å utnytte sekvensmodelleringsevnen til autoregressive transformatorer. Samtidig vant diffusjonsmodeller popularitet, og i dag har diffusjonsmodeller etablert seg som en dominant paradigme for høyoppløselig bilde-generering. Takket være diffusjonsmodellenes og forsterkingslæringens evner, forsøkes det å kombinere de to tilnærmingene, med målet å utnytte diffusjonsmodellenes fleksibilitet som trajektori-modeller, belønning-modeller, planleggere og som politikk for data-forbedring i offline-forsterkingslæring.

Verdensmodeller tilbyr en lovende metode for å trene forsterkingslæring-agenter trygt og effektivt. Tradisjonelt bruker disse modellene sekvenser av diskrete latente variabler for å simulere miljø-dynamikk. Imidlertid kan denne komprimeringen overse visuelle detaljer som er avgjørende for forsterkingslæring. Samtidig har diffusjonsmodeller steget i popularitet for bilde-generering, og utfordrer tradisjonelle metoder som bruker diskrete latenter. Inspirert av denne skiftet, skal vi i denne artikkelen snakke om DIAMOND (DIffusion As a Model Of eNvironment Dreams), en forsterkingslæring-agent som er trent innen en diffusjonsverdensmodell. Vi skal utforske de nødvendige design-valg for å gjøre diffusjon egnet for verdensmodellering og vise at forbedrede visuelle detaljer fører til bedre agent-ytelse. DIAMOND setter en ny standard på den konkurranse-Atari 100k-testen, med en gjennomsnittlig menneske-normalisert score på 1,46, den høyeste for agenter som er trent helt innen en verdensmodell.

DIAMOND: DIffusion As a Model Of eNvironment Dreams

Verdensmodeller eller generative modeller av miljøer oppstår som en av de viktigste komponentene for generative agenter til å planlegge og resonere om sine miljøer. Selv om bruk av forsterkingslæring har oppnådd betydelig suksess i de siste årene, er modeller som implementerer forsterkingslæring kjent for å være prøve-ineffektive, noe som begrenser deres virkelige verden-applikasjoner betydelig. På den andre siden har verdensmodeller demonstrert sin evne til å effektivt trene forsterkingslæring-agenter over forskjellige miljøer med en betydelig forbedret prøve-effektivitet, noe som tillater modellen å lære fra virkelige verden-erfaringer. Nylige verdensmodell-rammer modellerer vanligvis miljø-dynamikk som en sekvens av diskrete latente variabler, med modellen som diskretiserer den latente rommet for å unngå feil-akkumulering over multi-trinn-tids-horisonter. Selv om tilnærmingen kan levere betydelige resultater, er den også assosiert med en informasjons-tap, noe som fører til en reduksjon i rekonstruksjonskvalitet og generellhet. Informasjons-tapet kan bli en betydelig hindring for virkelige scenarioer som krever informasjonen å være godt definert, som trening av autonome kjøretøy. I slike oppgaver kan små endringer eller detaljer i den visuelle innmatningen, som fargen på trafikklys eller blinklys på kjøretøyet foran, kan endre agentens politikk. Selv om økning av antall diskrete latenter kan hjelpe med å unngå informasjons-tap, øker det beregningskostnadene betydelig.

Videre, i de siste årene, har diffusjonsmodeller oppstått som den dominerende tilnærmingen for høykvalitets bilde-genererings-rammer, siden rammer bygget på diffusjonsmodeller lærer å reversere en støy-prosess, og konkurrerer direkte med noen av de mer etablerte tilnærmingene som modellerer diskrete token, og tilbyr derfor en lovende alternativ for å eliminere behovet for diskretisering i verdensmodellering. Diffusjonsmodeller er kjent for sin evne til å være lett kondisjonert og å modellere komplekse, multi-modale distribusjoner uten mod-kollaps. Disse egenskapene er avgjørende for verdensmodellering, da kondisjonering tillater en verdensmodell å reflektere en agent’s handlinger nøyaktig, noe som fører til mer pålitelig kreditt-tildeling. I tillegg tilbyr modellering av multi-modale distribusjoner en større diversitet av trenings-scenarier for agenten, noe som forbedrer dens totale ytelse.

Bygget på disse egenskapene, er DIAMOND (DIffusion As a Model Of eNvironment Dreams) en forsterkingslæring-agent som er trent innen en diffusjonsverdensmodell. DIAMOND-rammen gjør omhyggelige design-valg for å sikre at dens diffusjonsverdensmodell forblir effektiv og stabil over lange tids-horisonter. Rammen tilbyr en kvalitativ analyse for å demonstrere viktigheten av disse design-valg. DIAMOND setter en ny standard med en gjennomsnittlig menneske-normalisert score på 1,46 på den etablerte Atari 100k-benchmark, den høyeste for agenter som er trent helt innen en verdensmodell. Å operere i bilde-rommet tillater DIAMOND’s diffusjonsverdensmodell å erstatter miljøet uten å påvirke agentens atferd. Notabelt er den forbedrede ytelsen i visse spill tilskrevet bedre modellering av kritiske visuelle detaljer. DIAMOND-rammen modellerer miljøet som en standard POMDP eller Partially Observable Markov Decision Process med en mengde tilstander, en mengde diskrete handlinger og en mengde bilde-observasjoner. Overgangs-funksjonene beskriver miljø-dynamikken, og belønning-funksjonen kartlegger overgangene til skalar-belønninger.

DIAMOND: Metodologi og Arkitektur

I kjernen er diffusjonsmodeller en klasse av generative modeller som genererer en prøve ved å reversere støy-prosessen, og trekker tungt inspirasjon fra ikke-likvid termodynamikk. DIAMOND-rammen betrakter en diffusjons-prosess indeksert av en kontinuerlig tid-variabel med tilhørende marginaler og grense-betingelser med en tractabel ustrukturert prior-distribusjon. Videre, for å oppnå en generativ modell som kartlegger fra støy til data, må DIAMOND-rammen reversere prosessen, med reversering-prosessen også være en diffusjons-prosess som kjører bakover i tid. Videre, på et gitt tidspunkt, er det ikke trivielt å estimere score-funksjonen siden DIAMOND-rammen ikke har tilgang til den sanne score-funksjonen, og modellen overvinner denne hindringen ved å implementere score-matching-objektet, en tilnærming som tillater en ramme å trene en score-modell uten å kjenne den underliggende score-funksjonen. Score-basert diffusjonsmodell tilbyr en ubetinget generativ modell. Imidlertid er en betinget generativ modell av miljø-dynamikk nødvendig for å fungere som en verdensmodell, og for å fungere dette formålet, ser DIAMOND-rammen på den generelle tilfelle av POMDP-tilnærmingen, hvor rammen kan bruke tidligere observasjoner og handlinger til å approksimere den ukjente markovianske tilstanden. Som demonstrert i Figur 1, bruker DIAMOND-rammen denne historien til å kondisjonere en diffusjonsmodell for å estimere og generere den neste observasjonen direkte. Selv om DIAMOND-rammen i teorien kan bruke hvilken som helst SDE eller ODE-løser, er det en avveining mellom NFE eller Antall funksjons-evalueringer og prøve-kvalitet som påvirker inferens-kostnadene til diffusjonsmodellene betydelig.

Bygget på disse lærdommene, la oss nå se på den praktiske realiseringen av DIAMOND-rammen av en diffusjons-basert verdensmodell, inkludert drift- og diffusjons-koeffisienter som tilhører en bestemt valg av diffusjons-tilnærming. I stedet for å velge DDPM, en naturlig egnet kandidat for oppgaven, bygger DIAMOND-rammen på EDM-formuleringen og betrakter en perturbasjons-kernel med en reell-verdi-funksjon av diffusjons-tid kalt støy-skjedulen. Rammen velger forhånds-betingelsene for å holde inn- og ut-gang-variansen for enhver stemme-nivå. Netverks-treningen blandes signal og støy adaptivt avhengig av degraderings-nivået, og når støyen er lav, og målet blir forskjellen mellom den rene og den perturberte signalen, dvs. den tilføyde gaussiske støyen. Intuitivt forhindrer dette at trening-objektet blir trivielt i lav-støy-regimet. I praksis er dette objektet høyt-varians på ytterpunktene av støy-skjedulen, så modellen sampler støy-nivået fra en log-normal-distribusjon valgt empirisk for å konsolidere treningen rundt medium-støy-regionene. DIAMOND-rammen bruker en standard U-Net 2D-komponent for vektor-feltet og holder en buffer av tidligere observasjoner og handlinger som rammen bruker til å kondisjonere seg selv. DIAMOND-rammen kobler så disse tidligere observasjonene til den neste støy-observasjonen og inn-aksjoner gjennom adaptive gruppe-normaliserings-lag i rest-lagene av U-Net.

DIAMOND: Eksperimenter og Resultater

For en omfattende evaluering, velger DIAMOND-rammen Atari 100k-benchmarken. Atari 100k-benchmarken består av 26 spill designet for å teste en rekke agent-egenskaper. I hvert spill er en agent begrenset til 100k handlinger i miljøet, som tilsvarer omtrent 2 timer av menneske-spill, for å lære spillet før evaluering. For sammenligning, ubegrensede Atari-agenter trener vanligvis i 50 millioner trinn, noe som representerer en 500-gangs økning i erfaring. Vi trente DIAMOND fra scratch ved å bruke 5 tilfeldige frø for hvert spill. Hver trening-runde krevde omtrent 12GB av VRAM og tok omtrent 2,9 dager på en enkelt Nvidia RTX 4090, noe som tilsvarer 1,03 GPU-år i alt. Følgende tabell gir scoren for alle spill, gjennomsnittet og IQM eller interkvartil-gjennomsnittet av menneske-normaliserte score.

Etter begrensningene av punkt-estimater, tilbyr DIAMOND-rammen stratifisert bootstrap-konfidens i gjennomsnittet og IQM eller interkvartil-gjennomsnittet av menneske-normaliserte score, samt ytelses-profiler og andre metrikker, som summeres opp i følgende figur.

Resultatene viser at DIAMOND utfører seg eksepsjonelt godt over hele benchmarken, overgår menneske-spillere i 11 spill og oppnår en overmenneskelig gjennomsnittlig HNS på 1,46, noe som setter en ny rekord for agenter som er trent helt innen en verdensmodell. I tillegg er DIAMOND’s IQM sammenlignbar med STORM og overgår alle andre referanse-verdier. DIAMOND utmerker seg i miljøer hvor fanget av små detaljer er avgjørende, som Asterix, Breakout og RoadRunner. Videre, som diskutert tidligere, har DIAMOND-rammen fleksibiliteten til å implementere noen som helst diffusjonsmodell i sin pipeline, selv om den velger EDM-tilnærmingen, ville det vært en naturlig valg å velge DDPM-modellen siden den allerede er implementert i mange bilde-genererings-applikasjoner. For å sammenligne EDM-tilnærmingen mot DDPM-implementeringen, trener DIAMOND-rammen begge variantene med samme nettverks-arkitektur på samme delte statiske datasett med over 100k-rammer samlet inn med en ekspert-politikk. Antall av-støy-trinn er direkte relatert til inferens-kostnadene til verdensmodellen, og færre trinn vil redusere kostnadene til å trene en agent på forestilte trajektorier. For å sikre at vår verdensmodell forblir komputasjonelt sammenlignbar med andre referanse-verdier, som IRIS som krever 16 NFE per tids-steg, sikter vi på å bruke ikke mer enn titalls av-støy-trinn, helst færre. Imidlertid kan å sette antall av-støy-trinn for lavt føre til en reduksjon i visuell kvalitet, noe som kan føre til feil-akkumulering. For å vurdere stabiliteten til forskjellige diffusjons-variante, viser vi forestilte trajektorier generert auto-regressivt opp til t = 1000 tids-steg i følgende figur, ved å bruke forskjellige antall av-støy-trinn n ≤ 10.

Vi observerer at å bruke DDPM (a) i dette regime resulterer i alvorlige feil-akkumuleringer, noe som fører til at verdensmodellen raskt drifter ut av distribusjonen. I kontrast forblir EDM-basert diffusjons-verdensmodell (b) mye mer stabil over lange tids-horisonter, selv med bare ett av-støy-trinn. Forestilte trajektorier med diffusjons-verdensmodeller basert på DDPM (venstre) og EDM (høyre) vises. Den innledende observasjonen ved t = 0 er den samme for begge, og hver rad korresponderer til et synkende antall av-støy-trinn n. Vi observerer at DDPM-basert generering lider av feil-akkumuleringer, med færre antall av-støy-trinn som fører til raskere feil-akkumulering. I kontrast forblir DIAMOND’s EDM-baserte verdensmodell mye mer stabil, selv for n = 1. Den optimale enkelt-trinns-prediksjonen er forventningen over mulige rekonstruksjoner for en gitt støy-innmatning, som kan være utenfor distribusjonen hvis den posteriøre distribusjonen er multi-modal. Mens noen spill, som Breakout, har deterministiske overganger som kan modelleres nøyaktig med ett enkelt av-støy-trinn, utviser andre spill partiell observabilitet, noe som resulterer i multi-modale observasjons-distribusjoner. I disse tilfellene er en iterativ løser nødvendig for å guide sampling-prosedyren mot en bestemt modus, som illustrert i spillet Boxing i følgende figur. Derfor setter DIAMOND-rammen n = 3 i alle våre eksperimenter.

Den ovenstående figuren sammenligner enkelt-trinns (øverste rad) og multi-trinns (nederste rad) sampling i Boxing. Bevegelsene til den svarte spilleren er uforutsigbare, noe som fører til at enkelt-trinns av-støyning interpolerer mellom mulige resultater, noe som resulterer i uklare prediksjoner. I kontrast produserer multi-trinns sampling en klar bilde ved å guide genereringen mot en bestemt modus. Interessant nok, siden politikken kontrollerer den hvite spilleren, er hans handlinger kjent for verdensmodellen, noe som eliminerer usikkerheten. Derfor prediker både enkelt-trinns og multi-trinns sampling korrekt den hvite spillerens posisjon.

I den ovenstående figuren viser trajektoriene forestilt av DIAMOND generelt høyere visuell kvalitet og er mer tro mot den sanne miljøet sammenlignet med de forestilte av IRIS. Trajektoriene generert av IRIS inneholder visuelle inkonsistenser mellom rammer (høydet av hvite bokser), som fiender som vises som belønninger og vice-versa. Selv om disse inkonsistensene kan bare påvirke noen få piksler, kan de ha en betydelig innvirkning på forsterkingslæring. For eksempel vil en agent vanligvis prøve å målrette belønninger og unngå fiender, så disse små visuelle diskrepanser kan gjøre det mer vanskelig å lære en optimal politikk. Figuren viser påfølgende rammer forestilt med IRIS (venstre) og DIAMOND (høyre). De hvite boksene høyder inkonsistenser mellom rammer, som bare oppstår i trajektorier generert med IRIS. I Asterix (øverste rad) blir en fiende (oransje) en belønning (rød) i den andre rammen, og så igjen en fiende i den tredje, og igjen en belønning i den fjerde. I Breakout (midterste rad) er mursteinene og poengsummen inkonsistente mellom rammer. I Road Runner (nederste rad) er belønningene (små blå punkter på veien) inkonsistente mellom rammer. Disse inkonsistensene oppstår ikke med DIAMOND. I Breakout er poengsummen pålitelig oppdatert med +7 når en rød murstein brytes.

Konklusjon

I denne artikkelen har vi snakket om DIAMOND, en forsterkingslæring-agent som er trent innen en diffusjonsverdensmodell. DIAMOND-rammen gjør omhyggelige design-valg for å sikre at dens diffusjonsverdensmodell forblir effektiv og stabil over lange tids-horisonter. Rammen tilbyr en kvalitativ analyse for å demonstrere viktigheten av disse design-valg. DIAMOND setter en ny standard med en gjennomsnittlig menneske-normalisert score på 1,46 på den etablerte Atari 100k-benchmark, den høyeste for agenter som er trent helt innen en verdensmodell. Å operere i bilde-rommet tillater DIAMOND’s diffusjonsverdensmodell å erstatter miljøet uten å påvirke agentens atferd. Notabelt er den forbedrede ytelsen i visse spill tilskrevet bedre modellering av kritiske visuelle detaljer. DIAMOND-rammen modellerer miljøet som en standard POMDP eller Partially Observable Markov Decision Process med en mengde tilstander, en mengde diskrete handlinger og en mengde bilde-observasjoner. Overgangs-funksjonene beskriver miljø-dynamikken, og belønning-funksjonen kartlegger overgangene til skalar-belønninger.