Resumen Técnico: Muestreo Adaptativo de Tareas Robusto Distribucionalmente (DRATS)

Enunciado del Problema

El Aprendizaje por Refuerzo Multi-Tarea (MTRL) tiene como objetivo entrenar un único agente para optimizar el rendimiento en múltiples tareas simultáneamente, aprovechando la transferencia de conocimientos para mejorar la eficiencia de los datos. Sin embargo, los enfoques estándar de MTRL a menudo sufren de aprendizaje desequilibrado: los agentes resuelven rápidamente las tareas fáciles mientras progresan lentamente en las más difíciles.

Trabajos anteriores han atribuido en gran medida este desequilibrio a gradientes de tareas en conflicto, proponiendo soluciones que involucran manipulación de gradientes (por ejemplo, cirugía de gradientes) o arquitecturas de red especializadas. Este artículo identifica una causa distinta y poco explorada del desequilibrio: asignación de datos desequilibrada. El MTRL estándar asigna un presupuesto de interacción igual a todas las tareas. Dado que las tareas difíciles requieren significativamente más experiencia para resolverse que las fáciles, la asignación uniforme resulta en una sobreasignación de datos a las tareas fáciles (que necesitan pocos datos adicionales) y una subasignación de datos a las tareas difíciles. Aunque los métodos de Aprendizaje Curricular (CL) adaptan las distribuciones de muestreo, generalmente priorizan primero las tareas fáciles para facilitar la transferencia, lo que puede exacerbar el desequilibrio de datos en escenarios donde la transferencia positiva es débil o las tareas entran en conflicto.

Metodología: Muestreo Adaptativo de Tareas Robusto Distribucionalmente (DRATS)

Los autores proponen DRATS, un algoritmo que prioriza adaptativamente el muestreo de tareas que están más lejos de ser resueltas. El método se deriva a través de los siguientes pasos:

1. Formalización como un Problema de Factibilidad

En lugar de maximizar la recompensa promedio en todas las tareas, el MTRL se formaliza como un problema de factibilidad: encontrar una política $\theta$ tal que la recompensa $J_i(\theta)$ en cada tarea $i$ cumpla con un objetivo de referencia $J^{ref}_i$.
El objetivo se transforma en minimizar la brecha de recompensa $g_i(\theta) = \max(J^{ref}_i - J_i(\theta), 0)$. Para manejar casos donde las restricciones no pueden satisfacerse simultáneamente (debido a límites de capacidad o conflictos), el objetivo pasa a ser minimizar la brecha de recompensa máxima en todas las tareas (un objetivo minimax):
$$\min_{\theta} \max_{i \in [k]} g_i(\theta)$$

2. Optimización Minimax y Robustez Distribucional

La maximización interna se reencuadra como una maximización sobre una distribución de tareas $q \in \Delta_k$ (el simplex de probabilidad). Resolver esto directamente conduce a una solución no suave y codiciosa que concentra toda la probabilidad en la única tarea con peor rendimiento, causando inestabilidad e ineficiencia de datos.

Para abordar esto, DRATS emplea Optimización Robusta Distribucionalmente (DRO) restringiendo la distribución de tareas $q$ para que se encuentre dentro de una bola de Kullback-Leibler (KL) de radio $\epsilon$ centrada en una distribución base $p_0$ (uniforme). Esto produce la optimización con restricciones:
$$\min_{\theta} \max_{q \in \mathcal{Q}} \mathbb{E}_{i \sim q}[g_i(\theta)]$$
donde $\mathcal{Q} = \{q \in \Delta_k \mid \text{KL}(q \| p_0) \leq \epsilon\}$.

3. Solución de Forma Cerrada y Algoritmo

Utilizando relajación lagrangiana, la distribución de tareas óptima $q^*$ se deriva en forma cerrada como un softmax de las brechas de recompensa:
$$q^*_i = \frac{\exp(\eta g_i(\theta))}{\sum_j \exp(\eta g_j(\theta))}$$
Aquí, $\eta$ actúa como un parámetro de temperatura inversa que controla la nitidez de la priorización.

Algoritmo DRATS (Algoritmo 1):

1. Recolección de Datos: Muestrear tareas desde la distribución actual $q_t$ y recolectar trayectorias.
2. Estimación de Brecha: Calcular estimaciones de Monte Carlo de las recompensas $\hat{J}_i$ y actualizar las recompensas de referencia $J^{ref}_i$ (estimadas en línea como la recompensa máxima observada).
3. Normalización de Brecha: Calcular las brechas de recompensa normalizadas $\hat{g}_i$.
4. Actualización de Distribución: Actualizar $q_t$ mediante ascenso de espejo (gradiente exponenciado) sobre el objetivo $L(\theta, q) = \mathbb{E}_{i \sim q}[g_i(\theta)] - \frac{1}{\eta}\text{KL}(q \| p_0)$.
5. Proyección: Imponer una probabilidad mínima de muestreo $\epsilon$ para asegurar que todas las tareas se muestreen lo suficiente para una estimación estable de la brecha.
6. Actualización de Política: Actualizar la política $\theta$ utilizando un algoritmo en política (por ejemplo, PPO) para maximizar la recompensa ponderada.

4. Garantías de Convergencia

Los autores demuestran que, bajo supuestos estándar (incluyendo una garantía de bajo arrepentimiento para la actualización de la política), DRATS converge a la brecha de recompensa en el peor de los casos óptima a una tasa de $O(1/\sqrt{T})$, donde $T$ es el número de iteraciones.

Contribuciones Clave

1. Algoritmo DRATS: Introducción de un algoritmo de RL multi-tarea que prioriza adaptativamente las tareas basándose en la brecha de recompensa, derivado de una formulación minimax del MTRL.
2. Convergencia Teórica: Demostración de que DRATS converge dentro de un error $\epsilon$ de la brecha de recompensa en el peor de los casos óptima a una tasa estándar de $O(1/\sqrt{T})$.
3. Validación Empírica: Demostración de que DRATS logra una mayor eficiencia de datos, una mayor recompensa agregada y una recompensa significativamente mayor en la tarea más difícil en comparación con las líneas base existentes de aprendizaje curricular (por ejemplo, Uniforme, Progreso de Aprendizaje, Potencial de Aprendizaje, Primero Difícil).
4. Ortogonalidad: Demostración de que DRATS es compatible con arquitecturas de red multi-tarea existentes (por ejemplo, Modularización Suave, MOORE) y mejora el rendimiento independientemente de si las tareas exhiben transferencia positiva o conflicto.

Resultados Experimentales

Los autores evaluaron DRATS en cuatro puntos de referencia: un Gridworld de 4 tareas, MuJoCo6 (6 tareas de control continuo), MetaWorld-MT10 y MetaWorld-MT50.

• Rendimiento Agregado: DRATS logró la mayor recompensa agregada en todos los puntos de referencia.
• Rendimiento en la Tarea Más Difícil: La mejora fue impulsada principalmente por un rendimiento superior en las tareas más difíciles. Por ejemplo, en MuJoCo6, DRATS fue el único método que resolvió con éxito las difíciles tareas Ant y Humanoid, mientras que las líneas base no lograron progresar en Ant.
• Eficiencia de Datos: DRATS alcanzó los niveles de rendimiento de las líneas base con significativamente menos interacciones. En MT10, superó la recompensa final de la línea base Uniforme en tareas difíciles (por ejemplo, pick-place, push) después de solo 6 millones de pasos de tiempo.
• Comparación con Líneas Base:
  • Uniforme: Sobreasigna a tareas fáciles; no logra resolver tareas difíciles.
  • Progreso de Aprendizaje/Potencial: Tienden a priorizar tareas con altas tasas de aprendizaje o alto error de estimación de valor, lo que a menudo se correlaciona con tareas fáciles, descuidando las tareas difíciles que están estancadas.
  • Primero Difícil (SMT): Sufre un colapso del rendimiento en tareas difíciles; una vez que una tarea se declara "resuelta" basándose en un umbral, la probabilidad de muestreo cae abruptamente, haciendo que el agente olvide la tarea y requiera reaprender.
• Ablaciones:
  • Muestreo Adaptativo vs. Reponderación: El muestreo adaptativo (DRATS) es significativamente más eficiente en datos que reponderar el objetivo de datos muestreados uniformemente. La reponderación conduce a estimaciones de gradiente de alta varianza para tareas difíciles debido a la recolección insuficiente de datos.
  • Hiperparámetros: DRATS es robusto ante la elección de la temperatura $\eta$ y el tamaño de paso $\alpha$.

Significado y Afirmaciones

El artículo afirma que DRATS aborda una ineficiencia fundamental en el MTRL: la mala asignación de recursos de datos. Al desplazar el enfoque del conflicto de gradientes hacia la asignación de datos, DRATS proporciona un mecanismo eficiente de datos para equilibrar el rendimiento entre tareas.

Los autores enfatizan que:

• DRATS es la primera aplicación de la Optimización Robusta Distribucionalmente al RL multi-tarea en entornos de control continuo.
• A diferencia de muchos métodos de Aprendizaje Curricular que asumen transferencia positiva de tareas fáciles a difíciles, DRATS apunta explícitamente a la brecha de recompensa en el peor de los casos, haciéndolo efectivo incluso cuando las tareas entran en conflicto o la transferencia es débil.
• El método es modular; modifica únicamente el proceso de recolección de datos y puede combinarse con arquitecturas especializadas o técnicas de manipulación de gradientes.

El artículo concluye que DRATS ofrece un enfoque principiado y teóricamente fundamentado para lograr un aprendizaje multi-tarea equilibrado, mejorando tanto la eficiencia del aprendizaje como la robustez de la política final en todas las tareas.