Steering Your Diffusion Policy with Latent Space Reinforcement Learning

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

Licenze

• Open access
• Articolo: CC BY 4.0

Tutto ciò che segue si basa sull'articolo completo secondo la sua licenza CC BY 4.0; la figura è riportata citando gli autori.

In breve

DSRL migliora una diffusion policy o una flow policy congelata addestrando una piccola policy RL che ne sceglie il rumore in ingresso, adattando robot reali e π0 in un numero di episodi che va da qualche decina a qualche centinaio.

Perché è importante

Un team può migliorare sul posto una policy di behavior cloning già funzionante, con qualche decina di episodi reali e una ricompensa di successo sparsa, senza toccarne i pesi, il che rende adattabili anche i modelli black-box o accessibili solo via API. Negli esperimenti con π0, DSRL è stato addestrato su una RTX 3070 da 8 GB, mentre il repository di π0 indica che il fine-tuning con LoRA richiede almeno 22,5 GB.

Problema

Le policy di behavior cloning, oggi spesso diffusion policy, non possono imparare dalla propria esperienza in esercizio, quindi correggere prestazioni scarse significa di solito raccogliere altre dimostrazioni umane. L'RL potrebbe migliorarle in modo autonomo, ma il fine-tuning di una diffusion policy con l'RL richiede la retropropagazione attraverso la sua catena di denoising in più passi, che è computazionalmente pesante e numericamente instabile, e l'RL standard richiede troppi campioni per i robot reali.

Soluzione proposta

DSRL tratta la diffusion policy preaddestrata (con campionamento deterministico DDIM o di flusso) come parte dell'ambiente e apprende una policy separata sul rumore latente che sceglie il rumore iniziale invece di estrarlo da N(0, I). Ciò trasforma l'adattamento in un normale RL su un 'MDP ad azioni latenti' e richiede solo passaggi in avanti della policy di base, non i suoi pesi. La variante con aliasing del rumore, DSRL-NA, addestra un critic sulle azioni reali e lo distilla in un critic nello spazio del rumore, così può usare dati offline e sfruttare i molti valori di rumore che corrispondono alla stessa azione.

Idee chiave

• Con il campionamento deterministico DDIM o di flusso, il rumore in ingresso determina completamente l'azione, quindi scegliere il rumore orienta la policy senza modificarla.
• Poiché la policy di base produce solo azioni interne alla distribuzione, ottimizzare sul suo rumore fornisce il conservatorismo dell'RL offline senza un termine di penalità esplicito.
• Molti valori di rumore corrispondono quasi alla stessa azione; distillare un critic nello spazio delle azioni nel critic dello spazio del rumore sfrutta questo aliasing e ha richiesto circa la metà dei campioni di un semplice SAC nello spazio del rumore.
• Per lo spazio delle azioni a chunk di 1.600 dimensioni di π0, ripetere un unico vettore di rumore a 32 dimensioni lungo il chunk è abbastanza espressivo da orientare la policy.
• L'orientamento ha funzionato con policy di base di dimensioni, durate di addestramento e qualità dei dati diverse, purché la policy di base mostri un qualche comportamento risolutivo del task.

Risultati principali

• L'adattamento online su Robomimic e OpenAI Gym ha raggiunto prestazioni quasi ottimali con almeno 5-10x campioni in meno rispetto a DPPO, IDQL, DQL, DIPO e QSM.
• Su 10 task offline di OGBench, DSRL è risultato allo stato dell'arte in circa la metà e ha migliorato sostanzialmente la flow policy non orientata.
• Franka reale, task singolo cubo nella ciotola: il successo è salito dal 20 % a quasi il 100 % in 3.500 passi online (circa 40 episodi); RLPD, con o senza interventi umani, non ha risolto il task.
• WidowX 250 reale con una policy multi-task addestrata su Bridge V2: miglioramenti significativi nel pick-and-place, nella chiusura di un cassetto e nell'impilamento di blocchi in 100-150 episodi, mentre RLPD non ha appreso.
• π0 su LIBERO: successo da circa il 20 % a circa il 100 % dopo circa 10.000 campioni online; ResiP e V-GPS non lo hanno migliorato in modo significativo.
• Franka reale con il checkpoint pubblico π0-DROID: migliorati i task del tostapane e del cucchiaio nel piatto in 80 e 65 episodi (circa 11.000 e 10.000 passi), quello che gli autori definiscono il primo fine-tuning RL di π0 nel mondo reale.

Limiti

• L'esplorazione è limitata dalla policy di base: non c'è garanzia che ogni diffusion policy sia orientabile, e una distribuzione delle azioni molto stretta può lasciare troppe poche opzioni (autori).
• Non esiste ancora un modo per quantificare in anticipo quanto possa essere migliorata una data policy di base (autori).
• Come ogni metodo RL richiede segnali di ricompensa, rollout online per la variante online e reset nel mondo reale, oltre a quanto richiesto dall'apprendimento per imitazione (autori).
• I task reali usano ricompense di successo 0-1 definite a mano e le valutazioni sono limitate (10-20 prove per task sul WidowX).

Impatto sul settore

Finalista ai premi CoRL 2025 Best Paper / Best Student Paper, con 242 citazioni finora. L'articolo rimanda a un sito web del progetto (diffusion-steering.github.io) e si basa sui checkpoint pubblici di π0; la copia con licenza su PMLR non indica un repository di codice.

Abstract

Le policy di controllo robotico apprese da dimostrazioni umane hanno ottenuto risultati notevoli in molte applicazioni reali. Tuttavia, negli scenari in cui le prestazioni iniziali non sono soddisfacenti, come accade spesso in nuovi contesti open-world, queste policy apprese tramite behavior cloning (BC) richiedono in genere la raccolta di ulteriori dimostrazioni umane per migliorare il proprio comportamento, un processo costoso e lungo. Al contrario, l'apprendimento per rinforzo (RL) promette di consentire il miglioramento autonomo e online delle policy, ma spesso non ci riesce a causa del gran numero di campioni che richiede di solito. In questo lavoro compiamo dei passi verso l'adattamento autonomo e rapido di policy addestrate con BC tramite RL efficiente nel mondo reale. Concentrandoci in particolare sulle diffusion policy, una metodologia di BC allo stato dell'arte, proponiamo il diffusion steering tramite apprendimento per rinforzo (DSRL): adattare la policy BC eseguendo l'RL sul suo spazio del rumore latente. Mostriamo che DSRL è molto efficiente in termini di campioni, richiede solo un accesso black-box alla policy BC e consente un miglioramento autonomo efficace delle policy nel mondo reale. Inoltre, DSRL evita molte delle difficoltà associate al fine-tuning delle diffusion policy, eliminando del tutto la necessità di modificare i pesi della policy di base. Dimostriamo DSRL su benchmark simulati, su task robotici reali e nell'adattamento di policy generaliste preaddestrate, illustrandone l'efficienza in termini di campioni e le prestazioni efficaci nel miglioramento delle policy nel mondo reale.

Link

• 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
• Pagina del progetto: https://diffusion-steering.github.io/