Resumo Técnico: Modelos de Mundo Sensoriomotores (SMWM)

Definição do Problema

O artigo aborda o desafio de aprender modelos de mundo latentes a partir de observações de alta dimensão (ex: pixels) sem sinais de recompensa ou rótulos de tarefas. Embora as Arquiteturas de Predição de Embedding Conjunto (JEPA) ofereçam um framework promissor ao prever estados futuros diretamente em um espaço de embedding (evitando a reconstrução de pixels), elas sofrem de colapso de representação. Quando um codificador e um modelo de dinâmica direta são treinados conjuntamente para minimizar o erro de predição, o codificador pode aprender a mapear todas as observações para um embedding constante. Essa solução trivial minimiza a perda de predição, mas torna o modelo inútutil para controle ou planejamento subsequente.

Soluções existentes para evitar o colapso frequentemente dependem de mecanismos complexos, tais como:

• Congelar codificadores pré-treinados (ex: DINO-WM).
• Adicionar regularizadores de distribuição (ex: alinhar embeddings a uma Gaussiana isotrópica via SIGReg).
• Utilizar redes alvo de Média Móvel Exponencial (EMA) com stop-gradients.
• Empregar funções de perda de múltiplos termos envolvendo restrições de variância-covariância.

Os autores postulam que essas abordagens podem não abordar plenamente a questão central do que uma representação útil de um modelo de mundo deve reter. Baseando-se no conceito de "percepção para a ação" e na percepção enativa, o artigo argumenta que as representações devem ser moldadas por sua relevância para as ações, preservando distinções que importam para dinâmicas controláveis enquanto descartam distratores incontroláveis.

Metodologia: Modelo de Mundo Sensoriomotor (SMWM)

Os autores propõem o SMWM, um modelo de mundo latente treinado de ponta a ponta a partir de trajetórias offline, livres de recompensa $(o_t, a_t, o_{t+1})$. A inovação central é o uso da regularização de dinâmica inversa como o único mecanismo para prevenir o colapso da representação.

Arquitetura

O modelo consiste em três componentes treinados conjuntamente:

1. Codificador ($f_\theta$): Mapeia observações $o_t$ para um embedding latente compacto $z_t$.
2. Modelo de Dinâmica Direta ($g_\phi$): Prediz o próximo estado latente $\hat{z}_{t+1}$ dado o estado atual $z_t$ e a ação $a_t$.
3. Modelo de Dinâmica Inversa ($h_\psi$): Prediz a ação $\hat{a}_t$ tomada entre dois embeddings consecutivos $z_t$ e $z_{t+1}$.

Objetivo de Treinamento

O modelo é otimizado minimizando uma função de perda conjunta:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Perda Direta ($\mathcal{L}_{fwd}$): Erro Quadrático Médio (MSE) entre o próximo embedding predito $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ e o próximo estado codificado real $z_{t+1} = f_\theta(o_{t+1})$.
• Perda Inversa ($\mathcal{L}_{inv}$): MSE entre a ação predita $\hat{a}_t = h_\psi(z_t, z_{t+1})$ e a ação real $a_t$.
• Peso de Regularização ($\lambda$): Um único hiperparâmetro que controla a força do sinal de dinâmica inversa.

Mecanismo de Ação

A perda inversa atua como um mecanismo anti-colapso. Se o codificador colapsar para um embedding constante, o modelo inverso receberá entradas idênticas para todas as transições e poderá apenas prever uma ação constante (a média da distribuição de ações). Para reduzir a perda inversa abaixo desse risco de preditor constante, o codificador deve preservar informações sobre a ação $a_t$ contida na transição $(o_t, o_{t+1})$. Isso força a representação a reter informações relevantes para a ação (graus de liberdade controláveis) enquanto filtra naturalmente distratores incontroláveis (ruído exógeno) que não influenciam a recuperação da ação.

Diferente dos regularizadores de distribuição que prescrevem a geometria do espaço de embedding (ex: forçando-o a ser uma Gaussiana), a dinâmica inversa ancora a representação a uma quantidade fundamentada na tarefa (a ação) sem impor priors geométricos específicos.

Principais Contribuições

1. Arquitetura SMWM: Um modelo de mundo sensoriomotor simples e treinável de ponta a ponta que utiliza a dinâmica inversa como o único mecanismo anti-colapso, exigindo apenas um hiperparâmetro adicional ($\lambda$).
2. Propriedades da Representação: Evidência empírica de que as representações aprendidas não são apenas não-colapsadas, mas especificamente rastreiam graus de liberdade controláveis, preservam a geometria espacial (tanto cartesiana quanto periódica) e filtram distratores incontroláveis.
3. Desempenho em Tarefas Subsequentes: Demonstração de desempenho competitivo de planejamento em tarefas de controle 2D e 3D, mostrando que as vantagens estruturais da representação se traduzem em um planejamento condicionado a objetivos eficaz.

Resultados Experimentais

Testbeds Controlados (Dot World & Sprite World)

• Recuperação de Dimensão Intrínseca: Em um ambiente "dot world", o modelo recuperou com sucesso a verdadeira dimensão intrínseca 2D do espaço de estados a partir de pixels, mesmo quando embutido em um espaço latente de 64 dimensões.
• Filtragem de Distratores: Quando objetos "distratores" incontroláveis foram introduzidos (movendo-se aleatoriamente), o codificador os ignorou. A dimensão latente efetiva correspondeu aos graus de liberdade controláveis, não à dimensão total da observação.
• Percepção Dependente de Controle: Em um experimento de "sprite world", o modelo aprendeu a preservar variáveis de pose (posição, orientação) apenas quando estas eram controláveis. Quando a orientação era incontrolável, a representação a mediava, resultando em uma reconstrução borrada da orientação.

Benchmarks de Planejamento

Os autores avaliaram o SMWM em quatro tarefas de controle offline: TwoRoom (navegação 2D), Reacher (controle contínuo), Push-T (manipulação 2D) e OGBench-Cube (manipulação 3D).

• Baselines: Comparado contra SIGReg (um forte baseline de JEPA usando regularizadores Gaussianos), Forward-only (sem perda inversa) e amostragem de ação Random.
• Desempenho:
  • O SMWM igualou o desempenho do SIGReg nas três tarefas 2D.
  • O SMWM superou o SIGReg na tarefa 3D OGBench-Cube (84% de sucesso vs. 59%).
  • Ambos os métodos regularizados superaram significativamente a ablação Forward-only, confirmando a necessidade de um mecanismo anti-colapso.
• Geometria Latente: Experimentos de sondagem revelaram que os embeddings do SMWM formam subespaços de baixa dimensão que espelham a geometria física do estado (ex: estruturas cilíndricas para ângulos de junta periódicos no Reacher, mapas cartesianos para navegação).

Significância e Alegações

O artigo afirma que o SMWM fornece uma instanciação principiada e simples de "percepção para a ação" no aprendizado de modelos de mundo latentes. Ao forçar os estados latentes a preservar informações suficientes para recuperar as ações do agente, o modelo naturalmente enviesa a representação para os graus de liberdade controláveis do ambiente.

Os pontos principais incluem:

• Simplicidade: A dinâmica inversa sozinha é suficiente para estabilizar o treinamento de JEPA de ponta a ponta sem reconstrução, sinais de recompensa, codificadores congelados ou complexos priors de distribuição.
• Viés Estrutural: O método induz representações que são interpretáveis e alinhadas com a estrutura causal do ambiente (ações e seus efeitos), em vez de apenas correlações preditivas.
• Utilidade Prática: As representações aprendidas suportam um planejamento latente eficaz, particularmente em ambientes 3D complexos onde outros baselines regularizados encontram dificuldades.

Os autores reconhecem limitações, como a suposição de que as ações são recuperáveis a partir de observações consecutivas e a incapacidade atual de capturar quantidades que não são identificáveis a partir de frames únicos (ex: velocidade) sem históricos de observação. No entanto, eles posicionam o SMWM como uma abordagem robusta e livre de recompensa para aprender modelos de mundo acionáveis a partir de dados offline.