Resumen Técnico: KAGE-Bench

Planteamiento del Problema

Los agentes de aprendizaje por refuerzo (RL) entrenados en observaciones de píxeles de alta dimensión son notoriamente frágiles cuando se enfrentan a cambios de distribución puramente visuales, incluso cuando la dinámica latente, los kernels de transición y las funciones de recompensa permanecen inalterados. Aunque los agentes pueden converger y desempeñarse bien en la distribución original (in-distribution), suelen fallar catastróficamente ante cambios en la apariencia, la iluminación o el ruido visual. Un obstáculo principal para diagnosticar estos fallos es la falta de bancos de pruebas (benchmarks) que puedan aislar fuentes específicas de cambio visual. Los bancos de pruebas existentes (por ejemplo, Procgen, Distracting Control Suite) suelen entrelazar múltiples fuentes de variación —como la apariencia del fondo, la geometría, la dinámica y los distractores— lo que hace imposible atribuir las brechas de rendimiento a factores visuales específicos. Además, muchos entornos de RL basados en píxeles son computacionalmente costosos, lo que dificulta los estudios de ablación a gran escala y las pruebas de hipótesis sistemáticas.

Metodología

Los autores presentan KAGE-Env y KAGE-Bench para abordar estas limitaciones mediante una filosofía de diseño de "ejes conocidos".

KAGE-Env: Una Plataforma de Plataformas Factorizada

KAGE-Env es un entorno de plataformas 2D nativo de JAX diseñado para factorizar el proceso de observación en ejes visuales controlables de forma independiente, manteniendo fijo el problema de control latente.

• Arquitectura: El entorno está implementado íntegramente en JAX con compilación JIT de extremo a extremo y ejecución vectorizada (usando vmap y lax.scan). Esto permite un paralelismo masivo, alcanzando hasta 33 millones de pasos de entorno por segundo en una sola GPU con $2^{16}$ entornos paralelos.
• Factorización: El kernel de observación $O_\xi(\cdot|s)$ está parametrizado por una configuración visual $\xi$. Crucialmente, el kernel de transición $P$ y la función de recompensa $r$ son invariantes a $\xi$. Esto asegura que cualquier diferencia de rendimiento entre el entrenamiento y la evaluación surja únicamente de cómo una política basada en píxeles $\pi(a|o)$ responde a diferentes renderizados del mismo estado latente $s$.
• Configuración: El entorno se controla mediante un único archivo YAML, que expone 93 parámetros controlables explícitamente en categorías como la imagen/color del fondo, la apariencia del agente/sprites, los personajes no jugadores (NPCs), los distractores, los filtros fotométricos (brillo, tono, ruido), los efectos de iluminación y el diseño (layout).

KAGE-Bench: El Protocolo de Evaluación

Construido sobre KKAGE-Env, KAGE-Bench define un protocolo de evaluación estandarizado que comprende seis suites de "ejes conocidos" y 34 pares de configuración de entrenamiento-evaluación.

• Aislamiento: En cada par, la configuración de entrenamiento ($\xi_{train}$) y la de evaluación ($\xi_{eval}$) difieren solo en un subconjunto designado de parámetros de renderizado (un eje específico), mientras que todos los demás parámetros y la dinámica subyacente permanecen idénticos.
• Fundamento Teórico: El artículo formaliza que la generalización visual bajo estas condiciones es equivalente a un cambio en la política de estado inducida $\pi_\xi$. Al marginalizar el kernel de observación, la política de píxeles $\pi(a|o)$ induce una distribución de acciones condicionada al estado $\pi_\xi(a|s)$. El banco de pruebas demuestra que las brechas de rendimiento son exactamente las diferencias entre estas políticas inducidas actuando sobre el mismo MDP latente.
• Métricas: Más allá del retorno episódico estándar, el banco de pruebas reporta métricas a nivel de trayectoria: distancia recorrida, progreso normalizado hacia la meta y éxito binario de la tarea. Esta distinción es vital, ya que el retorno por sí solo puede ocultar fallos donde el agente realiza un progreso hacia adelante pero no logra completar la tarea.

Contribuciones Clave

1. KAGE-Env: Un entorno de RL de alto rendimiento nativo de JAX, con 93 parámetros visuales controlables, capaz de simulación vectorizada de hasta 33M de pasos/segundo en una sola GPU.
2. KAGE-Bench: Un protocolo de evaluación que aísla los cambios de distribución visual por construcción, presentando seis suites (Agente, Fondo, Distractores, Efectos, Filtros, Diseño) y 34 pares específicos de entrenamiento-evaluación.
3. Diagnóstico Empírico: Una evaluación sistemática de la generalización visual utilizando un baseline estándar PPO-CNN, cuantificando cómo diferentes ejes visuales impactan el rendimiento e identificando clases específicas de cambios que inducen una degradación severa.

Resultados

Utilizando un baseline estándar PPO-CNN entrenado desde frames RGB únicos, los autores observaron fallos dependientes de los ejes:

• Degradación Severa: Los Filtros (ej. cambio de tono, saturación, brillo) y los Efectos (ej. intensidad de iluminación, cantidad) causaron las colapsos de rendimiento más significativos. Las tasas de éxito cayeron un 86.8% y un 80.5% respectivamente, a menudo mientras el agente mantenía cierto movimiento hacia adelante.
• Cambios de Fondo: Los cambios en la imagen o el color del fondo provocaron una caída del 53.3% en la tasa de éxito y una caída del 30.5% en la distancia recorrida, indicando que los cambios de fondo perjudican tanto la locomoción como la completitud de la tarea.
• Cambios Benignos: Los cambios en la apariencia del Agente (ej. cambiar el sprite de un esqueleto a un payaso) fueron comparativamente benignos, causando solo una caída del 21.1% en la tasa de éxito.
• La Brecha "Movimiento vs. Completitud": Varios cambios (notablemente Distractores y Diseño) preservaron el movimiento hacia adelante (brechas de distancia pequeñas de ~3-4%) pero causaron caídas masivas en la completitud de la tarea (brechas de Tasa de Éxito de 30-60%). Esto resalta que las métricas basadas en el retorno pueden enmascarar fallos de generalización donde un agente se mueve pero no logra resolver la tarea.
• Baselines de Robustez: Los autores probaron varios métodos de robustez (RAD-PPO con cutout-color, crop-resize, CURL, DrQ, SVEA). Aunque RAD-PPO con cutout-color mostró cierta mejora en configuraciones difíciles, los métodos que dependen de aumentos geométricos (crop-resize) fallaron en converger, probablemente porque interrumpieron las posiciones espaciales absolutas requeridas para la precisión en el tiempo de salto y la evasión de obstáculos en este entorno de plataformas.

Significado y Reivindicaciones

El artículo afirma que KAGE-Bench proporciona un marco rápido, reproducible y diagnóstico para evaluar el RL basado en píxeles bajo variación visual controlada.

• Atribución: Al factorizar el proceso de observación, este banco de pruebas permite la atribución precisa de los fallos de generalización a factores visuales específicos, eliminando los efectos de confusión de la dinámica entrelazada o los cambios en la estructura de la tarea encontrados en otros bancos de pruebas.
• Escalabilidad: La implementación nativa de JAX aborda el cuello de botella computacional de la evaluación de la generalización visual, haciendo factible la realización de barridos exhaustivos sobre los parámetros visuales.
• Utilidad Diagnóstica: Los resultados demuestran que la dificultad de la generalización visual no es uniforme; varía sustancialmente entre los ejes. El banco de pruebas sirve como una herramienta para identificar qué factores visuales específicos (ej. cambios fotométricos frente a la apariencia del agente) son más críticos para que una arquitectura de política aprenda las invarianzas necesarias.

Los autores posicionan a KAGE-Bench no como un proxy exhaustivo para todos los desafíos del mundo real del RL, sino como una herramienta diagnóstica controlada para apoyar un análisis más sistemático de la robustez visual y guiar el desarrollo de algoritmos capaces de manejar cambios visuales de "eje conocido" específicos.