WOMBET: World Model-based Experience Transfer for Robust and Sample-efficient Reinforcement Learning

1. Planteamiento del Problema

El Aprendizaje por Refuerzo (RL) en robótica enfrenta dos limitaciones críticas: el alto costo y el riesgo de la recolección de datos en el mundo real, y la ineficiencia de las muestras (sample inefficiency).

• RL Online: Requiere interacciones extensas y reseteos frecuentes, lo cual es peligroso o costoso en entornos físicos.
• RL Offline: Evita la exploración insegura pero su rendimiento se degrada cuando los datos son limitados o no coinciden con la distribución de la tarea objetivo (desplazamiento de distribución).
• RL Offline-a-Online: Combina preentrenamiento con datos históricos y ajuste fino (fine-tuning) en línea. Sin embargo, la mayoría de los métodos existentes asumen la existencia de un conjunto de datos offline de alta calidad y predefinido, sin abordar cómo generar datos fiables para la transferencia ni cómo gestionar la incertidumbre del modelo.

El objetivo de WOMBET es crear un marco unificado que no solo utilice datos previos, sino que genere activamente experiencia transferible y confiable a partir de una tarea fuente para acelerar el aprendizaje en una tarea objetivo.

2. Metodología: WOMBET

WOMBET (World Model-based Experience Transfer) es un marco que acopla la generación de datos basada en modelos con el ajuste fino en línea mediante una co-evolución iterativa del modelo y la política. El proceso se divide en las siguientes etapas clave:

A. Generación de Datos Basada en Modelos (Fuente)

En lugar de asumir un dataset fijo, WOMBET genera un dataset offline ($D_S$) en la tarea fuente utilizando un Modelo del Mundo (un ensemble probabilístico de dinámicas).

• Planificación con Penalización de Incertidumbre: Se utiliza Control Predictivo Basado en Modelos (MPC) para generar trayectorias. Para evitar la sobre-optimización (exploit de errores del modelo), se introduce una penalización de incertidumbre epistémica en la función de recompensa:
  $$\tilde{r}(s, a) = r(s, a) - \lambda u(s, a)$$
  Donde $u(s, a)$ es la incertidumbre estimada por el ensemble. Esto garantiza que la planificación maximice una cota inferior probada del retorno verdadero.

B. Filtrado de Doble Criterio (Dual-Criterion Filtering)

No todas las trayectorias generadas por el modelo son útiles. WOMBET aplica un filtro estricto para construir $D_S$:

1. Alto Retorno: La trayectoria debe tener un retorno acumulado superior a un umbral ($J_{th}$).
2. Baja Incertidumbre: La incertidumbre promedio a lo largo de la trayectoria debe estar por debajo de un umbral ($u_{th}$).
   Esto elimina trayectorias sesgadas o fuera de distribución (OOD), resultando en un dataset compacto y de alto valor. Las trayectorias aceptadas se reetiquetan con la recompensa de la tarea objetivo ($r_T$).

C. Entrenamiento Híbrido y Muestreo Adaptativo

En la tarea objetivo, se realiza un ajuste fino en línea utilizando una mezcla de datos offline ($D_S$) y datos online ($D_T$).

• Muestreo Adaptativo: La proporción de datos offline vs. online ($\alpha$) no es fija. Se ajusta dinámicamente en función del Error Temporal Diferencial (TD Error) del crítico.
  • Si el error TD es alto (indicando subajuste o inestabilidad), se aumenta la ponderación de datos offline para estabilizar el entrenamiento.
  • A medida que el error disminuye, se priorizan los datos online para adaptarse a la dinámica específica de la tarea objetivo.
• Regularización Implícita: Para evitar la explosión de valores en regiones OOD, el algoritmo utiliza:
  • LayerNorm: En las capas ocultas de los críticos para limitar la magnitud de las activaciones.
  • Mínimo del Ensemble: Se utiliza el mínimo de los valores predichos por un ensemble de críticos (pessimismo) para reducir la sobreestimación.

3. Contribuciones Clave

1. Generación y Transferencia Acopladas: WOMBET no asume un dataset offline fijo; genera y refina iterativamente los datos de transferencia utilizando un modelo del mundo que se actualiza constantemente con la experiencia acumulada.
2. Curación de Datos Confiables: Propone un filtro de doble criterio (retorno alto + incertidumbre baja) para seleccionar trayectorias, asegurando que el dataset inicial sea robusto y libre de sesgos del modelo.
3. Muestreo Adaptativo Teóricamente Fundado: Introduce una estrategia de mezcla de datos basada en el error del crítico, que equilibra la varianza y el sesgo.
4. Validación Teórica y Empírica:
   • Demuestra que la planificación con penalización de incertidumbre maximiza una cota inferior del retorno verdadero.
   • Deriva una descomposición de error de muestra finita que captura el desajuste de distribución y el error de aproximación.

4. Resultados Experimentales

Los experimentos se realizaron en diversos entornos de control continuo (MuJoCo):

• Eficiencia de Muestras: WOMBET supera significativamente a algoritmos de RL online puro (SAC, PPO, TD3) entrenados desde cero, logrando retornos asintóticos comparables o mejores con menos de la mitad del presupuesto de interacción.
• Necesidad de Ajuste Fino: Se demostró que los métodos puramente offline (sin ajuste en línea) fallan cuando hay un desplazamiento moderado entre las tareas fuente y objetivo. WOMBET, al combinar preentrenamiento con ajuste en línea, adapta la política a nuevas recompensas y dinámicas.
• Comparativa con Baselines: WOMBET supera a otros métodos de estado del arte en RL offline-a-online (como IQL, CQL, MOPO) que utilizan el mismo dataset generado. Esto confirma que la mejora proviene del diseño de transferencia (filtrado y muestreo adaptativo) y no solo de los datos.
• Ablación: La eliminación del filtrado de doble criterio o del muestreo adaptativo degrada el rendimiento, confirmando que ambos componentes son esenciales para la estabilidad y la eficiencia.

5. Significado e Impacto

WOMBET representa un avance significativo al cerrar la brecha entre el RL basado en modelos (eficiente en muestras pero propenso a sesgos) y el RL libre de modelos (robusto pero ineficiente).

• Robustez: Al integrar la incertidumbre en la generación de datos y el entrenamiento, el método mitiga el riesgo de políticas inseguras derivadas de errores de modelado.
• Aplicabilidad Real: Al no depender de demostraciones expertas costosas ni de datasets fijos de alta calidad, WOMBET ofrece un marco viable para la transferencia de aprendizaje en robótica real, donde la recolección de datos es limitada y riesgosa.
• Teoría: Proporciona garantías teóricas sobre la cota inferior del retorno y la descomposición de errores, ofreciendo una base sólida para el desarrollo futuro de algoritmos de transferencia de experiencia.

En resumen, WOMBET demuestra que la optimización conjunta de la generación de datos y la transferencia de aprendizaje es crucial para lograr un RL robusto y eficiente en muestras para aplicaciones robóticas complejas.