Improving Zero-Shot Offline RL via Behavioral Task Sampling

Aquí se presenta un resumen técnico detallado del artículo "Mejora del RL Offline Zero-Shot mediante Muestreo de Tareas Conductuales".

1. Enunciado del Problema

El Aprendizaje por Refuerzo Offline Zero-Shot (ZSRL) tiene como objetivo entrenar agentes que puedan optimizar funciones de recompensa no vistas sin interactuar adicionalmente con el entorno, confiando únicamente en un conjunto de datos offline pre-recopilado.

• Enfoque Estándar: Los métodos actuales más avanzados (por ejemplo, Características de Sucesor, representaciones Forward-Backward) aprenden una incrustación de estado $\phi(s)$ y entrenan una política condicionada a un vector de tarea $z$. La función de recompensa se define linealmente como $R_z(s) = \phi(s)^\top z$.
• El Defecto: Para entrenar estas políticas, los algoritmos existentes suelen muestrear vectores de tarea $z$ uniformemente desde una hipersfera unitaria de alta dimensión ($S^{d-1}$).
• El Problema Central: Los autores argumentan que este muestreo uniforme crea una discrepancia entre la geometría del espacio de tareas y la física del entorno. En espacios de alta dimensión, el fenómeno de "concentración de la medida" hace que los vectores muestreados uniformemente sean casi ortogonales a los comportamientos reales que el agente puede producir (el espacio conductual). Esto da lugar a una "dilución de la señal", donde la señal de recompensa se vuelve insignificante para casi todos los comportamientos, haciendo imposible que el agente distinga las acciones óptimas de las subóptimas. Esto resulta en una generalización zero-shot deficiente.

2. Metodología: Distribución de Tareas Conductuales (BTD)

Los autores proponen reemplazar el muestreo uniforme de tareas con un enfoque basado en datos llamado Muestreo de Tareas Conductuales.

A. Perspectiva Teórica

• Ocupación de Características: Definen la ocupación de características $\psi_\pi$ de una política como la suma descontada de las características de estado a lo largo de su trayectoria.
• Varianza Desvaneciente: Demuestran (Proposición 4.1) que a medida que aumenta la dimensión latente $d$, la varianza esperada de los retornos entre políticas bajo muestreo de tareas uniforme converge a cero. Esto significa que la señal de aprendizaje colapsa, impidiendo una optimización efectiva de la política.

B. El Algoritmo Propuesto

En lugar de muestrear $z$ desde una distribución uniforme, el método extrae vectores de tarea directamente del conjunto de datos offline:

1. Extracción: Del conjunto de datos offline $D$, los autores extraen sub-trayectorias $\tau$. Para cada sub-trayectoria, calculan la ocupación de características empírica $\tilde{\psi}_\tau = \sum \gamma^t \phi(s_t)$.
2. Definición del Vector de Tarea: El vector de tarea para una trayectoria se define como la ocupación normalizada: $z_\tau = \tilde{\psi}_\tau / \|\tilde{\psi}_\tau\|_2$. Este vector representa la dirección en el espacio de características que maximiza la recompensa para ese comportamiento observado específico.
3. Modelado de la Distribución: El conjunto de todos los vectores extraídos forma una distribución empírica $p_{data}$. Los autores ajustan un modelo de densidad paramétrico (específicamente un Modelo de Mezcla Gaussiana o GMM) a estos datos, creando la Distribución de Tareas Conductuales (BTD), denotada como $p_\theta(z)$.
4. Entrenamiento: Durante el entrenamiento de la política, los vectores de tarea se muestrean de la BTD aprendida ($z \sim p_\theta(z)$) en lugar de la hipersfera unitaria. Esto asegura que el agente se entrena en tareas que son físicamente realizables y están fundamentadas en la dinámica del entorno.

3. Contribuciones Clave

1. Identificación de la Dilución de la Señal: El artículo demuestra teórica y empíricamente que el muestreo uniforme de tareas en espacios de alta dimensión conduce al colapso de la señal de aprendizaje debido a la ortogonalidad con el espacio conductual.
2. Distribución de Tareas Conductuales (BTD): Un marco novedoso y agnóstico al método que extrae y modela la distribución de tareas realizables directamente desde datos offline.
3. Integración: El método es un reemplazo directo para el paso de muestreo de tareas en algoritmos ZSRL existentes (Características de Sucesor y Forward-Backward) sin modificar los objetivos subyacentes de aprendizaje de representaciones.
4. Evaluación Exhaustiva: Experimentos extensos a través de múltiples puntos de referencia, métodos de aprendizaje de representaciones y dimensiones latentes.

4. Resultados Experimentales

Los autores evaluaron su método en puntos de referencia estándar de ZSRL (Cheetah, Walker, Quadruped) utilizando conjuntos de datos recopilados mediante Distilación de Redes Aleatorias (RND). Compararon contra siete líneas base (incluyendo Autoencoders, Modelos de Transición, LRA y BYOL) en varias dimensiones latentes.

• Ganancia de Rendimiento: El muestreo BTD mejoró el rendimiento zero-shot en un promedio del 20% en todos los entornos y líneas base.
• Robustez ante la Dimensionalidad:
  • Fallo de la Línea Base: A medida que aumentaba la dimensión latente (hasta 1000), los métodos de muestreo uniforme estándar sufrieron una caída brusca en el rendimiento, colapsando a menudo a un rendimiento cercano al aleatorio.
  • Estabilidad de BTD: El muestreo BTD mantuvo un alto rendimiento incluso en dimensiones altas (por ejemplo, 1000), demostrando que mitiga eficazmente el problema de la dilución de la señal.
• Varianza Reducida: El muestreo BTD redujo significativamente la varianza en el rendimiento de prueba entre diferentes métodos de aprendizaje de representaciones, indicando un aprendizaje más consistente y fiable.
• Estudios de Ablación:
  • Estrategias de Mezcla: Mezclar el muestreo uniforme con el muestreo BTD (interpolando entre ambos) degradó el rendimiento, confirmando que introducir tareas arbitrarias e inalcanzables daña la señal de aprendizaje.
  • Componentes GMM: El rendimiento se estabilizó después de ~20 componentes GMM, sugiriendo que el espacio de tareas conductuales puede capturarse con una distribución relativamente simple.
  • Visualización: Las proyecciones UMAP mostraron que el muestreo uniforme cubre toda la hipersfera, mientras que las tareas conductuales reales ocupan una variedad pequeña y específica, validando la necesidad de un muestreo basado en datos.

5. Significado y Conclusión

Este trabajo cambia fundamentalmente el paradigma del RL offline ZSOL de la generación arbitraria de tareas al fundamento conductual.

• Impacto Práctico: Destaca que, en entornos offline, qué tareas se eligen para el entrenamiento es tan crítico como cómo la política las aprende.
• Escalabilidad: Al desacoplar la distribución de tareas de la maldición de la dimensionalidad inherente al muestreo uniforme, el método permite el entrenamiento de agentes generalistas en espacios latentes de alta dimensión.
• Generalizabilidad: El enfoque es compatible con diversas técnicas de aprendizaje de representaciones (SF, FB, Autoencoders, etc.), lo que lo convierte en un camino escalable hacia agentes de aprendizaje por refuerzo más robustos y generalizables para su implementación en el mundo real, donde las tareas están restringidas físicamente.

En resumen, el artículo argumenta que para lograr una verdadera generalización zero-shot en RL offline, los agentes deben entrenarse en tareas que reflejen las capacidades y dinámicas reales del entorno, en lugar de una distribución uniforme abstracta de posibilidades matemáticas.