Resumen Técnico: Dirección de Difusión por Perturbación Lagrangiana (LP-DS)

1. Planteamiento del Problema

Las políticas generativas de alta capacidad, como las basadas en modelos de difusión y de ajuste de flujo (flow-matching), se han consolidado como paradigmas potentes para el control continuo y la robótica. Estos modelos destacan por representar distribuciones de acción multimodales y expresivas aprendidas a partir de datos offline. Sin embargo, enfrentan dos limitaciones primordiales:

1. Restricciones de Clonación de Comportamiento (BC): Las políticas entrenadas mediante BC heredan las limitaciones de sus conjuntos de datos de demostración, incluyendo una cobertura incompleta del espacio de estados, la omisión de comportamientos de recuperación poco frecuentes y fragilidad ante el cambio de distribución.
2. Inestabilidad del Aprendizaje por Refuerzo (RL): Si bien el RL ofrece un mecanismo para la mejora online, el ajuste fino directo de decodificadores generativos de gran escala es frecuentemente inestable e ineficiente en términos de muestras. La interacción entre los objetivos de optimización y la dinámica de eliminación de ruido o integración de largo horizonte suele derivar en inestabilidad durante el entrenamiento.

Las alternativas existentes, como la Dirección de Difusión mediante Aprendizaje por Refuerzo (DSRL), intentan mejorar políticas congeladas aprendiendo una política de ruido latente. Sin embargo, DSRL reemplaza efectivamente el prior de preentrenamiento por una política latente aprendida. Este enfoque sufre de dos modos de fallo:

• Deriva Fuera de la Variedad (Off-Manifold Drift): El ruido dirigido puede desviarse del soporte Gaussiano estándar utilizado para entrenar el decodificador, lo que provoca una decodificación inestable.
• Colapso de Modos (Mode Collapse): La política latente aprendida puede colapsar el comportamiento multimodal del backbone hacia un subconjunto estrecho de modos, reduciendo la diversidad conductual.

2. Metodología: Dirección de Difusión por Perturbación Lagrangiana (LP-DS)

Los autores proponen LP-DS, un marco de adaptación ligero que mejora las políticas generativas congeladas mediante el aprendizaje de un residuo condicionado al estado en el espacio de ruido latente, en lugar de modificar los pesos del decodificador o reemplazar el prior por completo.

Mecanismo Central

LP-DS trata el backbone generativo congelado $\Phi$ como un decodificador de caja negra que mapea el ruido latente $w$ y el estado $s$ a acciones. En lugar de muestrear el ruido directamente de un prior $p_0$ (por ejemplo, $\mathcal{N}(0, I)$), LP-DS introduce una red de perturbación aprendible $\Delta_\theta(s)$:
$$w = \epsilon + \Delta_\theta(s), \quad \text{donde } \epsilon \sim \mathcal{N}(0, I)$$
Cuando $\Delta_\theta(s) = 0$, la política recupera exactamente la distribución original de clonación de comportamiento.

Optimización Restringida mediante Relajación Lagrangiana

Para evitar que la perturbación desplace el ruido fuera de la variedad donde se entrenó el decodificador, LP-DS formula el problema de aprendizaje como una optimización restringida. El objetivo es maximizar el valor Q esperado de las acciones generadas sujeto a una restricción de región de confianza sobre la magnitud de la perturbación.

Los autores aproximan la divergencia Kullback–Leibler (KL) esperada entre la distribución de ruido dirigida y el prior utilizando la magnitud al cuadrado del desplazamiento:
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
donde $\delta$ es un hiperparámetro que controla la desviación permitida.

Esto se resuelve mediante un enfoque de relajación Lagrangiana. Se introduce una variable dual (multiplicador de Lagrange) $\alpha$ para imponer la restricción de forma dinámica:
$$\mathcal{L}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
El algoritmo alterna entre:

1. Actualización del Actor: Maximizar $\mathcal{L}$ con respecto a $\theta$ (ascenso de gradiente).
2. Actualización Dual: Maximizar $\mathcal{L}$ con respecto a $\alpha$ (ascenso de gradiente proyectado) para asegurar que $\alpha \geq 0$. Si la perturbación viola la región de confianza, $\alpha$ aumenta, penalizando al actor y forzándolo a adherirse más al prior.

Componentes del Algoritmo

El procedimiento de entrenamiento (Algoritmo 1) implica:

• Recolección de Datos: Muestreo de acciones utilizando la política dirigida $w = \epsilon + \Delta_\theta(s)$.
• Actualizaciones del Crítico: Actualización de un crítico de valor de acción $Q^A$ mediante aprendizaje de Diferencia Temporal (TD) y un crítico de valor latente $Q^W$ mediante destilación desde $Q^A$ sobre la distribución de ruido base.
• Actualizaciones Primal-Dual: Actualización de la red de perturbación $\theta$ y del multiplicador $\alpha$ para equilibrar la maximización de la recompensa con la preservación del prior.

3. Contribuciones Clave

El artículo detalla cuatro contribuciones principales:

1. Marco LP-DS: Un método de adaptación ligero que mejora las políticas generativas congeladas mediante el control de ruido residual condicionado al estado, evitando la inestabilidad y el coste del ajuste fino completo del decodificador.
2. Objetivo de Región de Confianza Lagrangiana: Un objetivo novedoso que restringe dinámicamente la magnitud de la perturbación, proporcionando un mecanismo controlable para equilibrar la maximización de la recompensa frente a la preservación del prior latente preentrenado.
3. Mitigación del Colapso de Modos: Evidencia empírica que muestra que LP-DS preserva una mayor diversidad en el espacio de acciones y mitiga el colapso de modos en comparación con los métodos previos de dirección en el espacio de ruido que reemplazan el prior (como DSRL).
4. Generalizabilidad: Validación del método más allá de las políticas de difusión compactas, incluyendo backbones de ajuste de flujo, un modelo grande de visión-lenguaje-acción ($\pi_0$) y experimentos físicos con el robot Franka.

4. Resultados Experimentales

Los autores evaluaron LP-DS en los entornos de referencia RoboMimic (manipulación), OpenAI Gym (locomoción) y Adroit (manipulación diestra), comparándolo con DSRL, DPPO (ajuste fino de política de gradiente) y baselines de aprendizaje de difusión offline-to-online.

• Rendimiento: LP-DS mejoró consistentemente las tasas de éxito y los retornos episódicos. Notablemente, en el entorno WALKER2D, logró aproximadamente una mejora del 25% en el retorno sobre el baseline más fuerte.
• Eficiencia de Muestras: Al actualizar solo un módulo de perturbación ligero en lugar del decodificador completo, LP-DS demostró una eficiencia de muestras superior y una dinámica de aprendizaje más estable.
• Preservación de la Diversidad:
  • En una tarea de navegación multi-objetivo de juguete, LP-DS permitió el control explícito sobre el compromiso entre rendimiento y diversidad mediante el límite de la región de confianza $\delta$. Un $\delta$ pequeño preservó el comportamiento multimodal, mientras que un $\delta$ mayor permitió la especialización. En contraste, DSRL colapsó inmediatamente hacia un único modo.
  • En el entorno AVOIDING, LP-DS preservó un amplio conjunto de rutas factibles para la evasión de obstáculos (multimodalidad a nivel de trayectoria) con un $\delta$ pequeño, mientras que DSRL colapsó hacia una única trayectoria dominante.
  • Análisis de Entropía: LP-DS mantuvo consistentemente una mayor entropía en el espacio de acciones que DSRL durante la adaptación online, lo que indica que la restricción lagrangiana evita eficazmente el colapso prematuro mientras sigue mejorando el rendimiento.
• Estudios de Ablación: Eliminar la penalización lagrangiana condujo a un aprendizaje inestable y un rendimiento degradado. Eliminar tanto la penalización lagrangiana como el límite de ruido resultó en un comportamiento altamente inestable y bajos niveles de éxito, confirmando que la restricción de región de confianza es el mecanismo principal para prevenir la dirección fuera de la distribución.
• Sensibilidad: Se encontró que el método es robusto al valor específico del objetivo de la región de confianza $\delta$, con valores por encima de 0.1 que generalmente ofrecen un fuerte rendimiento.

5. Significado y Reivindicaciones

El artículo sostiene que LP-DS ofrece una alternativa estable y eficiente al ajuste fino completo para políticas generativas. Al tratar el backbone congelado como un prior fijo y aprender solo un desplazamiento residual acotado, el método:

• Previene la Deriva Fuera de la Variedad: Mantiene las consultas latentes dentro del rango operativo válido del decodificador, evitando el comportamiento errático observado en la dirección de ruido no restringida.
• Preserva la Multimodalidad: Mitiga el colapso de modos observado frecuentemente al reemplazar el prior de preentrenamiento, reteniendo así la diversidad conductual inherente al backbone congelado.
• Se Generaliza a través de Arquitecturas: El enfoque no se limita a modelos de difusión, sino que se aplica a cualquier política generativa con un mapeo de latente a acción, incluyendo modelos de ajuste de flujo y modelos grandes de visión-lenguaje-acción.
• Permite el Despliegue Físico: El método fue validado con éxito en experimentos físicos con el robot Franka utilizando adaptación latente basada en simulación, demostrando su aplicabilidad más allá de los entornos simulados.

Los autores posicionan a LP-DS como una solución práctica para la "mejora online" de políticas generativas de alta capacidad, equilibrando la necesidad de optimización específica de la tarea con la necesidad de mantener la estructura robusta y multimodal del modelo preentrenado.