Resumen Técnico: JEDI: Modelo de Mundo de Difusión de Incrustación Conjunta para Aprendizaje por Refuerzo Basado en Modelos en Línea

Enunciado del Problema

El Aprendizaje por Refuerzo Basado en Modelos en Línea (MBRL) depende de modelos de mundo aprendidos para predecir resultados futuros y mejorar la eficiencia de las muestras. Aunque los modelos de difusión han surgido como columnas vertebrales poderosas para el modelado de mundo en dominios visualmente complejos, los enfoques actuales enfrentan una tensión significativa:

1. Difusión de Píxeles: Modelos como DIAMOND operan directamente en el espacio de píxeles. Aunque son efectivos, son computacionalmente costosos en cuanto a memoria (VRAM), tiempo de muestreo y duración del entrenamiento.
2. Difusión Latente: Mover la difusión al espacio latente mejora la eficiencia, pero a menudo depende de latentes entrenados por separado (por ejemplo, mediante objetivos de reconstrucción) en lugar de objetivos de modelo de mundo de extremo a extremo. Estos enfoques, como Horizon Imagination (HI), aún no han igualado el rendimiento de los modelos de difusión basados en píxeles y carecen del aprendizaje de representación predictiva de extremo a extremo que ha impulsado el progreso moderno del MBRL.

La pregunta central abordada es si un modelo de mundo de difusión puede aprender su representación latente de extremo a extremo a partir del objetivo de eliminación de ruido utilizando un marco de Arquitectura de Incrustación Conjunta Predictiva (JEPA), combinando así la eficiencia de la difusión latente con el rendimiento del aprendizaje predictivo de extremo a extremo.

Metodología: JEDI

Los autores introducen JEDI (Joint Embedding Diffusion), el primer modelo de mundo de difusión latente en línea de extremo a extremo construido sobre un marco JEPA.

Arquitectura Central

JEDI aprende simultáneamente un codificador de espacio latente ($E_\phi$) y un modelo de dinámicas. El sistema consta de:

• Codificador: Comprime las observaciones $x_t$ en latentes limpios $z^0_t$.
• Dinámicas de Difusión Latente: Predice el siguiente estado latente mediante eliminación de ruido condicional. En lugar de una predicción de un solo paso, utiliza una cadena de difusión inversa condicionada al latente actual $z^0_t$ y a las acciones pasadas para predecir el latente futuro limpio $z^0_{t+1}$.
• Cabezas de Recompensa y Terminación: Predicen recompensas y señales de terminación directamente desde el estado latente.

Mecanismo de Entrenamiento

A diferencia de los modelos de difusión latente tradicionales que dependen de pérdidas de reconstrucción (por ejemplo, VAEs) para entrenar el codificador, JEDI entrena el codificador de extremo a extremo utilizando la pérdida de eliminación de ruido de difusión.

• Objetivo de Incrustación Conjunta: El codificador se actualiza mediante gradientes derivados del objetivo de eliminación de ruido condicional. El objetivo latente futuro se trata como una "etiqueta" derivada del entorno, y el modelo aprende a eliminar el ruido desde una versión ruidosa de este latente futuro hasta el estado limpio, condicionado al estado actual.
• Stop-Gradient: Para prevenir el colapso de la representación y estabilizar el entrenamiento, se aplica un stop-gradient al objetivo latente futuro durante las actualizaciones del codificador, siguiendo los principios de JEPA y TDMPC2.
• Técnicas de Estabilización: El modelo emplea un acotamiento diferenciable tanh en las activaciones latentes y un mecanismo de conmutación aleatoria entre la salida del eliminador de ruido y la salida del codificador para la condición del siguiente paso, equilibrando consistencia y supervisión.

Motivación Teórica

El artículo proporciona una justificación teórica que vincula a JEDI con el principio de Cuello de Botella de Información.

• JEPA como Cuello de Botella: Se demuestra que el objetivo estándar de JEPA induce un cuello de botella de información variacional, recompensando representaciones que retienen información útil para predecir el futuro mientras comprimen detalles irrelevantes.
• Descomposición Difusión-JEPA: Los autores demuestran que la difusión latente condicional admite una descomposición similar de cuello de botella de información variacional. La pérdida de eliminación de ruido condicional actúa como el término predictivo dentro de esta estructura de cuello de botella. Esto sugiere que la eliminación de ruido por difusión proporciona una señal de aprendizaje predictivo de extremo a extremo y fundamentada para el codificador latente, en lugar de servir meramente como un sustituto generativo.

Contribuciones Clave

1. Primer Modelo de Mundo de Difusión Latente de Extremo a Extremo: JEDI es el primer modelo MBRL en línea que entrena su codificador latente de extremo a extremo utilizando un objetivo de eliminación de ruido de difusión dentro de un marco JEPA, eliminando la necesidad de preentrenamiento separado o pérdidas de reconstrucción.
2. Conexión Teórica: El artículo establece un vínculo matemático entre los objetivos de difusión condicional y el cuello de botella de información variacional, explicando por qué la eliminación de ruido puede apoyar eficazmente el aprendizaje de representación predictiva.
3. Ganancias de Eficiencia: Al pasar del espacio de píxeles a un espacio latente aprendido compacto ($16 \times 8 \times 8$) y entrenar de extremo a extremo, JEDI reduce significativamente la sobrecarga computacional en comparación con la difusión basada en píxeles.

Resultados Experimentales

JEDI fue evaluado en la referencia Atari100k (26 tareas, 5 semillas) y en el entorno Craftium (toma de decisiones encarnada en 3D).

Rendimiento

• Atari100k: JEDI logra un rendimiento competitivo, igualando o superando a bases sólidas (DIAMOND, DreamerV3, STORM) en las métricas de Media Cuartílica Inter (IQM) y Brecha de Optimalidad. Logra resultados de vanguardia (SOTA) en 7 juegos y un rendimiento de segundo lugar en la cantidad de tareas superhumanas.
• Craftium: JEDI supera a Horizon Imagination (HI) en las cuatro tareas probadas, demostrando que los latentes predictivos de extremo a extremo pueden apoyar un rendimiento sólido en entornos 3D sin entrenamiento latente basado en reconstrucción.
• Estocasticidad: En experimentos con salto de cuadros estocástico, JEDI supera consistentemente a DreamerV3, lo que sugiere que los modelos de difusión pueden ser más robustos ante una alta incertidumbre aleatoria.

Eficiencia

En relación con la base de píxeles DIAMOND, JEDI demuestra mejoras sustanciales en eficiencia:

• Uso de VRAM: Utiliza un 43% menos de VRAM.
• Velocidad de Muestreo: Logra un muestreo de modelo de mundo >3 veces más rápido.
• Velocidad de Entrenamiento: Entrena >2.5 veces más rápido (38 horas frente a 98 horas en A100 para Atari100k).
• Tamaño del Modelo: Mantiene una cantidad de parámetros similar (13.5M) a DIAMOND, pero supera significativamente a HI (97M de parámetros) en eficiencia mientras mantiene un rendimiento competitivo.

Cambio en el Perfil de Rendimiento

Un hallazgo distintivo es que JEDI exhibe un perfil de rendimiento a nivel de tarea diferente en comparación con DIAMOND. Mientras que DIAMOND destaca en regímenes de alta HNS (Puntuación Normalizada Humana), JEDI muestra mejoras marcadas en regímenes de baja HNS (juegos difíciles). Esto se captura mediante la métrica Har-HNS, donde JEDI supera a DIAMOND, lo que indica que falla menos severamente en las tareas más difíciles. Los autores atribuyen esto a que la interfaz latente de extremo a extremo reduce la dimensionalidad de la entrada para el actor-crítico, permitiendo una mejor optimización en juegos complejos con espacios de acción grandes.

Significado y Afirmaciones

El artículo afirma que JEDI demuestra que la difusión latente predictiva es una alternativa viable y eficiente tanto a la difusión en espacio de píxeles como a los modelos latentes entrenados por separado para el MBRL en línea.

• Más allá del Cómputo: Los resultados sugieren que los latentes predictivos de extremo a extremo cambian el comportamiento del propio modelo de mundo, no solo el costo computacional. El cambio en los perfiles de rendimiento indica que la representación aprendida altera fundamentalmente cómo el agente interactúa con entornos complejos.
• Enfoque Fundamentado: Al fundamentar el método en la teoría del cuello de botella de información, el artículo posiciona la difusión latente predictiva como una dirección fundamentada para escalar el MBRL en línea.
• Limitaciones: Los autores notan modestamente que la evidencia está actualmente limitada a Atari100k y Craftium, y que el método aún no ha sido probado en dominios como DMControl o Procgen. También reconocen que el rendimiento depende de elecciones de diseño específicas (por ejemplo, acotamiento, conmutación aleatoria) y que la derivación teórica motiva propiedades de representación en lugar de garantizar la optimización del objetivo variacional completo.