Foi em 2018, quando surgiu a ideia de aprendizagem de reforço no contexto de um modelo mundial de rede neural foi introduzido pela primeira vez e logo esse princípio fundamental foi aplicado em modelos mundiais. Alguns dos modelos proeminentes que implementam a aprendizagem por reforço foram a estrutura Dreamer, que introduziu a aprendizagem por reforço a partir do espaço latente de um modelo de espaço de estados recorrente. O DreamerV2 demonstrou que o uso de latentes discretas pode resultar na redução de erros de composição, e a estrutura DreamerV3 foi capaz de alcançar desempenho semelhante ao humano em uma série de tarefas em diferentes domínios com hiperparâmetros fixos. 

Além disso, podem ser traçados paralelos entre os modelos de geração de imagens e os modelos mundiais, indicando que o progresso alcançado nos modelos de visão generativa poderia ser replicado para beneficiar os modelos mundiais. Desde o uso de transformadores em processamento de linguagem natural Os frameworks ganharam popularidade, surgiram os frameworks DALL-E e VQGAN. As estruturas implementaram codificadores automáticos discretos para converter imagens em tokens discretos e foram capazes de construir modelos geradores de texto para imagem altamente poderosos e eficientes, aproveitando as habilidades de modelagem de sequência dos transformadores autorregressivos. Ao mesmo tempo, os modelos de difusão ganharam força e hoje, os modelos de difusão se estabeleceram como um paradigma dominante para geração de imagens de alta resolução. Devido às capacidades oferecidas pelos modelos de difusão e aprendizagem 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 como política para aumento de dados em aprendizagem por reforço offline. 

Os modelos mundiais oferecem um método promissor para treinar agentes de aprendizagem por reforço com segurança e eficiência. Tradicionalmente, esses modelos utilizam sequências de variáveis ​​latentes discretas para simular a dinâmica do ambiente. No entanto, esta compressão pode ignorar detalhes visuais cruciais para a aprendizagem por reforço. Ao mesmo tempo, a popularidade dos modelos de difusão aumentou para geração de imagens, desafiando os métodos tradicionais que usam latentes discretas. Inspirados por essa mudança, neste artigo falaremos sobre DIAMOND (DIffusion As a Model Of eNvironment Dreams), um agente de aprendizagem por reforço treinado dentro de um modelo mundial de difusão. Exploraremos as opções de design necessárias para tornar a difusão adequada para a modelagem mundial e mostrar que detalhes visuais aprimorados levam a um melhor desempenho do agente. DIAMOND estabelece uma nova referência no teste competitivo do Atari 100k, alcançando uma pontuação humana normalizada média de 1.46, a mais alta para agentes treinados inteiramente dentro de um modelo mundial. 

DIAMANTE: DIFUSÃO COMO MODELO DE SONHOS AMBIENTAIS

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

Além disso, nos últimos anos, os modelos de difusão surgiram como a abordagem dominante para estruturas de geração de imagens de alta qualidade, uma vez que as estruturas construídas em modelos de difusão aprendem a reverter um processo de ruído e competem diretamente com algumas das abordagens mais bem estabelecidas que modelam tokens discretos. e, portanto, oferece uma alternativa promissora para eliminar a necessidade de discretização na modelagem mundial. Os modelos de difusão são conhecidos por sua capacidade de serem facilmente condicionados e de modelar com flexibilidade distribuições complexas e multimodais sem colapso de modo. Estes atributos são cruciais para a modelagem mundial, uma vez que o condicionamento permite que um modelo mundial reflita com precisão as ações de um agente, levando a uma atribuição de crédito mais confiável. Além disso, a modelagem de distribuições multimodais oferece uma maior diversidade de cenários de treinamento para o agente, melhorando seu desempenho geral. 

Com base nessas características, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), um agente de aprendizagem por reforço treinado dentro de um modelo mundial de difusão. A estrutura DIAMOND faz escolhas cuidadosas de design para garantir que seu modelo mundial de difusão permaneça eficiente e estável em horizontes de longo prazo. A estrutura fornece uma análise qualitativa para demonstrar a importância dessas escolhas de design. DIAMOND estabelece um novo estado da arte com uma pontuação humana normalizada média de 1.46 no bem estabelecido benchmark Atari 100k, a mais alta para agentes treinados inteiramente dentro de um modelo mundial. Operar no espaço de imagem permite que o modelo mundial de difusão do DIAMOND substitua perfeitamente o ambiente, oferecendo maiores insights sobre o modelo mundial e os comportamentos dos agentes. Notavelmente, o desempenho melhorado em certos jogos é atribuído a uma melhor modelagem de detalhes visuais críticos. A estrutura DIAMOND modela o ambiente como um POMDP padrão ou Processo de Decisão Markov 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 então utilizadas pelos agentes para ver os ambientes, uma vez que não podem acessar diretamente os estados. O objetivo principal da abordagem era obter uma política que mapeasse as observações em ações com a tentativa de maximizar o retorno esperado do desconto com um fator de desconto. Os modelos mundiais são modelos generativos do ambiente, e os modelos mundiais podem ser usados ​​para criar ambientes simulados para treinar agentes de aprendizagem por reforço no ambiente real e treinar agentes de aprendizagem por reforço no ambiente do modelo mundial. A Figura 1 demonstra o desenrolar da imaginação da estrutura DIAMOND ao longo do tempo. 

DIAMANTE: Metodologia e Arquitetura

Em sua essência, os modelos de difusão são uma classe de modelos generativos que geram uma amostra invertendo o processo de ruído e se inspiram fortemente na termodinâmica de não equilíbrio. A estrutura DIAMOND considera um processo de difusão indexado por uma variável de tempo contínua com marginais correspondentes e condições de contorno com uma distribuição anterior não estruturada tratável. Além disso, para obter um modelo generativo, que mapeie o ruído para os dados, a estrutura DIAMOND deve reverter o processo, sendo o processo de reversão também um processo de difusão, retrocedendo no tempo. Além disso, em qualquer momento, não é trivial estimar a função de pontuação, uma vez que a estrutura DIAMOND não acessa a verdadeira função de pontuação, e o modelo supera esse obstáculo implementando o objetivo de correspondência de pontuação, uma abordagem que facilita uma estrutura para treinar 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 generativo incondicional. No entanto, é necessário um modelo generativo condicional da dinâmica ambiental para servir como modelo mundial e, para servir este propósito, o quadro DIAMOND analisa o caso geral da abordagem POMDP, no qual o quadro pode fazer uso de observações e ações passadas para aproximar o estado Markoviano desconhecido. Conforme demonstrado na Figura 1., a estrutura DIAMOND faz uso desse histórico para condicionar um modelo de difusão, para estimar e gerar diretamente a próxima observação. Embora a estrutura DIAMOND possa recorrer a qualquer solucionador SDE ou EDO em teoria, há uma compensação entre NFE ou Número de Avaliações de Função e a qualidade da amostra que impacta significativamente o custo de inferência dos modelos de difusão. 

Com base nos aprendizados acima, vejamos agora a realização prática da estrutura DIAMOND de um modelo mundial baseado na difusão, incluindo os coeficientes de deriva e difusão correspondentes a uma escolha específica de abordagem de difusão. Em vez de optar pelo DDPM, um candidato naturalmente adequado para a tarefa, a estrutura DIAMOND baseia-se na formulação EDM e considera um kernel de perturbação com uma função de tempo de difusão com valor real chamada cronograma de ruído. A estrutura seleciona os pré-condicionadores para manter a variação de entrada e saída para qualquer nível de voz. O treinamento da rede mistura sinal e ruído de forma adaptativa dependendo do nível de degradação, e quando o ruído é baixo, o alvo passa a ser a diferença entre o sinal limpo e o perturbado, ou seja, o ruído gaussiano adicionado. Intuitivamente, isso evita que o objetivo do treinamento se torne trivial no regime de baixo ruído. Na prática, esse objetivo é a alta variância nos extremos do cronograma de ruído, de modo que 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. A estrutura DIAMOND faz uso de um componente U-Net 2D padrão para o campo vetorial e mantém um buffer de observações e ações anteriores que a estrutura usa para se condicionar. A estrutura DIAMOND então concatena essas observações anteriores com a próxima observação ruidosa e insere ações por meio de camadas de normalização de grupo adaptativas nos blocos residuais da U-Net. 

DIAMANTE: Experimentos e Resultados

Para uma avaliação abrangente, a estrutura DIAMOND opta pelo benchmark Atari 100k. O benchmark Atari 100k consiste em 26 jogos projetados para testar uma ampla gama de capacidades dos agentes. Em cada jogo, um agente está limitado a 100 mil ações no ambiente, o que equivale aproximadamente a 2 horas de jogo humano, para aprender o jogo antes da avaliação. Para efeito de comparação, os agentes irrestritos da Atari normalmente treinam 50 milhões de passos, representando um aumento de 500 vezes na experiência. Treinamos DIAMOND do zero usando 5 sementes aleatórias para cada jogo. Cada execução de treinamento exigiu cerca de 12 GB de VRAM e durou aproximadamente 2.9 dias em uma única Nvidia RTX 4090, totalizando 1.03 anos de GPU. A tabela a seguir fornece a pontuação de todos os jogos, a média e o IQM ou média interquartil das pontuações normalizadas por humanos. 

Seguindo as limitações das estimativas pontuais, a estrutura DIAMOND fornece confiança de inicialização estratificada na média e no IQM ou média interquartil de pontuações normalizadas por humanos, juntamente com perfis de desempenho e métricas adicionais, conforme resumido na figura a seguir. 

Os resultados mostram que o DIAMOND tem um desempenho excepcionalmente bom em todo o benchmark, superando jogadores humanos em 11 jogos e alcançando um HNS médio sobre-humano de 1.46, estabelecendo um novo recorde para agentes treinados inteiramente dentro de um modelo mundial. Além disso, o IQM do DIAMOND é comparável ao STORM e excede todas as outras linhas de base. DIAMOND se destaca em ambientes onde a captura de pequenos detalhes é crucial, como Asterix, Breakout e RoadRunner. Além disso, como discutido anteriormente, o quadro DIAMOND tem a flexibilidade de implementar qualquer modelo de difusão no seu pipeline, embora opte pela abordagem EDM, teria sido uma escolha natural optar pelo modelo DDPM, uma vez que já está a ser implementado em numerosos países. aplicações geradoras de imagens. Para comparar a abordagem EDM com a implementação DDPM, a estrutura DIAMOND treina ambas as variantes com a mesma arquitetura de rede no mesmo conjunto de dados estáticos compartilhados com mais de 100 mil quadros coletados com uma política especializada. O número de etapas de remoção de ruído está diretamente relacionado ao custo de inferência do modelo mundial e, portanto, menos etapas reduzirão o custo de treinamento de um agente em trajetórias imaginadas. Para garantir que o nosso modelo mundial permaneça computacionalmente comparável com outras linhas de base, como o IRIS, que requer 16 NFE por passo de tempo, pretendemos utilizar não mais do que dezenas de passos de eliminação de ruído, de preferência menos. No entanto, definir o número de etapas de eliminação de ruído muito baixo pode degradar a qualidade visual, levando a erros compostos. Para avaliar a estabilidade de diferentes variantes de difusão, exibimos trajetórias imaginadas geradas autoregressivamente até t = 1000 passos de tempo na figura a seguir, usando diferentes números de etapas de remoção de ruído n ≤ 10. 

Observamos que o uso do DDPM (a), neste regime, resulta em graves erros de composição, fazendo com que o modelo mundial saia rapidamente da distribuição. Em contraste, o modelo mundial de difusão baseado em EDM (b) permanece muito mais estável ao longo de longos horizontes de tempo, mesmo com uma única etapa de remoção de ruído. Trajetórias imaginadas com modelos mundiais de difusão baseados em DDPM (esquerda) e EDM (direita) são mostrados. A observação inicial em t = 0 é a mesma para ambos, e cada linha corresponde a um número decrescente de etapas de remoção de ruído n. Observamos que a geração baseada em DDPM sofre de erros compostos, com números menores de etapas de remoção de ruído levando a um acúmulo de erros mais rápido. Em contraste, o modelo mundial baseado em EDM do DIAMOND permanece muito mais estável, mesmo para n = 1. A previsão ideal de etapa única é a expectativa sobre possíveis reconstruções para uma determinada entrada ruidosa, que pode estar fora de distribuição se a distribuição posterior for multimodal . Embora alguns jogos, como Breakout, tenham transições determinísticas que podem ser modeladas com precisão com uma única etapa de remoção de ruído, outros jogos exibem observabilidade parcial, resultando em distribuições de observação multimodais. Nestes casos, é necessário um solucionador iterativo para orientar o procedimento de amostragem para um modo específico, conforme ilustrado no jogo Boxe na figura a seguir. Conseqüentemente, a estrutura DIAMOND definiu n = 3 em todos os nossos experimentos.

A figura acima compara a amostragem de etapa única (linha superior) e de várias etapas (linha inferior) no Boxe. Os movimentos do jogador preto são imprevisíveis, fazendo com que a eliminação de ruído em uma única etapa interpole entre os resultados possíveis, resultando em previsões borradas. Em contraste, a amostragem em várias etapas produz uma imagem nítida, orientando a geração para um modo específico. Curiosamente, uma vez que a política controla o jogador branco, as suas ações são conhecidas pelo modelo mundial, eliminando a ambiguidade. Assim, tanto a amostragem de etapa única quanto a de múltiplas etapas predizem corretamente a posição do jogador branco.

Na figura acima, as trajetórias imaginadas pelo DIAMOND geralmente apresentam maior qualidade visual e são mais fiéis ao verdadeiro ambiente em comparação com aquelas imaginadas pelo IRIS. As trajetórias geradas pelo IRIS contêm inconsistências visuais entre os frames (destacados 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 evitar inimigos, portanto, essas pequenas discrepâncias visuais podem tornar mais desafiador aprender uma política ideal. A figura mostra quadros consecutivos imaginados com IRIS (esquerda) e DIAMOND (direita). As caixas brancas destacam inconsistências entre quadros, que aparecem apenas em trajetórias geradas com IRIS. Em Asterix (linha superior), um inimigo (laranja) torna-se uma recompensa (vermelho) no segundo quadro, depois reverte para um inimigo no terceiro e novamente para uma recompensa no quarto. No Breakout (linha do meio), os tijolos e a pontuação são inconsistentes entre os quadros. Em Road Runner (linha inferior), as recompensas (pequenos pontos azuis na estrada) são renderizadas de forma inconsistente entre os quadros. Essas inconsistências não ocorrem com DIAMOND. No Breakout, a pontuação é atualizada de forma confiável em +7 quando um tijolo vermelho é quebrado. 

Conclusão

Neste artigo falamos sobre DIAMOND, um agente de aprendizagem por reforço treinado dentro de um modelo mundial de difusão. A estrutura DIAMOND faz escolhas cuidadosas de design para garantir que seu modelo mundial de difusão permaneça eficiente e estável em horizontes de longo prazo. A estrutura fornece uma análise qualitativa para demonstrar a importância dessas escolhas de design. DIAMOND estabelece um novo estado da arte com uma pontuação humana normalizada média de 1.46 no bem estabelecido benchmark Atari 100k, a mais alta para agentes treinados inteiramente dentro de um modelo mundial. Operar no espaço de imagem permite que o modelo mundial de difusão do DIAMOND substitua perfeitamente o ambiente, oferecendo maiores insights sobre o modelo mundial e os comportamentos dos agentes. Notavelmente, o desempenho melhorado em certos jogos é atribuído a uma melhor modelagem de detalhes visuais críticos. A estrutura DIAMOND modela o ambiente como um POMDP padrão ou Processo de Decisão Markov 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.