SkyJEPA: Learning Long-Horizon World Models for Zero-Shot Sim-to-Real Control of Quadrotors

Resumo Técnico: SkyJEPA: Aprendendo Modelos de Mundo de Longo Horizonte para Controle Sim-to-Real Zero-Shot de Quadricópteros

Declaração do Problema

Modelos de dinâmica precisos são essenciais para a tomada de decisão informada em veículos aéreos ágeis, particularmente para o controle baseado em modelo, que requer o raciocínio sobre a evolução futura do sistema. No entanto, os modelos de dinâmica de redes neurais existentes enfrentam limitações significativas:

1. Instabilidade de Longo Horizonte: A maioria das abordagens utiliza arquiteturas codificador-decodificador autorregressivas que alimentam previsões recursivamente de volta como entradas. Esse mecanismo amplifica pequenos erros de um passo ao longo do tempo, levando à deriva, instabilidade e trajetórias fisicamente implausíveis conforme o horizonte de previsão aumenta.
2. Falta de Interpretabilidade: Muitos modelos aprendidos operam em espaços latentes abstratos ou reconstroem observações completas, falhando em expor quantidades fisicamente significativas (ex: posição, velocidade, atitude) necessárias para impor restrições e limites de segurança em controladores.
3. Dependência de Dados e Generalização: Coletar dados do mundo real diversos e informativos é caro, arriscado e trabalhoso. Além disso, modelos treinados em plataformas ou tarefas específicas frequentemente falham ao generalizar para novas configurações (ex: mudanças de carga útil, troca de hélices) sem retreinamento.
4. Restrições de Tempo Real: O controle de alta frequência em hardware embarcado com recursos limitados exige modelos que possam ser avaliados repetidamente com baixa latência, um desafio para modelos complexos baseados em reconstrução.

O artigo aborda a questão central: Podemos aprender um modelo dinâmico eficiente para robôs aéreos que forneça previsão de longo horizonte precisa, rollouts interpretáveis fisicamente, capacidade de controle em tempo real e generalização zero-shot entre tarefas, reduzindo simultaneamente a dependência da coleta de dados perigosa no mundo real?

Metodologia

Os autores propõem o SkyJEPA, um framework baseado em Arquiteturas de Predição de Embedding Conjunto (JEPAs) adaptado para o controle de quadricópteros em tempo real. A abordagem consiste em quatro componentes principais:

1. Modelo de Dinâmica Latente Estilo JEPA:
   Em vez de prever estados ou observações futuras de forma autorregressiva, o modelo aprende a prever representações futuras em um espaço latente compacto. Ele recebe um histórico de estados ($X_t$) e ações ($A_t$) sobre uma janela $H$, codifica-os em representações latentes ($s_t, z_t$) e prevê o próximo estado latente ($\tilde{s}_{t+1}$) condicionado a ações futuras. Esta formulação evita o erro cumulativo inerente à previsão de estado autorregressiva ao operar em um espaço latente estruturado.

2. Objetivo de Treinamento (SIGReg):
   O modelo é treinado usando uma perda de dois termos:

   • Consistência Preditiva ($L_{pred}$): Minimiza a distância entre as previsões latentes projetadas (rollouts) e os estados latentes codificados.
   • Regularização Anti-Colapso (SIGReg): Utiliza a Regularização Gaussiana Isotrópica Esboçada (Sketched Isotropic Gaussian Regularization) para evitar o colapso da representação. Ela incentiva os embeddings latentes a corresponderem a uma distribuição Gaussiana isotrópica, comparando projeções aleatórias unidimensionais da distribuição latente a uma Gaussiana padrão. Isso promove a diversidade sem exigir mecanismos complexos de gradiente interrompido (stop-gradient) ou múltiplos termos de perda.

3. Prober Inspirado na Física (Physics-Inspired Prober):
   Para tornar as previsões latentes utilizáveis para controle, uma rede "prober" mapeia os rollouts latentes congelados para estados métricos fisicamente significativos. Este é um processo de dois estágios:

   • O modelo de dinâmica latente é treinado primeiro.
   • Os codificadores e o preditor são congelados, e uma rede de sondagem (prober) leve é treinada para prever termos de correção residual (aceleração translacional e angular) com base nos estados latentes.
   • Esses resíduos são integrados a um modelo cinemático diferenciável (usando mapas exponenciais SO(3) para atitude) para reconstruir a trajetória completa do estado. Isso garante que a saída seja interpretável e respeite as restrições físicas.

4. Síntese Automatizada de Dados e Randomização de Domínio:
   Para eliminar a necessidade de extensa coleta de dados no mundo real, os autores propõem um pipeline para gerar dados de treinamento diversos inteiramente em simulação:

   • Geração de Referência: Processos Gaussianos geram trajetórias de referência suaves e randomizadas que cobrem um amplo envelope de voo.
     z Rastreamento de Malha Fechada: Um controlador de rastreamento (combinando NMPC e MPPI) executa essas referências em simulação para produzir pares estado-ação fisicamente viáveis.
   • Randomização de Domínio: Cada rollout é simulado com parâmetros físicos randomizados (massa, inércia, arrasto, constantes de tempo do motor, etc.) para expor o modelo a uma família de sistemas, em vez de um único modelo nominal.

5. Controle Baseado em Amostragem (MPPI):
   O modelo aprendido é integrado a um controlador de Integral de Caminho de Predição de Modelo (MPPI). O controlador amostra sequências de ações, realiza os rollouts usando o modelo de dinâmica latente e o prober, avalia os custos da trajetória e atualiza a entrada de controle. Isso roda inteiramente em hardware embarcado (NVIDIA Orin NX) em altas frequências.

Principais Contribuições

1. JEPA para Controle em Tempo Real: Introdução do primeiro modelo de dinâmica latente estilo JEPA treinado de ponta a ponta para controle de quadricóptero em tempo real, indo além do planejamento cinemático de alto nível.
2. Probing Inspirado na Física: Um mecanismo inovador que mapeia representações latentes congeladas para estados métricos interpretáveis através de um integrador cinemático estruturado, permitindo a previsão de longo horizonte com fundamentação física.
3. Transferência Sim-to-Real Zero-Shot: Demonstração de que um modelo treinado inteiramente em dados de simulação com randomização de domínio pode alcançar controle robusto em hardware real sem ajuste fino específico para a tarefa.
4. Pipeline de Dados Automatizado: Um procedimento principiado para geração automatizada de datasets usando processos Gaussianos e randomização de domínio, reduzindo a dependência de coleta de dados insegura no mundo real.

Resultados

O framework foi avaliado através de previsões em malha aberta e experimentos em malha fechada ao ar livre em um quadricóptero de 1,3 kg.

• Precisão de Longo Horizonte: O SkyJEPA superou significativamente as linhas de base preditivas autorregressivas em rollouts de malha aberta. Exibiu uma razão de acúmulo menor (crescimento do erro em relação à previsão com teacher-forcing) e menores taxas de crescimento de erro em horizontos longos (até 60 passos).
• Interpretabilidade: O prober inspirado na física melhorou drasticamente a recuperação do estado métrico. Comparado a um prober MLP padrão, o PI Prober reduziu o RMSE de posição de ~5,56m para ~1,43m e o erro de atitude de ~40° para ~4,7° em testes de malha aberta.
• Transferência Sim-to-Real Zero-Shot: Em experimentos no mundo real, o controlador proposto alcançou menores erros de rastreamento de posição e atitude do que as linhas de base (MPPI com modelos preditivos) em diversas trajetórias (círculos, oito, lemniscatas). Por exemplo, em uma trajetória circular, o RMSE de posição foi reduzido de 0,39m (baseline) para 0,24m.
• Robustez: O sistema manteve o desempenho sob variações significativas de plataforma sem retreinamento, incluindo adições de carga útil (300g) e troca de hélices. Consistiu em superar as linhas de base nessas condições não nominais.
• Qualidade dos Dados: Uma análise usando uma pontuação de Qualidade de Distribuição de Trajetória (TDQ) mostrou uma forte correlação inversa entre a diversidade/cobertura do dataset e o erro de previsão, validando a abordagem de síntese de dados com randomização de domínio.
• Desempenho em Tempo Real: O modelo leve (~9K parâmetros) e a inferência otimizada permitiram a execução em tempo real em hardware embarcado dentro do orçamento de controle de 10ms.

Significância

O artigo afirma que o SkyJEPA satisfaz as quatro propriedades fundamentais exigidas para um modelo de mundo útil em robótica aérea: previsão de longo horizonte precisa, interpretabilidade, inferência em tempo real e generalização de tarefa zero-shot.

A principal significância reside em demonstrar que o treinamento apenas em simulação com randomização de domínio é suficiente para o controle robusto no mundo real quando pareado com um framework apropriado. Ao evitar o acúmulo de erro autorregressivo através da predição latente e fundamentar a saída com estrutura física, a abordagem alcança um nível de generalização e estabilidade que modelos puramente preditivos não conseguem. Os autores sugerem que este framework oferece uma alternativa escalável à coleta de dados dispendiosa no mundo real e à identificação manual de sistemas, permitindo que robôs aéreos ágeis operem de forma confiável em ambientes incertos e mutáveis. Trabalhos futuros são propostos para estender isso a entradas visuais de alta dimensão e incorporar restrições de segurança explícitas.