Resumo Técnico: AdaJEPA: Um Modelo de Mundo Latente Adaptativo

Problema

Modelos de mundo latentes, particularmente as Arquiteturas de Predição de Embedding Conjunto (JEPAs), permitem um planejamento eficiente ao prever estados futuros em um espaço latente compacto. Esses modelos são tipicamente treinados em trajetórias offline livres de recompensa e pareados com Controle Preditivo por Modelo (MPC) para tarefas condicionadas a objetivos. No entanto, uma limitação crítica existe: o modelo de mundo geralmente é congelado após o treinamento.

Durante a implantação, se as previsões do modelo se tornarem imprecisas devido a mudanças na distribuição no tempo de teste (ex: mudanças visuais como ruído ou iluminação, ou mudanças dinâmicas como fricção e massa), o planejador MPC otimiza um objetivo incorreto. Ações que parecem eficazes na "imaginação latente" podem falhar no ambiente real. Pequenos erros de um passo se acumulam ao longo do horizonte de planejamento, levando a uma degradação substancial de desempenho. As abordagens atuais frequentemente exigem demonstrações de especialistas adicionais, dados de ajuste fino no domínio alvo ou loops externos complexos para resolver isso, o que dificulta a aplicação no mundo real.

Metodologia: AdaJEPA

Os autores propõem o AdaJEPA, um modelo de mundo latente adaptativo que realiza adaptação no tempo de teste dentro do loop fechado do MPC. Em vez de tratar o modelo como estático, o AdaJEPA recalibra continuamente o modelo de mundo usando sinais autossupervisionados derivados das próprias interações do agente.

O Ciclo Planejar–Executar–Adaptar–Replanejar

O mecanismo central opera da seguinte forma em cada passo do MPC:

1. Planejar: O agente planeja uma sequência de ações usando o modelo de mundo atual (potencialmente desatualizado) para minimizar um custo de alcance de objetivo latente.
2. Executar: O agente executa o primeiro bloco de ações ($a_t$) no ambiente.
3. Observar: O agente recebe a próxima observação ($o_{t+1}$).
4. Adaptar: A transição $\{o_t, a_t, o_{t+1}\}$ é adicionada a um buffer de replay. Os parâmetros do modelo são atualizados via um pequeno número de passos de gradiente (frequentemente apenas um) para minimizar o erro de predição latente nesta nova transição.
5. Replanejar: O modelo atualizado é imediatamente utilizado para a próxima iteração de planejamento.

Componentes Técnicos

• Arquitetura: O modelo de mundo subjacente consiste em um codificador sensorial ($\mathcal{E}_s$), um codificador de ação ($\mathcal{E}_a$) e um preditor ($f_\theta$). Ele é treinado offline para prever representações latentes futuras ($z_{t+1}$) em vez de reconstruir pixels.
• Perda de Adaptação: A adaptação utiliza o mesmo objetivo de predição latente autossupervisionado do pré-treinamento:
  $$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$$
  onde $\mathcal{B}$ é um buffer de transições recentes, $\ell$ é uma perda de predição latente (ex: MSE) e $\text{sg}(\cdot)$ é um operador de stop-gradient para evitar o colapso.
• Atualizações de Parâmetros: Para garantir a eficiência, a adaptação é leve. Os autores atualizam apenas um pequeno subconjunto de parâmetros (ex: as camadas finais do codificador e do preditor), em vez de toda a rede.
• Estratégias de Buffer: O buffer de replay $\mathcal{B}$ pode ser gerenciado via "recent-N" (mantendo as $N$ transições mais recentes) ou "hard-N" (mantendo transições com os maiores erros de predição).

Principais Contribuições

1. Primeira Adaptação de JEPA durante o Planejamento: Tanto quanto os autores saibam, este é o primeiro trabalho a adaptar um modelo de mundo JEPA durante a fase de planejamento, em vez de depender apenas de pré-treinamento ou ajuste fino offline.
2. Autossupervisão em Loop Fechado: O método utiliza as consequências imediatas das ações executadas como um sinal autossupervisionado, eliminando a necessidade de rótulos de recompensa, demonstrações de especialistas ou fases separadas de coleta de dados.
3. Eficiência de Amostragem e Computação: O framework alcança ganhos significativos de desempenho com overhead computacional mínimo. Em muitos experimentos, um único passo de gradiente por passo de replanejamento do MPC é suficiente.
4. Robustez a Mudanças de Distribuição: A abordagem demonstra robustez através de mudanças visuais (ruído, desfoque, mudanças de cor), mudanças dinâmicas (massa, amortecimento) e mudanças estruturais (layouts de labirinto não vistos).

Resultados Experimentais

Os autores avaliaram o AdaJEPA nos benchmarks PushT (empurrar blocos) e PointMaze (navegação) sob várias mudanças de distribuição de treino-teste.

• Desempenho Dentro da Distribuição: Mesmo em tarefas similares aos dados de treinamento, o AdaJEPA melhora as taxas de sucesso (ex: ganho de >20% em formas específicas de PushObj) ou mantém o desempenho da linha de base sem degradação, provando que é seguro aplicá-lo dentro da distribuição.
• Desempenho Fora da Distribuição (OOD):
  • Mudanças de Forma: Em formas de blocos não vistas, o AdaJEPA quase dobrou a taxa de sucesso de planejamento comparado ao modelo congelado.
  • Mudanças Visuais: O método mostrou ganhos claros sob desfoque Gaussiano, ruído sal e pimenta e iluminação escura.
  • Mudanças de Dinâmica e Layout: No PointMaze, o AdaJEPA recuperou-se de falhas causadas por baixa massa ou alto amortecimento e melhorou significativamente as taxas de sucesso em layouts de labirinto não vistos (ex: ganho de +25,3% em layouts específicos).
• Latência: O passo de adaptação introduz latência desprezível (aprox. 0,01–0,03s por replan) porque atualiza apenas um pequeno subconjunto de parâmetros. Além disso, ao recuperar-se de erros de predição mais rapidamente, o agente frequentemente atinge o objetivo em menos replanamentos totais, reduzindo o tempo total de execução.
• Generalização para Escassez de Dados: Em regimes de baixos dados, a adaptação no tempo de teste compensou a cobertura limitada de treinamento. Por exemplo, com dados de treinamento limitados, o AdaJEPA superou um modelo congelado treinado com 16x mais trajetórias.

Significância e Alegações

O artigo argumenta que os modelos de mundo latentes não devem permanecer fixos após o treinamento, mas sim continuar a melhorar a partir da experiência encontrada durante a implantação.

• Inspiração Biológica: A abordagem é fundamentada na adaptação sensoriomotora biológica, onde sistemas (como o cerebelo) ajustam o comportamento com base em mudanças de inputs e dinâmicas.
• Complemento à Escala de Dados: Embora dados de treinamento maiores e mais diversos melhorem a generalização, o AdaJEPA fornece uma rota eficiente em termos de amostras para refinar modelos durante a implantação, atuando como um complemento à escala de dados.
• Direção Futura: Os autores sugerem que o próximo passo é combinar esta adaptação leve no tempo de teste com aprendizado contínuo e ativo para expandir a cobertura do modelo de mundo ao longo do tempo, permitindo uma percepção e planejamento mais resilientes em ambientes mutáveis.

Os autores enfatizam que seu método é agnóstico à implementação específica do modelo de mundo, mostrando melhorias consistentes através de diferentes arquiteturas (ResNet, Transformer, baseada em DINO) e planejadores (Gradiente Descendente, CEM).