Recovering Hidden Reward in Diffusion-Based Policies

1. Planteamiento del Problema

Las políticas basadas en difusión se han convertido en un paradigma de vanguardia para aprender habilidades de manipulación robótica a partir de demostraciones expertas, debido a su capacidad para modelar distribuciones de acciones complejas y multimodales. Sin embargo, estas políticas suelen depender del Clonamiento de Comportamiento (BC), que aprende a imitar acciones sin modelar explícitamente la intención subyacente ni la estructura de recompensa.

• Limitaciones de la Difusión Estándar: Carecen de una representación explícita de energía, lo que las hace propensas a fallar cuando las situaciones en el momento de la prueba se desvían de la distribución de entrenamiento (Fuera de Distribución o OOD). No pueden proporcionar fácilmente una señal de recompensa para el Aprendizaje por Refuerzo (RL) posterior sin recurrir a un entrenamiento adversarial inestable o a un Aprendizaje por Refuerzo Inverso (IRL) computacionalmente costoso.
• La Brecha: Existe la necesidad de un marco que unifique el modelado generativo de acciones con la recuperación de recompensas, permitiendo tanto una imitación de alta calidad como la extracción de una señal de recompensa válida para el refinamiento de la política, todo ello garantizando teóricamente la generalización.

2. Metodología: ENERGYFLOW

Los autores proponen ENERGYFLOW, un marco que parametriza una función de energía escalar $E_\theta(s, a)$ cuyo gradiente sirve como campo de desruido para la generación de acciones.

Arquitectura Central

• Parametrización de Energía Escalar: A diferencia de las políticas de difusión estándar que regresan directamente una predicción de ruido vectorial $\epsilon_\theta(s, a)$, ENERGYFLOW aprende una función de energía escalar $E_\theta(a, s, t)$.
• Campo Vectorial Conservativo: La función de puntuación (campo de desruido) se deriva mediante diferenciación automática:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  Por construcción, esto asegura que el campo vectorial sea conservativo (sin rotacional, $\nabla \times S = 0$), lo que significa que es el gradiente de un potencial escalar. Esto impone la integrabilidad, asegurando que las preferencias de acción sean transitivas y físicamente realizables.
• Entrenamiento: El modelo se entrena utilizando Emparejamiento de Puntuación de Desruido (DSM). El objetivo minimiza la diferencia entre el gradiente aprendido y la dirección real del ruido:
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• Inferencia: Las acciones se generan integrando la EDO de flujo de probabilidad utilizando el gradiente de la función de energía aprendida.

Extracción de Recompensa

El marco extrae una señal de recompensa directamente del paisaje de energía aprendido. Para manejar la constante de integración dependiente del estado $c(s)$ (que hace que los valores de energía crudos sean ambiguos entre diferentes estados), los autores proponen Moldeado Centrado:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
Esto resta una línea base dependiente del estado (estimada mediante muestreo de Monte Carlo), normalizando la recompensa para reflejar solo las preferencias de acción relativas dentro de un estado.

3. Contribuciones Teóricas Clave

El artículo proporciona demostraciones teóricas rigurosas que establecen la conexión entre el emparejamiento de puntuación de difusión y la recuperación de recompensas:

1. Equivalencia Puntuación-Recompensa: Bajo el supuesto de Optimalidad de Entropía Máxima, el artículo demuestra que la función de puntuación aprendida por los modelos de difusión es proporcional al gradiente de la función Q suave del experto ($\nabla_a Q^* \propto S^*$). Por lo tanto, el emparejamiento de puntuación recupera implícitamente el gradiente de la recompensa sin optimización min-max adversarial.
2. Reducción de Complejidad mediante Restricciones Conservativas: Los autores demuestran que restringir el espacio de hipótesis a campos conservativos (gradientes de escalares) reduce estrictamente la complejidad de Rademacher en comparación con los campos vectoriales sin restricciones. Específicamente, el límite de complejidad para campos conservativos escala con la suavidad de la representación ($L$), mientras que los campos sin restricciones escalan con la dimensión de salida ($\sqrt{d}$).
3. Límites de Generalización OOD: Aprovechando la complejidad reducida, el artículo deriva límites de generalización más ajustados para escenarios Fuera de Distribución (OOD). La restricción conservativa actúa como un regularizador geométrico, evitando que el modelo aprenda direcciones vectoriales arbitrarias y forzándolo a extrapolar la forma del paisaje de energía, lo que conduce a una mejor robustez bajo cambios de distribución.
4. Identificabilidad y Propagación de Errores: El artículo establece que, aunque el valor Q absoluto no puede identificarse de manera única (debido a la constante desconocida dependiente del estado), la clasificación de acciones dentro del estado es exacta. Además, acota la propagación de errores de estimación de puntuación a las preferencias de acción, mostrando que los errores escalan linealmente con la distancia entre acciones (continuidad Lipschitz).

4. Resultados Experimentales

Los autores evaluaron ENERGYFLOW en los puntos de referencia RoboMimic y Meta-World, así como en tareas de robots del mundo real.

• Rendimiento de Imitación (RQ1): ENERGYFLOW logró un rendimiento de vanguardia (SOTA).
  • RoboMimic: 93.8% de tasa de éxito promedio (vs. 91.2% para Política de Difusión).
  • Meta-World: 92.5% de tasa de éxito promedio (vs. 90.7% para Política de Difusión).
  • Superó significativamente a otros métodos basados en energía (como BC Implícito) y a las líneas base de IRL.
• Despliegue en el Mundo Real (RQ2): Desplegado en un robot AGIBOT G1 para las tareas "Botella" y "Cajón". ENERGYFLOW logró una tasa de éxito del 100% con trayectorias suaves y menos vacilaciones en comparación con las líneas base.
• Calidad de la Recompensa para RL Posterior (RQ3): Cuando se utilizó como señal de recompensa para entrenar un agente Soft Actor-Critic (SAC):
  • La recompensa "Energía Centrada" permitió que SAC alcanzara un rendimiento cercano al oráculo, superando significativamente a las recompensas dispersas y a las recompensas de energía cruda.
  • Proporcionó una guía densa paso a paso que evitó los estancamientos tempranos en el entrenamiento.
• Generalización OOD (RQ4): Bajo perturbaciones crecientes en la posición inicial, ENERGYFLOW mantuvo tasas de éxito robustas donde las políticas de difusión y flujo sin restricciones se degradaron significativamente. Esto confirma el beneficio teórico de la restricción conservativa.
• Eficiencia (RQ6): ENERGYFLOW alcanzó velocidades de inferencia comparables a las Políticas de Flujo no basadas en energía (~10ms de latencia) mientras proporcionaba los beneficios de una función de energía explícita, evitando los tiempos de muestreo lentos de los Modelos Basados en Energía (EBM) tradicionales.

5. Significado e Impacto

• Unificación de Imitación y RL: ENERGYFLOW cierra la brecha entre el Clonamiento de Comportamiento y el Aprendizaje por Refuerzo. Permite que una política se entrene mediante imitación mientras expone simultáneamente una señal de recompensa válida para un refinamiento posterior, eliminando la necesidad de diseño manual de recompensas o IRL adversarial inestable.
• Garantías Teóricas para la Difusión: El trabajo proporciona la primera justificación teórica de por qué imponer un campo conservativo (integrabilidad) en las políticas de difusión es beneficioso, vinculándolo a una complejidad de hipótesis reducida y una mejor generalización OOD.
• Utilidad Práctica: El método es computacionalmente eficiente, escalable a espacios de acción de alta dimensión y efectivo en entornos robóticos del mundo real, ofreciendo una solución robusta para aprender habilidades complejas de manipulación a partir de demostraciones limitadas.

En resumen, ENERGYFLOW demuestra que las restricciones estructurales requeridas para una extracción válida de recompensas (campos conservativos) sirven simultáneamente como sesgos inductivos poderosos que mejoran la generalización y la robustez de las políticas basadas en difusión.