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

Licenças

Acesso aberto
Artigo: CC BY 4.0

Tudo o que segue se baseia no artigo completo sob a licença CC BY 4.0; sua figura é mostrada com crédito aos autores.

Em resumo

O DSRL melhora uma política de difusão ou de fluxo congelada treinando uma pequena política de RL que escolhe seu ruído de entrada, adaptando robôs reais e o π0 em algumas dezenas a poucas centenas de episódios.

Por que importa

Uma equipe pode melhorar no próprio lugar uma política de clonagem de comportamento que já funciona, com algumas dezenas de episódios reais e uma recompensa de sucesso esparsa, sem tocar em seus pesos, o que também torna adaptáveis modelos de caixa-preta ou acessíveis só por API. Nos experimentos com o π0, o DSRL foi treinado em uma RTX 3070 de 8 GB, enquanto o repositório do π0 informa que o ajuste fino com LoRA precisa de pelo menos 22,5 GB.

Problema

Políticas de clonagem de comportamento, hoje muitas vezes políticas de difusão, não conseguem aprender com a própria experiência de implantação, então corrigir um desempenho fraco geralmente significa coletar mais demonstrações humanas. O RL poderia melhorá-las de forma autônoma, mas o ajuste fino de uma política de difusão com RL exige retropropagação pela sua cadeia de remoção de ruído de vários passos, o que é computacionalmente pesado e numericamente instável, e o RL padrão precisa de amostras demais para robôs reais.

Solução proposta

O DSRL trata a política de difusão pré-treinada (com amostragem determinística DDIM ou de fluxo) como parte do ambiente e aprende uma política de ruído latente separada que escolhe o ruído inicial em vez de sorteá-lo de N(0, I). Isso transforma a adaptação em RL comum sobre um 'MDP de ações latentes' e exige apenas passagens diretas da política base, não seus pesos. A variante com aliasing de ruído, DSRL-NA, treina um crítico sobre ações reais e o destila em um crítico do espaço de ruído, de modo que pode usar dados offline e explorar os muitos valores de ruído que levam à mesma ação.

Principais ideias

Com amostragem determinística DDIM ou de fluxo, o ruído de entrada determina totalmente a ação, então escolher o ruído direciona a política sem alterá-la.
Como a política base só produz ações dentro da distribuição, otimizar sobre seu ruído dá o conservadorismo do RL offline sem um termo de penalidade explícito.
Muitos valores de ruído levam quase à mesma ação; destilar um crítico no espaço de ações no crítico do espaço de ruído explora esse aliasing e precisou de cerca de metade das amostras do SAC simples no espaço de ruído.
Para o espaço de ações em chunks de 1.600 dimensões do π0, repetir um único vetor de ruído de 32 dimensões ao longo do chunk é expressivo o suficiente para direcionar a política.
O direcionamento funcionou com políticas base de diferentes tamanhos, durações de treinamento e qualidade de dados, desde que a política base mostre algum comportamento que resolva a tarefa.

Principais resultados

A adaptação online em Robomimic e OpenAI Gym chegou a desempenho quase ótimo com pelo menos 5-10x menos amostras que DPPO, IDQL, DQL, DIPO e QSM.
Em 10 tarefas offline do OGBench, o DSRL foi o estado da arte em cerca de metade e melhorou substancialmente em relação à política de fluxo sem direcionamento.
Franka real, tarefa única de colocar um cubo em uma tigela: o sucesso subiu de 20 % para quase 100 % em 3.500 passos online (cerca de 40 episódios); o RLPD, com ou sem intervenções humanas, não resolveu a tarefa.
WidowX 250 real com uma política multitarefa treinada no Bridge V2: ganhos significativos em pegar e colocar, fechar gaveta e empilhar blocos em 100-150 episódios, enquanto o RLPD não aprendeu.
π0 no LIBERO: sucesso de cerca de 20 % para cerca de 100 % após cerca de 10.000 amostras online; ResiP e V-GPS não o melhoraram significativamente.
Franka real com o checkpoint público π0-DROID: melhorou as tarefas da torradeira e da colher no prato em 80 e 65 episódios (cerca de 11.000 e 10.000 passos), o que os autores chamam de primeiro ajuste fino com RL do π0 no mundo real.

Limitações

A exploração é limitada pela política base: não há garantia de que toda política de difusão seja direcionável, e uma distribuição de ações muito estreita pode deixar poucas opções (autores).
Ainda não há como quantificar de antemão quanto uma dada política base pode ser melhorada (autores).
Como qualquer método de RL, precisa de sinais de recompensa, execuções online para a variante online e reinicializações no mundo real, além do que o aprendizado por imitação exige (autores).
As tarefas do mundo real usam recompensas de sucesso 0-1 definidas à mão, e as avaliações são pequenas (10-20 tentativas por tarefa no WidowX).

Impacto na área

Finalista dos prêmios CoRL 2025 Best Paper / Best Student Paper, com 242 citações até agora. O artigo traz o link de um site do projeto (diffusion-steering.github.io) e se baseia em checkpoints públicos do π0; a cópia licenciada do PMLR não cita um repositório de código.

Resumo

Políticas de controle robótico aprendidas a partir de demonstrações humanas alcançaram resultados impressionantes em muitas aplicações do mundo real. No entanto, em cenários em que o desempenho inicial não é satisfatório, como costuma acontecer em ambientes novos de mundo aberto, essas políticas aprendidas por clonagem de comportamento (BC) normalmente exigem a coleta de demonstrações humanas adicionais para melhorar ainda mais seu comportamento, um processo caro e demorado. Em contraste, o aprendizado por reforço (RL) promete permitir a melhoria autônoma de políticas online, mas muitas vezes não consegue isso devido ao grande número de amostras que normalmente exige. Neste trabalho, damos passos rumo à adaptação autônoma rápida de políticas treinadas por BC por meio de RL eficiente no mundo real. Com foco em particular nas políticas de difusão, uma metodologia de BC de ponta, propomos o direcionamento de difusão via aprendizado por reforço (DSRL): adaptar a política de BC executando RL sobre seu espaço de ruído latente. Mostramos que o DSRL é altamente eficiente em amostras, requer apenas acesso de caixa-preta à política de BC e permite uma melhoria autônoma eficaz de políticas no mundo real. Além disso, o DSRL evita muitos dos desafios associados ao ajuste fino de políticas de difusão, eliminando totalmente a necessidade de modificar os pesos da política base. Demonstramos o DSRL em benchmarks simulados, em tarefas robóticas do mundo real e na adaptação de políticas generalistas pré-treinadas, ilustrando sua eficiência em amostras e seu desempenho eficaz na melhoria de políticas no mundo real.

Open Source Robotics