Resumen Técnico: AdaJEPA: Un Modelo de Mundo Latente Adaptativo

Declaración del Problema

Los modelos de mundo latentes, particularmente las Arquitecturas de Predicción de Embebimiento Conjunto (JEPA, por sus siglas en inglés), permiten una planificación eficiente al predecir estados futuros en un espacio latente compacto. Estos modelos se entrenan típicamente en trayectorias offline libres de recompensas y se combinan con el Control Predictivo por Modelo (MPC) para tareas condicionadas por objetivos. Sin embargo, existe una limitación crítica: el modelo de mundo suele quedar congelado después del entrenamiento.

Al desplegarse, si las predicciones del modelo se vuelven inexactas debido a cambios en la distribución durante el tiempo de prueba (por ejemplo, cambios visuales como ruido o iluminación, o cambios dinámicos como fricción y masa), el planificador MPC optimiza un objetivo incorrecto. Las acciones que parecen efectivas en la "imaginación latente" pueden fallar en el entorno real. Los pequeños errores de un solo paso se acumulan a lo largo del horizonte de planificación, lo que provoca una degradación sustancial del rendimiento. Los enfoques actuales a menudo requieren demostraciones adicionales de expertos, datos de ajuste fino en el dominio objetivo o bucles externos complejos para abordar esto, lo que dificulta su aplicación en el mundo real.

Metodología: AdaJEPA

Los autores proponen AdaJEPA, un modelo de mundo latente adaptativo que realiza una adaptación en el tiempo de prueba dentro del bucle cerrado de MPC. En lugar de tratar al modelo como estático, AdaJEPA recalibra continuamente el modelo de mundo utilizando señales autosupervisadas derivadas de las propias interacciones del agente.

El Bucle Planificación–Ejecución–Adaptación–Replanificación

El mecanismo central opera de la siguiente manera en cada paso de MPC:

1. Planificar (Plan): El agente planifica una secuencia de acciones utilizando el modelo de mundo actual (potencialmente desactualizado) para minimizar un costo de alcance de objetivo latente.
2. Ejecutar (Execute): El agente ejecuta el primer bloque de acciones ($a_t$) en el entorno.
3. Observar (Observe): El agente recibe la siguiente observación ($o_{t+1}$).
4. Adaptar (Adapt): La transición $\{o_t, a_t, o_{t+1}\}$ se añade a un búfer de repetición (replay buffer). Los parámetros del modelo se actualizan mediante un pequeño número de pasos de gradiente (a menudo solo uno) para minimizar el error de predicción latente en esta nueva transición.
5. Replanificar (Replan): El modelo actualizado se utiliza inmediatamente para la siguiente iteración de planificación.

Componentes Técnicos

• Arquitectura: El modelo de mundo subyacente consiste en un codificador sensorial ($\mathcal{E}_s$), un codificador de acciones ($\mathcal{E}_a$) y un predictor ($f_\theta$). Se entrena offline para predecir representaciones latentes futuras ($z_{t+1}$) en lugar de reconstruir píxeles.
• Pérdida de Adaptación: La adaptación utiliza el mismo objetivo de predicción latente autosupervisado que el preentrenamiento:
  $$\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}))$$
  donde $\mathcal{B}$ es un búfer de transiciones recientes, $\ell$ es una pérdida de predicción latente (por ejemplo, MSE) y $\text{sg}(\cdot)$ es un operador de detención de gradiente (stop-gradient) para evitar el colapso.
• Actualizaciones de Parámetros: Para asegurar la eficiencia, la adaptación es ligera. Los autores actualizan solo un pequeño subconjunto de parámetros (por ejemplo, las capas finales del codificador y el predictor) en lugar de toda la red.
• Estrategias de Búfer: El búfer de repetición $\mathcal{B}$ puede gestionarse mediante "recientes-N" (manteniendo las $N$ transiciones más recientes) o "difíciles-N" (manteniendo las transiciones con los mayores errores de predicción).

Contribuciones Clave

1. Primera Adaptación de JEPA durante la Planificación: Según conocimiento de los autores, este es el primer trabajo en adaptar un modelo de mundo JEPA durante la fase de planificación, en lugar de depender únicamente del preentrenamiento o del ajuste fino offline.
2. Autosupervisión en Bucle Cerrado: El método utiliza las consecuencias inmediatas de las acciones ejecutadas como una señal autosupervisada, eliminando la necesidad de etiquetas de recompensa, demostraciones de expertos o fases de recolección de datos separadas.
3. Eficiencia de Muestreo y Cómputo: El marco logra mejoras significativas de rendimiento con un overhead computacional mínimo. En muchos experimentos, un solo paso de gradiente por cada paso de replanificación de MPC es suficiente.
4. Robustez a Cambios de Distribución: El enfoque demuestra robustez ante cambios visuales (ruido, desenfoque, cambios de color), cambios dinámicos (masa, amortiguación) y cambios estructurales (diseños de laberintos no vistos).

Resultados Experimentales

Los autores evaluaron AdaJEPA en los benchmarks PushT (empuje de bloques) y PointMaze (navegación) bajo diversas distribuciones de cambio entre entrenamiento y prueba.

• Rendimiento Intra-Distribución: Incluso en tareas similares a los datos de entrenamiento, AdaJEPA mejora las tasas de éxito (por ejemplo, >20% de ganancia en formas específicas de PushObj) o mantiene el rendimiento base sin degradación, demostando que es seguro aplicarlo intra-distribución.
• Rendimiento Fuera de la Distribución (OOD):
  • Cambios de Forma: En formas de bloques no vistas, AdaJEPA casi duplicó la tasa de éxito de planificación en comparación con el modelo congelado.
  • Cambios Visuales: El método mostró ganancias claras bajo desenfoque gaussiano, ruido de sal y pimienta y oscuridad.
  • Cambios de Dinámica y Diseño: En PointMaze, AdaJEPA se recuperó de los fallos causados por baja masa o alta amortiguación y mejoró significativamente las tasas de éxito en diseños de laberintos no vistos (por ejemplo, +25.3% de ganancia en cambios de diseño específicos).
• Latencia: El paso de adaptación introduce una latencia insignificante (aprox. 0.01–0.03s por replanificación) porque actualiza solo un pequeño subconjunto de parámetros. Además, al recuperarse de los errores de predicción más rápido, el agente a menudo alcanza el objetivo en menos replanificaciones totales, reduciendo el tiempo de ejecución global.
• Generalización a la Escasez de Datos: En regímenes de pocos datos, la adaptación en tiempo de prueba compensó la cobertura limitada de entrenamiento. Por ejemplo, con datos de entrenamiento limitados, AdaJEPA superó a un modelo congelado entrenado con 16 veces más trayectorias.

Significado y Reivindicaciones

El artículo argumenta que los modelos de mundo latentes no deben permanecer fijos tras el entrenamiento, sino que deben continuar mejorando a partir de la experiencia encontrada durante el despliegue.

• Inspiración Biológica: El enfoque se fundamenta en la adaptación sensoriomotora biológica, donde los sistemas (como el cerebelo) ajustan el comportamiento basándose en cambios de entradas y dinámicas.
• Complemento al Escalado de Datos: Si bien un conjunto de datos de entrenamiento más grande y diverso mejora la generalización, AdaJEPA proporciona una ruta de alta eficiencia de muestreo para refinar los modelos durante el despliegue, actuando como un complemento al escalado de datos.
• Dirección Futura: Los autores sugieren que el siguiente paso es combinar esta adaptación ligera en tiempo de prueba con el aprendizaje continuo y activo para expandir la cobertura del modelo de mundo a lo largo del tiempo, permitiendo una percepción y planificación más resilientes en entornos cambiantes.

Los autores enfatizan que su método es agnóstico a la implementación específica del modelo de mundo, mostrando mejoras consistentes a través de diferentes arquitecturas (ResNet, Transformer, basados en DINO) y planificadores (Descenso de Gradiente, CEM).