Resumen Técnico: Uso de la Incertidumbre de la Recompensa para Inducir Comportamientos Diversos en el Aprendizaje por Refuerzo

Planteamiento del Problema

El aprendizaje por refuerzo (RL) clásico suele buscar una política determinista que maximice la suma esperada de una recompensa escalar. Si bien este supuesto garantiza la existencia de una política determinista óptima, presenta limitaciones significativas en aplicaciones modernas como el ajuste fino de modelos de lenguaje de gran tamaño (LLM) y el descubrimiento científico. En estos dominios, la convergencia hacia un único comportamiento suele ser indeseable. Por ejemplo, preservar la diversidad de las políticas preentrenadas es crucial para tareas creativas y el razonamiento de cadena de pensamiento (chain-of-thought). Además, cuando los modelos de recompensa son aproximaciones imperfectas de las preferencias reales subyacentes, la sobreoptimización hacia una única política determinista puede conducir a la degradación de la calidad de la generación (reward hacking).

Los remedios existentes, como la regularización de entropía o los bonos de diversidad, intentan imponer estocasticidad pero sufren de desventajas críticas:

1. Compromisos Frágiles: La regularización de la entropía fuerza un compromiso donde el aumento de la estocasticidad tiene como costo la reducción de las recompensas esperadas.
2. Ordenamiento Incorrecto de Políticas: Los bonos de diversidad a menudo resultan en clasificaciones indeseables donde se prefieren políticas subóptimas con alta diversidad sobre políticas con mayores recompensas promedio.
3. Inexpresividad: Muchos objetivos, incluyendo el gradiente de política vanilla y los retornos escalares esperados (ESR), poseen inherentemente políticas deterministas óptimas, fallando al inducir la diversidad necesaria.
4. Desafíos de Optimización: Las formulaciones de retorno escalarizado (SER) a menudo carecen de estimadores de gradiente basados en muestras que sean insesgados para objetivos no lineales.

Los autores argumentan que la diversidad debe entenderse no como una restricza artificial, sino como la respuesta racional a la incertidumbre en la función de recompensa. Cuando la función de recompensa no se conoce perfectamente (debido a preferencias ambiguas o modelos de recompensa imperfectos), comprometerse con una única acción puede ser subóptimo.

Metodología: ROSA

El artículo propone una reformulación fundamental del objetivo de RL reemplazando la recompensa escalar por una distribución sobre funciones de recompensa y aplicando un objetivo no lineal sobre conjuntos de acciones. Este marco se denomina Randomized Objectives, Set Actions (ROSA).

El Criterio ROSA+Max

Centrándose en el entorno de bandidos contextuales, los autores definen el objetivo principal, ROSA+Max. Sea $\rho$ una distribución sobre funciones de recompensa $R$, y sea $\mathbf{Y} = (Y_i)_{i=1}^n$ un multiconjunto de $n$ acciones muestreadas i.i.d. de una política $\pi_\theta$. El objetivo es:

$$J_{\text{ROSA+Max}}(\pi_\theta) = \mathbb{E}_{X \sim \mu, \mathbf{Y} \sim \pi_\theta} \left[ \mathbb{E}_{R \sim \rho} \left[ \max_{1 \le i \le n} R(X, Y_i) \right] \right]$$

Intuitivamente, este criterio agrega el rendimiento tomando la recompensa máxima alcanzada dentro de un conjunto de acciones muestreadas para cada realización de la función de recompensa muestreada. Esto incentiva a la política a asegurar que al menos una acción en el conjunto funcione bien para cada realización posible de la función de recompensa, en lugar de optimizar para la recompensa promedio.

Estimación del Gradiente

Los autores derivan un estimador de gradiente de política que es un principio, insesgado y de varianza reducida para este objetivo. Dado $m$ funciones de recompensa muestreadas y $n$ acciones muestreadas, el estimador del gradiente es:

$$\hat{g}(\theta) = \frac{1}{m} \sum_{k=1}^m \left[ \sum_{i=1}^n \left( \max_{1 \le j \le n} R_k(X, Y_j) - \max_{1 \le j \le n, j \ne i} R_k(X, Y_j) \right) \nabla_\theta \log \pi_\theta(Y_i|X) \right]$$

Este estimador asigna crédito efectivamente solo a la acción que logra la recompensa máxima para una muestra específica de la función de recompensa, con el segundo término sirviendo como una covariable de control para la reducción de la varianza.

Generalización

El marco se generaliza a una familia de funciones de conjunto $f$. Los autores demuestran que para funciones de recompensa binarias, cualquier función de conjunto $f$ que sea estrictamente creciente y estrictamente cóncava con respecto a la suma de recompensas (por ejemplo, Softmax) produce una política óptima máximamente diversa.

Contribuciones Clave

1. Marco ROSA: Introducción de una nueva familia de objetivos de RL que proporcionan un control de diversidad calibrado a través de una distribución sobre funciones de recompensa.
2. Análisis Teórico:
   • Política Óptima: Se demostró que para una distribución uniforme sobre funciones de recompensa binarias, la política única óptima de ROSA+Max es aquella que asigna la misma masa de probabilidad a todas las acciones óptimas distintas.
   • Diversidad Controlable: Se derivó una fórmula que muestra cómo las distribuciones de funciones de recompensa no uniformes inducen políticas estocásticas específicas y controlables, permitiendo que la probabilidad de muestrear respuestas correctas específicas se ajuste mediante los pesos de la distribución de recompensa.
   • Análisis de Paisaje: Se demostró que ROSA+Max evita los paisajes "planos" de la regularización de entropía y los problemas de "reordenamiento" de los bonos de diversidad, proporcionando una dirección de gradiente que aumenta simultáneamente la recompensa y la diversidad.
3. Validación Empírica: Experimentos extensos en dominios de juguete y entornos de gran escala (incluyendo el post-entrenamiento de LLMs) demostrando que ROSA induce comportamientos diversos sin sacrificar la recompensa esperada, superando a los estándares de gradiente de política y a las líneas base de bonos de diversidad.

Resultados Experimentales

El artículo evalúa ROSA en tres escenarios principales:

1. Preferencias Conflictivas: En una tarea con dos funciones de recompensa contradictorias (ej. paridad vs. anti-paridad), el gradiente de política vanilla falla debido a la cancelación de la recompensa. ROSA aprende con éxito una política que equilibra ambas preferencias, generando salidas válidas diversas.
2. Múltiples Respuestas Correctas: En tareas de razonamiento matemático donde múltiples longitudes de respuesta son válidas, las políticas entrenadas con ROSA muestrean de todas las preferencias de longitud especificadas, mientras que la política de gradiente estándar colapsa a un único rango de longitud.
3. Recompensas de Verdad Fundamental Desconocidas: En entornos con modelos de recompensa ruidosos (simulados por un ensamble de jueces con precisión variable), ROSA demuestra una robustez superior. Captura la incertidumbre en la función de recompensa distribuyendo la masa de probabilidad a través de los modos de recompensa en proporción a sus pesos, mientras que la política de gradiente con regularización de entropía tiende a sobreajustarse a la recompensa más probable o a explorar demasiado ampliamente hacia acciones subóptimas.

En evaluaciones fuera de la distribución (AIME 2025), los modelos entrenados con ROSA en datos de MATH con ruido lograron tasas de pass@k más altas en comparación con el PG estándar y las líneas base de ensambles pesimistas.

Significado y Reivindicaciones

El artículo afirma que ROSA ofrece una alternativa robusta y teóricamente fundamentada para tareas complejas de RL donde las formulaciones tradicionales fallan al inducir la amplitud deseada del comportamiento del agente. Al replantear la diversidad como una respuesta racional a la incertidumbre de la recompensa, el marco resuelve la tensión entre la maximización de la recompensa y la entropía de la política.

Los autores enfatizan que ROSA:

• Generaliza naturalmente tanto el gradiente de política vanilla como los enfoques recientes de conjuntos de acciones (como best-of-$n$).
• Proporciona un mecanismo para controlar la diversidad a través de la distribución de la función de recompensa en lugar de recurrir al ajuste frágil de hiperparámetros de coeficientes de entropía.
• Asegura que la solución óptima sea una política estocástica que refleje la distribución de las posibles funciones de recompensa, evitando el colapso hacia un único comportamiento determinista incluso cuando el modelo de recompensa es imperfecto.

En última instancia, este trabajo proporciona una base unificada para entrenar políticas estocásticas que sean tan diversas y adaptables como los entornos que habitan, particularmente en dominios como el ajuste fino de LLMs y el descubrimiento científico donde existen múltiples soluciones válidas.