Diffusion Policies

Une politique de diffusion est une politique de contrôle robotique qui génère des actions comme un modèle de diffusion génère des images : elle part d’un bruit aléatoire qu’elle élimine progressivement, en fonction des observations actuelles du robot, jusqu’à obtenir une séquence de commandes motrices. Entraînée sur des démonstrations humaines, elle relève de l’apprentissage par imitation directe et peut représenter plusieurs façons distinctes d’accomplir une tâche au lieu de les moyenner en une seule. Cheng Chi, Shuran Song et leurs collègues de Columbia University, du Toyota Research Institute et du MIT l’ont présentée à RSS 2023 dans Diffusion Policy: Visuomotor Policy Learning via Action Diffusion ; elle a surpassé les méthodes d’apprentissage par imitation antérieures de 46,9 % en moyenne sur 12 tâches issues de 4 benchmarks de manipulation en robotique.

Pourquoi la régression échoue avec les démonstrations multimodales#

Les démonstrations humaines comportent souvent plusieurs actions valides pour une même observation. Dans une démonstration, une personne qui pousse un bloc peut l’approcher par la gauche et, dans la suivante, par la droite ; un conducteur confronté à un obstacle peut le contourner d’un côté ou de l’autre. Une politique entraînée par régression avec une fonction de perte d’erreur quadratique moyenne prédit la moyenne de ces choix, qui peut correspondre à une action qu’aucun démonstrateur n’a effectuée, comme foncer tout droit sur l’obstacle. Les premières solutions discrétisaient les actions en intervalles, ajustaient des mélanges gaussiens ou utilisaient des modèles fondés sur l’énergie, comme l’apprentissage par imitation directe implicite, qui nécessite des échantillons négatifs pendant l’entraînement et se révélait souvent instable. Un modèle de diffusion peut représenter n’importe quelle distribution d’actions, et sa fonction de perte de prédiction du bruit s’entraîne avec la même stabilité qu’une régression classique.

Cet exemple PyTorch entraîne une petite politique de diffusion sur des démonstrations où l’on contourne autant de fois l’obstacle par la gauche que par la droite :

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 politique de régression produit 0.00 et fonce droit sur l’obstacle. Les échantillons de la politique de diffusion se répartissent entre 44 % de virages à gauche et 55 % de virages à droite, avec seulement environ 1 % entre les deux modes ; chaque échantillon choisit donc une manœuvre valide.

Fonctionnement des politiques de diffusion#

Pendant l’entraînement, une séquence d’actions tirée d’une démonstration est mélangée à du bruit gaussien à une étape choisie au hasard parmi 100 ; un réseau neuronal apprend à prédire ce bruit à partir des actions bruitées et de l’observation. Au moment de l’exécution, le robot échantillonne une séquence de bruit pur, puis la débruite en fonction de sa dernière observation. L’article original utilisait 10 étapes de débruitage DDIM sur ses robots réels, pour une latence d’inférence de 0,1 seconde sur un GPU NVIDIA RTX 3080.

La politique planifie par séquences et replanifie dans une boucle à horizon glissant. Pour la plupart des tâches de l’article, elle observe les 2 dernières étapes temporelles, prédit les 16 actions suivantes, en exécute les 8 premières, puis débruite une nouvelle séquence à partir d’observations fraîches ; cela maintient la fluidité du mouvement tout en permettant de réagir aux changements :

Les auteurs ont créé deux réseaux de prédiction du bruit : un réseau neuronal convolutif temporel 1D qui intègre les caractéristiques d’observation par conditionnement FiLM, recommandé comme point de départ, et un transformer pour les tâches où les actions changent rapidement. Les images de caméra passent par un ResNet-18 entraîné de zéro, avec un regroupement softmax spatial à la place du regroupement moyen global et une normalisation par groupes à la place de la normalisation par lots. La prédiction des positions cibles plutôt que des vitesses a donné systématiquement de meilleurs résultats et rendu la politique plus robuste à la latence.

Les politiques peuvent aussi observer un état de faible dimension plutôt que des pixels. Dans le benchmark Push-T, où un robot pousse un bloc en forme de T sur une cible, la variante fondée sur l’état observe 9 points clés 2D du bloc ainsi que la position de l’effecteur terminal. Un modèle Ultralytics YOLO26 d’estimation de pose entraîné avec une disposition de points clés personnalisée, définie par kpt_shape dans le YAML du jeu de données, renvoie ces points clés à partir des images de caméra via result.keypoints.xy.

Successeurs plus rapides et plus grands#

Le débruitage nécessite plusieurs passages du réseau pour chaque séquence d’actions ; de nombreux travaux ultérieurs visent donc l’inférence en temps réel. En 2024, Aaditya Prasad et ses collègues ont proposé Consistency Policy, qui distille une politique de diffusion en un modèle de cohérence nécessitant une ou quelques étapes et accélérant l’inférence d’un ordre de grandeur. En 2024 également, Yanjie Ze et ses collègues ont proposé 3D Diffusion Policy (DP3), qui s’appuie sur des nuages de points compacts plutôt que sur des images : elle a traité la plupart de ses 72 tâches simulées avec seulement 10 démonstrations et atteint 85 % de réussite sur 4 tâches robotiques réelles avec 40 démonstrations chacune.

La même approche alimente désormais les modèles fondamentaux pour la robotique. π0 (2024) de Physical Intelligence associe un expert en actions entraîné avec le flow matching, proche parent de la diffusion, à un modèle vision-langage de 3 milliards de paramètres et génère des séquences de 50 actions pour contrôler des robots jusqu’à 50 Hz. La bibliothèque open source LeRobot de Hugging Face implémente Diffusion Policy ainsi que d’autres politiques d’apprentissage par imitation pour la recherche en IA incarnée.

Politiques de diffusion et notions apparentées#

• Découpage en séquences d’actions : prédire une séquence d’actions futures plutôt qu’une seule. Les politiques de diffusion génèrent leurs séquences par débruitage, tandis que la politique Action Chunking with Transformers (ACT) les génère à l’aide d’un autoencodeur variationnel conditionnel.
• Diffusion forcing : attribue un niveau de bruit propre à chaque token d’une séquence afin qu’un seul modèle causal puisse dérouler les étapes une à une et planifier vers des objectifs. Dans une tâche robotique nécessitant de la mémoire, cette méthode a réussi dans 80 % des cas, alors que Diffusion Policy, dépourvue de mémoire, a échoué.
• Modèles de diffusion : famille générale de modèles génératifs qui apprennent à inverser un processus graduel de bruitage et servent à générer des images, des vidéos et de l’audio. Une politique de diffusion est un modèle de diffusion portant sur les actions et conditionné par les observations.