Resumo Técnico: Correspondência de Tokens Identificáveis para Modelos de Mundo

1. Declaração do Problema

Modelos de mundo baseados em Transformers demonstraram forte desempenho em aprendizado por reforço (RL) visual ao tratar sequências de estados e ações passadas como fluxos de tokens para prever estados futuros. No entanto, esses modelos frequentemente sofrem de inconsistência temporal durante rollouts de longo horizonte. Modos de falha comuns incluem duplicação de objetos, desaparecimento e transmutação em objetos diferentes.

Os autores identificam uma causa primária para esses erros: as abordagens existentes tratam a previsão do próximo quadro puramente como um problema de geração de tokens. Elas falham em modelar explicitamente a correspondência entre tokens ao longo do tempo. Em ambientes visuais realistas, tokens em quadros sucessivos frequentemente correspondem às mesmas entidades subjacentes que persistem e se movem. Quando um modelo é forçado a regenerar cada token do zero sem rastrear a identidade, ele luta para preservar estruturas persistentes, levando a uma preservação de identidade frágil e a erros cumulativos ao longo do tempo.

2. Metodologia: Correspondência de Tokens Identificáveis (ITC)

O artigo propõe a Correspondência de Tokens Identificáveis (ITC), um modelo de mundo Transformer que formula a previsão do próximo quadro como um problema de inferência probabilística estruturada com variáveis latentes de correspondência de tokens.

Mecanismo Central

Em vez de gerar cada token do próximo quadro do zero, a ITC determina, para cada token do próximo quadro, se ele deve ser:

1. Copiado de um token no quadro anterior (preservando a identidade).
2. Gerado como um novo token (lidando com mudanças ou novo conteúdo).

Essa atribuição é modelada como um problema de Transporte Ótimo (OT).

Arquitetura Técnica

1. Tokenização e Incorporação:

   • Observações visuais são convertidas em tokens discretos usando um tokenizador de busca de patches por vizinho mais próximo.
   • As incorporações de tokens são aumentadas com Incorporações de Posição Rotacional 3D (RoPE). Diferentemente da RoPE 1D padrão, esta formulação codifica coordenadas temporais, verticais e horizontais independentemente, permitindo que o modelo capture estruturas relacionais localizadas tanto no espaço quanto no tempo.

2. Solver de Transporte Ótimo:

   • O modelo constrói um grafo onde os vértices de origem representam tokens do estado anterior ($u$) e os tokens candidatos do próximo estado previstos pelo Transformer ($\tilde{p}$). Os vértices de destino representam os tokens finais do próximo estado ($\hat{u}$).
   • Uma matriz de afinidade é computada para pontuar a relação entre tokens anteriores e tokens previstos. Esta matriz inclui:
     • $A^{(prev)}$: Pontua a afinidade entre tokens anteriores e previsões, penalizada pela distância espacial para impor plausibilidade física (por exemplo, objetos não podem se teleportar arbitrariamente para longe).
     • $A^{(gen)}$: Pontua o "bônus" de admitir um token recém-gerado (curinga) em vez de reutilizar um anterior, sujeito a uma penalidade constante.
   • O solver (usando o algoritmo Sinkhorn) computa um plano de transporte que minimiza o custo de mover massa das origens para os destinos.

3. Binarização e Decodificação:

   • O plano de transporte contínuo é convertido em uma matriz de atribuição binária usando um procedimento de binarização gananciosa (Algoritmo 2). Isso garante um mapeamento estrito um para um.
   • Regra de Decodificação: Para cada posição no próximo quadro, se o plano de transporte atribuir um token anterior, esse token é copiado. Se o plano atribuir um token de "gerar", o modelo amostra da distribuição prevista pelo Transformer.

Loop de Treinamento

O modelo de mundo e a política são treinados concorrentemente. O modelo de mundo gera "imaginações" (trajetórias sintéticas) que são usadas junto com dados reais do ambiente para atualizar a política via Otimização de Política Proximal (PPO).

3. Contribuições Chave

• Formulação da Correspondência de Tokens: O artigo reformula a previsão do próximo quadro ao modelar explicitamente a correspondência de tokens como um problema de transporte ótimo, separando as tarefas de "reutilizar" entidades persistentes e "gerar" novo conteúdo.
• Integração de Transporte Ótimo: Introduz um mecanismo diferenciável (via Sinkhorn) porém discreto (via binarização) para selecionar entre copiar tokens anteriores e gerar novos, reduzindo efetivamente alucinações.
• Codificação Posicional 3D: A implementação da RoPE 3D permite que o modelo distinga entre diferenças temporais e espaciais, um componente necessário para rastrear objetos em movimento em sequências de vídeo.

4. Resultados Experimentais

Os autores avaliaram a ITC em quatro benchmarks desafiadores: Craftax-classic, Craftax, MinAtar e Atari 100K.

• Craftax-classic: A ITC alcançou um retorno de 72,5% e uma pontuação de 35,6%, superando significativamente o estado da arte anterior (Dedieu et al., 2025), que alcançou 67,4% de retorno e 27,9% de pontuação. Também demonstrou convergência mais rápida.
• Craftax: Neste ambiente mais complexo, a ITC alcançou um retorno de 7,09% e uma pontuação de 2,40%, superando baselines incluindo Simulus e Dedieu et al.
• MinAtar: A ITC superou o estado da arte anterior para RL baseado em modelo em todos os quatro jogos (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Aplicado a um conjunto de 26 jogos, a ITC alcançou novo desempenho de estado da arte na Média Interquartil (IQM) de 1,092 e na Lacuna de Otimalidade de 0,376, superando o modelo de mundo baseado em tokens Simulus.
• Precisão de Previsão: Em um conjunto de dados de 10.000 transições, a ITC melhorou a precisão geral de previsão em 3,39% em comparação com a baseline, com uma melhoria notável de 3,45% em transições envolvendo criaturas em movimento.
• Melhorias Qualitativas: A análise visual confirmou que a ITC elimina artefatos comuns, como duplicação e desaparecimento de objetos, vistos em modelos baseline como DreamerV3 e Dedieu et al.
• Eficiência: O método introduziu sobrecarga negligenciável, aumentando o tempo total de treinamento ponta a ponta em apenas 2,8%.

5. Significado e Alegações

O artigo afirma que, ao modelar explicitamente a correspondência entre tokens ao longo do tempo, a ITC resolve um problema generalizado em modelos de mundo onde agentes falham em manter a identidade de objetos ao longo de horizontes longos. Os autores afirmam que essa abordagem permite que o modelo "reutilize seletivamente tokens que têm uma forte correspondência com o quadro anterior", aproveitando assim as semelhanças quadro a quadro em vez de depender exclusivamente do Transformer para regenerar cada token.

O significado reside em alcançar desempenho de estado da arte em ambientes diversos (de mundos abertos baseados em grade a jogos padrão do Atari) mantendo eficiência computacional. Os autores concluem que a ITC previne efetivamente inconsistências temporais como duplicação e desaparecimento, levando a imaginações de maior qualidade e melhor desempenho da política.