A fost în 2018 când ideea de învățare prin întărire în contextul unui model de lume neural a fost introdusă pentru prima dată, și curând, acest principiu fundamental a fost aplicat pe modele de lume. Unele dintre modelele proeminente care implementează învățarea prin întărire au fost framework-ul Dreamer, care a introdus învățarea prin întărire din spațiul latent al unui model de stare recurentă. DreamerV2 a demonstrat că utilizarea latentelor discrete poate rezulta în reducerea erorilor de compunere, și framework-ul DreamerV3 a fost capabil să atingă performanțe umane pe o serie de sarcini din diferite domenii cu hiperparametri fixi.

Mai mult, pot fi trasate paralele între modelele de generare a imaginilor și modelele de lume, indicând că progresul realizat în modelele generative de viziune poate fi replicat pentru a beneficia modelele de lume. De când utilizarea transformatorilor în framework-urile de procesare a limbajului natural a câștigat popularitate, framework-urile DALL-E și VQGAN au apărut. Framework-urile au implementat autoencoder-uri discrete pentru a converti imagini în token-uri discrete și au fost capabile să construiască modele generative de text-la-imagini foarte puternice și eficiente prin exploatarea capacităților de modelare secvențială ale transformatorilor autoregresivi. În același timp, modelele de difuzie au câștigat popularitate, și astăzi, modelele de difuzie s-au stabilit ca un paradigm dominant pentru generarea de imagini de înaltă rezoluție. Datorită capacităților oferite de modelele de difuzie și învățarea prin întărire, se fac încercări de a combina cele două abordări, cu scopul de a profita de flexibilitatea modelelor de difuzie ca modele de traiectorie, modele de recompensă, planificatori și ca politici de augmentare a datelor în învățarea prin întărire offline.

Modelele de lume oferă o metodă promițătoare pentru antrenarea agenților de învățare prin întărire în siguranță și eficient. În mod tradițional, aceste modele utilizează secvențe de variabile latente discrete pentru a simula dinamica mediului. Cu toate acestea, această compresie poate ignora detaliile vizuale cruciale pentru învățarea prin întărire. În același timp, modelele de difuzie au câștigat popularitate pentru generarea de imagini, provocând metodele tradiționale care utilizează latente discrete. Inspirat de acest schimb, în acest articol, vom discuta despre DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agent de învățare prin întărire antrenat într-un model de lume de difuzie. Vom explora alegerile de proiectare necesare pentru a face difuzia potrivită pentru modelarea lumii și vom arăta că detaliile vizuale îmbunătățite duc la o performanță mai bună a agentului. DIAMOND stabilește un nou punct de referință pe testul competitiv Atari 100k, atingând un scor mediu normalizat uman de 1,46, cel mai mare pentru agenții antrenați complet într-un model de lume.

DIAMOND : DIffusion As a Model Of eNvironment Dreams

Modelele de lume sau modelele generative de mediu sunt în curs de a deveni una dintre componentele mai importante pentru agenții generativi pentru a planifica și a raționa despre mediile lor. Deși utilizarea învățării prin întărire a obținut succese considerabile în ultimii ani, modelele care implementează învățarea prin întărire sunt cunoscute pentru a fi ineficiente din punct de vedere al eşantionării, ceea ce limitează semnificativ aplicațiile lor în lumea reală. Pe de altă parte, modelele de lume au demonstrat capacitatea lor de a antrena agenții de învățare prin întărire în mod eficient în diverse medii, cu o eficiență a eşantionării semnificativ îmbunătățită, permițând modelului să învețe din experiențe reale. Framework-urile recente de modelare a lumii modelează în mod obișnuit dinamica mediului ca o secvență de variabile latente discrete, cu modelul care discretizează spațiul latent pentru a evita erorile de compunere pe orizonturi de timp multiplu. Deși abordarea poate livra rezultate substanțiale, este asociată și cu o pierdere de informații, ceea ce duce la o pierdere a calității de reconstrucție și a generalității. Pierderea de informații poate deveni un obstacol semnificativ pentru scenariile din lumea reală care necesită informații bine definite, cum ar fi antrenarea vehiculelor autonome. În astfel de sarcini, schimbări mici sau detalii în intrarea vizuală, cum ar fi culoarea semaforului sau indicatorul de direcție al vehiculului din față, pot schimba politica unui agent. Deși creșterea numărului de latente discrete poate ajuta la evitarea pierderii de informații, crește semnificativ costurile de calcul.

Mai mult, în ultimii ani, modelele de difuzie au apărut ca abordarea dominantă pentru framework-urile de generare de imagini de înaltă calitate, deoarece framework-urile construite pe modele de difuzie învață să inverseze un proces de zgomot și concurează direct cu unele dintre abordările mai bine stabilite care modelează token-uri discrete, oferind astfel o alternativă promițătoare pentru a elimina nevoia de discretizare în modelarea lumii. Modelele de difuzie sunt cunoscute pentru capacitatea lor de a fi condiționate ușor și de a modela distribuții complexe, multimodale, fără colaps de mod, ceea ce este crucial pentru modelarea lumii, deoarece condiționarea permite unui model de lume să reflecte cu acuratețe acțiunile unui agent, conducând la o atribuire mai fiabilă a creditului. Mai mult, modelarea distribuțiilor multimodale oferă o diversitate mai mare de scenarii de antrenament pentru agent, îmbunătățindu-i performanța generală.

Pe baza acestor caracteristici, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), un agent de învățare prin întărire antrenat într-un model de lume de difuzie. Framework-ul DIAMOND face alegeri de proiectare atente pentru a se asigura că modelul său de lume de difuzie rămâne eficient și stabil pe orizonturi de timp lungi. Framework-ul oferă o analiză calitativă pentru a demonstra importanța acestor alegeri de proiectare. DIAMOND stabilește un nou punct de referință cu un scor mediu normalizat uman de 1,46 pe binecunoscutul benchmark Atari 100k, cel mai mare pentru agenții antrenați complet într-un model de lume. Funcționarea în spațiul de imagine permite modelului de lume de difuzie al lui DIAMOND să înlocuiască în mod transparent mediul, oferind insights mai bune în comportamentul modelului de lume și al agentului. În mod semnificativ, performanța îmbunătățită în anumite jocuri este atribuită unei modelări mai bune a detaliilor vizuale critice. Framework-ul DIAMOND modelează mediul ca un proces decizional Markov parțial observabil standard, cu un set de stări, un set de acțiuni discrete și un set de observații de imagine. Funcțiile de tranziție descriu dinamica mediului, iar funcția de recompensă asociază tranzițiile cu recompense scalare.

DIAMOND : Metodologia și Arhitectura

La nivel fundamental, modelele de difuzie sunt o clasă de modele generative care generează un exemplar prin inversarea procesului de zgomot și se inspiră puternic din termodinamica neechilibrată. Framework-ul DIAMOND consideră un proces de difuzie indexat de o variabilă temporală continuă, cu marginale și condiții de frontieră corespunzătoare, cu o distribuție prioritară neordonată tratabilă. Mai mult, pentru a obține un model generativ, care mapă de la zgomot la date, framework-ul DIAMOND trebuie să inverseze procesul, procesul de inversare fiind și el un proces de difuzie, care rulează înapoi în timp. Mai mult, la un moment dat în timp, nu este trivial să estimezi funcția de scor, deoarece framework-ul DIAMOND nu are acces la funcția de scor reală, iar modelul depășește acest obstacol prin implementarea obiectivului de potrivire a scorului, o abordare care facilitează un framework să antreneze un model de scor fără a cunoaște funcția de scor subiacentă. Modelul de difuzie bazat pe scor oferă un model generativ necondiționat. Cu toate acestea, un model generativ condiționat de dinamică de mediu este necesar pentru a servi ca model de lume și pentru a servi acestui scop, framework-ul DIAMOND examinează cazul general al abordării POMDP, în care framework-ul poate utiliza observații și acțiuni trecute pentru a aproxima starea Markoviană necunoscută. Așa cum se arată în Figura 1, framework-ul DIAMOND utilizează această istorie pentru a condiționa un model de difuzie, pentru a estima și genera următoarea observație direct. Deși framework-ul DIAMOND poate recurge în teorie la orice solver SDE sau ODE, există un compromis între NFE sau Numărul de Evaluări ale Funcției și calitatea eşantionului, care afectează semnificativ costul de inferență al modelelor de difuzie.

Pe baza acestor învățăminte, să examinăm acum realizarea practică a framework-ului DIAMOND al unui model de lume bazat pe difuzie, inclusiv coeficienții de derivă și difuzie corespunzători unei anumite alegeri de abordare de difuzie. În loc să opteze pentru DDPM, un candidat natural potrivit pentru sarcină, framework-ul DIAMOND se bazează pe formularea EDM și consideră un kernel de perturbare cu o funcție cu valori reale de timp de difuzie numită programul de zgomot. Framework-ul selectează preconditionările pentru a menține varianța de intrare și ieșire pentru orice nivel de voce. Antrenarea rețelei combină semnalul și zgomotul adaptiv, în funcție de nivelul de degradare, iar atunci când zgomotul este scăzut, ținta devine diferența dintre semnalul curat și semnalul perturbat, adică zgomotul Gaussian adăugat. În mod intuitiv, acest lucru previne ca obiectivul de antrenare să devină trivial în regimul de zgomot scăzut. În practică, acest obiectiv este de variație ridicată la extremele programului de zgomot, astfel încât modelul eşantionează nivelul de zgomot dintr-o distribuție log-normală aleasă empiric pentru a concatena antrenarea în jurul regiunilor medii de zgomot. Framework-ul DIAMOND utilizează un component U-Net 2D standard pentru câmpul vectorial și păstrează un buffer de observații și acțiuni trecute pe care le utilizează pentru a se condiționa. Framework-ul DIAMOND concatenează apoi aceste observații trecute la următoarea observație zgomotoasă și acțiunile de intrare prin straturi de normalizare adaptivă de grup în blocurile reziduale ale U-Net.

DIAMOND: Experimente și Rezultate

Pentru o evaluare cuprinzătoare, framework-ul DIAMOND optează pentru benchmark-ul Atari 100k. Benchmark-ul Atari 100k constă în 26 de jocuri proiectate pentru a testa o gamă largă de capacități ale agentului. În fiecare joc, un agent este limitat la 100k de acțiuni în mediu, ceea ce este aproximativ echivalent cu 2 ore de joc uman, pentru a învăța jocul înainte de evaluare. Pentru comparație, agenții Atari neconstrânși se antrenează de obicei timp de 50 de milioane de pași, reprezentând o creștere de 500 de ori a experienței. Am antrenat DIAMOND de la zero, utilizând 5 semințe aleatorii pentru fiecare joc. Fiecare rulare de antrenare a necesitat aproximativ 12 GB de VRAM și a durat aproximativ 2,9 zile pe un singur Nvidia RTX 4090, ceea ce înseamnă 1,03 ani de GPU în total. Tabelul următor oferă scorul pentru toate jocurile, media și IQM sau media intercuartilă a scorurilor normalizate uman.

Urmând limitările estimărilor punctuale, framework-ul DIAMOND oferă intervalul de încredere bootstrap stratificat pentru media și IQM sau media intercuartilă a scorurilor normalizate uman, împreună cu profilurile de performanță și metrice suplimentare, astfel cum se rezumă în figura următoare.

Rezultatele arată că DIAMOND performează excepțional de bine pe întregul benchmark, depășind jucătorii umani în 11 jocuri și atingând un scor mediu normalizat uman de 1,46, stabilind un nou record pentru agenții antrenați complet într-un model de lume. În plus, IQM-ul lui DIAMOND este comparabil cu STORM și depășește toate celelalte linii de bază. DIAMOND excelează în medii în care capturarea detaliilor mici este crucială, cum ar fi Asterix, Breakout și RoadRunner. Mai mult, așa cum s-a discutat anterior, framework-ul DIAMOND are flexibilitatea de a implementa orice model de difuzie în pipeline-ul său, deși optează pentru abordarea EDM, ar fi fost o alegere naturală să opteze pentru modelul DDPM, deoarece este deja implementat în numeroase aplicații generative de imagini. Pentru a compara abordarea EDM cu implementarea DDPM, framework-ul DIAMOND antrenează ambele variante cu aceeași arhitectură de rețea pe același set de date static împărtășit, cu peste 100k de cadre colectate cu o politică expertă. Numărul de pași de denoising este direct legat de costul de inferență al modelului de lume, și astfel, mai puțini pași vor reduce costul de antrenare a unui agent pe traiectorii imaginate. Pentru a ne asigura că modelul nostru de lume rămâne comparabil din punct de vedere computațional cu alte linii de bază, cum ar fi IRIS, care necesită 16 NFE pe timp, ne propunem să utilizăm nu mai mult de zeci de pași de denoising, preferabil mai puțini. Cu toate acestea, stabilirea numărului de pași de denoising prea mic poate deteriora calitatea vizuală, conducând la erori de compunere. Pentru a evalua stabilitatea diferitelor variante de difuzie, afișăm traiectorii imaginate generative autoregresiv până la t = 1000 de pași de timp, utilizând diferite numere de pași de denoising n ≤ 10.

Observăm că utilizarea DDPM (a) în acest regim duce la erori de compunere severe, cauzând modelul de lume să se deplaseze rapid din distribuție. În contrast, modelul de lume de difuzie bazat pe EDM (b) rămâne mult mai stabil pe orizonturi de timp lungi, chiar și cu un singur pas de denoising. Traiectoriile imaginate cu modele de difuzie bazate pe DDPM (stânga) și EDM (dreapta) sunt arătate. Observația inițială la t = 0 este aceeași pentru ambele, iar fiecare rând corespunde unui număr descrescător de pași de denoising. Observăm că generația bazată pe DDPM suferă de erori de compunere, cu numere mai mici de pași de denoizing conducând la o acumulare mai rapidă a erorilor. În contrast, modelul de lume de difuzie al lui DIAMOND rămâne mult mai stabil, chiar și pentru n = 1. Predicția optimă cu un singur pas este așteptarea peste reconstrucții posibile pentru o intrare zgomotoasă dată, care poate fi în afara distribuției dacă distribuția posterioară este multimodală. Deși unele jocuri, cum ar fi Breakout, au tranziții deterministe care pot fi modelate cu acuratețe cu un singur pas de denoising, alte jocuri prezintă parțialitatea observației, rezultând în distribuții de observație multimodale. În aceste cazuri, un solver iterativ este necesar pentru a ghida procedura de eşantionare către un mod specific, așa cum se ilustrează în jocul Boxing în figura următoare. Prin urmare, framework-ul DIAMOND setează n = 3 în toate experimentele noastre.

Figura de mai sus compară eşantionarea cu un singur pas (rândul de sus) și eşantionarea multi-pas (rândul de jos) în Boxing. Mișcările jucătorului negru sunt imprevizibile, cauzând eşantionarea cu un singur pas să interpoleze între rezultate posibile, rezultând în predicții blurate. În contrast, eşantionarea multi-pas produce o imagine clară, ghidând generația către un mod specific. În mod interesant, deoarece politica controlează jucătorul alb, acțiunile sale sunt cunoscute de modelul de lume, eliminând ambiguitatea. Astfel, atât eşantionarea cu un singur pas, cât și cea multi-pas, prevăd corect poziția jucătorului alb.

În figura de mai sus, traiectoriile imaginate de DIAMOND prezintă în general o calitate vizuală mai ridicată și sunt mai fidele față de mediul real în comparație cu cele imaginate de IRIS. Traiectoriile generate de IRIS conțin incoerențe vizuale între cadre (evidențiate de cutii albe), cum ar fi dușmanii care sunt afișați ca recompense și viceversa. Deși aceste incoerențe pot afecta doar câteva pixeli, ele pot avea un impact semnificativ asupra învățării prin întărire. De exemplu, un agent tinde de obicei să vizeze recompensele și să evite dușmanii, astfel încât aceste mici discrepanțe vizuale pot face mai dificilă învățarea unei politici optime. Figura arată cadre consecutive imaginate cu IRIS (stânga) și DIAMOND (dreapta). Cutiile albe evidențiază incoerențe între cadre, care apar doar în traiectoriile generate de IRIS. În Asterix (rândul de sus), un dușman (portocaliu) devine o recompensă (roșu) în al doilea cadru, apoi revine la un dușman în al treilea cadru și din nou la o recompensă în al patrulea cadru. În Breakout (rândul de mijloc), zidurile și scorul sunt inconstante între cadre. În Road Runner (rândul de jos), recompensele (puncte albastre mici pe drum) sunt renderizate inconstant între cadre. Aceste incoerențe nu apar cu DIAMOND. În Breakout, scorul este actualizat în mod fiabil cu +7 atunci când o cărămidă roșie este spartă.

Concluzie

În acest articol, am discutat despre DIAMOND, un agent de învățare prin întărire antrenat într-un model de lume de difuzie. Framework-ul DIAMOND face alegeri de proiectare atente pentru a se asigura că modelul său de lume de difuzie rămâne eficient și stabil pe orizonturi de timp lungi. Framework-ul oferă o analiză calitativă pentru a demonstra importanța acestor alegeri de proiectare. DIAMOND stabilește un nou punct de referință cu un scor mediu normalizat uman de 1,46 pe binecunoscutul benchmark Atari 100k, cel mai mare pentru agenții antrenați complet într-un model de lume. Funcționarea în spațiul de imagine permite modelului de lume de difuzie al lui DIAMOND să înlocuiască în mod transparent mediul, oferind insights mai bune în comportamentul modelului de lume și al agentului. În mod semnificativ, performanța îmbunătățită în anumite jocuri este atribuită unei modelări mai bune a detaliilor vizuale critice. Framework-ul DIAMOND modelează mediul ca un proces decizional Markov parțial observabil standard, cu un set de stări, un set de acțiuni discrete și un set de observații de imagine. Funcțiile de tranziție descriu dinamica mediului, iar funcția de recompensă asociază tranzițiile cu recompense scalare.