Resumo Técnico: FF-JEPA: Planejamento de Longo Horizonte em Modelos de Mundo com Planejadores Latentes

Declaração do Problema

Avanços recentes em modelos de mundo visuais, particularmente as Arquiteturas de Predição de Embedding Conjunto (JEPAs) como o LeWorldModel (LeWM), demonstraram a capacidade de prever observações futuras e suportar controle baseado em modelo via planejadores baseados em amostragem, como o Método de Entropia Cruzada (CEM). No entanto, os autores identificam três limitações críticas que impedem a implantação desses modelos de mundo "planos" em tarefas de longo horizonte no mundo real:

1. Proibições Computacionais: Planejar sobre horizontes longos requer rollouts repetidos e sequências de ações cada vez maiores, tornando a otimização por CEM computacionalmente intratável.
2. Acúmulo de Erro: Erros de predição autoregressiva acumulam-se ao longo de muitos passos, fazendo com que o planejador divirja para trajetórias complexas.
3. Dependência de Objetivo: Métodos existentes tipicamente requerem uma imagem ou estado de objetivo explícito ($o_g$) para conduzir o planejamento. Em muitos cenários do mundo real, tal especificação de objetivo concreta é indisponível ou impraticável.

Embora métodos hierárquicos e abordagens generativas tenham tentado abordar esses problemas, eles frequentemente ainda operam no espaço de pixels ou dependem de imagens/trajetórias de objetivo especificadas externamente, limitando sua aplicabilidade em ambientes de código aberto.

Metodologia: Forward-Forward JEPA (FF-JEPA)

O artigo propõe o FF-JEPA, um framework hierárquico que estende os modelos de mundo do tipo JEPA introduzindo um planejador latente ($G$) que opera ao lado de um modelo de dinâmica direta condicionado a ação ($P$) padrão. Esta abordagem decompõe o planejamento de longo horizonte em uma sequência de problemas de otimização de curto prazo tratáveis, sem exigir imagens de objetivo explícitas.

Arquitetura Central

O sistema consiste em dois componentes primários treinados dentro do espaço latente definido por um codificador de modelo de mundo pré-treinado e congelado ($E$):

1. O Modelo de Mundo ($P$): Um modelo de dinâmica direta padrão condicionado a ação (baseado em LeWM) que prevê o próximo estado latente dado um histórico de estados e uma ação: $\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$.
2. O Planejador Latente ($G$): Um modelo de dinâmica direta livre de ação que prevê o próximo sub-objetivo latente ($\hat{z}_{sg}$) baseando-se apenas em um histórico de sub-objetivos latentes anteriores. Ele opera em uma resolução temporal superior (a cada $H$ passos do ambiente).
   • Planejador Determinístico ($G_{Det}$): Uma arquitetura transformer (semelhante ao preditor do modelo de mundo, mas sem condicionamento de ação) treinada para minimizar o erro quadrático médio entre os sub-objetivos previstos e os reais.
   • Planejador de Difusão ($G_{DM}$): Um backbone DiT (Diffusion Transformer) treinado com um objetivo de correspondência de score de denoising para prever uma sequência de sub-objetivos futuros.

Inferência de Planejamento

O processo de inferência é hierárquico:

1. Predição de Sub-objetivo: O planejador latente $G$ prevê o próximo estado de sub-objetivo $\hat{z}_{sg, m+1}$ dado o histórico de sub-objetivos atual.
2. Otimização de Ação: O modelo de mundo $P$ usa CEM para encontrar uma sequência de ações ($a_{1:H}$) que leve o estado previsto em direção ao sub-objetivo previsto $\hat{z}_{sg, m+1}$ dentro de um horizonte curto $H$.
3. Iteração: Este ciclo se repete, efetivamente encadeando planos de curto horizonte para alcançar objetivos de longo horizonte.

Crucialmente, este método remove a necessidade de uma imagem de objetivo alvo no momento da inferência; o planejador gera a trajetória em direção a um objetivo implicitamente definido.

Principais Contribuições

• Planejamento de Longo Horizonte Livre de Objetivo: O framework permite planejar em modelos de mundo sem imagens de objetivo explícitas, usando um planejador latente para gerar sub-objetivos intermediários.
• Decomposição Hierárquica: Aborda o "colapso de longo horizonte" dos modelos de mundo planos, quebrando trajetórias complexas em subproblemas curtos e solucionáveis, mitigando o acúmulo de erro e os custos computacionais.
• Treinamento Não Rotulado: O planejador latente pode ser treinado em demonstrações bem-sucedidas não rotuladas (subamostradas do espaço latente do modelo de mundo), reduzindo a dependência de dados rotulados com ações.
• Framework Unificado: Reinterpreta o modelo de mundo como um módulo de dinâmica inversa que opera sobre trajetórias latentes imaginadas, unificando modelagem preditiva e controle dentro de um único espaço latente coerente.

Avaliação Empírica

Os autores avaliaram o FF-JEPA na tarefa PushT usando três configurações experimentais: Curto horizonte (25 passos), Longo horizonte (75 passos) e Inicialização Aleatória (300 passos com posições iniciais aleatórias).

Resultados

• Desempenho de Longo Horizonte: O baseline plano LeWM colapsou em tarefas de longo horizonte (taxa de sucesso de 3,52% em $t=75$) e falhou completamente sob inicialização aleatória (0,00%). Em contraste, o FF-JEPA (Difusão) alcançou 91,80% de sucesso em $t=75$ e 82,42% sob inicialização aleatória.
• Vantagem Livre de Objetivo: Ao contrário dos baselines (LeWM, DINO-WM) que requerem uma imagem de objetivo para inferência, o FF-JEPA não requer. Apesar disso, superou os baselines hierárquicos DINO-WM tanto em curto horizonte (96,09% vs. 89,0%) quanto em longo horizonte (91,80% vs. 61,0%) nas configurações.
• Eficiência de Dados: Estudos de ablação mostraram que treinar o planejador de difusão em apenas 200 demonstrações de alta qualidade (uma redução de 40x nos dados) ainda produziu um desempenho competitivo (76,17% vs. 82,42% com dados totais), sugerindo alta eficiência de dados.
• Sobrecarga Computacional: O planejador determinístico adicionou um overhead de tempo de inferência negligível (2,1ms) comparado ao CEM (926ms), enquanto o planejador de difusão adicionou ~242ms. Ambos são significativamente mais rápidos do que os rollouts repetidos exigidos para o planejamento plano de longo horizonte.

Análise de Falhas

Os casos de falha ocorreram tipicamente quando o agente se desviava demais da posição prevista do sub-objetivo ou quando o planejador latente gerava um estado de sub-objetivo de baixa qualidade (por exemplo, omitindo o agente).

Significância e Alegações

O artigo posiciona o FF-JEPA como uma direção promissora para o aprendizado robótico escalável e livre de objetivos. Ao demonstrar que modelos de mundo condicionados a ação pareados com planejadores latentes treinados em dados não rotulados podem superar as limitações do planejamento plano, os autores sugerem que esta abordagem permite que agentes operem em ambientes de código aberto onde especificações de objetivos explícitos não estão disponíveis. O trabalho destaca que decompor o planejamento em sub-objetivos de curto prazo é uma estratégia viável para mitigar os problemas de computação e acúmulo de erro inerentes às abordagens atuais de modelos de mundo.