Steering Your Diffusion Policy with Latent Space Reinforcement Learning

Autores: Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, Sergey Levine · CoRL 2025

Licencias

• Acceso abierto
• Artículo: CC BY 4.0

Todo lo que sigue se basa en el artículo completo bajo su licencia CC BY 4.0; su figura se muestra con crédito a los autores.

En resumen

DSRL mejora una política de difusión o de flujo congelada entrenando una pequeña política de RL que elige su ruido de entrada, adaptando robots reales y π0 en entre decenas y unos pocos cientos de episodios.

Por qué importa

Un equipo puede mejorar en su sitio una política de clonación del comportamiento que ya funciona, con unas pocas decenas de episodios reales y una recompensa de éxito dispersa, sin tocar sus pesos, lo que también hace adaptables los modelos de caja negra o accesibles solo por API. En los experimentos con π0, DSRL se entrenó en una RTX 3070 de 8 GB, mientras que el repositorio de π0 indica que el ajuste fino con LoRA necesita al menos 22,5 GB.

Problema

Las políticas de clonación del comportamiento, hoy a menudo políticas de difusión, no pueden aprender de su propia experiencia de despliegue, por lo que corregir un rendimiento flojo suele significar recopilar más demostraciones humanas. El RL podría mejorarlas de forma autónoma, pero el ajuste fino de una política de difusión con RL requiere retropropagación a través de su cadena de eliminación de ruido de varios pasos, lo que es computacionalmente pesado y numéricamente inestable, y el RL estándar necesita demasiadas muestras para robots reales.

Solución propuesta

DSRL trata la política de difusión preentrenada (con muestreo determinista DDIM o de flujo) como parte del entorno y aprende una política de ruido latente separada que elige el ruido inicial en lugar de extraerlo de N(0, I). Esto convierte la adaptación en RL ordinario sobre un 'MDP de acciones latentes' y solo necesita pasadas hacia delante de la política base, no sus pesos. La variante con aliasing de ruido, DSRL-NA, entrena un crítico sobre acciones reales y lo destila en un crítico del espacio de ruido, de modo que puede usar datos offline y aprovechar los muchos valores de ruido que corresponden a la misma acción.

Ideas clave

• Con muestreo determinista DDIM o de flujo, el ruido de entrada determina por completo la acción, así que elegir el ruido dirige la política sin cambiarla.
• Como la política base solo produce acciones dentro de la distribución, optimizar sobre su ruido aporta el conservadurismo del RL offline sin un término de penalización explícito.
• Muchos valores de ruido corresponden casi a la misma acción; destilar un crítico en el espacio de acciones en el crítico del espacio de ruido aprovecha este aliasing y necesitó aproximadamente la mitad de muestras que SAC simple en el espacio de ruido.
• Para el espacio de acciones en chunks de 1600 dimensiones de π0, repetir un único vector de ruido de 32 dimensiones a lo largo del chunk es lo bastante expresivo para dirigir la política.
• El direccionamiento funcionó con políticas base de distintos tamaños, duraciones de entrenamiento y calidad de datos, siempre que la política base muestre algún comportamiento que resuelva la tarea.

Resultados principales

• La adaptación online en Robomimic y OpenAI Gym alcanzó un rendimiento casi óptimo con al menos 5-10x menos muestras que DPPO, IDQL, DQL, DIPO y QSM.
• En 10 tareas offline de OGBench, DSRL fue el estado del arte en aproximadamente la mitad y mejoró sustancialmente respecto a la política de flujo sin direccionar.
• Franka real, tarea única de meter un cubo en un cuenco: el éxito pasó del 20 % a casi el 100 % en 3500 pasos online (unos 40 episodios); RLPD, con o sin intervenciones humanas, no resolvió la tarea.
• WidowX 250 real con una política multitarea entrenada en Bridge V2: mejoras significativas en coger y colocar, cerrar un cajón y apilar bloques en 100-150 episodios, mientras que RLPD no aprendió.
• π0 en LIBERO: el éxito pasó de alrededor del 20 % a alrededor del 100 % tras unas 10 000 muestras online; ResiP y V-GPS no lo mejoraron significativamente.
• Franka real con el checkpoint público π0-DROID: mejoró las tareas de la tostadora y de la cuchara en el plato en 80 y 65 episodios (unos 11 000 y 10 000 pasos), lo que los autores llaman el primer ajuste fino con RL de π0 en el mundo real.

Limitaciones

• La exploración está acotada por la política base: no hay garantía de que toda política de difusión sea direccionable, y una distribución de acciones muy estrecha puede dejar muy pocas opciones (autores).
• Todavía no hay forma de cuantificar de antemano cuánto puede mejorarse una política base dada (autores).
• Como cualquier método de RL, necesita señales de recompensa, ejecuciones online para la variante online y reinicios en el mundo real, además de lo que requiere el aprendizaje por imitación (autores).
• Las tareas del mundo real usan recompensas de éxito 0-1 definidas a mano, y las evaluaciones son pequeñas (10-20 ensayos por tarea en el WidowX).

Impacto en el campo

Finalista de los premios CoRL 2025 Best Paper / Best Student Paper, con 242 citas hasta ahora. El artículo enlaza un sitio web del proyecto (diffusion-steering.github.io) y se basa en checkpoints públicos de π0; la copia con licencia de PMLR no menciona ningún repositorio de código.

Resumen

Las políticas de control robótico aprendidas a partir de demostraciones humanas han logrado resultados impresionantes en muchas aplicaciones del mundo real. Sin embargo, en escenarios en los que el rendimiento inicial no es satisfactorio, como suele ocurrir en entornos nuevos de mundo abierto, estas políticas aprendidas por clonación del comportamiento (BC) suelen requerir la recopilación de demostraciones humanas adicionales para seguir mejorando su comportamiento, un proceso caro y que consume mucho tiempo. En cambio, el aprendizaje por refuerzo (RL) promete permitir la mejora autónoma de políticas en línea, pero a menudo no lo consigue debido al gran número de muestras que suele requerir. En este trabajo damos pasos hacia la adaptación autónoma rápida de políticas entrenadas por BC mediante RL eficiente en el mundo real. Centrándonos en particular en las políticas de difusión, una metodología de BC de última generación, proponemos el direccionamiento de la difusión mediante aprendizaje por refuerzo (DSRL): adaptar la política de BC ejecutando RL sobre su espacio de ruido latente. Mostramos que DSRL es muy eficiente en muestras, requiere solo acceso de caja negra a la política de BC y permite una mejora autónoma eficaz de políticas en el mundo real. Además, DSRL evita muchos de los retos asociados al ajuste fino de políticas de difusión, eliminando por completo la necesidad de modificar los pesos de la política base. Demostramos DSRL en benchmarks simulados, en tareas robóticas del mundo real y en la adaptación de políticas generalistas preentrenadas, ilustrando su eficiencia en muestras y su rendimiento eficaz en la mejora de políticas en el mundo real.

Enlaces

• arXiv 2506.15799: https://arxiv.org/abs/2506.15799
• DOI 10.48550/arXiv.2506.15799: https://doi.org/10.48550/arXiv.2506.15799
• PDF: https://raw.githubusercontent.com/mlresearch/v305/main/assets/wagenmaker25a/wagenmaker25a.pdf
• Página del proyecto: https://diffusion-steering.github.io/