Diffusion Policies

Eine Diffusionsrichtlinie ist eine Robotersteuerungsrichtlinie, die Aktionen ähnlich wie ein Diffusionsmodell Bilder erzeugt: Sie beginnt mit zufälligem Rauschen und entfernt dieses schrittweise. Dabei orientiert sie sich an den aktuellen Beobachtungen des Roboters, bis eine Folge von Motorbefehlen übrig bleibt. Auf menschlichen Demonstrationen trainiert, ist sie eine Form des Verhaltensklonens, die mehrere unterschiedliche Möglichkeiten zur Ausführung einer Aufgabe darstellen kann, statt sie zu einer einzigen zu mitteln. Cheng Chi, Shuran Song und Kollegen von der Columbia University, dem Toyota Research Institute und dem MIT stellten die Methode 2023 auf der RSS in Diffusion Policy: Visuomotor Policy Learning via Action Diffusion vor. Über 12 Aufgaben aus 4 Benchmarks zur robotischen Manipulation hinweg übertraf sie frühere Imitationslernverfahren im Durchschnitt um 46,9 %.

Warum Regression bei multimodalen Demonstrationen versagt#

Menschliche Demonstrationen enthalten oft mehrere gültige Aktionen für dieselbe Beobachtung. Eine Person kann sich in einer Demonstration einem Block von links und in der nächsten von rechts nähern. Ein Fahrer kann auf beiden Seiten an einem Hindernis vorbeifahren. Eine mit Regression und einem mittleren quadratischen Fehler als Verlustfunktion trainierte Richtlinie sagt den Durchschnitt dieser Möglichkeiten voraus. Dieser kann einer Aktion entsprechen, die kein Demonstrator ausgeführt hat, etwa geradewegs in das Hindernis zu fahren. Frühere Ansätze unterteilten Aktionen in Wertebereiche, passten Gaußsche Mischverteilungen an oder verwendeten energiebasierte Modelle wie Implicit Behavioral Cloning, das während des Trainings negative Beispiele benötigt und oft instabil war. Ein Diffusionsmodell kann jede Aktionsverteilung darstellen, und sein Rauschvorhersageverlust lässt sich ebenso stabil trainieren wie eine gewöhnliche Regression.

Dieses PyTorch-Beispiel trainiert eine winzige Diffusionsrichtlinie anhand von Demonstrationen, bei denen ein Hindernis gleich oft links wie rechts umfahren wird:

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"
)

Die Regressionsrichtlinie gibt 0.00 aus und fährt damit geradewegs in das Hindernis. Bei den Stichproben der Diffusionsrichtlinie entfallen 44 % auf das Lenken nach links und 55 % auf das Lenken nach rechts. Nur etwa 1 % liegt zwischen diesen beiden Varianten. Jede Stichprobe legt sich also auf ein gültiges Manöver fest.

So funktionieren Diffusionsstrategien#

Beim Training wird eine Aktionssequenz aus einer Demonstration genommen, in einem zufällig ausgewählten Schritt von 100 mit gaußschem Rauschen vermischt und ein neuronales Netz darauf trainiert, dieses Rauschen anhand der verrauschten Aktionen und der Beobachtung vorherzusagen. Zur Laufzeit erstellt der Roboter eine Sequenz aus reinem Rauschen und entrauscht sie, während er die neueste Beobachtung berücksichtigt. Im Originalartikel wurden auf realen Robotern 10 Entrauschungsschritte mit DDIM eingesetzt. Auf einer NVIDIA RTX 3080 GPU ergab sich dadurch eine Inferenzlatenz von 0,1 Sekunden.

Die Richtlinie plant in Abschnitten und plant in einer Regelkreisschleife mit rückschreitendem Horizont neu. Bei den meisten Aufgaben des Artikels beobachtet sie die letzten 2 Zeitschritte, sagt die nächsten 16 Aktionen voraus, führt die ersten 8 aus und entrauscht dann anhand neuer Beobachtungen eine neue Sequenz. Dadurch bleiben die Bewegungen flüssig, während der Roboter trotzdem auf Änderungen reagiert:

Die Autoren entwickelten zwei Rauschvorhersagenetze: ein zeitliches eindimensionales Faltungsnetz, das Beobachtungsmerkmale über eine FiLM-Konditionierung einbindet und das sie als Ausgangspunkt empfehlen, sowie einen Transformer für Aufgaben mit schnell wechselnden Aktionen. Kamerabilder werden durch ein von Grund auf trainiertes ResNet-18 verarbeitet. Dabei ersetzt räumliches Softmax-Pooling das globale Durchschnitts-Pooling und Gruppen-Normalisierung die Batch-Normalisierung. Zielpositionen vorherzusagen statt Geschwindigkeiten funktionierte durchgehend besser und machte die Richtlinie robuster gegenüber Latenz.

Die Richtlinien können statt Pixeln auch einen niedrigdimensionalen Zustand beobachten. Beim Push-T-Benchmark, bei dem ein Roboter einen T-förmigen Block auf ein Ziel schiebt, beobachtet die zustandsbasierte Variante 9 zweidimensionale Schlüsselpunkte des Blocks sowie die Position des Endeffektors. Ein mit einer benutzerdefinierten Schlüsselpunktekonfiguration trainiertes Ultralytics-Pose-Estimationsmodell YOLO26, die in der Datensatz-YAML-Datei durch kpt_shape festgelegt wird, gibt solche Schlüsselpunkte aus Kamerabildern über result.keypoints.xy aus.

Schnellere und größere Nachfolgemodelle#

Für das Entrauschen sind mehrere Durchläufe des Netzes pro Aktionsabschnitt nötig. Viele Folgearbeiten konzentrieren sich daher auf die Echtzeitinferenz. Die von Aaditya Prasad und Kollegen 2024 vorgestellte Consistency Policy destilliert eine Diffusionsrichtlinie in ein Konsistenzmodell, das einen oder wenige Schritte benötigt und die Inferenz um eine Größenordnung beschleunigt. Die 2024 von Yanjie Ze und Kollegen vorgestellte 3D Diffusion Policy (DP3) nutzt kompakte Punktwolken statt Bilder als Bedingung. Sie bewältigte die meisten ihrer 72 simulierten Aufgaben mit nur 10 Demonstrationen und erreichte bei 4 Aufgaben mit realen Robotern und jeweils 40 Demonstrationen eine Erfolgsquote von 85 %.

Dasselbe Prinzip treibt inzwischen die Grundlagenmodelle für Robotik an. π0 von Physical Intelligence (2024) ergänzt ein mit Flow Matching trainiertes Aktionsmodell, das eng mit Diffusion verwandt ist, zu einem Bild-Sprach-Modell mit 3 Milliarden Parametern. Es erzeugt Aktionsabschnitte mit 50 Aktionen, um Roboter mit bis zu 50 Hz zu steuern. Die quelloffene Bibliothek LeRobot von Hugging Face implementiert Diffusion Policy und weitere Imitationslernrichtlinien für die Forschung zur verkörperten KI.

Diffusionsrichtlinien im Vergleich zu verwandten Begriffen#

• Bündeln von Aktionen: Vorhersage einer Folge zukünftiger Aktionen statt nur einer einzelnen Aktion. Diffusionsrichtlinien erzeugen ihre Abschnitte durch Entrauschen. Die Richtlinie Action Chunking with Transformers (ACT) erzeugt sie mithilfe eines bedingten Variations-Autoencoders.
• Diffusion Forcing: Weist jedem Token einer Sequenz einen eigenen Rauschpegel zu, damit ein einziges kausales Modell die Sequenz Schritt für Schritt durchlaufen und auf Ziele hin planen kann. Bei einer Roboteraufgabe, die Gedächtnis erforderte, war das Verfahren in 80 % der Fälle erfolgreich, während die Diffusion Policy ohne Gedächtnis versagte.
• Diffusionsmodelle: Die allgemeine Familie generativer Modelle, die lernen, einen schrittweisen Rauschprozess umzukehren, und für die Erzeugung von Bildern, Videos und Audio verwendet werden. Eine Diffusionsrichtlinie ist ein Diffusionsmodell über Aktionen, das sich an Beobachtungen orientiert.