Resumo Técnico: Fast LeWorldModel

Declaração do Problema

Arquiteturas de Predição de Embedding Conjunto (JEPAs), como o LeWorldModel (LeWM), estabeleceram uma base para modelos de mundo visuais livres de reconstrução capazes de planejamento condicionado a objetivos e livre de recompensas. No entanto, os planejadores do tipo LeWM enfrentam duas limitações críticas devido à sua dependência de transições latentes autorregressivas de um passo:

1. Ineficiência Computacional: Para avaliar uma sequência candidata de ações, o planejador deve aplicar repetidamente um modelo de transição local passo a passo. Este rollout autorregressivo exige a geração de toda a trajetória latente sequencialmente, invocando repetidamente módulos de codificação de ação e de predição latente. Este processo torna-se um gargalo durante o planejamento baseado em amostragem (ex: Método de Cruzamento de Entropia, CEM), onde milhares de sequências candidatas são avaliadas.
2. Erro de Predição Acumulado: Como cada estado previsto $\hat{z}_{t+k}$ depende do estado previamente imaginado $\hat{z}_{t+k-1}$, os erros introduzidos em passos iniciais ou intermediários propagam-se e acumulam-se à medida que o horizonte de planejamento aumenta. Isso torna as predições de longo horizonte cada vez mais não confiáveis.

Metodologia: Fast LeWorldModel (Fast-LeWM)

Os autores propõem o Fast-LeWorldModel (Fast-LeWM), um modelo de mundo latente que substitui o rollout autorregressivo sequencial pela predição de prefixo de ação. Em vez de prever o próximo estado a partir do estado atual, o Fast-LeWM prevê o estado latente futuro alcançado após a execução de um prefixo específico de uma sequência de ações, condicionado diretamente ao estado latente observado atual.

Arquitetura Central

O método consiste em dois componentes primários:

1. Codificador de Prefixo de Ação:

   • Recebe uma sequência de ações candidata $(a_t, \dots, a_{t+H-1})$ e o estado latente atual $z_t$.
   • Mapeia $z_t$ para um "token de estado" e o concatena ao início dos tokens de ação.
   • Utiliza um Transformer causal para processar esta sequência. Devido à máscara causal, a saída na posição $k$ (o token de prefixo $p_{t,k}$) resume apenas as ações $(a_t, \dots, a_{t+k-1})$ e o contexto do estado atual, evitando o vazamento de informações de ações futuras.
   • Isso gera um conjunto de tokens de prefixo $\{p_{t,1}, \dots, p_{t,H}\}$, onde cada token representa o efeito acumulado de um prefixo de ação específico.

2. Preditor Latente Paralelo:

   • Recebe o âncora latente atual $z_t$ e o conjunto de tokens de prefixo $\{p_{t,1}, \dots, p_{t,H}\}$.
   • Prevê todos os futuros latentes $\{\hat{z}_{t+1}, \dots, \hat{z}_{t+H}\}$ em uma única passagem direta (forward pass) usando uma rede preditora paralela.
   • A predição para o horizonte $k$ é definida como $\hat{z}_{t+k} = G_\phi(z_t, p_{t,k})$.

Objetivo de Treinamento

O Fast-LeWM emprega supervisão densa em nível de prefixo. Diferente dos modelos padrão que supervisionam apenas o estado final ou transições de um único passo, o Fast-LeWM minimiza a perda entre os latentes previstos e os reais para cada horizonte $k$:
$$\mathcal{L}_{prefix} = \frac{1}{H} \sum_{k=1}^{H} \|\hat{z}_{t+k} - z_{t+k}\|_2^2$$
Isso força o modelo a aprender como os estados evoluem continuamente sob diferentes prefixos de ação, em vez de apenas ajustar transições locais de um passo. Um regularizador SIGReg é mantido para evitar o colapso latente.

Protocolo de Planejamento

Durante o planejamento, o Fast-LeWM mantém o mesmo protocolo CEM condicionado a objetivos do LeWM. No entanto, em vez de realizar o rollout de uma sequência passo a passo, ele codifica a sequência de ações uma única vez e consulta todos os horizontes futuros em paralelo. Um termo de autoconsistência opcional pode ser adicionado à função de custo, penalizando discrepâncias entre uma predição terminal direta e uma predição derivada via um passo de prefixo intermediário, o que estabiliza ainda mais a seleção de candidatos.

Principais Contribuições

1. Identificação de Gargalos: O artigo identifica a interface de transição local de um passo do LeWM como a causa primária de rollouts autorregressivos lentos e do acúmulo de erro ao longo de horizontes extensos.
2. Reformulação por Prefixo de Ação: Os autores propõem reformular a modelagem da dinâmica latente de transições de um único passo para a predição de prefixo de ação. Isso permite a supervisão densa em nível de prefixo e transforma os efeitos de ação acumulados em alvos diretos de treinamento dinâmico.
3. Predição de Múltiplos Horizontos em Paralelo: Ao desacoplar as predições de estados futuros das dependências sequenciais, o modelo alcança a predição de evolução de estado paralela, reduzindo significativamente o custo computacional e a propagação de erro.

Resultados Experimentais

Os autores avaliaram o Fast-LeWM em quatro tarefas de planejamento condicionado a objetivos (Two-Room, Reacher, Push-T, OGBench-Cube) utilizando os mesmos protocolos e conjuntos de dados do LeWM.

• Sucesso de Planejamento: O Fast-LeWM melhorou a taxa média de sucesso de 85,8% (LeWM) para 90,5%. Com o termo de autoconsistência opcional, este valor aumentou para 92,0%.
• Eficiência:
  • Tempo do Módulo de Dinâmica: Reduzido de 31,4s para 8,0s (um aceleração de 3,92×) na tarefa Two-Room ao substituir o rollout sequencial pela predição paralela.
  • Tempo Total de Resolução do CEM: Reduzido de 54,4s para 28,3s (uma redução de 48,0%).
• Precisão de Predição:
  • Erro em Malha Aberta (Open-Loop): O Fast-LeWM alcançou um erro de predição latente inicial menor e uma taxa de crescimento de erro significativamente mais lenta à medida que o horizonte aumentava em comparação ao LeWM.
  • Sondagem Física (Physical Probing): Na tarefa Push-T, o Fast-LeWM demonstrou desempenho superior em sondagens MLP para localização do agente, localização do bloco e ângulo do bloco, sugerindo que o espaço latente aprendido retém informações de estado físico mais ricas do que o LeWM ou PLDM.

Significância e Alegações

O artigo afirma que a predição de prefixo de ação serve como uma interface de dinâmica eficiente e eficaz para o planejamento de modelos de mundo latentes. Ao mudar de rollouts autorregressivos sequenciais para consultas paralelas baseadas em prefixos, o Fast-LeWM aborda o compromisso fundamental entre velocidade de planejamento e confiabilidade de predição.

Os autores enfatizam que esta abordagem não requer um novo objetivo de representação visual (permanecendo dentro do framework JEPA), mas sim otimiza como as sequências de ações candidatas são avaliadas. Os resultados sugerem que modelar a evolução do estado através de prefixos de ação acumulados permite que o modelo de dinâmica aprenda transições de estado contínuas e mais precisas, mitigando os erros de composição que tipicamente prejudicam o planejamento visual de longo horizonte.