Resumo Técnico: JEDI: Modelo de Mundo de Difusão de Incorporação Conjunta para Aprendizado por Reforço Baseado em Modelo Online

Declaração do Problema

O Aprendizado por Reforço Baseado em Modelo Online (MBRL) depende de modelos de mundo aprendidos para prever resultados futuros e melhorar a eficiência de amostragem. Embora os modelos de difusão tenham surgido como backbones poderosos para modelagem de mundo em domínios visualmente complexos, as abordagens atuais enfrentam uma tensão significativa:

1. Difusão de Pixels: Modelos como DIAMOND operam diretamente no espaço de pixels. Embora eficazes, são computacionalmente caros em termos de memória (VRAM), tempo de amostragem e duração do treinamento.
2. Difusão Latente: Mover a difusão para o espaço latente melhora a eficiência, mas frequentemente depende de latentes treinados separadamente (por exemplo, via objetivos de reconstrução) em vez de objetivos de modelo de mundo de ponta a ponta. Essas abordagens, como Horizon Imagination (HI), ainda não igualaram o desempenho dos modelos de difusão baseados em pixels e carecem da aprendizagem de representação preditiva de ponta a ponta que impulsionou o progresso moderno do MBRL.

A questão central abordada é se um modelo de mundo de difusão pode aprender sua representação latente de ponta a ponta a partir do objetivo de remoção de ruído usando uma estrutura de Arquitetura Preditiva de Incorporação Conjunta (JEPA), combinando assim a eficiência da difusão latente com o desempenho da aprendizagem preditiva de ponta a ponta.

Metodologia: JEDI

Os autores apresentam o JEDI (Joint Embedding Diffusion), o primeiro modelo de mundo de difusão latente online de ponta a ponta construído sobre uma estrutura JEPA.

Arquitetura Central

O JEDI aprende simultaneamente um codificador de espaço latente ($E_\phi$) e um modelo de dinâmica. O sistema consiste em:

• Codificador: Comprime as observações $x_t$ em latentes limpos $z^0_t$.
• Dinâmica de Difusão Latente: Prevê o próximo estado latente através de remoção de ruído condicional. Em vez de uma previsão de passo único, utiliza uma cadeia de difusão reversa condicionada ao latente atual $z^0_t$ e às ações passadas para prever o latente futuro limpo $z^0_{t+1}$.
• Cabeças de Recompensa e Terminação: Preveem recompensas e sinais de terminação diretamente a partir do estado latente.

Mecanismo de Treinamento

Diferentemente dos modelos de difusão latente tradicionais que dependem de perdas de reconstrução (por exemplo, VAEs) para treinar o codificador, o JEDI treina o codificador de ponta a ponta usando a perda de remoção de ruído de difusão.

• Objetivo de Incorporação Conjunta: O codificador é atualizado via gradientes do objetivo de remoção de ruído condicional. O alvo latente futuro é tratado como um "rótulo" derivado do ambiente, e o modelo aprende a remover o ruído de uma versão ruidosa desse latente futuro de volta ao estado limpo, condicionado ao estado atual.
• Stop-Gradient: Para evitar o colapso da representação e estabilizar o treinamento, um stop-gradient é aplicado ao alvo latente futuro durante as atualizações do codificador, seguindo os princípios JEPA e TDMPC2.
• Técnicas de Estabilização: O modelo emprega clamping tanh diferenciável nas ativações latentes e um mecanismo de comutação aleatória entre a saída do removedor de ruído e a saída do codificador para condicionamento do próximo passo, equilibrando consistência e supervisão.

Motivação Teórica

O artigo fornece uma justificativa teórica que liga o JEDI ao princípio do Gargalo de Informação.

• JEPA como Gargalo: O objetivo padrão JEPA é mostrado como induzindo um gargalo de informação variacional, recompensando representações que retêm informações úteis para prever o futuro enquanto comprimem detalhes irrelevantes.
• Decomposição Difusão-JEPA: Os autores demonstram que a difusão condicional latente admite uma decomposição variacional de gargalo de informação similar. A perda de remoção de ruído condicional atua como o termo preditivo dentro dessa estrutura de gargalo. Isso sugere que a remoção de ruído de difusão fornece um sinal de aprendizagem preditiva de ponta a ponta principiado para o codificador latente, em vez de servir meramente como um substituto generativo.

Contribuições Principais

1. Primeiro Modelo de Mundo de Difusão Latente de Ponta a Ponta: O JEDI é o primeiro modelo MBRL online a treinar seu codificador latente de ponta a ponta usando um objetivo de remoção de ruído de difusão dentro de uma estrutura JEPA, eliminando a necessidade de pré-treinamento separado ou perdas de reconstrução.
2. Conexão Teórica: O artigo estabelece um link matemático entre objetivos de difusão condicional e o gargalo de informação variacional, explicando por que a remoção de ruído pode suportar efetivamente a aprendizagem de representação preditiva.
3. Ganhos de Eficiência: Ao migrar do espaço de pixels para um espaço latente aprendido compacto ($16 \times 8 \times 8$) e treinar de ponta a ponta, o JEDI reduz significativamente a sobrecarga computacional em comparação com a difusão baseada em pixels.

Resultados Experimentais

O JEDI foi avaliado no benchmark Atari100k (26 tarefas, 5 sementes) e no ambiente Craftium (tomada de decisão corporificada 3D).

Desempenho

• Atari100k: O JEDI alcança desempenho competitivo, igualando ou superando bases fortes (DIAMOND, DreamerV3, STORM) nas métricas Média Interquartil (IQM) e Lacuna de Otimalidade. Alcança resultados de Estado da Arte (SOTA) em 7 jogos e desempenho de vice-campeão no número de tarefas super-humanas.
• Craftium: O JEDI supera o Horizon Imagination (HI) em todas as quatro tarefas testadas, demonstrando que latentes preditivos de ponta a ponta podem suportar alto desempenho em ambientes 3D sem treinamento latente baseado em reconstrução.
• Estocasticidade: Em experimentos com pulo de quadros estocástico, o JEDI supera consistentemente o DreamerV3, sugerindo que modelos de difusão podem ser mais robustos a alta incerteza aleatória.

Eficiência

Em relação à base baseada em pixels DIAMOND, o JEDI demonstra melhorias substanciais de eficiência:

• Uso de VRAM: Utiliza 43% menos VRAM.
• Velocidade de Amostragem: Alcança amostragem do modelo de mundo >3× mais rápida.
• Velocidade de Treinamento: Treina >2,5× mais rápido (38 horas vs. 98 horas em A100 para Atari100k).
• Tamanho do Modelo: Mantém uma contagem de parâmetros semelhante (13,5M) ao DIAMOND, mas supera significativamente o HI (97M parâmetros) em eficiência, mantendo desempenho competitivo.

Mudança no Perfil de Desempenho

Uma descoberta distinta é que o JEDI exibe um perfil de desempenho em nível de tarefa diferente em comparação com o DIAMOND. Enquanto o DIAMOND se destaca em regimes de alto HNS (Human Normalized Score), o JEDI mostra melhorias marcantes em regimes de baixo HNS (jogos difíceis). Isso é capturado pela métrica Har-HNS, onde o JEDI supera o DIAMOND, indicando que ele falha menos severamente nas tarefas mais difíceis. Os autores atribuem isso à interface latente de ponta a ponta reduzir a dimensionalidade da entrada para o ator-crítico, permitindo melhor otimização em jogos complexos com alto espaço de ações.

Significado e Alegações

O artigo alega que o JEDI demonstra que a difusão latente preditiva é uma alternativa viável e eficiente tanto à difusão no espaço de pixels quanto a modelos latentes treinados separadamente para MBRL online.

• Além do Cálculo: Os resultados sugerem que latentes preditivos de ponta a ponta alteram o comportamento do próprio modelo de mundo, não apenas o custo computacional. A mudança nos perfis de desempenho indica que a representação aprendida altera fundamentalmente como o agente interage com ambientes complexos.
• Abordagem Principiada: Ao fundamentar o método na teoria do gargalo de informação, o artigo posiciona a difusão latente preditiva como uma direção principiada para escalar o MBRL online.
• Limitações: Os autores notam modestamente que as evidências atualmente se limitam ao Atari100k e ao Craftium, e o método ainda não foi testado em domínios como DMControl ou Procgen. Eles também reconhecem que o desempenho depende de escolhas de design específicas (por exemplo, clamping, comutação aleatória) e que a derivação teórica motiva propriedades de representação em vez de garantir a otimização do objetivo variacional completo.