Resumo Técnico: RepWAM – Modelagem de Ação de Mundo com Representação

1. Declaração do Problema

Os Modelos de Ação de Mundo (WAMs - World Action Models) visam preencher a lacuna entre a previsão passiva do futuro e o controle robótico incorporado, aprendendo um espaço latente único que prevê a dinâmica visual e infere as ações robóticas necessárias para realizá-las. No entanto, os atuais WAMs enfrentam duas limitações críticas em seu design de representação:

1. Deficiências de Representação Visual: Os WAMs existentes tipicamente herdam tokenizadores de vídeo de modelos de geração de vídeo pré-treinados (ex: WAN, VQ-VAE). Esses tokenizadores são otimizados para reconstrução de pixels, dedicando capacidade latente à aparência de baixo nível (texturas de fundo) em vez de semântica relevante para a tarefa (identidade de objetos, relações espaciais, pistas de interação). Consequentmente, a modelagem de mundo condicionada por instruções é mal fundamentada, pois os objetivos de linguagem devem ser inferidos de um espaço latente dominado pela variação de aparência.
2. Desacoplamento Ação-Visual: Os latentes visuais e os comandos motores residem em espaços disjuntos. Isso força um Modelo de Dinâmica Inversa (IDM) a realizar a ponte de uma lacuna modal significativa em cada etapa, impedindo que o especialista de ação fundamente efetivamente os comandos motores no conteúdo da cena.

Os autores postulam que a fidelidade dos WAMs é limitada por esses espaços latentes e questionam: Que tipo de representação um modelo de ação de mundo realmente precisa?

2. Metodologia: RepWAM

O artigo propõe o RepWAM, um WAM centrado na representação construído sobre um Tokenizador de Ação Visual-Representação (RepViTok). O framework consiste em três componentes principais:

A. Tokenizador de Ação Visual-Representação (RepViTok)

O RepViTok alinha os latentes visuais com um modelo de fundação visual congelado e induz tokens de ação latentes como transições dentro deste espaço semântico compartilhado.

• Tokenização Visual: Um Autoencoder de Vision Transformer (ViT) processa observações visuais. Diferente dos tokenizadores padrão, ele é supervisionado por um objetivo duplo:
  1. Perda de Reconstrução: Combina perdas $L_1$, perceptuais e adversariais para garantir a fidelidade de pixels.
  2. Perda de Alinhamento de Características: Puxa os latentes de vídeo em direção a um modelo de fundação visual congelado (ex: Codificador de Percepção) via uma projeção linear. Isso garante que o espaço latente capture semântica de alto nível (identidade de objetos, interações) em vez de apenas aparência de baixo nível.
• Tokenização de Ação Latente: Baseando-se nos tokens visuais alinhados, um Tokenizador de Ação Latente (LAT) aprende tokens de ação que representam transições entre estados visuais. Ele acopla um Modelo de Dinâmica Inversa (IDM) e um Modelo de Dinâmica Direta (FDM):
  • O IDM comprime a transição entre frames latentes consecutivos ($z_t, z_{t+1}$) em uma ação latente compacta $\ell_t$.
  • O FDM realiza $\ell_t$ como um mapa de transporte (roteamento de conteúdo visual) e um resíduo, reconstruindo o próximo latente visual.
  • Este design garante que as ações latentes sejam definidas como transformações entre estados visuais semânticos, tornando-as transferíveis entre diferentes embodiments.

B. Modelo de Ação de Mundo Causal

O RepWAM pré-treina um transformer de difusão causal para modelar conjuntamente estados visuais futuros e ações latentes sob condicionamento de linguagem.

• Arquitetura: O modelo utiliza uma máscara de causalidade de bloco onde os tokens visuais e de ação compartilham pesos de atenção, mas utilizam redes feed-forward (FFNs) específicas para cada modalidade.
• Objetivo de Treinamento: O modelo é treinado usando um objetivo de correspondência de fluxo condicional (conditional flow-matching). Ele amostra um bloco de latentes visuais e ações latentes pareadas, interpolando-os com ruído Gaussiano. A rede regride a velocidade alvo para prever simultaneamente o próximo estado visual e a ação latente correspondente.
• Adaptação: Após o pré-treinamento em dados de vídeo-ação em larga escala, o modelo é adaptado para trajetórias robóticas reais para manipulação de loop fechado.

3. Principais Contribuições

1. Tokenização Semântica Visual-Ação: A introdução do RepViTok, que alinha latentes visuais com um modelo de fundação e aprende ações latentes como transições dentro deste espaço semântico, abordando a natureza "rasa" e "desacoplada" das representações de WAM existentes.
2. Estrutura de Modelagem Conjunta: Uma arquitetura unificada que pré-treina um transformer de difusão causal para prever conjuntamente estados visuais futuros e ações latentes, seguido de adaptação para controle robótico real.
3. Validação Empírica: Avaliação abrangente em tarefas de manipulação de braço duplo no mundo real e no benchmark de simulação RoboTwin 2.0, demonstrando que o alinhamento semântico melhora tanto a previsão da dinâmica visual quanto a execução em loop fechado.

4. Resultados Experimentais

Os autores avaliaram o RepWAM (escalas de 1.3B e 5B de parâmetros) contra baselines incluindo $\pi_0.5$, Motus e Lingbot-VA.

• Manipulação no Mundo Real (Franka Dual-Arm):
  • Em tarefas como "Pegar o fruto", "Empurrar a gaveta" e "Inserir o tubo", o RepWAM-5B alcançou taxas de sucesso de 60%, 80% e 60%, respectivamente.
  • Estes resultados superaram o $\pi_0.5$ (por 10–50 pontos percentuais) e o Lingbot-VA (por 10–20 pontos percentuais), particularmente em tarefas de longo horizonte e de grão fino onde dinâmicas visuais coerentes são cruciais.
• Benchmarks de Simulação (RoboTwin 2.0):
  • O RepWAM-5B alcançou 89,3% em tarefas Fáceis (Easy) e 88,4% em tarefas Difíceis (Hard), superando a média de 50 tarefas do $\pi_0.5$ e Motus.
  • Notavelmente, o RepWAM foi treinado do zero, sem herdar pesos de geração de vídeo pré-treinados, mas permaneceu competitivo com modelos como o Lingbot-VA que utilizam tal pré-treinamento.
• Estudos de Ablação:
  • Design do Tokenizador: Substituir o VAE padrão do WAN2.2 pelo RepViTok melhorou o Score de Loop Aberto (OLS) e as taxas de sucesso em robôs reais, confirmando que o alinhamento semântico é crítico para o seguimento de instruções.
  • Treinamento de Ação Latente: Uma abordagem de "Dois Estágios" (pré-treino em ações latentes, seguido de adaptação ao controle robótico) superou significativamente a "Previsão Conjunta" ou o treinamento sem ações latentes, validando a utilidade das ações latentes como uma representação intermediária.
  • Guia de Classificador Livre (CFG): Ao contrário dos modelos de geração de vídeo, o RepWAM alcançou desempenho ideal com um fator de escala de CFG de 1.0 (sem extrapolação), sugerindo que o espaço semântico é robusto o suficiente para reduzir a dependência de CFG.

5. Significância e Alegações

O artigo afirma que a tokenização de ação visual-representação é uma base promissora para modelos de ação de mundo. Ao fundamentar os latentes visuais e de ação em uma representação semântica compartilhada, o RepWAM demonstra que:

• A geração explícita do futuro não é estritamente necessária para o aprendizado de políticas se a representação capturar semântica relevante para a tarefa.
• O alinhamento semântico permite que os WAMs alcancem forte generalização e desempenho em loop fechado mesmo sem herdar backbones de geração de vídeo pré-treinados.
• Esta abordagem caminha em direção a políticas robóticas generalistas ao criar um espaço de representação onde a ação é entendida como uma transformação de estados semânticos, em vez de um comando motor disjunto.

Os autores concluem que o design do espaço latente é fundamental para a capacidade de agentes incorporados, e o RepWAM oferece um passo em direção a uma manipulação robótica mais interpretável, transferível e eficaz. Sugere-se como trabalho futuro escalar o pré-treino além de vídeos específicos de robótica para incluir vídeos da internet e de perspectiva egocêntrica humana em larga escala.