Diffusion Policies

Una policy di diffusione è una policy di controllo robotico che genera azioni come un modello di diffusione genera immagini: parte da rumore casuale e lo rimuove gradualmente, condizionata da ciò che il robot osserva in quel momento, finché resta una sequenza di comandi motori. Addestrata su dimostrazioni umane, è una forma di apprendimento per imitazione comportamentale in grado di rappresentare diversi modi di eseguire un'attività, anziché farne la media. Cheng Chi, Shuran Song e colleghi della Columbia University, del Toyota Research Institute e del MIT l'hanno presentata a RSS 2023 in Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, superando i precedenti metodi di apprendimento per imitazione in media del 46,9% su 12 attività tratte da 4 benchmark di manipolazione robotica.

Perché la regressione non funziona con dimostrazioni multimodali#

Le dimostrazioni umane spesso includono più di un'azione valida per la stessa osservazione. In una dimostrazione, una persona può spingere un blocco avvicinandosi da sinistra e, in quella successiva, da destra; un conducente che incontra un ostacolo può passargli da entrambi i lati. Una policy addestrata con la regressione e una funzione di perdita dell'errore quadratico medio prevede la media di queste scelte, che può corrispondere a un'azione mai eseguita da alcun dimostratore, come guidare dritto contro l'ostacolo. Le soluzioni precedenti discretizzavano le azioni in intervalli, adattavano miscele gaussiane o usavano modelli basati sull'energia come Implicit Behavioral Cloning, che richiede campioni negativi durante l'addestramento ed era spesso instabile. Un modello di diffusione può rappresentare qualsiasi distribuzione di azioni e la sua funzione di perdita di previsione del rumore si addestra con la stessa stabilità della regressione ordinaria.

Questo esempio PyTorch addestra una piccola policy di diffusione su dimostrazioni che sterzano a sinistra o a destra per aggirare un ostacolo in egual numero:

import torch
from torch import nn

torch.manual_seed(0)
# Demonstrations at one observation: facing an obstacle, experts steer left (-1) or right (+1) equally often
actions = torch.cat([torch.randn(500, 1) * 0.05 - 1, torch.randn(500, 1) * 0.05 + 1])
print(f"Regression policy (mean squared error) outputs: {actions.mean():.2f}")

# Diffusion policy: train a network to predict the noise mixed into an action at denoising step k
K, betas = 50, torch.linspace(1e-4, 0.2, 50)
alpha_bar = torch.cumprod(1 - betas, 0)
net = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for _ in range(3000):
    k, noise = torch.randint(0, K, (1000, 1)), torch.randn(1000, 1)
    noisy = alpha_bar[k].sqrt() * actions + (1 - alpha_bar[k]).sqrt() * noise
    loss = (net(torch.cat([noisy, k / K], 1)) - noise).pow(2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Sampling: start from pure noise and remove the predicted noise one step at a time
a = torch.randn(1000, 1)
with torch.no_grad():
    for k in reversed(range(K)):
        eps = net(torch.cat([a, torch.full_like(a, k / K)], 1))
        a = (a - betas[k] / (1 - alpha_bar[k]).sqrt() * eps) / (1 - betas[k]).sqrt()
        a += betas[k].sqrt() * torch.randn_like(a) if k > 0 else 0
print(
    f"Diffusion policy samples: {(a < -0.5).float().mean():.0%} steer left, {(a > 0.5).float().mean():.0%} steer right"
)

La policy di regressione produce 0.00, cioè punta dritta contro l'ostacolo. I campioni della policy di diffusione si dividono tra il 44% di sterzate a sinistra e il 55% a destra, con solo circa l'1% tra le due modalità, quindi ogni campione sceglie una manovra valida.

Come funzionano le policy di diffusione#

Durante l'addestramento, si prende una sequenza di azioni da una dimostrazione, vi si mescola rumore gaussiano a un passo scelto casualmente su 100 e si addestra una rete neurale a prevedere quel rumore a partire dalle azioni rumorose e dall'osservazione. In esecuzione, il robot campiona una sequenza di puro rumore e la denoisa condizionandola sulla sua osservazione più recente. L'articolo originale usava 10 passi di denoising DDIM sui robot reali, ottenendo una latenza di inferenza di 0,1 secondi su una GPU NVIDIA RTX 3080.

La policy pianifica in blocchi e ripianifica con un ciclo a orizzonte scorrevole. Nella maggior parte delle attività dell'articolo osserva gli ultimi 2 passi temporali, prevede le 16 azioni successive, ne esegue le prime 8 e poi denoisa una nuova sequenza a partire da nuove osservazioni: così il movimento resta fluido e reagisce comunque ai cambiamenti.

Gli autori hanno realizzato due reti di previsione del rumore: una rete neurale convoluzionale temporale 1D che integra le caratteristiche delle osservazioni tramite condizionamento FiLM, consigliata come punto di partenza, e un Transformer per le attività in cui le azioni cambiano rapidamente. Le immagini della telecamera passano attraverso una ResNet-18 addestrata da zero, con pooling softmax spaziale al posto del pooling medio globale e normalizzazione a gruppi al posto della normalizzazione del batch. Prevedere le posizioni obiettivo anziché le velocità ha funzionato costantemente meglio e ha reso la policy più robusta alla latenza.

Le policy possono osservare anche uno stato a bassa dimensionalità anziché i pixel. Nel benchmark Push-T, in cui un robot spinge un blocco a forma di T verso un obiettivo, la variante basata sullo stato osserva 9 punti chiave 2D del blocco, oltre alla posizione dell'end effector. Un modello Ultralytics YOLO26 di stima della posa addestrato con una configurazione personalizzata dei punti chiave, definita con kpt_shape nel file YAML del dataset, restituisce questi punti chiave dai fotogrammi della telecamera tramite result.keypoints.xy.

Successori più veloci e più grandi#

Il denoising richiede diversi passaggi della rete per ogni blocco di azioni, quindi molti lavori successivi puntano all'inferenza in tempo reale. Consistency Policy, presentata da Aaditya Prasad e colleghi nel 2024, distilla una policy di diffusione in un modello di consistenza che richiede uno o pochi passi, accelerando l'inferenza di un ordine di grandezza. 3D Diffusion Policy (DP3), presentata da Yanjie Ze e colleghi nel 2024, usa nuvole di punti compatte anziché immagini per il condizionamento; gestisce la maggior parte delle sue 72 attività simulate con sole 10 dimostrazioni e raggiunge l'85% di successo in 4 attività con robot reali, con 40 dimostrazioni ciascuna.

La stessa ricetta è ora alla base dei modelli fondazionali per la robotica. π0 (2024) di Physical Intelligence aggiunge un esperto di azioni addestrato con il flow matching, strettamente imparentato con la diffusione, a un modello visione-linguaggio da 3 miliardi di parametri e genera blocchi di 50 azioni per controllare robot fino a 50 Hz. La libreria open source LeRobot di Hugging Face implementa Diffusion Policy insieme ad altre policy di apprendimento per imitazione per la ricerca sull'IA incarnata.

Policy di diffusione e termini correlati#

• Raggruppamento delle azioni: prevedere una sequenza di azioni future anziché una sola. Le policy di diffusione generano i blocchi tramite denoising; la policy Action Chunking with Transformers (ACT) li genera con un autoencoder variazionale condizionale.
• Diffusion forcing: assegna a ogni token di una sequenza un livello di rumore distinto, così un singolo modello causale può generare la sequenza passo dopo passo e pianificare verso gli obiettivi. In un'attività robotica che richiedeva memoria è riuscito nell'80% dei casi, mentre Diffusion Policy, priva di memoria, non ci è riuscita.
• Modelli di diffusione: la famiglia generale di modelli generativi che imparano a invertire un processo graduale di aggiunta del rumore, usata per generare immagini, video e audio. Una policy di diffusione è un modello di diffusione applicato alle azioni e condizionato sulle osservazioni.