Simulus: Combining Improvements in Sample-Efficient World Model Agents

Resumo Técnico: Simulus – Combinando Melhorias em Agentes de Modelo de Mundo Eficientes em Amostras

Declaração do Problema
Os modelos de mundo (WMs) representam uma fronteira no aprendizado por reforço (RL) eficiente em amostras, mas seu potencial frequentemente não é realizado devido à complexidade de integrar melhorias promissoras. Embora avanços individuais existam, combiná-los exige expertise significativa e esforço de engenharia. Além disso, os modelos de mundo baseados em tokens existentes (TBWMs) são amplamente restritos a observações de imagem e ações discretas, limitando sua aplicabilidade a controle contínuo e configurações multimodais. Adicionalmente, componentes como motivação intrínseca (para redução da incerteza epistêmica) e replay priorizado são frequentemente negligenciados em agentes eficientes em amostras, em parte devido ao alto custo de integração e ao risco de que a exploração intrínseca desperdice orçamentos de interação escassos em experiências irrelevantes para a tarefa. Esse cenário espelha a era pré-Rainbow no RL profundo, onde melhorias validadas individualmente não haviam sido combinadas sistematicamente.

Metodologia
Os autores introduzem o Simulus, um agente modular de modelo de mundo baseado em tokens projetado para integrar e validar um conjunto de melhorias complementares. O agente é construído sobre a arquitetura REM e compreende quatro módulos distintos: um módulo de representação ($V$), um modelo de mundo ($M$), um controlador ($C$) e um buffer de replay.

1. Framework de Tokenização Flexível ($V$): O Simulus emprega um sistema de tokenização modular que suporta combinações arbitrárias de modalidades de observação e ação. Ele lida nativamente com imagens, vetores contínuos, variáveis categóricas e "categóricos 2D" (grades de variáveis categóricas semelhantes a imagens). Os codificadores produzem sequências de tokens de comprimento fixo, que são concatenadas em uma representação unificada. Isso permite que o agente opere em ambientes diversos, incluindo controle contínuo e tarefas simbólicas multimodais.
2. Modelo de Mundo ($M$): O modelo de dinâmica utiliza uma Rede Retentiva (RetNet) aumentada com um mecanismo de Previsão Paralela de Observação (POP) para modelar a dinâmica do ambiente como modelagem de sequência sobre tokens discretos.
   • Motivação Intrínseca: Para reduzir a incerteza epistêmica, o Simulus emprega um ensemble de cabeças de previsão da próxima observação. A discordância entre essas cabeças, medida via Divergência de Jensen-Shannon (JSD), gera recompensas intrínsecas ($r_{int}$). Estas são combinadas com recompensas extrínsecas para guiar o controlador em direção a regiões de alta incerteza durante a "imaginação" (planejamento dentro do modelo).
   • Replay Priorizado: O buffer de replay mantém valores de perda do modelo de mundo para exemplos armazenados. O treinamento utiliza uma mistura de amostragem uniforme e priorizada (controlada por um único parâmetro $\alpha$), onde as amostras são selecionadas proporcionalmente às suas perdas transformadas por softmax.
3. Controlador ($C$): O controlador é uma arquitetura estendida de ator-crítico que opera inteiramente sobre experiência simulada gerada pelo modelo de mundo.
   • Regressão como Classificação (RaC): Para lidar robustamente com sinais de recompensa e retorno ilimitados, o Simulus adota uma rede de classificação com bins exponenciais para previsão de valor. Utiliza o método de perda HL-Gauss (baseado em entropia cruzada) em vez do erro quadrático médio tradicional, melhorando a estabilidade na estimativa de valor.
4. Ciclo de Treinamento: O agente segue um ciclo de coleta de dados, aprendizado de representação, aprendizado de modelo de mundo e aprendizado de controle na imaginação. O controlador é treinado usando retornos $\lambda$ e um objetivo REINFORCE com uma linha de base de valor.

Principais Contribuições

• Integração de Componentes Negligenciados: O artigo identifica e integra motivação intrínseca, replay de modelo de mundo priorizado e uma abordagem híbrida RaC. Demonstra que esses componentes não são apenas individualmente eficazes, mas produzem ganhos sinérgicos quando combinados.
• Agentes Baseados em Tokens Versáteis: Ao adaptar a tokenização multimodal ao regime eficiente em amostras, o Simulus estabelece os TBWMs como agentes versáteis capazes de lidar com entradas visuais, contínuas e multimodais, superando as limitações dos modelos de ação discreta baseados apenas em imagem.
• Validação da Exploração Induzida por Modelo: O trabalho mostra que a exploração induzida por modelo de mundo melhora significativamente o desempenho mesmo sob orçamentos de interação restritos, desafiando a noção de que a motivação intrínseca é muito custosa para configurações eficientes em amostras.
• Fundação de Código Aberto: Os autores fornecem um ponto de partida sólido para a pesquisa em TBWMs ao disponibilizar o código e os pesos do modelo, reduzindo a barreira de entrada para o desenvolvimento de agentes de modelo de mundo competitivos.

Resultados
O Simulus foi avaliado em três benchmarks diversos com orçamentos de interação estritos:

• Visual Atari 100K: O Simulus alcançou eficiência de amostra state-of-the-art (SOTA) entre métodos sem planejamento, superando baselines como DreamerV3, TWM, STORM e DIAMOND. Notavelmente, é o primeiro modelo de mundo sem planejamento a atingir pontuações de Média Interquartil (IQM) e mediana no nível humano, alcançando desempenho super-humano em 13 de 26 jogos.
• Controle Contínuo (DMC Proprioception 500K): Em tarefas de controle contínuo, o Simulus consistentemente igualou ou superou ligeiramente o DreamerV3, demonstrando a eficácia de arquiteturas baseadas em tokens em domínios contínuos com vocabulários compactos.
• Simbólico Multimodal (Craftax-1M): Em um ambiente de sobrevivência 2D parcialmente observável que combina grades simbólicas e características contínuas, o Simulus superou métodos de modelo de mundo concorrentes e todas as baselines sem modelo (incluindo RND, PPO-RNN e E3B), lidando eficientemente com sequências longas de tokens.

Estudos de ablação confirmaram que cada componente contribui individualmente, com recompensas intrínsecas provando-se particularmente cruciais. A combinação de todos os componentes produziu o desempenho mais forte, indicando sinergia significativa.

Significado
O artigo argumenta que o potencial não realizado de agentes de modelo de mundo decorre da dificuldade de combinar melhorias modulares. Ao se inspirar no framework Rainbow no DQN, o Simulus demonstra que uma combinação sistemática de flexibilidade de tokenização, motivação intrínseca, replay priorizado e técnicas avançadas de previsão de valor leva à eficiência de amostra state-of-the-art. O trabalho estabelece que os modelos de mundo baseados em tokens são uma abordagem versátil e escalável para diversos domínios de RL, desde que os componentes sejam cuidadosamente integrados para evitar a interferência frequentemente observada em arquiteturas monolíticas. Os autores mantêm que, embora os agentes baseados em tokens atualmente enfrentem limitações de velocidade de treinamento em comparação com algumas baselines, seu design modular oferece vantagens distintas para inferência de política e escalabilidade de pesquisa futura.