Det var 2018, när idén om förstärkning lärande i samband med ett neuralt nätverk introducerades världsmodellen först, och snart tillämpades denna grundläggande princip på världsmodeller. Några av de framträdande modellerna som implementerar förstärkningsinlärning var Dreamer-ramverket, som introducerade förstärkningsinlärning från det latenta rummet av en återkommande tillståndsmodell. DreamerV2 visade att användningen av diskreta latenter kan resultera i minskade kompounderingsfel, och DreamerV3-ramverket kunde uppnå människoliknande prestanda på en rad uppgifter över olika domäner med fasta hyperparametrar. 

Dessutom kan paralleller dras mellan bildgenereringsmodeller och världsmodeller, vilket indikerar att de framsteg som gjorts i generativa visionsmodeller skulle kunna replikeras till nytta för världsmodellerna. Ända sedan användningen av transformatorer i naturlig språkbehandling ramverk blev populära, DALL-E och VQGAN ramverk uppstod. Ramverken implementerade diskreta autokodare för att konvertera bilder till diskreta tokens och kunde bygga mycket kraftfulla och effektiva text-till-bild-generativa modeller genom att utnyttja sekvensmodelleringsförmågan hos de autoregressiva transformatorerna. Samtidigt fick diffusionsmodeller draghjälp och idag har diffusionsmodeller etablerat sig som ett dominerande paradigm för högupplöst bildgenerering. På grund av de möjligheter som spridningsmodeller och förstärkningsinlärning erbjuder, görs försök att kombinera de två tillvägagångssätten, i syfte att dra fördel av diffusionsmodellernas flexibilitet som banamodeller, belöningsmodeller, planerare och som policy för dataökning i offline förstärkningsinlärning. 

Världsmodeller erbjuder en lovande metod för att träna förstärkningsinlä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. Denna komprimering kan dock förbise visuella detaljer som är avgörande för förstärkningsinlärning. Samtidigt har diffusionsmodeller ökat i popularitet för bildgenerering, vilket utmanar traditionella metoder som använder diskreta latenter. Inspirerade av detta skifte kommer vi i den här artikeln att prata om DIAMOND (DIffusion As a Model Of Environment Dreams), en förstärkningsinlärningsagent utbildad inom en diffusionsvärldsmodell. Vi kommer att utforska de nödvändiga designvalen för att göra diffusion lämplig för världsmodellering och visa att förbättrade visuella detaljer leder till bättre agentprestanda. DIAMOND sätter ett nytt riktmärke för det konkurrenskraftiga Atari 100k-testet, och uppnår ett medelvärde för mänskligt normaliserat resultat på 1.46, det högsta för agenter utbildade helt inom en världsmodell. 

DIAMOND: Diffusion som en modell av miljödrömmar

Världsmodeller eller Generativa modeller av miljöer växer fram som en av de viktigare komponenterna för generativa agenter att planera och resonera kring sina miljöer. Även om användningen av förstärkningsinlärning har nått avsevärd framgång de senaste åren, är modeller som implementerar förstärkningsinlärning kända för att vara ineffektiva med prov, vilket avsevärt begränsar deras verkliga tillämpningar. Å andra sidan har världsmodeller visat sin förmåga att effektivt träna förstärkningsinlärningsagenter i olika miljöer med en avsevärt förbättrad proveffektivitet, vilket gör att modellen kan lära sig av verkliga erfarenheter. Nya ramverk för världsmodellering modellerar vanligtvis miljödynamik som en sekvens av diskreta latenta variabler, där modellen diskretiserar det latenta utrymmet för att undvika förvärrade fel över flerstegs tidshorisonter. Även om tillvägagångssättet kan ge betydande resultat, är det också förknippat med förlust av information, vilket leder till förlust av rekonstruktionskvalitet och förlust av allmänhet. Förlusten av information kan bli en betydande vägspärr för verkliga scenarier som kräver att informationen är väldefinierad, som att träna autonoma fordon. I sådana uppgifter kan små förändringar eller detaljer i den visuella inmatningen som färgen på trafikljuset eller blinkersen på det framförvarande fordonet ändra en agents policy. Även om en ökning av antalet diskreta latenter kan hjälpa till att undvika informationsförlust, ökar beräkningskostnaderna avsevärt. 

Vidare, under de senaste åren har diffusionsmodeller dykt upp som det dominerande tillvägagångssättet för ramverk för bildgenerering av hög kvalitet, eftersom ramverk byggda på diffusionsmodeller lär sig att vända en brusprocess och direkt konkurrerar med några av de mer väletablerade metoderna som modellerar diskreta tokens , 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 lätt konditioneras och att flexibelt modellera komplexa, multimodala distributioner utan mod kollaps. Dessa attribut är avgörande för världsmodellering, eftersom konditionering gör det möjligt för en världsmodell att exakt återspegla en agents agerande, vilket leder till mer tillförlitlig kreditgivning. Dessutom erbjuder modellering av multimodala distributioner en större mångfald av träningsscenarier för agenten, vilket förbättrar dess övergripande prestanda. 

Bygger på dessa egenskaper, DIAMOND, (DIffusion As a Model Of Environment Dreams), en förstärkningsinlärningsagent utbildad inom en diffusionsvärldsmodell. Ramverket DIAMOND gör noggranna designval för att säkerställa att dess spridningsvärldsmodell förblir effektiv och stabil över långa tidshorisonter. Ramverket tillhandahåller en kvalitativ analys för att visa vikten av dessa designval. DIAMOND sätter en ny state-of-the-art med ett medelvärde för mänskligt normaliserat resultat på 1.46 på det väletablerade Atari 100k-riktmärket, det högsta för agenter som helt utbildats inom en världsmodell. Genom att arbeta i bildrymden kan DIAMONDs diffusionsvärldsmodell sömlöst ersätta miljön, vilket ger större insikter i världsmodeller och agentbeteenden. Särskilt den förbättrade prestandan i vissa spel tillskrivs bättre modellering av kritiska visuella detaljer. DIAMOND-ramverket modellerar miljön som en standard POMDP eller Partially Observable Markov Decision Process med en uppsättning tillstånd, en uppsättning diskreta åtgärder och en uppsättning bildobservationer. Övergångsfunktionerna beskriver omgivningens dynamik, och belöningsfunktionen kartlägger övergångarna till skalära belöningar. Observationsfunktionen beskriver observationssannolikheterna, och avger bildobservationer, som sedan används av agenterna för att se miljöerna, eftersom de inte direkt kan komma åt tillstånden. Det primära syftet med tillvägagångssättet var att få fram en policy som kartlägger observationer till åtgärder med försök att maximera den förväntade rabattavkastningen med en diskonteringsfaktor. Världsmodeller är generativa modeller av miljön, och världsmodeller kan användas för att skapa simulerade miljöer för att träna förstärkningsinlärningsagenter i den verkliga miljön, och träna förstärkningsinlärningsagenter i världsmodellmiljön. Figur 1 visar den rullande fantasin hos DIAMOND-ramverket över tiden. 

DIAMOND: metodik och arkitektur

I sin kärna är diffusionsmodeller en klass av generativa modeller som genererar ett prov genom att vända brusprocessen och hämtar tung inspiration från termodynamik som inte är i jämvikt. DIAMOND-ramverket betraktar en diffusionsprocess som indexeras av en kontinuerlig tidsvariabel med motsvarande marginaler och gränsvillkor med en hanterbar ostrukturerad tidigare fördelning. Dessutom, för att erhålla en generativ modell, som kartlägger från brus till data, måste DIAMOND-ramverket vända processen, där återgångsprocessen också är en diffusionsprocess som går bakåt i tiden. Vidare, vid en given tidpunkt, är det inte trivialt att uppskatta poängfunktionen eftersom DIAMOND-ramverket inte har tillgång till den sanna poängfunktionen, och modellen övervinner detta hinder genom att implementera poängmatchningsmål, ett tillvägagångssätt som underlättar ett ramverk för att träna en poängmodell utan att känna till den underliggande poängfunktionen. Den poängbaserade diffusionsmodellen ger en ovillkorlig generativ modell. Emellertid krävs en villkorad generativ modell av miljödynamik för att fungera som en världsmodell, och för att tjäna detta syfte tittar DIAMOND-ramverket på det allmänna fallet med POMDP-metoden, där ramverket kan använda sig av tidigare observationer och åtgärder för att ungefär den okända markoviska staten. 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 tillgripa vilken SDE- eller ODE-lösare som helst, finns det en avvägning mellan NFE- eller antal funktionsutvärderingar och provkvaliteten som avsevärt påverkar slutsatskostnaden för diffusionsmodeller. 

Med utgångspunkt i ovanstående lärdomar, låt oss nu titta på det praktiska förverkligandet av DIAMOND-ramverket för en diffusionsbaserad världsmodell inklusive drift- och diffusionskoefficienterna som motsvarar ett särskilt val av diffusionsmetod. Istället för att välja DDPM, en naturligt lämplig kandidat för uppgiften, bygger DIAMOND-ramverket på EDM-formuleringen och överväger en störningskärna med en verkligt värderad funktion av diffusionstid som kallas brusschemat. Ramverket väljer förkonditionerarna för att behålla ingångs- och utgångsvariansen för alla röstnivåer. 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, det vill säga det adderade gaussiska bruset. Intuitivt förhindrar detta träningsmålet från att bli trivialt i det låga bullersystemet. I praktiken är detta mål hög varians vid extremerna av bullerschemat, så modellen samplar brusnivån från en log-normalfördelning vald empiriskt för att sammanfoga träningen runt medelbrusregionerna. DIAMOND-ramverket använder en standard U-Net 2D-komponent för vektorfältet och håller en buffert av tidigare observationer och åtgärder som ramverket använder för att konditionera sig själv. Ramverket DIAMOND sammanfogar sedan dessa tidigare observationer till nästa bullriga observation, och matar in åtgärder genom adaptiva gruppnormaliseringsskikt i restblocken av U-Net. 

DIAMANT: Experiment och resultat

För en omfattande utvärdering väljer DIAMOND-ramverket Atari 100k-riktmärket. Atari 100k-riktmärket består av 26 spel utformade för att testa ett brett utbud av agentfunktioner. I varje spel är en agent begränsad till 100 2 handlingar i miljön, vilket ungefär motsvarar 50 timmars mänskligt spelande, för att lära sig spelet innan utvärdering. Som jämförelse tränar obegränsade Atari-agenter vanligtvis i 500 miljoner steg, vilket motsvarar en 5-faldig ökning av erfarenhet. Vi tränade DIAMOND från grunden med 12 slumpmässiga frön för varje spel. Varje träningskörning krävde cirka 2.9 GB VRAM och tog cirka 4090 dagar på en enda Nvidia RTX 1.03, vilket motsvarar XNUMX GPU-år totalt. Följande tabell visar poängen för alla spel, medelvärdet och IQM eller interkvartilmedelvärde för mänskligt normaliserade poäng. 

Efter begränsningarna för punktuppskattningar ger DIAMOND-ramverket stratifierat bootstrap-förtroende för medelvärdet och IQM eller interkvartilmedelvärde för mänskligt normaliserade poäng tillsammans med prestationsprofiler och ytterligare mätvärden, som sammanfattas i följande figur. 

Resultaten visar att DIAMOND presterar exceptionellt bra över hela riktmärket, överträffar mänskliga spelare i 11 spel och uppnår ett övermänskligt genomsnittligt HNS på 1.46, vilket sätter ett nytt rekord för agenter som är helt utbildade inom en världsmodell. Dessutom är DIAMONDs IQM jämförbar med STORM och överträffar alla andra baslinjer. DIAMOND utmärker sig i miljöer där fånga små detaljer är avgörande, som Asterix, Breakout och RoadRunner. Dessutom, som diskuterats tidigare, har DIAMOND-ramverket flexibiliteten att implementera vilken diffusionsmodell som helst i sin pipeline, även om den väljer EDM-metoden, skulle det ha varit ett naturligt val att välja DDPM-modellen eftersom den redan implementeras i många bildgenerativa applikationer. För att jämföra EDM-metoden mot DDPM-implementering tränar DIAMOND-ramverket båda varianterna med samma nätverksarkitektur på samma delade statiska datauppsättning med över 100 16 ramar som samlats in med en expertpolicy. Antalet denoising-steg är direkt relaterat till slutsatskostnaden för världsmodellen, och så färre steg kommer att minska kostnaden för att träna en agent på tänkta banor. För att säkerställa att vår världsmodell förblir beräkningsmässigt jämförbar med andra baslinjer, såsom IRIS som kräver 1000 NFE per tidssteg, siktar vi på att inte använda mer än tiotals nedtoningssteg, helst färre. Men om du ställer in antalet avbrutningssteg för lågt kan det försämra den visuella kvaliteten, vilket leder till förvärrade fel. För att bedöma stabiliteten hos olika diffusionsvarianter visar vi föreställda banor genererade autoregressivt upp till t = 10 tidssteg i följande figur, med olika antal avbrusningssteg n ≤ XNUMX. 

Vi observerar att användning av DDPM (a) i denna regim resulterar i allvarliga sammansättningsfel, vilket gör att världsmodellen snabbt glider ur distributionen. Däremot förblir den EDM-baserade diffusionsvärldsmodellen (b) mycket mer stabil över långa tidshorisonter, även med ett enda avbrutningssteg. Tänkte banor med diffusionsvärldsmodeller baserat 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 nedtoningssteg n. Vi observerar att DDPM-baserad generation lider av sammansättningsfel, med ett mindre antal avbrusningssteg som leder till snabbare felackumulering. Däremot förblir DIAMONDs EDM-baserade världsmodell mycket mer stabil, även för n = 1. Den optimala enstegsförutsägelsen är förväntningen över möjliga rekonstruktioner för en given bullrig ingång, som kan vara ur distribution om den bakre fördelningen är multimodal . Medan vissa spel, som Breakout, har deterministiska övergångar som kan modelleras exakt med ett enda avbrutningssteg, 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 styra samplingsproceduren mot ett specifikt läge, som illustreras i spelet Boxning i följande figur. Följaktligen satte DIAMOND-ramverket n = 3 i alla våra experiment.

Ovanstående figur jämför enstegs (övre raden) och flerstegs (nedre raden) provtagning i boxning. Den svarta spelarens rörelser är oförutsägbara, vilket gör att enkelstegs denoising interpolerar mellan möjliga utfall, vilket resulterar i suddiga förutsägelser. Däremot ger flerstegssampling en tydlig bild genom att styra genereringen mot ett specifikt läge. Intressant nog, eftersom policyn kontrollerar den vita spelaren, är hans handlingar kända för världsmodellen, vilket eliminerar tvetydighet. Således förutsäger både enstegs- och flerstegssampling korrekt den vita spelarens position.

I figuren ovan uppvisar banorna som föreställts av DIAMOND generellt högre visuell kvalitet och är mer trogna den verkliga miljön jämfört med de som IRIS föreställer sig. Banorna som genereras av IRIS innehåller visuella inkonsekvenser mellan bildrutor (markerade med vita rutor), som att fiender visas som belöningar och vice versa. Även om dessa inkonsekvenser bara kan påverka ett fåtal pixlar, kan de avsevärt påverka förstärkningsinlärningen. Till exempel strävar en agent vanligtvis efter belöningar och undvika fiender, så dessa små visuella avvikelser kan göra det mer utmanande att lära sig en optimal policy. Bilden visar på varandra följande ramar föreställda med IRIS (vänster) och DIAMOND (höger). De vita rutorna markerar inkonsekvenser mellan ramar, som endast visas i banor som genereras med IRIS. I Asterix (översta raden) blir en fiende (orange) en belöning (röd) i den andra bilden, för att sedan återgå 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 ramarna. I Road Runner (nedre raden) återges belöningarna (små blåa prickar på vägen) inkonsekvent mellan bildrutor. Dessa inkonsekvenser förekommer inte med DIAMOND. I Breakout uppdateras poängen tillförlitligt med +7 när en röd tegelsten är bruten. 

Slutsats

I den här artikeln har vi pratat om DIAMOND, en förstärkningslärande agent utbildad inom en diffusionsvärldsmodell. Ramverket DIAMOND gör noggranna designval för att säkerställa att dess spridningsvärldsmodell förblir effektiv och stabil över långa tidshorisonter. Ramverket tillhandahåller en kvalitativ analys för att visa vikten av dessa designval. DIAMOND sätter en ny state-of-the-art med ett medelvärde för mänskligt normaliserat resultat på 1.46 på det väletablerade Atari 100k-riktmärket, det högsta för agenter som helt utbildats inom en världsmodell. Genom att arbeta i bildrymden kan DIAMONDs diffusionsvärldsmodell sömlöst ersätta miljön, vilket ger större insikter i världsmodeller och agentbeteenden. Särskilt den förbättrade prestandan i vissa spel tillskrivs bättre modellering av kritiska visuella detaljer. DIAMOND-ramverket modellerar miljön som en standard POMDP eller Partially Observable Markov Decision Process med en uppsättning tillstånd, en uppsättning diskreta åtgärder och en uppsättning bildobservationer. Övergångsfunktionerna beskriver omgivningens dynamik, och belöningsfunktionen kartlägger övergångarna till skalära belöningar.