Det var 2018 när idén om förstärkt inlärning i sammanhanget med en neural nätverksvärldsmodell först introducerades, och snart tillämpades denna grundläggande princip på världsmodeller. Några av de framstående modellerna som implementerar förstärkt inlärning var Dreamer-ramverket, som introducerade förstärkt inlärning från den latenta utrymmet för en återkommande tillståndsrummodell. DreamerV2 visade att användningen av diskreta latenter kan resultera i minskade ackumulerade fel, och DreamerV3-ramverket kunde uppnå mänsklig prestanda på en serie uppgifter över olika domäner med fasta hyperparametrar.

Dessutom kan paralleller dras mellan bildgenereringsmodeller och världsmodeller, vilket indikerar att framstegen inom generativa synmodeller kan återges för att gynna världsmodellerna. Sedan användningen av transformer i naturligt språkbehandlingsramverk blev populär, uppstod DALL-E och VQGAN-ramverken. Ramverken implementerade diskreta autoencoders för att konvertera bilder till diskreta token, och kunde bygga mycket kraftfulla och effektiva text-till-bild-genereringsmodeller genom att utnyttja sekvensmodelleringsförmågan hos de autoregressiva transformerarna. Samtidigt fick diffusionsmodeller drag, och idag har diffusionsmodeller etablerat sig som en dominant paradigm för högupplösningsbildgenerering. Tack vare diffusionsmodellernas och förstärkt inlärnings förmågor görs försök att kombinera de två tillvägagångssätten, i syfte att dra nytta av diffusionsmodellernas flexibilitet som bana-, belönings-, planerings- och policy-modeller för dataförstärkning i offline-förstärkt inlärning.

Världsmodeller erbjuder ett lovande sätt att träna förstärkt inlärningsagenter på ett säkert och effektivt sätt. Traditionellt använder dessa modeller sekvenser av diskreta latenta variabler för att simulera miljödynamik. Men denna komprimering kan försumma visuella detaljer som är avgörande för förstärkt inlärning. Samtidigt har diffusionsmodeller blivit populära för bildgenerering, och utmanar traditionella metoder som använder diskreta latenter. Inspirerad av denna förändring, kommer vi i den här artikeln att tala om DIAMOND (Diffusion som en modell för miljödrömmar), en förstärkt inlärningsagent som tränas inom en diffusionsvärldsmodell. Vi kommer att undersöka de nödvändiga designvalen för att göra diffusionslämplig för världsmodellering och visa att förbättrade visuella detaljer leder till bättre agentprestanda. DIAMOND sätter en ny benchmark på den konkurrenskraftiga Atari 100k-testen, med en genomsnittlig mänsklig normaliserad poäng på 1,46, den högsta för agenter som tränats helt inom en världsmodell.

DIAMOND : Diffusion som en modell för miljödrömmar

Världsmodeller eller generativa modeller av miljöer är på väg att bli en av de viktigaste komponenterna för generativa agenter att planera och resonera om sina miljöer. Även om användningen av förstärkt inlärning har uppnått betydande framgång under de senaste åren, är modeller som implementerar förstärkt inlärning kända för att vara exempel på ineffektivitet, vilket begränsar deras verkliga världstillämpningar avsevärt. Å andra sidan har världsmodeller visat sin förmåga att effektivt träna förstärkt inlärningsagenter i olika miljöer med avsevärt förbättrad exempelhastighet, vilket möjliggör för modellen att lära sig från verkliga världsupplevelser. Nya världsmodellramverk modellerar vanligtvis miljödynamik som en sekvens av diskreta latenta variabler, med modellen som diskretiserar det latenta utrymmet för att undvika ackumulerade fel över multi-stegs tidshorisonter. Även om tillvägagångssättet kan leverera betydande resultat, är det också associerat med en informationsförlust, vilket leder till en förlust av rekonstruktionskvalitet och en förlust av allmänhet. Informationsförlusten kan bli en betydande hinder för verkliga scenarier som kräver att informationen är väldefinierad, som t.ex. utbildning av autonoma fordon. I sådana uppgifter kan små förändringar eller detaljer i den visuella inmatningen, som färgen på trafikljuset eller vändindikatorn på fordonet framför, ändra agentens policy. Även om en ökning av antalet diskreta latenter kan hjälpa till att undvika informationsförlust, ökar det beräkningskostnaderna avsevärt.

Dessutom har diffusionsmodeller under de senaste åren uppstått som den dominerande metoden för högkvalitativa bildgenereringsramverk, eftersom ramverk byggda på diffusionsmodeller lär sig att reversera en störningsprocess och direkt konkurrerar med några av de mer etablerade tillvägagångssätten som modellerar diskreta token, och erbjuder därför ett lovande alternativ för att eliminera behovet av diskretisering i världsmodellering. Diffusionsmodeller är kända för sin förmåga att konditioneras lätt och för att flexibelt modellera komplexa, multimodala distributioner utan modkollaps. Dessa attribut är avgörande för världsmodellering, eftersom konditionering möjliggör för en världsmodell att exakt återspegla en agents handlingar, vilket leder till en mer tillförlitlig kredit tilldelning. Dessutom erbjuder modellering av multimodala distributioner en större mångfald av tränings scenarier för agenten, vilket förbättrar dess övergripande prestanda.

Byggande på dessa egenskaper, är DIAMOND (Diffusion som en modell för miljödrömmar) en förstärkt inlärningsagent som tränas inom en diffusionsvärldsmodell. DIAMOND-ramverket gör noggranna designval för att säkerställa att dess diffusionsvärldsmodell förblir effektiv och stabil över långa tidshorisonter. Ramverket tillhandahåller en kvalitativ analys för att demonstrera vikten av dessa designval. DIAMOND sätter en ny standard på den etablerade Atari 100k-benchmarken, med en genomsnittlig mänsklig normaliserad poäng på 1,46, den högsta för agenter som tränats helt inom en världsmodell. Att verka i bildutrymme möjliggör för DIAMONDs diffusionsvärldsmodell att smidigt ersätta miljön, vilket ger en djupare inblick i världsmodell- och agentbeteende. Noterbart är att den förbättrade prestandan i vissa spel tillskrivs en bättre modellering av kritiska visuella detaljer. DIAMOND-ramverket modellerar miljön som en standard-POMDP eller delvis observerbar Markovbeslutsprocess med en uppsättning tillstånd, en uppsättning diskreta handlingar och en uppsättning bildobservationer. Övergångsfunktionerna beskriver miljödynamiken, och belöningsfunktionen kartar övergångarna till skalära belöningar.

DIAMOND : Metodik och Arkitektur

I sin kärna är diffusionsmodeller en klass av generativa modeller som genererar ett exempel genom att reversera störningsprocessen, och drar tung inspiration från icke-jämvikts termodynamik. DIAMOND-ramverket överväger en diffusionsprocess indexeras av en kontinuerlig tidvariabel med motsvarande marginaler och gränsvillkor med en hanterbar ostrukturerad priorfördelning. Dessutom, för att få en generativ modell, som kartar från brus till data, måste DIAMOND-ramverket reversera processen, med reversionprocessen också vara en diffusionsprocess, som körs bakåt i tiden. Dessutom, vid en given tidpunkt, är det inte trivialt att uppskatta score-funktionen, eftersom DIAMOND-ramverket inte har tillgång till den sanna score-funktionen, och modellen övervinner detta hinder genom att implementera en score-matchningsobjektiv, ett tillvägagångssätt som möjliggör för en ram att träna en score-modell utan att känna till den underliggande score-funktionen. Score-baserad diffusionsmodell tillhandahåller en ovillkorlig generativ modell. Men en villkorlig generativ modell av miljödynamik krävs för att fungera som en världsmodell, och för att uppfylla detta syfte, tittar DIAMOND-ramverket på det allmänna fallet av POMDP-tillvägagångssättet, där ramverket kan utnyttja tidigare observationer och handlingar för att approximera den okända markoviska tillståndet. Som visas i figur 1, använder DIAMOND-ramverket denna historia för att konditionera en diffusionsmodell, för att uppskatta och generera nästa observation direkt. Även om DIAMOND-ramverket i teorin kan använda valfri SDE eller ODE-lösare, finns det en avvägning mellan NFE eller antalet funktionsutvärderingar och exempelkvalitet som påverkar inferenskostnaden för diffusionsmodeller avsevärt.

Byggande på ovanstående lärdomar, låt oss nu titta på den praktiska realiseringen av DIAMOND-ramverket för en diffusionsbaserad världsmodell, inklusive drifthastighets- och diffusionskoefficienter som motsvarar ett visst val av diffusionsmetod. I stället för att välja DDPM, en naturlig kandidat för uppgiften, bygger DIAMOND-ramverket på EDM-formuleringen, och överväger en störningskärna med en reellvärdesfunktion av diffusions tid kallad brus schemat. Ramverket väljer förkonditionerare för att hålla ingångs- och utgångsvariansen för varje ljudnivå. Nätverksträningen blandar signal och brus adaptivt beroende på nedbrytningsnivån, och när bruset är lågt, och målet blir skillnaden mellan den rena och den störda signalen, dvs. det tillagda gaussiska bruset. Intuitivt förhindrar detta att träningsobjektivet blir trivialt i lågbrusregimen. I praktiken är detta objekt högvariations i ytterkanterna av brusschemat, så modellen provar brusnivån från en log-normalfördelning vald empiriskt för att konkatenera träningsomgången runt medelbrusområdena. DIAMOND-ramverket använder en standard U-Net 2D-komponent för vektorfältet, och håller ett buffert av tidigare observationer och handlingar som ramverket använder för att konditionera sig själv. DIAMOND-ramverket konkatenerar sedan dessa tidigare observationer till nästa bullrig observation, och ingångshandlingar genom adaptiva gruppnormaliseringslager i de residuala blocken i U-Net.

DIAMOND: Experiment och Resultat

För en omfattande utvärdering, väljer DIAMOND-ramverket Atari 100k-benchmarken. Atari 100k-benchmarken består av 26 spel som är utformade för att testa en bred uppsättning agentförmågor. I varje spel är en agent begränsad till 100 000 handlingar i miljön, vilket är ungefär lika med 2 timmars mänskligt spel, för att lära sig spelet innan utvärdering. För jämförelse, oberoende Atari-agenter tränas vanligtvis i 50 miljoner steg, vilket representerar en 500-faldig ökning av erfarenhet. Vi tränade DIAMOND från scratch med 5 slumpmässiga frön för varje spel. Varje träningskörning krävde cirka 12 GB VRAM och tog ungefär 2,9 dagar på en enda Nvidia RTX 4090, vilket motsvarar 1,03 GPU-år i totalt. Följande tabell tillhandahåller poängen för alla spel, medelvärdet och IQM eller interkvartilmedelvärdet av mänsklig-normaliserade poäng.

Följande begränsningar för punktskattningar, tillhandahåller DIAMOND-ramverket stratifierad bootstrap-förtroende för medelvärdet och IQM eller interkvartilmedelvärdet av mänsklig-normaliserade poäng, tillsammans med prestandaprofiler och ytterligare mått, som sammanfattas i följande figur.

Resultaten visar att DIAMOND presterar exceptionellt bra över hela benchmarken, och överträffar mänskliga spelare i 11 spel och uppnår en övermänsklig medel-HNS på 1,46, vilket sätter en ny rekord för agenter som tränats helt inom en världsmodell. Dessutom är DIAMONDS IQM jämförbar med STORM och överträffar alla andra baseline. DIAMOND excellerar i miljöer där fångst av små detaljer är avgörande, som Asterix, Breakout och RoadRunner. Dessutom, som diskuterats tidigare, har DIAMOND-ramverket flexibiliteten att implementera valfri diffusionsmodell i sin pipeline, även om det väljer EDM-tillvägagångssättet, skulle det ha varit ett naturligt val att välja DDPM-modellen, eftersom den redan implementeras i många bildgenereringsapplikationer. För att jämföra EDM-tillvägagångssättet mot DDPM-implementering, tränar DIAMOND-ramverket båda varianterna med samma nätverksarkitektur på samma delade statiska dataset med över 100 000 ramar insamlade med en expertpolicy. Antalet avbrusningssteg är direkt relaterat till inferenskostnaden för världsmodellen, och färre steg kommer att minska kostnaden för att träna en agent på föreställda banor. För att säkerställa att vår världsmodell förblir beräkningsmässigt jämförbar med andra baseline, som IRIS som kräver 16 NFE per tidssteg, syftar vi till att använda högst tiotals avbrusningssteg, helst färre. Men att ställa in antalet avbrusningssteg för lågt kan försämra den visuella kvaliteten, vilket leder till ackumulerade fel. För att bedöma stabiliteten hos olika diffusionsvarianter, visar vi föreställda banor som genereras autoregressivt upp till t = 1000 tidssteg i följande figur, med hjälp av olika antal avbrusningssteg n ≤ 10.

Vi observerar att användningen av DDPM (a) i detta regime resulterar i allvarliga ackumulerade fel, vilket orsakar att världsmodellen snabbt drar sig ur fördelningen. I kontrast, förblir den EDM-baserade diffusionsvärldsmodellen (b) mycket mer stabil över långa tidshorisonter, även med ett enda avbrusningssteg. Föreställda banor med diffusionsvärldsmodeller baserade på DDPM (vänster) och EDM (höger) visas. Den initiala observationen vid t = 0 är densamma för båda, och varje rad motsvarar ett minskande antal avbrusningssteg n. Vi observerar att DDPM-baserad generering lider av ackumulerade fel, med färre antal avbrusningssteg som leder till snabbare felackumulering. I kontrast, förblir DIAMONDS EDM-baserade världsmodell mycket mer stabil, även för n = 1. Den optimala enstegsprediktionen är förväntningen över möjliga rekonstruktioner för en given bullrig inmatning, som kan vara utanför fördelningen om den posteriorda fördelningen är multimodal. Medan vissa spel, som Breakout, har deterministiska övergångar som kan modelleras exakt med ett enda avbrusningssteg, uppvisar andra spel partiell observerbarhet, vilket resulterar i multimodala observationsfördelningar. I dessa fall är en iterativ lösare nödvändig för att vägleda urvalsproceduren mot en specifik mod, som illustreras i spelet Boxing i följande figur. Följaktligen, satte DIAMOND-ramverket n = 3 i alla våra experiment.

Figuren ovan jämför enstegs- (översta raden) och flerstegs- (nedre raden) sampling i Boxing. Rörelserna hos den svarta spelaren är oförutsägbara, vilket orsakar enstegsavbrusning att interpolera mellan möjliga resultat, vilket resulterar i suddiga förutsägelser. I kontrast, producerar flerstegssampling en tydlig bild genom att vägleda genereringen mot en specifik mod. Intressant nog, eftersom policyn kontrollerar den vita spelaren, är deras handlingar kända för världsmodellen, vilket eliminerar tvetydighet. Således, både enstegs- och flerstegssampling förutsäger korrekt den vita spelarens position.

I ovanstående figur, visar de föreställda banorna som genereras av DIAMOND i allmänhet en högre visuell kvalitet och är mer trogna den verkliga miljön jämfört med de som föreställs av IRIS. De banor som genereras av IRIS innehåller visuella inkonsekvenser mellan ramar (markerade med vita rutor), såsom fiender som visas som belöningar och vice versa. Även om dessa inkonsekvenser kan påverka endast ett fåtal pixlar, kan de ha en betydande inverkan på förstärkt inlärning. Till exempel, syftar en agent vanligtvis till att målmedvetet rikta sig mot belöningar och undvika fiender, så dessa små visuella diskrepanser kan göra det svårare att lära sig en optimal policy. Figuren visar på varandra följande ramar som föreställs av IRIS (vänster) och DIAMOND (höger). De vita rutorna markerar inkonsekvenser mellan ramar, som endast förekommer i banor som genereras av IRIS. I Asterix (översta raden), blir en fiende (orange) en belöning (röd) i den andra ramen, och återgår till en fiende i den tredje, och igen till en belöning i den fjärde. I Breakout (mittenraden), är tegelstenarna och poängen inkonsekventa mellan ramar. I Road Runner (nedre raden), är belöningarna (små blå prickar på vägen) inkonsekvent återgivna mellan ramar. Dessa inkonsekvenser förekommer inte med DIAMOND. I Breakout, uppdateras poängen tillförlitligt med +7 när en röd tegelsten bryts.

Slutsats

I denna artikel, har vi talat om DIAMOND, en förstärkt inlärningsagent som tränas inom en diffusionsvärldsmodell. DIAMOND-ramverket gör noggranna designval för att säkerställa att dess diffusionsvärldsmodell förblir effektiv och stabil över långa tidshorisonter. Ramverket tillhandahåller en kvalitativ analys för att demonstrera vikten av dessa designval. DIAMOND sätter en ny standard på den etablerade Atari 100k-benchmarken, med en genomsnittlig mänsklig normaliserad poäng på 1,46, den högsta för agenter som tränats helt inom en världsmodell. Att verka i bildutrymme möjliggör för DIAMONDs diffusionsvärldsmodell att smidigt ersätta miljön, vilket ger en djupare inblick i världsmodell- och agentbeteende. Noterbart är att den förbättrade prestandan i vissa spel tillskrivs en bättre modellering av kritiska visuella detaljer. DIAMOND-ramverket modellerar miljön som en standard-POMDP eller delvis observerbar Markovbeslutsprocess med en uppsättning tillstånd, en uppsättning diskreta handlingar och en uppsättning bildobservationer. Övergångsfunktionerna beskriver miljödynamiken, och belöningsfunktionen kartar övergångarna till skalära belöningar.