Steering Your Diffusion Policy with Latent Space Reinforcement Learning

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

Licences

• Accès libre
• Article: CC BY 4.0

Tout ce qui suit s'appuie sur l'article complet sous sa licence CC BY 4.0 ; sa figure est montrée en citant les auteurs.

En bref

DSRL améliore une politique de diffusion ou de flow figée en entraînant une petite politique RL qui choisit son bruit d'entrée, adaptant des robots réels et π0 en quelques dizaines à quelques centaines d'épisodes.

Pourquoi c'est important

Une équipe peut améliorer sur place une politique de clonage comportemental fonctionnelle, avec quelques dizaines d'épisodes réels et une récompense de succès éparse, sans toucher à ses poids, ce qui rend aussi adaptables les modèles en boîte noire ou accessibles uniquement par API. Dans les expériences sur π0, DSRL s'est entraîné sur une RTX 3070 de 8 GB, alors que le dépôt de π0 indique que l'affinage LoRA nécessite au moins 22,5 GB.

Problème

Les politiques de clonage comportemental, aujourd'hui souvent des politiques de diffusion, ne peuvent pas apprendre de leur propre expérience de déploiement ; corriger de faibles performances implique donc généralement de collecter davantage de démonstrations humaines. Le RL pourrait les améliorer de façon autonome, mais affiner une politique de diffusion par RL exige une rétropropagation à travers sa chaîne de débruitage en plusieurs étapes, coûteuse en calcul et numériquement instable, et le RL standard demande trop d'échantillons pour des robots réels.

Solution proposée

DSRL traite la politique de diffusion préentraînée (avec échantillonnage DDIM ou de flow déterministe) comme une partie de l'environnement et apprend une politique de bruit latent distincte qui choisit le bruit initial au lieu de le tirer de N(0, I). L'adaptation devient ainsi un RL ordinaire sur un 'MDP à actions latentes' et ne nécessite que des passes avant de la politique de base, pas ses poids. La variante à repliement du bruit, DSRL-NA, entraîne un critique sur les actions réelles et le distille dans un critique dans l'espace du bruit, ce qui lui permet d'utiliser des données hors ligne et d'exploiter les nombreuses valeurs de bruit correspondant à la même action.

Idées clés

• Avec un échantillonnage DDIM ou de flow déterministe, le bruit d'entrée détermine entièrement l'action ; choisir le bruit pilote donc la politique sans la modifier.
• Comme la politique de base ne produit que des actions dans la distribution, optimiser sur son bruit apporte le conservatisme du RL hors ligne sans terme de pénalité explicite.
• De nombreuses valeurs de bruit correspondent presque à la même action ; distiller un critique dans l'espace des actions vers le critique dans l'espace du bruit exploite ce repliement (aliasing) et a nécessité environ deux fois moins d'échantillons qu'un SAC simple sur l'espace du bruit.
• Pour l'espace d'actions par segments (chunks) de π0, à 1 600 dimensions, répéter un seul vecteur de bruit à 32 dimensions sur tout le segment est suffisamment expressif pour piloter la politique.
• Le pilotage a fonctionné sur des politiques de base de tailles, de durées d'entraînement et de qualités de données différentes, dès lors que la politique de base montre un certain comportement de résolution de la tâche.

Principaux résultats

• L'adaptation en ligne sur Robomimic et OpenAI Gym a atteint des performances quasi optimales avec au moins 5 à 10 fois moins d'échantillons que DPPO, IDQL, DQL, DIPO et QSM.
• Sur 10 tâches hors ligne d'OGBench, DSRL était à l'état de l'art sur environ la moitié et s'est nettement amélioré par rapport à la politique de flow non pilotée.
• Franka réel, tâche unique cube dans un bol : la réussite est passée de 20 % à près de 100 % en 3 500 pas en ligne (environ 40 épisodes) ; RLPD, avec ou sans interventions humaines, n'a pas résolu la tâche.
• WidowX 250 réel avec une politique multitâche entraînée sur Bridge V2 : gains significatifs en saisie-dépose, fermeture de tiroir et empilement de blocs en 100 à 150 épisodes, tandis que RLPD n'a pas appris.
• π0 sur LIBERO : réussite passant d'environ 20 % à environ 100 % après environ 10 000 échantillons en ligne ; ResiP et V-GPS ne l'ont pas améliorée de façon significative.
• Franka réel avec le checkpoint public π0-DROID : amélioration des tâches grille-pain et cuillère sur assiette en 80 et 65 épisodes (environ 11 000 et 10 000 pas), ce que les auteurs présentent comme le premier affinage par RL de π0 en conditions réelles.

Limites

• L'exploration est bornée par la politique de base : rien ne garantit que toute politique de diffusion soit pilotable, et une distribution d'actions très étroite peut laisser trop peu d'options (auteurs).
• Il n'existe pas encore de moyen de quantifier à l'avance dans quelle mesure une politique de base donnée peut être améliorée (auteurs).
• Comme toute méthode de RL, elle requiert des signaux de récompense, des rollouts en ligne pour la variante en ligne et des réinitialisations en conditions réelles, au-delà de ce qu'exige l'apprentissage par imitation (auteurs).
• Les tâches réelles utilisent des récompenses de succès 0-1 définies à la main, et les évaluations sont restreintes (10 à 20 essais par tâche sur le WidowX).

Impact sur le domaine

Finaliste des prix CoRL 2025 Best Paper / Best Student Paper, avec 242 citations à ce jour. L'article renvoie vers un site web du projet (diffusion-steering.github.io) et s'appuie sur des checkpoints publics de π0 ; la copie PMLR sous licence ne mentionne aucun dépôt de code.

Résumé

Les politiques de commande robotique apprises à partir de démonstrations humaines ont obtenu des résultats impressionnants dans de nombreuses applications réelles. Cependant, dans les scénarios où la performance initiale n'est pas satisfaisante, comme c'est souvent le cas dans de nouveaux contextes en monde ouvert, ces politiques apprises par clonage comportemental (BC) nécessitent généralement la collecte de démonstrations humaines supplémentaires pour améliorer davantage leur comportement – un processus coûteux et chronophage. À l'inverse, l'apprentissage par renforcement (RL) promet de permettre une amélioration autonome des politiques en ligne, mais n'y parvient souvent pas en raison du grand nombre d'échantillons qu'il requiert généralement. Dans ce travail, nous progressons vers une adaptation autonome rapide des politiques entraînées par BC grâce à un RL efficace en conditions réelles. En nous concentrant en particulier sur les politiques de diffusion – une méthodologie de BC de pointe –, nous proposons le pilotage de diffusion par apprentissage par renforcement (diffusion steering via reinforcement learning, DSRL) : adapter la politique BC en exécutant du RL sur son espace de bruit latent. Nous montrons que DSRL est très efficace en échantillons, ne requiert qu'un accès en boîte noire à la politique BC et permet une amélioration autonome efficace des politiques en conditions réelles. En outre, DSRL évite bon nombre des difficultés liées à l'affinage des politiques de diffusion, rendant totalement inutile la modification des poids de la politique de base. Nous démontrons DSRL sur des benchmarks simulés, des tâches robotiques réelles et pour l'adaptation de politiques généralistes préentraînées, illustrant son efficacité en échantillons et ses bonnes performances pour l'amélioration de politiques en conditions réelles.

Liens

• 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
• Page du projet: https://diffusion-steering.github.io/