EXPO: Stable Reinforcement Learning with Expressive Policies

1. Planteamiento del Problema

El artículo aborda el desafío de ajustar finamente clases de políticas expresivas (como Políticas de Difusión y Políticas de Ajuste de Flujo) utilizando Aprendizaje por Refuerzo (RL) en línea dado un conjunto de datos offline.

• El Conflicto Central: Si bien las políticas expresivas sobresalen en capturar distribuciones de comportamiento complejas y multimodales mediante aprendizaje por imitación (IL), son difíciles de optimizar directamente con RL. A diferencia de las políticas Gaussianas simples, las políticas expresivas se parametrizan mediante cadenas de eliminación de ruido largas. La retropropagación de gradientes de valor a través de estas cadenas para actualizar los parámetros de la política es computacionalmente costosa e inestable, lo que a menudo conduce al colapso del entrenamiento o a una convergencia deficiente.
• La Brecha: Los métodos existentes de RL en línea están diseñados para políticas Gaussianas simples y no logran aprovechar la capacidad de los modelos expresivos preentrenados. Por el contrario, los métodos que intentan ajustar finamente políticas expresivas a menudo luchan con la estabilidad o requieren modificaciones complejas en el proceso de eliminación de ruido.
• Objetivo: Desarrollar un algoritmo de RL en línea eficiente en muestras y estable que pueda ajustar finamente políticas expresivas preentrenadas utilizando datos offline sin optimizar directamente los parámetros de la política expresiva contra la función de valor.

2. Metodología: EXPO (Optimización de Política Expresiva)

Los autores proponen EXPO, un marco que desacopla la política expresiva de la maximización directa del valor. En lugar de retropropagar gradientes a través de la política expresiva, EXPO construye una política "sobre la marcha" (OTF) que maximiza los valores Q utilizando dos componentes distintos:

A. Arquitectura de Política de Dos Componentes

1. Política Expresiva Base ($\pi_{base}$):
   • Función: Genera candidatos de acción iniciales basados en la distribución de comportamiento.
   • Entrenamiento: Entrenada utilizando un objetivo estable de Aprendizaje por Imitación (IL) (por ejemplo, Modelos Probabilísticos de Difusión de Eliminación de Ruido). Nunca se entrena explícitamente para maximizar el valor Q, preservando su estabilidad y capacidad para capturar modos complejos.
2. Política de Edición Ligera ($\pi_{edit}$):
   • Función: Una pequeña política Gaussiana que toma una acción $a$ muestreada de $\pi_{base}$ y produce una edición residual $\hat{a}$ para generar una acción refinada $\tilde{a} = a + \hat{a}$.
   • Entrenamiento: Entrenada utilizando gradientes de política estándar de RL (maximizando $Q(s, a + \hat{a})$) con regularización de entropía.
   • Restricción: La edición está restringida (escalada por un hiperparámetro $\beta$) para mantenerse cerca de la acción original, asegurando que la política de edición resuelva un problema de optimización local en lugar de uno global. Esto permite un entrenamiento eficiente y estable.

B. Extracción de Política Sobre la Marcha (OTF)

La política real utilizada para el muestreo y la copia de seguridad de Diferencia Temporal (TD) no es una sola red neuronal, sino un proceso de selección dinámico:

1. Muestreo: Para un estado dado $s$, el sistema muestrea $N$ acciones de $\pi_{base}$.
2. Edición: Cada acción muestreada es refinada por $\pi_{edit}$.
3. Selección: El sistema evalúa el valor Q tanto para las acciones originales ($a_i$) como para las acciones editadas ($\tilde{a}_i$).
4. Maximización: La acción con el valor Q más alto entre todos los candidatos se selecciona para su ejecución y se utiliza como objetivo para la copia de seguridad TD.

C. Copia de Seguridad de Entropía para Regímenes con Datos Limitados

En escenarios con conjuntos de datos offline pequeños, los autores introducen una copia de seguridad de entropía. Dado que las políticas expresivas (como la difusión) carecen de una entropía de forma cerrada, construyen una distribución de muestreo suave sobre los $N$ candidatos de acción base + editados, ponderados por sus valores Q. Esto permite incluir una bonificación de entropía en la pérdida de RL para fomentar la exploración sin requerir cálculos analíticos de entropía.

3. Contribuciones Clave

1. Estrategia Estable de Maximización de Valor: El artículo introduce un enfoque novedoso para evitar la propagación directa de gradientes a través de cadenas de eliminación de ruido expresivas. Al delegar la maximización de valor a una política de edición ligera y un mecanismo de selección, EXPO logra estabilidad donde la optimización directa falla.
2. Eficiencia en Muestras: EXPO logra una mejora de 2 a 3 veces en la eficiencia de muestras en comparación con los métodos anteriores más avanzados tanto en configuraciones de RL en línea como de ajuste fino de offline a online.
3. Agnosticismo de la Política: El marco es agnóstico a la parametrización específica de la política base (funciona con Difusión, Ajuste de Flujo, etc.) y no requiere una función de valor preentrenada, lo que lo hace aplicable a cualquier política robótica preentrenada.
4. Copia de Seguridad TD Sobre la Marcha: Los autores demuestran que utilizar la acción que maximiza el valor para el objetivo TD (no solo para el muestreo) es crítico para la velocidad de aprendizaje, cerrando efectivamente la brecha entre actualizaciones tipo SARSA y tipo Q-learning en el contexto de políticas expresivas.

4. Resultados Experimentales

El método fue evaluado en 12 tareas de control continuo desafiantes a través de 4 dominios (AntMaze, Adroit, Robomimic, MimicGen), todos con recompensas escasas.

• RL en Línea (Sin Preentrenamiento): EXPO superó significativamente a las líneas base (incluyendo RLPD, IDQL y QSM) en eficiencia de muestras, convergiendo a menudo más rápido incluso sin aprovechar el conjunto de datos offline para preentrenamiento.
• Ajuste Fino de Offline a Online:
  • EXPO ajustó finamente con éxito políticas preentrenadas sin el colapso de rendimiento que a menudo se observa en otros métodos (por ejemplo, DAC colapsó, IDQL se estancó).
  • Mantuvo un alto rendimiento en tareas de manipulación de alta precisión (por ejemplo, enhebrar una aguja, apilar bloques) donde otros métodos lucharon por mejorar más allá del rendimiento inicial del aprendizaje por imitación.
• Estudios de Ablación:
  • Copia de Seguridad TD OTF: Eliminar la selección de maximización de valor para el objetivo TD (usando solo muestreo) redujo drásticamente el rendimiento, confirmando la importancia de la alineación del valor Q en la copia de seguridad.
  • Ediciones de Acción: Eliminar la política de edición resultó en un rendimiento deficiente, especialmente en tareas que requerían exploración más allá de la estrecha distribución de comportamiento.
  • Calidad de los Datos Offline: El rendimiento se correlacionó fuertemente con la calidad del conjunto de datos offline. Sin embargo, la variante de copia de seguridad de entropía permitió que EXPO tuviera éxito incluso con datos offline muy pobres (donde el aprendizaje por imitación por sí solo logró <10% de éxito).

5. Significado

EXPO representa un paso significativo hacia adelante en Robótica y RL al resolver el cuello de botella del "ajuste fino de políticas expresivas".

• Impacto Práctico: Permite el despliegue de grandes modelos generativos preentrenados (que son excelentes imitando datos humanos) en entornos del mundo real donde deben adaptarse y mejorar mediante interacción en línea.
• Perspectiva Teórica: Cuestiona la necesidad de retropropagar a través de cadenas generativas complejas para RL, mostrando que una estrategia de "edición local + selección global" es más estable y eficiente.
• Dirección Futura: El trabajo abre la puerta al ajuste fino de cualquier política generativa preentrenada (difusión, flujo, etc.) en configuraciones en línea, acelerando potencialmente la adopción de modelos fundacionales en robótica.

Limitaciones: El método requiere muestrear múltiples acciones para cada actualización de lote, lo cual es computacionalmente costoso. Además, asume que existe un prior razonable (conjunto de datos offline o política preentrenada); no está diseñado para aprender desde cero sin conocimiento previo.