Performance Asymmetry in Model-Based Reinforcement Learning

1. El Problema: Asimetría de Rendimiento y el Enmascaramiento de Métricas

Aunque los agentes de Aprendizaje por Refuerzo Basado en Modelos (MBRL) han alcanzado un rendimiento promedio "superhumano" en el benchmark Atari100k (superando un puntaje normalizado humano medio > 1.0), los autores descubren un problema crítico oculto: la Asimetría de Rendimiento.

• Definición: Los agentes MBRL actuales superan drásticamente a los humanos en ciertas tareas (denominadas Agent-Optimal), pero fracasan estrepitosamente en otras (denominadas Human-Optimal).
• La Trampa de las Métricas Actuales: Las métricas agregadas convencionales, como la Media Aritmética de los Puntajes Normalizados Humanos (HNS), enmascaran esta disparidad. Debido al gran rango de valores del HNS (desde negativos hasta valores muy altos) y la invarianza de la media aritmética ante valores extremos, un rendimiento excepcional en pocas tareas puede compensar matemáticamente un rendimiento pésimo en la mayoría de las tareas.
• Evidencia: El agente más avanzado (SOTA) basado en píxeles, DIAMOND, logró el mejor HNS medio general, pero obtuvo el peor rendimiento entre todas las líneas base en las tareas Human-Optimal, con una brecha de rendimiento de 21x entre los subconjuntos Agent-Optimal y Human-Optimal.
• Causas Identificadas:
  1. Complejidad de la Tarea: Las tareas Human-Optimal tienen espacios de acción más grandes y una mayor prevalencia de tareas tipo "Shooter" (disparos), lo que aumenta exponencialmente la complejidad de la búsqueda.
  2. Maldición de la Dimensionalidad: Los agentes basados en píxeles (que toman imágenes 64x64x3 como entrada) sufren de ineficiencia de entrenamiento en tareas complejas debido a la necesidad de aprender representaciones y políticas simultáneamente en un espacio de alta dimensión.
  3. Sesgo de los Modelos de Difusión: Los modelos de difusión basados en píxeles (como DIAMOND) excelen en tareas con alto detalle visual (ej. Breakout), lo que infla artificialmente la media general, pero fallan en tareas que requieren razonamiento complejo sobre acciones de alto riesgo (como el disparo en BankHeist o Hero).

2. Metodología y Propuesta: JEDI

Para abordar este problema, los autores proponen un nuevo marco de evaluación y un nuevo modelo de mundo.

A. Nueva Métrica de Evaluación: Sym-HNS

Los autores proponen Sym-HNS (Symmetric Human-Normalized Score) para revelar la asimetría:

• Partición: Dividen las 26 tareas de Atari100k en dos conjuntos equilibrados de 13 tareas cada uno: Agent-Optimal (donde los agentes suelen ganar) y Human-Optimal (donde los humanos suelen ganar).
• Cálculo: En lugar de la media aritmética, Sym-HNS utiliza la Media Armónica de los puntajes medios de ambos subconjuntos.
  $$\text{Sym-HNS} = \frac{2 \cdot \mu_{AO} \cdot \mu_{HO}}{\mu_{AO} + \mu_{HO}}$$
• Ventaja: La media armónica penaliza fuertemente los desequilibrios. Si un agente tiene un puntaje muy alto en un conjunto y muy bajo en el otro, el Sym-HNS será bajo, obligando a los modelos a ser equilibrados.

B. El Modelo Propuesto: JEDI (Joint Embedding DIffusion)

Para resolver la asimetría, proponen JEDI, un modelo de mundo basado en espacio latente que integra difusión y arquitectura JEPA (Joint Embedding Predictive Architecture).

• Arquitectura:
  • Codificador (Encoder): Mapea las observaciones de píxeles ($x_t$) a un espacio latente comprimido ($z_t$).
  • Modelo de Dinámica Latente: Utiliza un modelo de difusión (basado en EDM - Elucidating the Design Space of Diffusion-based Generative Models) para predecir el siguiente estado latente ($z_{t+1}$) dado el estado actual y la acción.
  • Entrenamiento End-to-End: A diferencia de los enfoques anteriores que requieren reconstrucción de imágenes o extracción de estados ocultos, JEDI entrena el codificador y el modelo de difusión conjuntamente utilizando una función de pérdida de denoising (ruido) directa.
• Innovación Clave (JEPA + Difusión):
  • JEDI es el primer método que demuestra que una pérdida de difusión "fuera de la caja" (off-the-shelf) es compatible directamente con el paradigma JEPA.
  • El codificador aprende representaciones útiles para la tarea de RL directamente a partir de los gradientes de la pérdida de denoising del modelo de difusión, sin necesidad de una cabeza de reconstrucción de imágenes.
  • Esto permite una compresión eficiente (12x) y evita la maldición de la dimensionalidad asociada a los inputs de píxeles.

3. Resultados Clave

El modelo JEDI fue evaluado en Atari100k y comparado con SOTA (DIAMOND, DreamerV3, STORM, etc.):

1. Reversión de la Asimetría: JEDI invierte la tendencia de asimetría. Logra el SOTA en Sym-HNS, demostrando un rendimiento equilibrado.
2. Rendimiento en Tareas Human-Optimal: JEDI supera a todos los baselines en las tareas Human-Optimal, incluyendo el logro de rendimiento superhumano en BankHeist (donde otros agentes fallan al no entender el uso estratégico de la dinamita).
3. Rendimiento en Tareas Visuales (Agent-Optimal): Contrario a la creencia de que los modelos latentes fallan en tareas visuales, JEDI logra el SOTA en Breakout (una tarea Agent-Optimal con alto detalle visual), demostrando que la difusión en espacio latente puede capturar detalles visuales necesarios.
4. Eficiencia Computacional:
   • Velocidad de Entrenamiento: 2x más rápido que DIAMOND.
   • Velocidad de Inferencia: 3x más rápido.
   • Uso de Memoria: Reduce el uso de memoria GPU en un 43% gracias a la compresión latente.
5. Robustez Estocástica: JEDI supera significativamente a los agentes latentes no difusivos en entornos con saltos de cuadros aleatorios (frameskipping), demostrando una mejor capacidad para modelar distribuciones multimodales estocásticas.

4. Contribuciones Principales

1. Identificación del Problema: Revelan la "Asimetría de Rendimiento" en MBRL, mostrando que el progreso actual es en realidad un sobreajuste a tareas específicas (Agent-Optimal) mientras se ignoran otras críticas.
2. Nueva Métrica: Introducen Sym-HNS y la partición de tareas, proporcionando una herramienta para evaluar el equilibrio real de los agentes en lugar de promedios engañosos.
3. Análisis de Causa Raíz: Vinculan el fracaso en tareas Human-Optimal a la maldición de la dimensionalidad de los inputs de píxeles y la complejidad de las tareas tipo "Shooter".
4. JEDI: Presentan un nuevo modelo de mundo latente end-to-end que combina difusión y JEPA, logrando un rendimiento equilibrado, eficiente y superior en tareas complejas sin sacrificar el rendimiento en tareas visuales.

5. Significado e Impacto

Este trabajo es fundamental porque desafía la narrativa de que los agentes MBRL están "resolviendo" el problema de Atari100k de manera general. Demuestra que las métricas actuales pueden ocultar fallos críticos en la generalización.

• Para la Investigación: Cambia el enfoque de maximizar el puntaje promedio a maximizar el rendimiento equilibrado (robustez) a través de tareas diversas.
• Para la Arquitectura de Modelos: Valida que los modelos de difusión pueden ser efectivos en espacios latentes comprimidos, superando las limitaciones de los modelos basados en píxeles en tareas complejas, y establece un nuevo estándar para el aprendizaje de representaciones sin reconstrucción explícita (JEPA-style).
• Implicación Futura: Sugiere que para lograr una inteligencia artificial verdaderamente robusta, los sistemas deben ser capaces de manejar tanto tareas de alta complejidad de acción (Human-Optimal) como tareas de alta fidelidad visual, evitando el sobreajuste a dominios específicos.