Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching

Resumen Técnico: Aprendizaje por Refuerzo de Máxima Entropía Escalable para Políticas de Difusión mediante el Emparejamiento de Adjuntos

Declaración del Problema

Las políticas de difusión han emergido como un paradigma poderoso para representar distribuciones de acciones multimodales complejas en el Aprendizaje por Refuerzo (RL). Sin embargo, su aplicación al RL online sigue limitada por el desafío de un entrenamiento escalable ante la ausencia de datos de verdad fundamental (ground-truth). Los objetivos de entrenamiento estándar, como el emparejamiento de puntuación (score matching) o el emparejamiento de puentes (bridge matching), no son directamente aplicables porque el RL online carece de un conjunto estático de muestras óptimas.

Los enfoques existentes para entrenar políticas de difusión en entornos online enfrentan tres limitaciones primarias:

1. Divergencia KL Inversa: Los métodos que minimizan la divergencia KL inversa requieren la retropropagación (backpropagation) a través de todo el proceso de difusión. Esto incurre en una sobrecarga de memoria y una ineficiencia computacional masivas, lo que impide escalar a problemas de mayor envergadura.
2. Emparejamiento con Pesado de Importancia: Los enfoques que utilizan el muestreo de importancia para minimizar los errores entre las acciones predichas y las de objetivo tienen dificultades para escalar a espacios de acción de alta dimensión.
3. Métodos de Q-Score: Las técnicas que utilizan el gradiente de la función Q ($\nabla_a Q$) a menudo dependen de la dinámica de Langevin sobreamortuada. Estos procesos suelen requerir truncamiento tras pasos finitos, ofreciendo pocas garantías teóricas en el régimen de tiempo finito, o dependen de técnicas de muestreo heurísticas que carecen de garantías para generar muestras de la verdadera distribución objetivo.

Metodología: Política de Difusión por Emparejamiento de Adjuntos (AMDP)

Los autores proponen la Política de Difusión por Emparejamiento de Adjuntos (AMDP), un marco que formula la optimización de la política de difusión para el RL de Máxima Entropía (MaxEnt-RL) como un problema de Control Óptimo Estocástico (SOC). La innovación central es la adaptación del Emparejamiento de Adjuntos (AM), específicamente el Emparejamiento de Adjuntos Recíproco (RAM), para permitir un entrenamiento libre de simulación.

1. Fundamento Teórico

El método enmarca el paso de mejora de la política como la búsqueda de una función de deriva $u$ tal que la distribución marginal de un proceso de difusión en el tiempo terminal $\tau=1$ coincida con la política óptima de MaxEnt $\pi^*(a|s) \propto \exp(Q^\pi(s,a)/\alpha)$.

• Puente de Schrödinger: El problema se plantea como un problema de puente de Schrödinger, minimizando el costo cinético del control $u$ sujeto a condiciones de contorno.
• Emparejamiento de Adjuntos Recíproco: En lugar de realizar la retropropagación a través de la trayectoria de difusión (lo que requiere almacenar estados intermedios), AMDP utiliza la proyección recíproca del proceso controlado sobre un proceso de referencia no controlado. Esto permite que la pérdida se minimice mediante el muestreo de estados intermedios $X_\tau$ directamente desde el estado terminal $X_1$ utilizando un proceso recíproco condicional $\Pi^0_{\tau|1}$.
• Entrenamiento Libre de Simulación: Al muestrear $X_\tau$ desde el estado terminal $X_1$ (que se almacena en un replay buffer), el algoritmo evita simular la SDE hacia adelante durante la actualización del gradiente. Esto elude la necesidad de una estimación explícita de la verosimilitud y la costosa retropropagación a través de la cadena de difusión.

2. Componentes Algorítmicos Clave

• Adaptación del Espacio de Acción (Aplastamiento de Función de Error):
  Para manejar espacios de acción acotados (por ejemplo, $A = [-1, 1]^d$), los autores introducen una transformación diferenciable $f: \mathbb{R}^d \to [-1, 1]^d$. A diferencia del aplastamiento estándar mediante $\tanh$ elemento a elemento, AMDP emplea un aplastamiento por función de error (erf) escalada:
  $$f_{\text{erf}}(x) = (\text{erf}(k x_i))_{i=1}^d$$
  El factor de escala $k$ se elige de tal manera que el determinante del Jacobiano de la transformación coincida exactamente con la densidad marginal Gaussiana del proceso de referencia. Esta elección específica hace que la razón de logaritmos de la densidad de referencia y el determinante del Jacobiano se convierta en una constante, simplificando significativamente el objetivo de regresión y mejorando la estabilidad numérica.

• Iteración de la Política:

  • Evaluación: Se utiliza un backup de Bellman "suave" modificado. Dado que la entropía marginal de los modelos de difusión es intratable, los autores proponen un límite inferior de entropya tratable derivado del costo de control, lo cual preserva las propiedades de convergencia.
  • Mejora: La política se actualiza minimizando la pérdida AMDP, que es una iteración de punto fijo. El punto fijo único de esta pérdida corresponde al control óptimo $u^*$.

• Actualizaciones de Región de Confianza:
  Para asegurar la estabilidad en entornos online donde la distribución objetivo cambia rápidamente, se añade una restricción de región de confianza. Esta penaliza la desviación del nuevo arrastre $u$ respecto a la iteración anterior $u_{\text{old}}$, evitando actualizaciones agresivas que podrían conducir al colapso de la distribución o a la divergencia de la función Q. Esto se resuelve mediante un formalismo Lagrangiano relajado.

Contribuciones Clave

1. Algoritmo de Entrenamiento Escalable: AMDP introduce un algoritmo altamente eficiente que optimiza políticas de difusión sin retropropagar a través del proceso de difusión, reduciendo drásticamente los requisitos de memoria.
2. Optimización Libre de Simulación: Al aprovechar el emparejamiento de adjuntos recíproco, el método permite la optimización utilizando únicamente las acciones terminales almacenadas en un replay buffer, permitiendo múltiples actualizaciones de gradiente por muestra sin re-simular la SDE.
3. Aplastamiento de Acción Numéricamente Estable: La propuesta del aplastamiento por función de error proporciona una alternativa estable a $\tanh$, simplificando el objetivo de entrenamiento al cancelar términos específicos en la función de pérdida.
4. Actualizaciones de Política Estables: La integración de una pérdida de región de confianza asegura la estabilidad de la iteración de la política en entornos online sin alterar las garantías teóricas del marco de emparejamiento de adjuntos.

Resultados Experimentales

Los autores evaluaron AMDP en 63 entornos a través de MuJoCo Playground, ManiSkill y HumanoidBench en el entorno on-policy, y en 7 entornos de alta dimensión de perros y humanoides (suite DMC) en el entorno off-policy.

• Rendimiento On-Policy:

  • AMDP iguala o supera a los fuertes baselines Gaussianos (REPPO, PPO, SPO) y a los métodos de difusión de vanguardia (DIME, DPPO, FPO).
  • En tareas complejas y de alta dimensión (HumanoidBench), AMDP supera significativamente a todos los baselines, incluso sin el muestreo Best-of-N (BoN).
  • AMDP converge más rápido que REPPO en eficiencia de muestras en tareas de locomoción de MuJoCo DMC y humanoides.
  • Cuando se combina con el muestreo BoN (seleccionando la mejor de 16 muestras), AMDP logra un rendimiento superior en todos los benchmarks.

• Rendimiento Off-Policy:

  • En las tareas de perro y humanoide de DMC, AMDP iguala el rendimiento de DIME (que optimiza la KL inversa) y converge ligeramente más rápido, demostando la competitividad del objetivo de emparejamiento de adjuntos en entornos off-policy.

• Eficiencia Computacional:

  • Tiempo de ejecución (Wall-clock Time): AMDP incurre solo en una pequeña sobrecarga constante (~10%) en comparación con métodos Gaussianos altamente eficientes como REPPO, independientemente del número de pasos de difusión.
  • Escalabilidad: En contraste, los métodos que utilizan retropropagación de la KL inversa son 10 veces más lentos con 16 pasos de difusión y hasta 72-80 veces más lentos con 128 pasos.
  • Escalamiento del Modelo: Los autores entrenaron con éxito una política de difusión grande (106M de parámetros, ~300x más grande que los modelos estándar) usando AMDP, demostrando su capacidad para escalar a modelos complejos donde los métodos basados en retropropagación fallarían debido a restricciones de memoria.

Significado y Reivindicaciones

El artículo afirma que AMDP proporciona un camino teóricamente fundamentado, computacionalmente eficiente y robusto para el despliegue de políticas generativas expresivas en aplicaciones de RL en el mundo real.

• Superación de Cuellos de Botella de Memoria: Al eliminar la necesidad de retropropagación a través de la cadena de difusión, AMDP hace que el entrenamiento de políticas de difusión sea viable para escenarios de RL online donde los recursos de memoria y computación están limitados.
• Garantías Teóricas: A diferencia de los métodos heurísticos de Q-score, AMDP se deriva del Control Óptimo Estocástico con garantías de punto fijo, asegurando que la política converja a la distribución de Máxima Entropía óptima.
• Viabilidad Práctica: La combinación de entrenamiento libre de simulación, estabilidad numérica mediante el aplastamiento por erf y actualizaciones de región de confianza aborda las barreras primarias (escalabilidad, estabilidad y eficiencia) que anteriormente habían limitado la adopción de las políticas de difusión en RL online.

Los autores concluyen que su enfoque reduce significativamente la sobrecarga computacional mientras mantiene un rendimiento competitivo, haciendo que las políticas de difusión sean una alternativa más viable a las políticas Gaussianas tradicionales para tareas de control continuo complejas.