Resumo Técnico: Aprendizado de Modelos Mundiais Baseados em Recursos Visuais via Ação Latente Residual

Declaração do Problema

Modelos mundiais visam prever transições de estado futuras a partir de observações e ações, servindo como base para aprendizado de políticas e raciocínio. O paradigma dominante atual depende da geração de vídeo em espaços latentes alinhados a pixels (por exemplo, VAEs). Embora visualmente atraentes, essas abordagens sofrem com alta sobrecarga computacional e uma propensão a alucinações (geração de estados fisicamente implausíveis).

Por outro lado, modelos mundiais baseados em recursos visuais (por exemplo, DINO-WM) preveem recursos visuais futuros (como tokens DINO) em vez de pixels brutos. Essa abordagem é mais eficiente e menos propensa a alucinações. No entanto, os métodos baseados em recursos existentes dependem de regressão direta para prever tokens futuros. Em interações 3D complexas, isso leva à "regressão à média", resultando em previsões desfocadas ou colapsadas. Uma alternativa aparentemente viável — usar modelos generativos (como difusão) diretamente no espaço de recursos — é impedida pela maldição da dimensionalidade; os tokens DINO têm dimensionalidade significativamente maior (por exemplo, 1 milhão de dimensões para uma imagem de 512×512) em comparação com os latentes VAE padrão (16 mil), tornando a modelagem generativa computacionalmente proibitiva e frustrando os objetivos de eficiência dos modelos baseados em recursos.

O artigo aborda duas questões-chave:

1. Como desenvolver um modelo mundial eficiente e preciso em um espaço de recursos visuais que escale para manipulação 3D complexa?
2. Como aproveitar tais modelos mundiais para melhorar políticas de robôs a jusante?

Metodologia

1. Ação Latente Residual (RLA)

Os autores propõem uma nova representação de ação latente chamada Ação Latente Residual (RLA). Em vez de prever o estado futuro absoluto $s_{t+h}$ ou gerá-lo via difusão pesada, a RLA codifica o resíduo entre os tokens DINO atuais ($s_t$) e os tokens DINO futuros ($s_{t+h}$) em um vetor latente compacto $z$.

• Arquitetura: Um autoencoder é treinado onde um codificador $f_{enc}$ processa o resíduo $(s_{t+h} - s_t)$ juntamente com consultas aprendíveis para produzir um latente compacto $z$. Um decodificador $f_{dec}$ reconstrói $s_{t+h}$ usando $s_t$ e $z$.
• Propriedades Chave:
  • Suficiência Preditiva: Ao contrário de ações latentes anteriores usadas como condicionamento fraco para geração iterativa, a RLA permite a reconstrução precisa de tokens futuros em um único passe feedforward.
  • Generalizabilidade: O codificador RLA generaliza para cenas e padrões de movimento novos, mesmo quando treinado em dados limitados e agnósticos a tarefas.
  • Topologia Temporal: O espaço latente codifica naturalmente a progressão temporal. A interpolação linear entre ruído gaussiano e um vetor RLA produz representações correspondentes a estados temporais intermediários.

2. Modelo Mundial RLA (RLA-WM)

Construindo sobre a RLA, os autores propõem o RLA-WM, um modelo mundial que prevê o vetor RLA compacto $z$ em vez de $s_{t+h}$ de alta dimensionalidade diretamente.

• Correspondência de Fluxo: O modelo usa correspondência de fluxo para prever a velocidade necessária para transformar ruído gaussiano no RLA alvo $z$.
  • Entrada: Estado atual $s_t$ e um bloco de ações $a_{t:t+h}$.
  • Processo: Uma rede de condição processa $s_t$ e ações para gerar tokens de condição. Estes são concatenados com um latente ruidoso $z_\tau$. Uma rede de fluxo prevê a velocidade $\hat{v}$ para transformar iterativamente o ruído no RLA final $z_1$.
  • Saída: O RLA previsto $z_1$ é decodificado com $s_t$ para reconstruir o estado futuro $s_{t+h}$.
• Eficiência: Ao realizar correspondência de fluxo no espaço RLA compacto (por exemplo, 2048 dimensões) em vez do espaço completo de tokens DINO, o modelo alcança alta fidelidade com custo computacional significativamente menor (FLOPs) em comparação com modelos de difusão de vídeo.

3. Aplicações em Aprendizado de Robôs

O artigo demonstra duas aplicações de RLA e RLA-WM:

• Modelo de Ação Mundial Minimalista (WAM) para Vídeos sem Ações:

  • Estende uma política padrão de Clonagem de Comportamento (BC) (ResNet-18) adicionando uma única camada linear para prever RLA $z$ a partir de observações.
  • Treinado em uma mistura de vídeos com e sem rótulos de ação. Para vídeos sem ações, a perda de ação é mascarada, e a rede aprende a prever o alvo RLA (derivado do par de vídeos) para guiar a política.
  • Isso permite aprendizado a partir de vídeos em grande escala e não rotulados sem acoplar a política a backbones pesados de geração de vídeo.

• Aprendizado por Reforço Visual (WMRL) dentro do RLA-WM:

  • A primeira demonstração de RL visual treinado inteiramente dentro de um modelo mundial aprendido apenas a partir de vídeos offline.
  • Nenhuma interação online, recompensas artesanais ou perda auxiliar BC são usadas durante a fase de RL.
  • Recompensa Alinhada a Vídeo (VAR): A recompensa é definida como a distância L1 negativa entre os tokens DINO da simulação do modelo e os tokens DINO de verdadeiros de um vídeo de referência offline (sincronizado no tempo).
  • A política (inicializada a partir de BC) é otimizada via PPO usando simulações geradas inteiramente dentro do RLA-WM.

Contribuições Principais

1. Ação Latente Residual (RLA): Uma representação de ação latente simples e compacta aprendida a partir de resíduos DINO que é preditiva, generalizável e codifica topologia temporal.
2. RLA-WM: Um modelo mundial baseado em recursos visuais de última geração que prevê RLA via correspondência de fluxo. Supera modelos baseados em recursos existentes e modelos de difusão de vídeo em precisão, sendo ordens de magnitude mais rápido.
3. Aplicações Inovadoras:
   • Um WAM minimalista que melhora o aprendizado por imitação a partir de vídeos sem ações.
   • Um framework de RL visual (WMRL) que otimiza políticas via simulações em um modelo mundial aprendido sem interação online ou engenharia manual de recompensas.

Resultados Experimentais

Qualidade de Previsão

Avaliado nos conjuntos de dados ManiSkill (simulação) e IWS (manipulação bimanual do mundo real):

• Métricas: O RLA-WM alcançou desempenho superior em LPIPS, SSIM e distância L1 DINO em comparação com baselines incluindo DINO-WM (regressão), RAE (difusão no espaço de recursos), FM-WM (correspondência de fluxo no espaço de recursos) e Vid2World (difusão de vídeo).
• Eficiência: O RLA-WM exigiu 3,5T FLOPs por inferência, significativamente menor que Vid2World (1,1P FLOPs) e RAE/FM-WM (14,3T FLOPs).
• Qualitativo: O RLA-WM produziu previsões fundamentadas fisicamente com alucinação mínima, enquanto modelos de difusão de vídeo (Vid2World) frequentemente geravam trajetórias nítidas, mas fisicamente inconsistentes.

Desempenho em Aprendizado de Robôs

• Aprendizado de Vídeo sem Ações: Em tarefas de aprendizado por imitação usando apenas 5% de dados rotulados com ações, o WAM baseado em RLA alcançou uma taxa de sucesso média de 35,6%, superando a baseline BC (27,2%) e outros métodos de ação latente (por exemplo, AdaWorld em 33,7%). A melhoria foi mais pronunciada na tarefa desafiadora Push-T (+15,2% sobre a baseline).
• RL Visual (WMRL): Através de 1500 sementes de avaliação, o WMRL melhorou significativamente o desempenho da política nos robôs XArm e UR10e em comparação com a política BC inicial. Em média, alcançou um ganho de +1,1% sobre BC em cinco tarefas sem usar dados adicionais de interação com o mundo real.

Significado e Alegações

O artigo alega que o RLA-WM representa um grande passo à frente para modelos mundiais baseados em recursos visuais. Ao mudar da regressão direta ou modelagem generativa pesada para modelagem de ação latente residual, os autores alcançam um equilíbrio de eficiência, precisão e escalabilidade que era anteriormente inatingível em espaços de recursos de alta dimensionalidade.

O significado reside em:

1. Superar a Maldição da Dimensionalidade: Demonstrar que a modelagem generativa no espaço de recursos é viável se o alvo for uma representação residual compacta em vez do vetor de recursos completo.
2. Habilitar RL Apenas Offline: Provar que RL visual de alta qualidade é possível usando apenas dados de vídeo offline e um modelo mundial aprendido, eliminando a necessidade de exploração online cara ou funções de recompensa artesanais.
3. Escalabilidade para Dados Não Rotulados: Fornecer um caminho para aproveitar a vasta quantidade de dados de vídeo de robôs não rotulados para melhoria de políticas através da abordagem WAM minimalista.

Os autores reconhecem limitações, como a dificuldade de lidar com movimento de fundo irrelevante para a tarefa (por exemplo, tremor de câmera) e observabilidade parcial (oclusões), sugerindo trabalho futuro em RLA independente de visão 3D e condicionamento de múltiplos quadros. No entanto, eles enfatizam que seus resultados em conjuntos de dados de pequena escala isolam os ganhos impulsionados pelo método do mero escalonamento de dados, estabelecendo uma base sólida para futuras aplicações em mundo aberto.