Resumo Técnico: Dueling World Models

Declaração do Problema

Modelos de mundo latentes, particularmente aqueles que seguem a Arquitetura de Predição de Embedding Conjunto (JEPA), planejam prevendo embeddings de estados futuros a partir de estados e ações atuais. No entanto, quando uma cena contém movimento que o agente não controla (distratores), esses modelos sofrem um modo de falha "silencioso": as previsões para diferentes ações tornam-se indistinguíveis, mesmo enquanto o erro de perda total de treinamento (erro de reconstrução ou de predição) continua a melhorar.

Remédios existentes tentam separar a dinâmica controlável de distratores exógenos por meio da adição de maquinário externo, como perdas de reconstrução, recompensas de tarefas, objetivos auxiliares de dinâmica inversa ou ramos de rede separados. Essas abordagens frequentemente requerem suposições específicas sobre o processo do distrator ou supervisão adicional. Os autores identificam que, em preditores monolíticos padrão $P(z, a)$, o canal de ação torna-se emaranhado com a dinâmica do distrator, a menos que seja explicitamente fundamentado por rótulos de ação verdadeiros ou purgado via máscaras de segmentação. Crucialmente, a seleção de modelo padrão baseada em perda de validação frequentemente falha em detectar essa "cegueira de ação", selecionando ativamente checkpoints onde o sinal da ação colapsou.

Metodologia: AD-JEPA

O artigo propõe o AD-JEPA (Action-Decomposed JEPA), uma reestruturação mínima de preditores latentes condicionados à ação que isola o canal de ação sem exigir recompensas, reconstrução ou perdas auxiliares específicas de distratores.

Mecanismo Central: Centralização da Média da Ação (Action-Mean Centering)

O método decompõe o preditor em dois fluxos:

1. Fluxo Passivo ($B(z)$): Prediz como a cena evolui com a ação marginalizada.
2. Fluxo de Ação Centralizado ($u(z, a)$): Representa o deslocamento específico da ação, calculado subtraindo a média da predição sobre as ações do deslocamento bruto da ação.

A regra de atualização é definida como:
$$\hat{z}_{t+1} = B(z_t) + \underbrace{[\Delta(z_t, a_t) - \bar{\Delta}(z_t)]}_{u(z_t, a_t)}$$
Onde $\bar{\Delta}(z) = \mathbb{E}_{a' \sim \nu}[\Delta(z, a')]$ é a média do deslocamento sobre uma distribuição de proposta $\nu$ de ações.

Fundamentação Teórica

Esta abordagem adapta a decomposição Dueling DQN ($Q = V + (A - \bar{A})$) de valores escalares para dinâmicas latentes de valores vetoriais. Os autores provam que esta subtração é uma identidade algébrica em vez de uma tendência aprendida:

• Cancelamento de Modo Comum: Qualquer variação que entra nos deslocamentos de ação identicamente em todas as ações (variação de modo comum), incluindo distratores independentes da ação, é cancelada exatamente.
• Exatidão: O cancelamento mantém-se tanto para conjuntos de ações discretas (via enumeração exata) quanto para conjuntos de ações contínuas (via amostragem de Monte-Carlo), desde que a média seja estimada sobre um conjunto finito de candidatos.
• Sem Suposição de Exogeneidade: O método não requer suposições de exogeneidade de nível de processo (isto é, que os distratores sejam independentes da política do agente no nível do processo); ele apenas requer que a contribuição do distrator para a transição seja de modo comum entre as ações.

Objetivo de Treinamento

O modelo é treinado com uma perda de predição JEPA padrão contra um alvo de EMA (Média Móvel Exponencial) com interrupção de gradiente, aumentada com:

• InfoNCE ($L_{act}$): Incentiva a predição da ação executada a ser a mais próxima do alvo entre as ações contrafatuais.
• Penalidade de Compactação ($L_{off}$): Desencoraja o canal de ação de absorver a persistência do estado.
• Regularização VICReg: Previne o colapso da representação.

Notavelmente, os autores relatam que adicionar um mecanismo de portão (gating) aprendido para selecionar dimensões foi inerte, degenerando para um escalonamento uniforme, confirmando que a centralização sozinha é o mecanismo operante.

Principais Contribuições

1. Centralização da Média da Ação: Uma reestruturação de uma linha de preditores condicionados à ação que torna o canal de ação invariante ao modo comum. Os autores fornecem provas (Proposições 1 e 2) mostrando o cancelamento exato de componentes de modo comum para conjuntos de ações discretos e amostrados.
2. Validação Empírica: O método foi validado em quatro configurações distintas:
   • Gridworld (FourRooms): Com células distratoras binárias rotativas.
   • Geradores Sintéticos: Com fatores controláveis e incômodos conhecidos.
   • Controle Contínuo (DMC): Com ocluidores móveis em nível de pixel.
   • Atari (Freeway): Com distratores de tráfego naturais.
     Em todos os casos, o canal centralizado preservou o efeito do agente enquanto reduzia o vazamento de incômodo para um nível indistinguível de zero, enquanto os preditores emaranhados falharam.
3. Aplicabilidade Apenas na Leitura (Readout-Only): Como o mecanismo é uma propriedade da leitura e não da arquitetura, ele pode ser aplicado post-hoc a modelos de mundo pré-treinados e congelados (especificamente RePo e TIA). Os autores demonstraram que subtrair a média da ação das leituras brutas desses modelos congelados recuperou com sucesso um canal de ação que estava de outra forma latente e não roteado.
4. Limite Definido: O artigo identifica explicitamente a limitação: o método falha quando a dinâmica do distrator é correlacionada com a ação (ex: distratores reafferentes ou agentes que reagem ao usuário). Nesses casos, o distrator não é de modo comum e não é cancelado.

Resultados

• Recuperação vs. Rejeição: Em testes sintéticos, o canal centralizado alcançou alta recuperação do fator controlável ($R^2 \approx 0.87\text{--}0.93$) enquanto rejeitou o fator incômodo ($R^2 \approx -0.00$). Variantes não centralizadas recuperaram o sinal, mas vazaram pesadamente o incômodo.
• Desempenho de Controle: No gridworld, um controlador de predição de modelo usando o canal centralizado manteve alto sucesso de alcance de meta ($\approx 0.92$), mesmo com 30 células de distração, enquanto o preditor padrão colapsou para $\approx 0.57$.
• Sucesso Post-Hoc: Aplicar a subtração aos modelos congelados RePo e TIA recuperou canais de ação com valores de $R^2$ entre 0.19 e 0.52, enquanto as leituras de delta de ação brutas estavam próximas de zero.
• Modos de Falha: O método falhou corretamente (colapsando para o acaso) em testes de estresse de ação correlacionada, alinhando-se com o limite teórico.

Significância e Alegações

O artigo afirma que a centralização é um plug-in universal para modelos de mundo condicionados à ação. Sua significância reside na:

• Simplicidade: Alcança a rejeição de distratores sem a complexidade de reconstrução, fatoração baseada em recompensa ou perdas auxiliares de dinâmica inversa.
• Garantia Estrutural: A separação é uma propriedade algébrica da parametrização, não uma tendência aprendida dependente de sinais de treinamento específicos.
• Capacidade de Retroajuste: Permite que modelos de mundo existentes e congelados sejam "desdistraídos" no tempo de inferência sem retreinamento, trazendo à tona canais de ação que anteriormente estavam ocultos por dinâmicas emaranhadas.

Os autores são modestos quanto à escalabilidade do controle: embora o canal permita o controle em gridworlds, uma tentativa registrada de escalar isso para tarefas de controle contínuo de DMC falhou em superar baselines aleatórios. Eles atribuem isso a gargalos de precisão do modelo no planejador, em vez de um defeito no próprio canal. A principal contribuição é, portanto, de nível de representação: fornecer um canal de ação limpo e invariante ao modo comum, onde preditores emaranhados tornam-se "cegos para a ação".