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

Resumen Técnico: SkyJEPA: Aprendizaje de Modelos de Mundo de Largo Horizonte para el Control Sim-to-Real Zero-Shot de Cuadricópteros

Declaración del Problema

Los modelos de dinámica precisos son esenciales para la toma de decisiones informada en vehículos aéreos ágiles, particularmente para el control basado en modelos, que requiere razonar sobre la evolución futura del sistema. Sin embargo, los modelos de dinámica de redes neuronales existentes enfrentan limitaciones significativas:

1. Inestabilidad de Largo Horizonte: La mayoría de los enfoques utilizan arquitecturas codificador-decodificador autorregresivas que retroalimentan las predicciones de forma recursiva como entradas. Este mecanismo amplifica los pequeños errores de un solo paso a lo largo del tiempo, lo que conduce al desvío, la inestabilidad y trayectorias físicamente implausibles a medida que aumenta el horizonte de predicción.
2. Falta de Interpretabilidad: Muchos modelos aprendidos operan en espacios latentes abstractos o reconstruyen observaciones completas, fallando en exponer cantidades físicamente significativas (por ejemplo, posición, velocidad, actitud) necesarias para imponer restricciones y límites de seguridad en los controladores.
3. Dependencia de Datos y Generalización: La recolección de datos diversos e informativos en el mundo real es costosa, riesgosa y laboriosa. Además, los modelos entrenados en plataformas o tareas específicas suelen fallar al generalizarse a nuevas configuraciones (por ejemplo, cambios en la carga útil o en las hélices) sin necesidad de reentrenamiento.
4. Restricciones de Tiempo Real: El control de alta frecuencia en hardware embebido con recursos limitados requiere modelos que puedan evaluarse repetidamente con baja latencia, un desafío para los modelos complejos basados en reconstrucción.

El artículo aborda la pregunta central: ¿Podemos aprender un modelo dinámico eficiente para robots aéreos que proporcione una predicción de largo horizonte precisa, rollouts físicamente interpretables, capacidad de control en tiempo real y generalización zero-shot entre tareas, reduciendo al mismo tiempo la dependencia de la costosa recolección de datos en el mundo real?

Metodología

Los autores proponen SkyJEPA, un marco basado en Arquitecturas de Predicción de Incrustación Conjunta (JEPAs, por sus siglas en inglés) adaptado para el control de cuadricópteros en tiempo real. El enfoque consta de cuatro componentes principales:

1. Modelo de Dinámica Latente Estilo JEPA:
   En lugar de predecir de forma autorregresiva estados u observaciones futuras, el modelo aprende a predecir representaciones futuras en un espacio latente compacto. Toma un historial de estados ($X_t$) y acciones ($A_t$) sobre una ventana $H$, los codifica en representaciones latentes ($s_t, z_t$), y predice el siguiente estado latente ($\tilde{s}_{t+1}$) condicionado a acciones futuras. Esta formulación evita el error acumulativo inherente a la predicción de estado autorregresiva al operar en un espacio latente estructurado.

2. Objetivo de Entrenamiento (SIGReg):
   El modelo se entrena utilizando una pérdida de dos términos:

   • Consistencia Predictiva ($L_{pred}$): Minimiza la distancia entre las predicciones latentes proyectadas (roll-out) y los estados latentes codificados.
   • Regularización Anti-Colapso (SIGReg): Utiliza la Regularización Gaussiana Isotrópica Esbozada (Sketched Isotropic Gaussian Regularization) para prevenir el colapso de la representación. Fomenta que las incrustaciones latentes coincidan con una distribución Gaussiana isotrópica comparando proyecciones unidimensionales aleatorias de la distribución latente con una Gaussiana estándar. Esto promueve la diversidad sin requerir mecanismos complejos de gradiente detenido (stop-gradient) o múltiples términos de pérdida.

3. Sonda Inspirada en la Física (Physics-Inspired Prober):
   Para hacer que las predicciones latentes sean utilizables para el control, una red de "sondeo" (prober) mapea los rollouts latentes congelados a estados métricos físicamente significativos. Este es un proceso de dos etapas:

   • Primero se entrena el modelo de dinámica latente.
   • Los codificadores y el predictor se congelan, y se entrena una red de sondeo ligera para predecir términos de corrección residuales (aceleración traslacional y angular) basados en los estados latentes.
   • Estos residuales se integran en un modelo cinemático diferenciable (usando mapas exponenciales de SO(3) para la actitud) para reconstruir la trayectoria completa del estado. Esto asegura que la salida sea interpretable y respete las restricciones físicas.

4. Síntesis Automatizada de Datos y Randomización de Dominio:
   Para eliminar la necesidad de una extensa recolección de datos en el mundo real, los autores proponen un flujo de trabajo para generar datos de entrenamiento diversos enteramente en simulación:

   • Generación de Referencia: Los procesos gaussianos generan trayectorias de referencia suaves y aleatorizadas que cubren un amplio entorno de vuelo.
   • Seguimiento de Bucle Cerrado: Un controlador de seguimiento (que combina NMPC y MPPI) ejecuta estas referencias en simulación para producir pares estado-acción físicamente factibles.
   • Randomización de Dominio: Cada ejecución (rollout) se simula con parámetros físicos aleatorizados (masa, inercia, arrastre, constantes de tiempo del motor, etc.) para exponer al modelo a una familia de sistemas en lugar de un único modelo nominal.

5. Control Basado en Muestreo (MPPI):
   El modelo aprendido se integra en un controlador de Integral de Trayectoria Predictiva de Camino (MPPI, por sus siglas en inglés). El controlador toma muestras de secuencias de acciones, realiza sus rollouts usando el modelo de dinámica latente y la sonda, evalúa los costos de la trayectoria y actualiza la entrada de control. Esto se ejecuta enteramente en hardware embebido (NVIDIA Orin NX) a altas frecuencias.

Contribuciones Clave

1. JEPA para Control en Tiempo Real: Introducción del primer modelo de dinámica latente de estilo JEPA entrenado de extremo a extremo para el control de cuadricópteros en tiempo real, yendo más allá de la planificación cinemática de alto nivel.
2. Sondeo Inspirado en la Física: Un mecanismo novedoso que mapea representaciones latentes congeladas a estados métricos interpretables mediante un integrador cinemático estructurado, permitiendo la predicción de largo horizonte físicamente fundamentada.
3. Transferencia Sim-to-Real Zero-Shot: Demostración de que un modelo entrenado enteramente con datos de simulación con randomización de dominio puede lograr un control robusto en hardware real sin necesidad de ajuste fino (fine-tuning) específico de la tarea.
4. Flujo de Trabajo de Generación de Datasets Automatizado: Un procedimiento basado en principios para la generación automatizada de conjuntos de datos utilizando procesos gaussianos y randomización de dominio, reduciendo la dependencia de la recolección insegura de datos en el mundo real.

Resultados

El marco fue evaluado mediante predicciones de bucle abierto y experimentos de bucle cerrado en exteriores en un cuadricóptero de 1.3 kg.

• Precisión de Largo Horizonte: SkyJEPA superó significativamente a las líneas base predictivas autorregresivas en rollouts de bucle abierto. Exhibió una menor relación de error acumulado (crecimiento del error respecto a la predicción con teacher-forcing) y menores tasas de crecimiento de error en horizontes largos (hasta 60 pasos).
• Interpretabilidad: La sonda inspirada en la física mejoró drásticamente la recuperación del estado métrico. En comparación con una sonda MLP estándar, la PI Prober redujo el RMSE de posición de ~5.56 m a ~1.43 m y el error de actitud de ~40° a ~4.7° en pruebas de bucle abierto.
• Transferencia Sim-to-Real Zero-Shot: En experimentos en el mundo real, el controlador propuesto logró errores de seguimiento de posición y actitud más bajos que las líneas base (MPPI con modelos predictivos) a través de diversas trayectorias (círculos, ochos, lemniscatas). Por ejemplo, en una trayectoria circular, el RMSE de posición se redujo de 0.39 m (línea base) a 0.24 m.
• Robustez: El sistema mantuvo el rendimiento bajo variaciones significativas de la plataforma sin reentrenamiento, incluyendo adiciones de carga útil (300 g) y cambio de hélices. Superó consistentemente a las líneas base en estas condiciones no nominales.
• Calidad de los Datos: El análisis mediante una puntuación de Calidad de Distribución de Trayectoria (TDQ) mostró una fuerte correlación inversa entre la diversidad/cobertura del conjunto de datos y el error de predicción, validando el enfoque de síntesis de datos con randomización de dominio.
• Desempeño en Tiempo Real: El modelo ligero (~9K parámetros) y la inferencia optimizada permitieron la ejecución en tiempo real en hardware embebido dentro del presupuesto de control de 10 ms.

Significado

El artículo afirma que SkyJEPA satisface las cuatro propiedades clave requeridas para un modelo de mundo útil en robótica aérea: predicción precisa de largo horizonte, interpretabilidad, inferencia en tiempo real y generalización de tareas zero-shot.

La importancia principal radica en demostrar que el entrenamiento únicamente en simulación con randomización de dominio es suficiente para un control robusto en el mundo real cuando se combina con un marco de aprendizaje apropiado. Al evitar la acumulación de errores autorregresivos mediante la predicción latente y fundamentar la salida con estructura física, el enfoque logra un nivel de generalización y estabilidad que los modelos puramente predictivos no pueden alcanzar. Los autores sugieren que este marco ofrece una alternativa escalable a la costosa recolección de datos en el mundo real y a la identificación manual de sistemas, permitiendo que los robots aéreos ágiles operen de manera confiable en entornos inciertos y cambiantes. Se propone como trabajo futuro extender esto a entradas visuales de alta dimensión e incorporar restricciones de seguridad explícitas.