Policy-Guided World Model Planning for Language-Conditioned Visual Navigation

1. Problema Abordado

O artigo aborda o desafio fundamental na IA corporificada (embodied AI) de navegar um agente robótico até um objetivo visualmente especificado, seguindo instruções em linguagem natural.

• Limitações das Abordagens Atuais:
  • Políticas Reativas (VLA): Modelos como o Octo são eficazes para controle de curto horizonte e compreensão de instruções, mas falham no planejamento de longo horizonte, pois preveem ações em uma única passagem sem raciocinar sobre consequências futuras.
  • Modelos de Mundo (World Models): Modelos baseados em JEPA (Joint-Embedding Predictive Architecture) permitem planejamento explícito no espaço latente, evitando a reconstrução de pixels. No entanto, otimizadores de amostragem como MPPI (Model Predictive Path Integral) frequentemente sofrem de inicialização pobre. Começar a busca a partir de uma distribuição Gaussiana não informada em espaços de ação de alta dimensão leva a uma convergência lenta e a mínimos locais.

2. Metodologia: O Framework PiJEPA

Os autores propõem o PiJEPA, um framework de duas etapas que combina as forças de uma política aprendida com o planejamento de modelo de mundo latente. A ideia central é usar a saída de uma política VLA (Visão-Linguagem-Ação) para "aquecer" (warm-start) o processo de planejamento do modelo de mundo.

Componentes Principais:

1. Codificador de Visão Congelado: Utiliza um encoder pré-treinado e congelado (DINOv2 ou V-JEPA-2) para mapear observações visuais para um espaço latente consistente ($z_t$).
2. Estágio 1: Prioridade de Política (Policy Prior):
   • Um modelo Octo (baseado em transformadores com cabeçalho de difusão) é ajustado (fine-tuned) no conjunto de dados CAST.
   • Dada a observação atual e a instrução de linguagem, o Octo gera uma distribuição de ações informada.
   • Amostras de ações são extraídas e transformadas para o sistema de coordenadas local do modelo de mundo.
   • Calculam-se a média ($\mu_\pi$) e o desvio padrão ($\sigma_\pi$) dessas amostras para formar uma distribuição inicial informada.
3. Estágio 2: Planejamento MPPI Guiado por Modelo de Mundo:
   • Um Modelo de Mundo JEPA (treinado com o mesmo encoder congelado) prevê estados latentes futuros com base em ações candidatas.
   • O algoritmo MPPI é executado, mas em vez de iniciar com uma Gaussiana não informada, ele inicia com a distribuição $(\mu_\pi, \sigma_\pi)$ derivada da política.
   • O planejador otimiza a sequência de ações minimizando a distância no espaço latente entre o estado futuro previsto e o objetivo codificado ($z_g$).
   • A instrução de linguagem influencia apenas a política inicial; o modelo de mundo é agnóstico à linguagem, focando apenas na dinâmica visual-ação.

3. Contribuições Chave

1. Framework Unificado (PiJEPA): Integração de uma política VLA ajustada com planejamento baseado em MPPI sobre um modelo de mundo JEPA, utilizando a política para inicializar o planejador.
2. Análise Sistemática de Encoders: Estudo comparativo de dois tipos de representações visuais congeladas:
   • DINOv2: Encoder de imagem (foco em características espaciais e de segmentação).
   • V-JEPA-2: Encoder de vídeo (foco em dinâmicas temporais).
   • O trabalho testa combinações cruzadas (ex: política DINOv2 + modelo de mundo V-JEPA-2) para entender quais representações são melhores para a prioridade de ação versus a dinâmica latente.
3. Validação em Tarefas do Mundo Real: Demonstração experimental de que o planejamento guiado por política supera significativamente tanto a execução pura da política reativa quanto o planejamento de modelo de mundo não informado.

4. Resultados Experimentais

Os experimentos foram realizados no conjunto de dados CAST (navegação visual em ambientes reais).

• Desempenho de Precisão Posicional:
  • O PiJEPA alcançou a melhor precisão posicional (menor erro de trajetória absoluta - ATE) em ambas as configurações de encoder.
  • Com o encoder V-JEPA-2, o PiJEPA obteve o melhor resultado geral (ATE XY RMSE de 1.65m), superando a política Octo pura e o MPPI não informado.
• Comparação de Baselines:
  • MPPI Não Informado: Teve desempenho ruim na posição (especialmente com V-JEPA-2), confirmando que a inicialização aleatória desperdiça o orçamento de amostragem. Curiosamente, teve bom desempenho em orientação (heading), sugerindo que o objetivo no espaço latente regulariza a direção mesmo quando a posição é imprecisa.
  • Octo-WM (Apenas Pontuação): Usar o modelo de mundo apenas para pontuar amostras da política (sem otimização iterativa) já trouxe grandes ganhos, indicando que filtrar propostas ruins é uma estratégia poderosa.
  • Octo Puro: Competitivo em métricas de passo único (controle local), mas acumula erros em horizontes longos.
• Latência: O overhead do planejamento MPPI é mínimo (0.35s) comparado ao tempo de amostragem da difusão da política (2.13s), tornando o sistema viável para inferência em tempo real.
• Análise de Falhas: O sistema ainda enfrenta desafios com instruções ambíguas (ex: "Siga o prédio" quando há vários), onde o modelo de mundo pode ficar estagnado (prevendo estados latentes idênticos). A prioridade da política ajuda a mitigar isso, mas não resolve completamente o problema de ambiguidade.

5. Significado e Conclusão

O trabalho PiJEPA estabelece uma divisão de trabalho eficiente entre aprendizado por reforço/imitação e planejamento:

• A política reativa fornece propostas de ação semanticamente informadas e eficientes para controle local.
• O modelo de mundo avalia e refina essas trajetórias ao longo de um horizonte estendido, corrigindo erros cumulativos.

A principal inovação é demonstrar que inicializar um planejador de modelo de mundo com uma prioridade de política aprendida resolve o gargalo de convergência lenta típico de otimizadores baseados em amostragem em espaços de alta dimensão. Isso permite que agentes robóticos naveguem com maior fidelidade a instruções complexas e maior precisão de destino, superando as limitações de abordagens puramente reativas ou puramente baseadas em planejamento cego.

O código está disponível publicamente, e o trabalho aponta para futuras direções como arquiteturas de dinâmica mais robustas para evitar estagnação e mecanismos de replanejamento em malha fechada.