Foi em 2018 que a ideia de aprendizado por reforço em um modelo de mundo neural foi introduzida pela primeira vez, e logo, esse princípio fundamental foi aplicado em modelos de mundo. Alguns dos modelos proeminentes que implementam aprendizado por reforço foram o framework Dreamer, que introduziu aprendizado por reforço a partir do espaço latente de um modelo de espaço de estado recorrente. O DreamerV2 demonstrou que o uso de latentes discretos pode resultar em erros de compounding reduzidos, e o framework DreamerV3 foi capaz de alcançar desempenho humano em uma série de tarefas em diferentes domínios com hiperparâmetros fixos.

Além disso, paralelos podem ser traçados entre modelos de geração de imagens e modelos de mundo, indicando que o progresso feito em modelos de visão generativos pode ser replicado para beneficiar os modelos de mundo. Desde que o uso de transformadores em frameworks de processamento de linguagem natural ganhou popularidade, os frameworks DALL-E e VQGAN surgiram. Os frameworks implementaram autoencoders discretos para converter imagens em tokens discretos e foram capazes de construir modelos generativos de texto para imagem altamente poderosos e eficientes, aproveitando as habilidades de modelagem de sequência dos transformadores autoregressivos. Ao mesmo tempo, os modelos de difusão ganharam tração, e hoje, os modelos de difusão estabeleceram-se como um paradigma dominante para geração de imagens de alta resolução. Devido às capacidades oferecidas pelos modelos de difusão e aprendizado por reforço, estão sendo feitas tentativas de combinar as duas abordagens, com o objetivo de aproveitar a flexibilidade dos modelos de difusão como modelos de trajetória, modelos de recompensa, planejadores e políticas para aumento de dados em aprendizado por reforço offline.

Os modelos de mundo oferecem um método promissor para treinar agentes de aprendizado por reforço de forma segura e eficiente. Tradicionalmente, esses modelos usam sequências de variáveis latentes discretas para simular a dinâmica do ambiente. No entanto, essa compressão pode ignorar detalhes visuais cruciais para o aprendizado por reforço. Ao mesmo tempo, os modelos de difusão ganharam popularidade para geração de imagens, desafiando métodos tradicionais que usam latentes discretos. Inspirados por essa mudança, neste artigo, vamos falar sobre DIAMOND (DIffusion As a Model Of eNvironment Dreams), um agente de aprendizado por reforço treinado dentro de um modelo de mundo de difusão. Vamos explorar as escolhas de design necessárias para tornar a difusão adequada para modelagem de mundo e mostrar que detalhes visuais aprimorados levam a um melhor desempenho do agente. O DIAMOND estabelece um novo padrão no competitivo teste Atari 100k, alcançando uma pontuação média normalizada humana de 1,46, a mais alta para agentes treinados inteiramente dentro de um modelo de mundo.

DIAMOND: DIffusion As a Model Of eNvironment Dreams

Os modelos de mundo ou modelos gerativos de ambientes estão surgindo como um dos componentes mais importantes para agentes gerativos planejarem e raciocinarem sobre seus ambientes. Embora o uso de aprendizado por reforço tenha alcançado um sucesso considerável nos últimos anos, os modelos que implementam aprendizado por reforço são conhecidos por serem ineficientes em termos de amostras, o que limita significativamente suas aplicações no mundo real. Por outro lado, os modelos de mundo demonstraram sua capacidade de treinar agentes de aprendizado por reforço de forma eficiente em diferentes ambientes, com uma eficiência de amostra significativamente melhorada, permitindo que o modelo aprenda com experiências do mundo real. Os frameworks de modelagem de mundo recentes geralmente modelam a dinâmica do ambiente como uma sequência de variáveis latentes discretas, com o modelo discretizando o espaço latente para evitar erros de compounding em horizontes de tempo multi-etapas. Embora a abordagem possa entregar resultados substanciais, também está associada à perda de informações, levando à perda de qualidade de reconstrução e perda de generalidade. A perda de informações pode se tornar um obstáculo significativo para cenários do mundo real que requerem que as informações sejam bem definidas, como o treinamento de veículos autônomos. Nesses tarefas, pequenas mudanças ou detalhes na entrada visual, como a cor do semáforo ou o indicador de direção do veículo à frente, podem alterar a política do agente. Embora aumentar o número de latentes discretos possa ajudar a evitar a perda de informações, isso aumenta significativamente os custos de computação.

Além disso, nos últimos anos, os modelos de difusão surgiram como a abordagem dominante para frameworks de geração de imagens de alta qualidade, desde que os frameworks construídos com base em modelos de difusão aprendem a reverter um processo de ruído e competem diretamente com algumas das abordagens mais estabelecidas que modelam tokens discretos, oferecendo assim uma alternativa promissora para eliminar a necessidade de discretização na modelagem de mundo. Os modelos de difusão são conhecidos por sua capacidade de ser facilmente condicionados e modelar distribuições complexas e multi-modais sem colapso de modos. Esses atributos são cruciais para a modelagem de mundo, pois a condicionação permite que um modelo de mundo reflita com precisão as ações de um agente, levando a uma atribuição de crédito mais confiável. Além disso, modelar distribuições multi-modais oferece uma maior diversidade de cenários de treinamento para o agente, melhorando seu desempenho geral.

Com base nesses atributos, o DIAMOND, (DIffusion As a Model Of eNvironment Dreams), um agente de aprendizado por reforço treinado dentro de um modelo de mundo de difusão. O framework DIAMOND faz escolhas de design cuidadosas para garantir que seu modelo de mundo de difusão permaneça eficiente e estável em longos horizontes de tempo. O framework fornece uma análise qualitativa para demonstrar a importância dessas escolhas de design. O DIAMOND estabelece um novo padrão com uma pontuação média normalizada humana de 1,46 no bem estabelecido benchmark Atari 100k, o mais alto para agentes treinados inteiramente dentro de um modelo de mundo. Operar no espaço de imagem permite que o modelo de mundo de difusão do DIAMOND substitua perfeitamente o ambiente, oferecendo insights mais profundos sobre o comportamento do modelo de mundo e do agente. Notavelmente, o desempenho melhorado em certos jogos é atribuído a uma melhor modelagem de detalhes visuais críticos. O framework DIAMOND modela o ambiente como um POMDP padrão ou um Processo de Decisão Markoviano Parcialmente Observável com um conjunto de estados, um conjunto de ações discretas e um conjunto de observações de imagem. As funções de transição descrevem a dinâmica do ambiente, e a função de recompensa mapeia as transições para recompensas escalares. A função de observação descreve as probabilidades de observação e emite observações de imagem, que são usadas pelos agentes para ver os ambientes, pois eles não podem acessar diretamente os estados. O objetivo principal da abordagem foi obter uma política que mapeie observações para ações com o objetivo de maximizar a recompensa esperada com um fator de desconto. Os modelos de mundo são modelos gerativos do ambiente, e os modelos de mundo podem ser usados para criar ambientes simulados para treinar agentes de aprendizado por reforço no ambiente real e treinar agentes de aprendizado por reforço no ambiente do modelo de mundo. A Figura 1 demonstra a imaginação desenrolada do framework DIAMOND ao longo do tempo.

DIAMOND: Metodologia e Arquitetura

Em sua essência, os modelos de difusão são uma classe de modelos gerativos que geram uma amostra revertendo o processo de ruído e tiram inspiração pesada da termodinâmica não equilibrada. O framework DIAMOND considera um processo de difusão indexado por uma variável de tempo contínua com marginais e condições de contorno correspondentes e uma distribuição de priori não estruturada traçável. Além disso, para obter um modelo gerativo, que mapeia do ruído para os dados, o framework DIAMOND deve reverter o processo, com o processo de reversão também sendo um processo de difusão, executado em sentido inverso ao tempo. Além disso, em qualquer ponto dado no tempo, não é trivial estimar a função de pontuação, pois o framework DIAMOND não tem acesso à função de pontuação real, e o modelo supera esse obstáculo implementando o objetivo de correspondência de pontuação, uma abordagem que facilita que um framework treine um modelo de pontuação sem conhecer a função de pontuação subjacente. O modelo de difusão baseado em pontuação fornece um modelo gerativo incondicional. No entanto, um modelo gerativo condicional de dinâmica do ambiente é necessário para servir como um modelo de mundo, e para atender a esse propósito, o framework DIAMOND olha para o caso geral da abordagem POMDP, na qual o framework pode usar observações e ações passadas para aproximar o estado markoviano desconhecido. Como demonstrado na Figura 1, o framework DIAMOND usa essa história para condicionar um modelo de difusão para estimar e gerar a próxima observação diretamente. Embora o framework DIAMOND possa recorrer a qualquer solucionador de EDO ou ODE em teoria, há um trade-off entre NFE ou Número de Avaliações de Função e qualidade da amostra que impacta significativamente o custo de inferência dos modelos de difusão.

Com base nas aprendizagens acima, vamos agora olhar para a realização prática do framework DIAMOND de um modelo de mundo baseado em difusão, incluindo os coeficientes de deriva e difusão correspondentes a uma escolha particular de abordagem de difusão. Em vez de optar pelo DDPM, um candidato naturalmente adequado para a tarefa, o framework DIAMOND se baseia na formulação EDM e considera um kernel de perturbação com uma função de valor real de tempo de difusão chamada de cronograma de ruído. O framework seleciona os pré-condicionadores para manter a variância de entrada e saída para qualquer nível de voz. O treinamento da rede mistura sinal e ruído adaptativamente dependendo do nível de degradação, e quando o ruído é baixo, e o alvo se torna a diferença entre o sinal limpo e o sinal perturbado, ou seja, o ruído gaussiano adicionado. Intuitivamente, isso impede que o objetivo de treinamento se torne trivial no regime de baixo ruído. Na prática, esse objetivo é de alta variância nos extremos do cronograma de ruído, então o modelo amostra o nível de ruído a partir de uma distribuição log-normal escolhida empiricamente para concatenar o treinamento em torno das regiões de ruído médio. O framework DIAMOND usa um componente U-Net 2D padrão para o campo vetorial e mantém um buffer de observações e ações passadas que o framework usa para se condicionar. O framework DIAMOND, em seguida, concatena essas observações passadas à próxima observação ruidosa e ações de entrada por meio de camadas de normalização de grupo adaptativas nos blocos residuais do U-Net.

DIAMOND: Experimentos e Resultados

Para uma avaliação abrangente, o framework DIAMOND opta pelo benchmark Atari 100k. O benchmark Atari 100k consiste em 26 jogos projetados para testar uma ampla gama de capacidades de agente. Em cada jogo, um agente é limitado a 100k ações no ambiente, o que é aproximadamente equivalente a 2 horas de jogo humano, para aprender o jogo antes da avaliação. Para comparação, agentes Atari não constrangidos normalmente treinam por 50 milhões de etapas, representando um aumento de 500 vezes na experiência. Treinamos o DIAMOND do zero usando 5 sementes aleatórias para cada jogo. Cada execução de treinamento exigiu cerca de 12GB de VRAM e levou aproximadamente 2,9 dias em um único Nvidia RTX 4090, totalizando 1,03 anos de GPU no total. A tabela a seguir fornece a pontuação para todos os jogos, a média e a IQM ou média interquartil de pontuações normalizadas humanas.

Seguindo as limitações das estimativas pontuais, o framework DIAMOND fornece confiança de bootstrap estratificada na média e na IQM ou média interquartil de pontuações normalizadas humanas, juntamente com perfis de desempenho e métricas adicionais, como resumido na figura a seguir.

Os resultados mostram que o DIAMOND se sai excepcionalmente bem em todo o benchmark, superando os jogadores humanos em 11 jogos e alcançando uma pontuação média normalizada humana super-humana de 1,46, estabelecendo um novo recorde para agentes treinados inteiramente dentro de um modelo de mundo. Além disso, a IQM do DIAMOND é comparável à do STORM e supera todas as outras linhas de base. O DIAMOND se destaca em ambientes onde a captura de pequenos detalhes é crucial, como Asterix, Breakout e RoadRunner. Além disso, como discutido anteriormente, o framework DIAMOND tem a flexibilidade de implementar qualquer modelo de difusão em sua pipeline, embora opte pela abordagem EDM, teria sido uma escolha natural optar pelo modelo DDPM, pois já está sendo implementado em inúmeras aplicações gerativas de imagens. Para comparar a abordagem EDM contra a implementação DDPM, o framework DIAMOND treina ambos os variantes com a mesma arquitetura de rede no mesmo conjunto de dados estático compartilhado com mais de 100k frames coletados com uma política de especialista. O número de etapas de desruído está diretamente relacionado ao custo de inferência do modelo de mundo, e assim, menos etapas reduzirão o custo de treinamento de um agente em trajetórias imaginadas. Para garantir que nosso modelo de mundo permaneça computacionalmente comparável com outras linhas de base, como o IRIS, que requer 16 NFE por etapa, visamos usar não mais de dezenas de etapas de desruído, preferencialmente menos. No entanto, definir o número de etapas de desruído muito baixo pode degradar a qualidade visual, levando a erros de compounding. Para avaliar a estabilidade de diferentes variantes de difusão, exibimos trajetórias imaginadas geradas autoregressivamente até t = 1000 etapas na figura a seguir, usando diferentes números de etapas de desruído n ≤ 10.

Observamos que usar o DDPM (a) nesse regime resulta em erros de compounding graves, causando o modelo de mundo a se afastar rapidamente da distribuição. Em contraste, o modelo de mundo de difusão baseado em EDM (b) permanece muito mais estável em longos horizontes de tempo, mesmo com uma única etapa de desruído. As trajetórias imaginadas com modelos de mundo de difusão baseados em DDPM (esquerda) e EDM (direita) são mostradas. A observação inicial em t = 0 é a mesma para ambos, e cada linha corresponde a um número decrescente de etapas de desruído n. Observamos que a geração baseada em DDPM sofre de erros de compounding, com números menores de etapas de desruído levando a uma acumulação de erros mais rápida. Em contraste, o modelo de mundo de difusão do DIAMOND permanece muito mais estável, mesmo para n = 1. A previsão de um passo ótimo é a expectativa sobre possíveis reconstruções para uma entrada ruidosa dada, que pode estar fora da distribuição se a distribuição posterior for multi-modal. Embora algumas tarefas, como Breakout, tenham transições determinísticas que podem ser modeladas com precisão com uma única etapa de desruído, outras tarefas exibem observabilidade parcial, resultando em distribuições de observação multi-modais. Nesses casos, um solucionador iterativo é necessário para guiar o procedimento de amostragem em direção a um modo específico, como ilustrado no jogo Boxing na figura a seguir. Consequentemente, o framework DIAMOND define n = 3 em todos os nossos experimentos.

A figura acima compara a amostragem de um passo (linha superior) e a amostragem de multi-passos (linha inferior) no Boxing. Os movimentos do jogador preto são imprevisíveis, causando a desruído de um passo a interpolar entre resultados possíveis, resultando em previsões borradas. Em contraste, a amostragem de multi-passos produz uma imagem clara, guiando a geração em direção a um modo específico. Interessantemente, desde que a política controla o jogador branco, suas ações são conhecidas pelo modelo de mundo, eliminando a ambiguidade. Assim, tanto a amostragem de um passo quanto a amostragem de multi-passos prevêem corretamente a posição do jogador branco.

Na figura acima, as trajetórias imaginadas pelo DIAMOND geralmente exibem uma qualidade visual mais alta e são mais fiéis ao ambiente real em comparação com as imaginadas pelo IRIS. As trajetórias geradas pelo IRIS contêm inconsistências visuais entre frames (realçadas por caixas brancas), como inimigos sendo exibidos como recompensas e vice-versa. Embora essas inconsistências possam afetar apenas alguns pixels, elas podem impactar significativamente o aprendizado por reforço. Por exemplo, um agente normalmente visa recompensas e evita inimigos, então essas pequenas discrepâncias visuais podem tornar mais difícil aprender uma política ótima. A figura mostra frames consecutivos imaginados com o IRIS (esquerda) e o DIAMOND (direita). As caixas brancas realçam inconsistências entre frames, que só aparecem em trajetórias geradas com o IRIS. Em Asterix (linha superior), um inimigo (laranja) se torna uma recompensa (vermelha) no segundo frame, então reverte para um inimigo no terceiro, e novamente para uma recompensa no quarto. Em Breakout (linha do meio), os tijolos e a pontuação são inconsistentes entre frames. Em Road Runner (linha inferior), as recompensas (pequenos pontos azuis na estrada) são renderizadas de forma inconsistente entre frames. Essas inconsistências não ocorrem com o DIAMOND. Em Breakout, a pontuação é atualizada de forma confiável por +7 quando um tijolo vermelho é quebrado.

Conclusão

Neste artigo, falamos sobre o DIAMOND, um agente de aprendizado por reforço treinado dentro de um modelo de mundo de difusão. O framework DIAMOND faz escolhas de design cuidadosas para garantir que seu modelo de mundo de difusão permaneça eficiente e estável em longos horizontes de tempo. O framework fornece uma análise qualitativa para demonstrar a importância dessas escolhas de design. O DIAMOND estabelece um novo padrão com uma pontuação média normalizada humana de 1,46 no bem estabelecido benchmark Atari 100k, o mais alto para agentes treinados inteiramente dentro de um modelo de mundo. Operar no espaço de imagem permite que o modelo de mundo de difusão do DIAMOND substitua perfeitamente o ambiente, oferecendo insights mais profundos sobre o comportamento do modelo de mundo e do agente. Notavelmente, o desempenho melhorado em certos jogos é atribuído a uma melhor modelagem de detalhes visuais críticos. O framework DIAMOND modela o ambiente como um POMDP padrão ou um Processo de Decisão Markoviano Parcialmente Observável com um conjunto de estados, um conjunto de ações discretas e um conjunto de observações de imagem. As funções de transição descrevem a dinâmica do ambiente, e a função de recompensa mapeia as transições para recompensas escalares.