Resumen Técnico: FF-JEPA: Planificación de Largo Horizonte en Modelos de Mundo con Planificadores Latentes

Planteamiento del Problema

Los avances recientes en modelos de mundo visuales, particularmente las Arquitecturas de Predicción de Incrustación Conjunta (JEPA, por sus siglas en inglés) como LeWorldModel (LeWM), han demostrado la capacidad de predecir observaciones futuras y soportar el control basado en modelos mediante planificadores basados en muestreo como el Método de Entropía Cruzada (CEM). Sin embargo, los autores identifican tres limitaciones críticas que impiden el despliegue de estos modelos de mundo "planos" en tareas de largo horizonte en el mundo real:

1. Prohibiciones Computacionales: La planificación sobre horizontes largos requiere rollouts repetidos y secuencias de acciones cada vez más grandes, lo que hace que la optimización mediante CEM sea computacionalmente intratable.
2. Acumulación de Errores: Los errores de predicción autorregresiva se acumulan a lo largo de muchos pasos, causando que el planificador diverja para trayectorias complejas.
3. Dependencia de Objetivos: Los métodos existentes suelen requerir una imagen o estado de objetivo explícito ($o_g$) para dirigir la planificación. En muchos escenarios del mundo real, tal especificación de objetivo concreto es indisponible o poco práctica.

Aunque los métodos jerárquicos y los enfoques generativos han intentado abordar estos problemas, a menudo siguen operando en el espacio de píxeles o dependen de imágenes/trayectorias de objetivos especificadas externamente, lo que limita su aplicabilidad en entornos abiertos.

Metodología: Forward-Forward JEPA (FF-JEPA)

El artículo propone FF-JEPA, un marco jerárquico que extiende los modelos de mundo de estilo JEPA introduciendo un planificador latente ($G$) que opera junto a un modelo de avance estándar condicionado por acciones ($P$). Este enfoque descompone la planificación de largo horizonte en una secuencia de problemas de optimización de corto plazo tratables sin requerir imágenes de objetivo explícitas.

Arquitectura Central

El sistema consta de dos componentes primarios entrenados dentro del espacio latente definido por un codificador de modelo de mundo preentrenado y congelado ($E$):

1. El Modelo de Mundo ($P$): Un modelo de dinámica de avance estándar condicionado por acciones (basado en LeWM) que predice el siguiente estado latente dado un historial de estados y una acción: $\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$.
2. El Planificador Latente ($G$): Un modelo de dinámica de avance libre de acciones que predice el siguiente subobjetivo latente ($\hat{z}_{sg}$) basándose únicamente en un historial de subobjetivos previos. Opera a una resolución temporal superior (cada $H$ pasos del entorno).
   • Planificador Determinista ($G_{Det}$): Una arquitectura transformer (similar al predictor del modelo de mundo pero sin condicionamiento por acción) entrenada para minimizar el error cuadrático medio entre los subobjetivos predichos y los reales.
   • Planificador de Difusión ($G_{DM}$): Un backbone DiT (Diffusion Transformer) entrenado con un objetivo de coincidencia de puntuación de eliminación de ruido para predecir una secuencia de subobjetivos futuros.

Inferencia de la Planificación

El proceso de inferencia es jerárquico:

1. Predicción de Subobjetivos: El planificador latente $G$ predice el siguiente estado de subobjetivo $\hat{z}_{sg, m+1}$ dado el historial actual de subobjetivos.
2. Optimización de Acciones: El modelo de mundo $P$ utiliza CEM para encontrar una secuencia de acciones ($a_{1:H}$) que dirija el estado predicho hacia el subobjetivo predicho $\hat{z}_{sg, m+1}$ dentro de un horizonte corto $H$.
3. Iteración: Este ciclo se repite, encadenando efectivamente planes de corto horizonte para lograr objetivos de largo horizonte.

Crucialmente, este método elimina la necesidad de una imagen de objetivo en el tiempo de inferencia; el planificador genera la trayectoria hacia un objetivo definido implícitamente.

Contribuciones Clave

• Planificación de Largo Horizonte Libre de Objetivos: El marco permite planificar en modelos de mundo sin imágenes de objetivo explícitas mediante el uso de un planificador latente para generar subobjetivos intermedios.
• Descomposición Jerárquica: Aborda el "colapso de largo horizonte" de los modelos de mundo planos descomponiendo trayectorias complejas en subproblemas cortos y solubles, mitigando la acumulación de errores y los costos computacionales.
• Entrenamiento sin Etiquetas: El planificador latente puede entrenarse en demostraciones exitosas no etiquetadas (submuestreadas del espacio latente del modelo de mundo), reduciendo la dependencia de datos etiquetados con acciones.
• Marco Unificado: Reinterpreta el modelo de mundo como un módulo de dinámica inversa que opera sobre trayectorias latentes imaginadas, unificando el modelado predictivo y el control dentro de un único espacio latente coherente.

Evaluación Empírica

Los autores evaluaron FF-JEPA en la tarea PushT utilizando tres configuraciones experimentales: Horizonte corto (25 pasos), Largo horizonte (75 pasos) y Inicialización Aleatoria (300 pasos con posiciones de inicio aleatorias).

Resultados

• Rendimiento de Largo Horizonte: El modelo base plano LeWM colapsó en tareas de largo horizonte (tasa de éxito del 3.52% en $t=75$) y falló completamente bajo inicialización aleatoria (0.00%). En contraste, FF-JEPA (Difusión) logró un 91.80% de éxito en $t=75$ y un 82.42% bajo inicialización aleatoria.
• Ventaja de la Ausencia de Objetivos: A diferencia de los baselines (LeWM, DINO-WM) que requieren una imagen de objetivo para la inferencia, FF-JEPA no la requiere. A pesar de esto, superó a los baselines jerárquicos de DINO-WM tanto en configuración de corto horizonte (96.09% vs. 89.0%) como de largo horizonte (91.80% vs. 61.0%).
• Eficiencia de Datos: Los estudios de ablación mostraron que entrenar el planificador de difusión con solo 200 demostraciones de alta calidad (una reducción de 40 veces en los datos) aún producía un rendimiento competitivo (76.17% vs. 82.42% con datos completos), lo que sugiere una alta eficiencia de datos.
• Sobrecarga Computacional: El planificador determinista añadió una sobrecarga de tiempo de inferencia insignificante (2.1ms) comparado con CEM (926ms), mientras que el planificador de difusión añadió ~242ms. Ambos son significativamente más rápidos que los rollouts repetidos requeridos para la planificación plana de largo horizonte.

Análisis de Fallos

Los casos de fallo ocurrieron típicamente cuando el agente se desviaba demasiado de la posición del subobjetivo predicho o cuando el planificador latente generaba un estado de subobjetivo de baja calidad (por ejemplo, omitiendo al agente).

Significancia y Reivindicaciones

El artículo posiciona a FF-JEPA como una dirección prometedora para el aprendizaje robótico escalable y libre de objetivos. Al demostrar que los modelos de mundo condicionados por acciones, combinados con planificadores latentes entrenados en datos no etiquetados, pueden superar las limitaciones de la planificación plana, los autores sugieren que este enfoque permite a los agentes operar en entornos de mundo abierto donde las especificaciones de objetivos explícitos no están disponibles. El trabajo destaca que descomponer la planificación en subobjetivos de corto plazo es una estrategia viable para mitigar los problemas computacionales y de acumulación de errores inherentes a los enfoques actuales de modelos de mundo.