DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

Resumen Técnico: DAWM – Modelos Mundiales de Acción por Difusión para Aprendizaje por Refuerzo Offline mediante Transiciones Inferidas por Acción

1. Planteamiento del Problema

Los modelos mundiales basados en difusión han demostrado capacidades sólidas para sintetizar trayectorias realistas de largo horizonte para el aprendizaje por refuerzo (RL) offline. Sin embargo, los métodos existentes enfrentan un problema crítico de compatibilidad con los algoritmos estándar de RL offline basados en valores (por ejemplo, TD3BC, IQL) que dependen del aprendizaje de diferencia temporal (TD) de un paso.

• Limitación de los Modelos Libres de Acción: Métodos como el Modelo Mundial de Difusión (DWM) generan estados y recompensas futuras, pero no generan acciones explícitamente. En consecuencia, no pueden construir transiciones completas $(s, a, r, s')$, lo que los hace incompatibles con el aprendizaje TD de un paso, el cual requiere consistencia entre estados, acciones y recompensas para una regularización efectiva del crítico.
• Limitación del Modelado Conjunto: Los enfoques que intentan modelar conjuntamente estados, recompensas y acciones a menudo sufren un aumento en la complejidad y la inestabilidad del entrenamiento, lo que conduce a un rendimiento práctico reducido.
• Limitación de los Métodos Basados en Planificación: Los planificadores basados en difusión (por ejemplo, Diffuser, Decision Diffuser) que recuperan acciones mediante Modelos de Dinámica Inversa (IDM) o modelado conjunto generalmente dependen de muestreo iterativo y selección de recompensas post-hoc. Esto resulta en altos costos de inferencia y elude la eficiencia del aprendizaje TD de un paso.

El desafío central es generar datos sintéticos que apoyen un aprendizaje TD de un paso estable y eficiente, sin incurrir en la inestabilidad del entrenamiento del modelado conjunto ni en los costos de inferencia de la planificación iterativa.

2. Metodología: DAWM

Los autores proponen DAWM (Modelo Mundial de Acción por Difusión), un marco modular que desacopla la generación de trayectorias de la inferencia de acciones para producir transiciones completas y etiquetadas adecuadas para RL offline. El marco consta de dos fases principales:

Fase 1: Generación Condicional de Trayectorias por Difusión

• Se entrena un modelo de difusión $p_\theta$ para generar secuencias futuras de estados y recompensas condicionadas al estado actual $s_t$, la acción $a_t$ y un retorno hasta el final ($R_{tg}$) objetivo.
• El modelo define la distribución condicional $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$.
• El objetivo de entrenamiento minimiza el error esperado entre el ruido predicho y el real utilizando un objetivo de eliminación de ruido con puntuación y guía libre de clasificadores.
• Crucialmente, este modelo genera trayectorias libres de acciones (solo estados y recompensas), simplificando la tarea generativa al evitar las complejidades de modelar acciones conjuntamente dentro del proceso de difusión.

Fase 2: Completado de Acciones mediante Modelo de Dinámica Inversa (IDM)

• Para convertir las trayectorias libres de acciones en transiciones completas, se emplea un Modelo de Dinámica Inversa (IDM), $f_\phi$, entrenado por separado.
• El IDM reconstruye las acciones faltantes modelando la distribución $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ como una distribución gaussiana diagonal.
• El IDM se entrena sobre datos offline reales para maximizar la verosimilitud de las acciones observadas dadas las transiciones de estado.
• En la inferencia, el IDM se aplica a las trayectorias generadas por difusión para inferir acciones, produciendo tuplas completas $(s, a, r, s')$.

Entrenamiento de la Política Offline

• Las transiciones completas sintetizadas se utilizan para entrenar agentes de RL offline utilizando algoritmos estándar basados en valores, específicamente TD3BC (que añade un regularizador de clonación de comportamiento a TD3) e IQL (que realiza regresión supervisada ponderada por ventaja).
• Las funciones Q se actualizan utilizando objetivos de Bellman calculados a partir de las transiciones reconstruidas, permitiendo actualizaciones TD de un paso eficientes.

3. Contribuciones Clave

1. Marco Modular: Los autores proponen un marco eficiente que combina la generación de estados-recompensa basada en difusión con un IDM entrenado por separado para la inferencia de acciones. Este diseño evita la inestabilidad del entrenamiento del modelado conjunto mientras produce transiciones completas.
2. Compatibilidad con el Aprendizaje TD de Un Paso: Al generar transiciones completas $(s, a, r, s')$, DAWM habilita el uso de algoritmos conservadores de RL offline (TD3BC, IQL) que dependen del aprendizaje TD de un paso, ofreciendo una alternativa computacionalmente eficiente a los planificadores de difusión iterativos.
3. Rendimiento Empírico: El artículo demuestra que los algoritmos conservadores de RL offline se benefician significativamente al entrenarse con estas trayectorias aumentadas, superando consistentemente a las líneas base basadas en difusión anteriores (DWM, DD) en múltiples tareas de locomoción de D4RL.
4. Viabilidad de Datos Sintéticos: Los resultados indican que los agentes entrenados con datos sintéticos generados por DAWM pueden lograr un rendimiento comparable o superior al de los agentes entrenados en conjuntos de datos offline reales, destacando el potencial de las trayectorias sintéticas como una alternativa práctica.

4. Resultados Experimentales

Los experimentos se realizaron en 9 benchmarks estándar de locomoción de D4RL (Hopper, Walker2d, HalfCheetah) con calidad de datos variable (Medium, Medium-Replay, Medium-Expert).

• Comparación con DWM: DAWM superó consistentemente a la línea base DWM libre de acciones. Con TD3BC, DAWM logró una mejora promedio del 9.3%, y con IQL, del 9.5%. Las ganancias fueron particularmente pronunciadas en tareas de control complejas como Walker2d.
• Comparación con Métodos Basados en Planificación (DD): DAWM superó a Decision Diffuser (DD) en términos de retorno normalizado. Además, DAWM demostró una eficiencia significativamente mayor, logrando una inferencia aproximadamente 4.0 veces más rápida en comparación con DD, ya que evita el muestreo iterativo durante la inferencia.
• Comparación con Datos Reales: Los agentes entrenados con datos generados por DAWM lograron retornos normalizados comparables a los agentes entrenados con datos reales utilizando TD3BC/IQL estándar, y en algunos casos (por ejemplo, Hopper-MR con TD3BC), los superaron. Aunque SSORL (una línea base que utiliza IDM en datos reales) funcionó bien, los autores señalan que SSORL se beneficia de filtrar las trayectorias más largas, un factor no presente en la generación de longitud fija de DAWM.
• Estudios de Ablación:
  • Volumen de Datos (DAWM-T): Reducir el número de pares de condicionamiento en un factor de 8 (DAWM-T) resultó en caídas de rendimiento solo modestas, lo que sugiere que el mecanismo de completado de acciones es un factor más crítico que el volumen de datos en sí.
  • Horizonte de Generación: DAWM mantuvo un rendimiento robusto a través de diferentes horizontes de generación ($H=1, 3, 7$), lo que indica que el marco de aprendizaje TD de un paso es estable independientemente de la longitud del despliegue.

5. Significado y Afirmaciones

El artículo afirma que DAWM cierra la brecha entre el modelado generativo eficiente y el aprendizaje de políticas basado en TD. Su significado principal radica en:

• Habilitar el Aprendizaje TD de Un Paso: Proporciona un mecanismo para utilizar datos generados por difusión en algoritmos estándar y estables de RL offline que anteriormente eran incompatibles con modelos mundiales libres de acciones.
• Eficiencia Computacional: Ofrece una solución práctica para el RL generativo offline a gran escala al reducir los costos de inferencia en comparación con los métodos de planificación iterativa.
• Alternativa Práctica: Los resultados sugieren que las trayectorias sintéticas generadas mediante este enfoque modular pueden servir como una alternativa viable y de alta calidad a los conjuntos de datos offline reales, mitigando potencialmente los problemas de escasez de datos en aplicaciones del mundo real donde la exploración es costosa o insegura.

Los autores concluyen que, aunque el método muestra promesa, el trabajo futuro investigará su extensión al control visual, la robótica de alta dimensión y la robustez bajo dinámicas imperfectas.