W 2018 roku po raz pierwszy wprowadzono ideę uczenia wzmacniania w kontekście środowiska modelu neuronowej sieci, a niedługo potem ten fundamentalny princip został zastosowany w modelach świata. Niektóre z najbardziej znaczących modeli, które zaimplementowały uczenie wzmacniania, to ramy Dreamer, które wprowadziły uczenie wzmacniania z latentnego przestrzeni modelu przestrzeni stanu. DreamerV2 wykazał, że użycie dyskretnych latentów może skutkować zmniejszeniem błędów składanych, a ramy DreamerV3 osiągnęły wyniki porównywalne z ludzkimi na serii zadań w różnych dziedzinach z ustalonymi hiperparametrami.

Ponadto można wysnuć paralele między modelami generowania obrazów a modelami świata, wskazując, że postęp w modelach widzenia generatywnego może być powielony, aby przynieść korzyści modelom świata. Odkąd użycie transformatorów w ramach przetwarzania języka naturalnego zyskało na popularności, pojawiły się ramy DALL-E i VQGAN. Ramy te zaimplementowały dyskretne autoenkodery do konwersji obrazów na dyskretne tokeny i były w stanie zbudować bardzo potężne i wydajne modele generatywne tekstu do obrazu, wykorzystując sekwencyjne możliwości modelowania autoregresyjnych transformatorów. W tym samym czasie modele dyfuzji zyskały na popularności, a dziś modele dyfuzji ustanowiły się jako dominujący paradygmat generowania obrazów o wysokiej rozdzielczości. Dzięki możliwościom oferowanym przez modele dyfuzji i uczenie wzmacniania, podejmowane są próby połączenia tych dwóch podejść, z celem wykorzystania elastyczności modeli dyfuzji jako modeli trajektorii, modeli nagród, planistów i jako polityki dla augmentacji danych w trybie offline uczenia wzmacniania.

Modele świata oferują obiecującą metodę szkolenia agentów uczenia wzmacniania w sposób bezpieczny i wydajny. Tradycyjnie te modele wykorzystują sekwencje dyskretnych zmiennych latentnych do symulowania dynamiki środowiska. Jednakże ta kompresja może zignorować wizualne szczegóły niezbędne dla uczenia wzmacniania. W tym samym czasie modele dyfuzji zyskały na popularności w generowaniu obrazów, wyzwając tradycyjne metody, które wykorzystują dyskretne latenty. Zainspirowani tym przesunięciem, w tym artykule będziemy rozmawiać o DIAMOND (DIffusion As a Model Of eNvironment Dreams), agencie uczenia wzmacniania szkolonym w ramach modelu świata dyfuzji. Będziemy eksplorować niezbędne wybory projektowe, aby uczynić dyfuzję odpowiednią dla modelowania świata, i pokażemy, że udoskonalone wizualne szczegóły prowadzą do lepszej wydajności agenta. DIAMOND ustanawia nowy benchmark na konkurencyjnym teście Atari 100k, osiągając średni wynik ludzki znormalizowany na poziomie 1,46, najwyższy dla agentów szkolonych całkowicie w ramach modelu świata.

DIAMOND: Wizualne szczegóły mają znaczenie w Atari i dyfuzji dla modelowania świata

Modele świata lub generatywne modele środowiska stają się jednym z najważniejszych komponentów dla agentów generatywnych do planowania i rozumowania o ich środowiskach. Chociaż użycie uczenia wzmacniania osiągnęło znaczny sukces w ostatnich latach, modele zaimplementowane z uczeniem wzmacniania są znane z tego, że są nieefektywne pod względem próbek, co znacznie ogranicza ich zastosowanie w świecie rzeczywistym. Z drugiej strony, modele świata wykazały swoją zdolność do efektywnego szkolenia agentów uczenia wzmacniania w różnych środowiskach z znacznie poprawioną efektywnością próbek, pozwalając modelowi uczyć się z doświadczeń świata rzeczywistego. Niedawne ramy modelowania świata zwykle modelują dynamikę środowiska jako sekwencję dyskretnych zmiennych latentnych, z modelem dyskretnym w przestrzeni latentnej, aby uniknąć błędów składanych w wieloetapowych horyzontach czasowych. Chociaż podejście to może dostarczyć znaczne wyniki, jest również związane ze stratą informacji, prowadzącą do utraty jakości rekonstrukcji i utraty ogólności. Utrata informacji może stać się znaczącą przeszkodą w scenariuszach świata rzeczywistego, które wymagają, aby informacje były dobrze zdefiniowane, takie jak szkolenie pojazdów autonomicznych. W takich zadaniach małe zmiany lub szczegóły wejścia wizualnego, takie jak kolor światła drogowego lub wskaźnik skrętu pojazdu przed, mogą zmienić politykę agenta. Chociaż zwiększenie liczby dyskretnych latentów może pomóc uniknąć utraty informacji, znacznie zwiększa koszty obliczeniowe.

Ponadto w ostatnich latach modele dyfuzji wyłoniły się jako dominujące podejście do generowania obrazów o wysokiej jakości, ponieważ ramy zbudowane na modelach dyfuzji uczą się odwracać proces hałasowania i bezpośrednio konkurują z niektórymi bardziej ugruntowanymi podejściami modelowania tokenów dyskretnych, oferując obiecującą alternatywę dla eliminacji potrzeby dyskretyzacji w modelowaniu świata. Modele dyfuzji są znane z ich zdolności do łatwego warunkowania i elastycznego modelowania złożonych, wielomodalnych rozkładów bez kolapsu trybu. Atrybuty te są kluczowe dla modelowania świata, ponieważ warunkowanie umożliwia modelowi świata dokładnie odzwierciedlić działania agenta, prowadząc do bardziej niezawodnego przypisania kredytu. Ponadto modelowanie rozkładów wielomodalnych oferuje większą różnorodność scenariuszy szkoleniowych dla agenta, poprawiając jego ogólną wydajność.

Budując na tych cechach, DIAMOND (DIffusion As a Model Of eNvironment Dreams) to agent uczenia wzmacniania szkolony w ramach modelu świata dyfuzji. Ramy DIAMOND dokonują starannych wyborów projektowych, aby zapewnić, że ich model świata dyfuzji pozostaje wydajny i stabilny w długich horyzontach czasowych. Ramy DIAMOND zapewniają jakościową analizę, aby zademonstrować znaczenie tych wyborów projektowych. DIAMOND ustanawia nowy stan sztuki z średnim wynikiem ludzki znormalizowanym na poziomie 1,46 na dobrze ugruntowanym teście Atari 100k, najwyższym dla agentów szkolonych całkowicie w ramach modelu świata. Działanie w przestrzeni obrazu pozwala modelowi świata DIAMOND na bezproblemową substytucję środowiska, oferując większe wglądy w zachowania modelu świata i agenta. Godne uwagi jest to, że poprawiona wydajność w pewnych grach jest przypisywana lepszemu modelowaniu krytycznych szczegółów wizualnych. Ramy DIAMOND modelują środowisko jako standardowy proces decyzyjny Markowa z częściowo obserwowalnym stanem z zestawem stanów, zestawem dyskretnych działań i zestawem obserwacji obrazu. Funkcje przejścia opisują dynamikę środowiska, a funkcja nagrody mapuje przejścia na skalarną nagrodę.

DIAMOND: Metodologia i Architektura

W swojej istocie modele dyfuzji są klasą modeli generatywnych, które generują próbkę przez odwrócenie procesu hałasowania i czerpią inspirację z termodynamiki nieequlibrium. Ramy DIAMOND rozważają proces dyfuzji indeksowany przez ciągłą zmienną czasową z odpowiednimi marginalnymi i warunkami brzegowymi z traktowalną nieustrukturalizowaną dystrybucją a priori. Ponadto, aby uzyskać model generatywny, który mapuje od hałasu do danych, ramy DIAMOND muszą odwrócić proces, a proces odwrócenia jest również modelem dyfuzji, który działa wstecz w czasie. Ponadto, w dowolnym punkcie czasu nie jest trywialne oszacowanie funkcji skory, ponieważ ramy DIAMOND nie mają dostępu do prawdziwej funkcji skory, a model pokonuje tę przeszkodę, implementując cel dopasowania skory, podejście, które umożliwia ramom szkolenia modelu skory bez znajomości podstawowej funkcji skory. Model dyfuzji oparty na skory zapewnia niewarunkowy model generatywny. Jednak warunkowy model generatywny dynamiki środowiska jest wymagany do służenia jako model świata, a aby służyć temu celowi, ramy DIAMOND rozważają ogólny przypadek podejścia POMDP, w którym ramy mogą wykorzystywać poprzednie obserwacje i działania, aby przybliżyć nieznany stan Markowa. Jak to zostało zademonstrowane na Rysunku 1, ramy DIAMOND wykorzystują tę historię do warunkowania modelu dyfuzji, aby oszacować i wygenerować następną obserwację bezpośrednio. Chociaż ramy DIAMOND mogą teoretycznie wykorzystywać dowolny rozwiązujący SDE lub ODE, istnieje kompromis między NFE (liczbą ocen funkcji) a jakością próbki, który znacznie wpływa na koszt wnioskowania modeli dyfuzji.

Budując na tych wnioskach, spójrzmy teraz na praktyczną realizację ram DIAMOND modelu świata opartego na dyfuzji, w tym współczynniki dryftu i dyfuzji odpowiadające określonemu wyborowi podejścia dyfuzji. Zamiast wybrać DDPM, naturalnie odpowiedniego kandydata do tego zadania, ramy DIAMOND budują na sformalizowaniu EDM i rozważają jądro perturbacji z funkcją wartości rzeczywistej czasu dyfuzji, zwaną harmonogramem hałasu. Ramy wybierają prekondycjonery, aby utrzymać wejściową i wyjściową wariancję dla dowolnego poziomu głosu. Szkolenie sieci łączy sygnał i hałas w zależności od poziomu degradacji, a gdy hałas jest niski, a celem staje się różnica między czystym a zakłóconym sygnałem, czyli dodanym hałasem Gaussa. Intuicyjnie, to zapobiega temu, aby cel szkolenia stał się trywialny w reżimie niskiego hałasu. W praktyce ten cel jest wysoko zmienny na krańcach harmonogramu hałasu, więc model próbkuj hałas z log-normalnej dystrybucji wybranej empirycznie, aby połączyć szkolenie wokół średnich obszarów hałasu. Ramy DIAMOND wykorzystują standardowy komponent U-Net 2D dla pola wektorowego i utrzymują bufor poprzednich obserwacji i działań, które ramy wykorzystują do warunkowania siebie. Następnie ramy DIAMOND łączą te poprzednie obserwacje z następną zakłóconą obserwacją i wejściowymi działaniami przez warstwy normalizacji grupowej w blokach resztowych U-Net.

DIAMOND: Eksperymenty i Wyniki

Do kompleksowej oceny ramy DIAMOND wybierają benchmark Atari 100k. Benchmark Atari 100k składa się z 26 gier zaprojektowanych do testowania szerokiego zakresu możliwości agenta. W każdej grze agent jest ograniczony do 100k akcji w środowisku, co jest równoważne z około 2 godzinami gry ludzkiej, aby nauczyć się gry przed oceną. Dla porównania, nieograniczone agenty Atari zwykle trenują przez 50 milionów kroków, reprezentując 500-krotny wzrost doświadczenia. Trenowaliśmy DIAMOND od podstaw przy użyciu 5 losowych nasion dla każdej gry. Każdy przebieg treningu wymagał około 12GB pamięci VRAM i trwał około 2,9 dni na jednej karcie graficznej Nvidia RTX 4090, co stanowi 1,03 roku GPU w sumie. Poniższa tabela zawiera wynik dla wszystkich gier, średnią i IQM (średnią międzykwartylową) wyników znormalizowanych do ludzkich.

Po ograniczeniach estymacji punktowych ramy DIAMOND zapewniają stratyfikowaną ufność bootstrap w średniej i IQM wyników znormalizowanych do ludzkich wraz z profilami wydajności i dodatkowymi metrykami, jak to zostało podsumowane w poniższym rysunku.

Wyniki pokazują, że DIAMOND wykonuje się wyjątkowo dobrze w całym benchmarku, przewyższając graczy ludzkich w 11 grach i osiągając nadludzki średni wynik znormalizowany do 1,46, ustanawiając nowy rekord dla agentów szkolonych całkowicie w ramach modelu świata. Ponadto, DIAMOND wyróżnia się w środowiskach, w których przechwycenie małych szczegółów jest kluczowe, takich jak Asterix, Breakout i RoadRunner. Ponadto, jak omówiono wcześniej, ramy DIAMOND mają elastyczność implementowania dowolnego modelu dyfuzji w swoim potoku, chociaż wybierają podejście EDM, byłoby to naturalne wybranie modelu DDPM, ponieważ jest on już wdrożony w licznych aplikacjach generowania obrazów. Aby porównać podejście EDM z implementacją DDPM, ramy DIAMOND trenują oba warianty z tą samą architekturą sieci na tym samym współdzielonym zestawie danych statycznych z ponad 100k klatkami zebranymi z ekspertem. Liczba kroków denoisingu jest bezpośrednio związana z kosztem wnioskowania modelu świata, a więc mniej kroków zmniejszy koszt treningu agenta na wyobrażonych trajektoriach. Aby upewnić się, że nasz model świata pozostaje komputacyjnie porównywalny z innymi benchmarkami, takimi jak IRIS, które wymagają 16 NFE na krok, staramy się wykorzystywać nie więcej niż dziesiątki kroków denoisingu, a najlepiej mniej. Jednak ustawienie liczby kroków denoisingu zbyt niskiego może pogorszyć jakość wizualną, prowadząc do błędów składanych. Aby ocenić stabilność różnych wariantów dyfuzji, wyświetlamy wyobrażone trajektorie generowane autoregresyjnie do t = 1000 kroków w poniższym rysunku, przy użyciu różnej liczby kroków denoisingu n ≤ 10.

Obserwujemy, że użycie DDPM (a) w tym reżimie skutkuje poważnymi błędami składanymi, powodując, że model świata szybko wykracza poza dystrybucję. W przeciwieństwie do tego, model świata dyfuzji oparty na EDM (b) pozostaje znacznie bardziej stabilny w długich horyzontach czasowych, nawet z jednym krokiem denoisingu. Wyobrażone trajektorie z modelem dyfuzji opartym na DDPM (po lewej) i EDM (po prawej) są pokazane. Początkowa obserwacja w t = 0 jest taka sama dla obu, a każdy wiersz odpowiada malejącej liczbie kroków denoisingu n. Obserwujemy, że generacja oparta na DDPM cierpi na błędy składane, a mniejsza liczba kroków denoisingu prowadzi do szybszego nagromadzenia błędów. W przeciwieństwie do tego, model świata DIAMOND oparty na EDM pozostaje znacznie bardziej stabilny, nawet dla n = 1. Optymalna predykcja jednego kroku jest oczekiwaniem nad możliwymi rekonstrukcjami dla danego wejścia zakłóconego, które może być poza dystrybucją, jeśli rozkład a posteriori jest wielomodalny. Chociaż niektóre gry, takie jak Breakout, mają deterministyczne przejścia, które mogą być dokładnie modelowane za pomocą jednego kroku denoisingu, inne gry wykazują częściową obserwowalność, skutkując wielomodalnymi rozkładami obserwacji. W takich przypadkach rozwiązanie iteracyjne jest konieczne do kierowania procedurą próbkowania w kierunku określonego trybu, jak to zostało zademonstrowane w grze Boxing w poniższym rysunku. W związku z tym ramy DIAMOND ustawiają n = 3 we wszystkich eksperymentach.

Powyższy rysunek porównuje próbkowanie jednego kroku (górny wiersz) i próbkowanie wielokroku (dolny wiersz) w Boxing. Ruchy czarnego gracza są nieprzewidywalne, powodując, że próbkowanie jednego kroku interpoluje między możliwymi wynikami, skutkując rozmytymi predykcjami. W przeciwieństwie do tego, próbkowanie wielokroku produkuje wyraźny obraz, kierując generację w kierunku określonego trybu. Co ciekawe, ponieważ polityka kontroluje białego gracza, jego działania są znane modelowi świata, eliminując niepewność. Zatem zarówno próbkowanie jednego kroku, jak i próbkowanie wielokroku poprawnie przewidują pozycję białego gracza.

W powyższym rysunku trajektorie wyobrażone przez DIAMOND ogólnie wykazują wyższą jakość wizualną i są bardziej wiernymi odwzorowaniem środowiska rzeczywistego w porównaniu z tymi wyobrażonymi przez IRIS. Trajektorie wygenerowane przez IRIS zawierają nieścisłości wizualne między klatkami (wyróżnione białymi prostokątami), takie jak wrogowie wyświetlani jako nagrody i odwrotnie. Chociaż te nieścisłości mogą wpływać tylko na kilka pikseli, mogą one znacznie wpłynąć na uczenie wzmacniania. Na przykład, agent zwykle stara się celować w nagrody i unikać wrogów, więc te małe nieścisłości wizualne mogą utrudnić naukę optymalnej polityki. Rysunek pokazuje kolejne klatki wyobrażone przez IRIS (po lewej) i DIAMOND (po prawej). Białe prostokąty wyróżniają nieścisłości między klatkami, które pojawiają się tylko w trajektoriach wygenerowanych przez IRIS. W Asterix (górny wiersz) wróg (pomarańczowy) staje się nagrodą (czerwony) w drugiej klatce, a następnie wraca do wroga w trzeciej, i ponownie do nagrody w czwartej. W Breakout (środkowy wiersz) cegły i wynik są nieścisłe między klatkami. W Road Runner (dolny wiersz) nagrody (małe niebieskie kropki na drodze) są nieścisłe między klatkami. Te nieścisłości nie występują w DIAMOND. W Breakout wynik jest niezawodnie aktualizowany o +7, gdy czerwona cegła jest zniszczona.

Wnioski

W tym artykule omówiliśmy DIAMOND, agenta uczenia wzmacniania szkolonego w ramach modelu świata dyfuzji. Ramy DIAMOND dokonują starannych wyborów projektowych, aby zapewnić, że ich model świata dyfuzji pozostaje wydajny i stabilny w długich horyzontach czasowych. Ramy DIAMOND zapewniają jakościową analizę, aby zademonstrować znaczenie tych wyborów projektowych. DIAMOND ustanawia nowy stan sztuki z średnim wynikiem ludzki znormalizowanym na poziomie 1,46 na dobrze ugruntowanym teście Atari 100k, najwyższym dla agentów szkolonych całkowicie w ramach modelu świata. Działanie w przestrzeni obrazu pozwala modelowi świata DIAMOND na bezproblemową substytucję środowiska, oferując większe wglądy w zachowania modelu świata i agenta. Godne uwagi jest to, że poprawiona wydajność w pewnych grach jest przypisywana lepszemu modelowaniu krytycznych szczegółów wizualnych. Ramy DIAMOND modelują środowisko jako standardowy proces decyzyjny Markowa z częściowo obserwowalnym stanem z zestawem stanów, zestawem dyskretnych działań i zestawem obserwacji obrazu. Funkcje przejścia opisują dynamikę środowiska, a funkcja nagrody mapuje przejścia na skalarną nagrodę.