Diffusion Policies

拡散方策は、拡散モデルが画像を生成するのと同じ方法で行動を生成するロボット制御方策です。ランダムノイズから始め、ロボットが現在観測している内容を条件として段階的にノイズを除去し、一連のモーター指令を得ます。人間のデモンストレーションで訓練する行動クローニングの一種であり、異なる方法を一つに平均化せず、複数の異なるタスク実行方法を表現できます。Columbia University、Toyota Research Institute、MITのCheng Chi、Shuran Songらは、2023年のRSSでDiffusion Policy: Visuomotor Policy Learning via Action Diffusionを発表しました。4つのロボティクス操作ベンチマークの12タスク全体で、従来の模倣学習手法を平均46.9%上回りました。

マルチモーダルなデモンストレーションで回帰が失敗する理由#

同じ観測に対して、人間のデモンストレーションには複数の有効な行動が含まれることがよくあります。ブロックを押す人は、あるデモンストレーションでは左から近づき、次のデモンストレーションでは右から近づく場合があり、障害物に直面したドライバーは左右どちらにも回避できます。平均二乗誤差の損失関数を使った回帰で訓練した方策は、これらの選択肢の平均を予測しますが、それは障害物に向かって直進するなど、デモンストレーターの誰も取らなかった行動になる可能性があります。従来の対策では、行動をビンに離散化したり、ガウス混合を適合させたり、Implicit Behavioral Cloningのようなエネルギーベースモデルを使ったりしましたが、後者は学習中に負例サンプルが必要で、不安定になることもよくありました。拡散モデルは任意の行動分布を表現でき、そのノイズ予測損失は通常の回帰と同じように安定して学習できます。

このPyTorchの例では、障害物を左右に回避するデモンストレーションを同数ずつ使って、小さな拡散方策を訓練します。

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

回帰方策は障害物に向かって直進する0.00を出力します。拡散方策のサンプルは左操舵が44%、右操舵が55%で、その2つのモードの間は約1%だけです。したがって、各サンプルは有効な操作方法のいずれかを選びます。

拡散ポリシーの仕組み#

学習では、デモンストレーションから行動系列を取り出し、100段階からランダムに選んだ時点でガウスノイズを混ぜます。そして、ノイズを加えた行動と観測からそのノイズを予測するようニューラルネットワークを訓練します。実行時には、ロボットが純粋なノイズから行動系列をサンプリングし、最新の観測を条件としてノイズを除去します。原論文では実機ロボットに10段階のDDIMノイズ除去を使い、NVIDIA RTX 3080 GPUで推論レイテンシ0.1秒を達成しました。

方策はチャンク単位で計画し、後退ホライズンループで再計画します。論文のほとんどのタスクでは、直近2タイムステップを観測して次の16行動を予測し、最初の8行動を実行してから、新しい観測を使って新たな系列のノイズを除去します。これにより、変化に対応しながら動作を滑らかに保ちます。

著者らはノイズ予測ネットワークを2種類構築しました。1つはFiLM条件付けで観測特徴を注入する1次元の時間方向畳み込みニューラルネットワークで、出発点として推奨されています。もう1つは、行動が急速に変化するタスク向けのTransformerです。カメラ画像はスクラッチから訓練したResNet-18に渡し、グローバル平均プーリングの代わりに空間ソフトマックスプーリングを、バッチ正規化の代わりにグループ正規化を使います。速度ではなく目標位置を予測する方が一貫して良い結果となり、レイテンシに対する方策の頑健性も高まりました。

方策はピクセルではなく低次元の状態を観測することもできます。ロボットがT字型ブロックを押して目標位置に移動させるPush-Tベンチマークでは、状態ベースのバリアントがブロック上の9個の2D キーポイントとエンドエフェクターの位置を観測します。カスタムキーポイントレイアウトで訓練したUltralytics YOLO26の姿勢推定モデルは、データセットYAMLのkpt_shapeで設定することで、カメラフレームからそのようなキーポイントをresult.keypoints.xy経由で返します。

より高速で大規模な後継手法#

ノイズ除去では行動チャンクごとにネットワークを複数回実行するため、後続研究の多くはリアルタイム推論を目指しています。Aaditya Prasadらが2024年に発表したConsistency Policyは、拡散方策を1回または数回のステップで実行できる一貫性モデルに蒸留し、推論を1桁高速化します。Yanjie Zeらが2024年に発表した3D Diffusion Policy (DP3)は、画像の代わりにコンパクトな点群を条件として使います。シミュレーションの72タスクの大半をデモンストレーション10件だけで処理し、実ロボットの4タスクでは各40件のデモンストレーションから85%の成功率を達成しました。

同じ手法は現在、ロボット基盤モデルにも使われています。Physical Intelligenceのπ0 (2024)は、拡散と近縁のフローマッチングで訓練した行動エキスパートを、30億パラメーターの視覚言語モデルに追加し、最大50 Hzでロボットを制御するための50行動のチャンクを生成します。Hugging FaceのオープンソースLeRobotライブラリは、身体性AI研究向けに、拡散方策とほかの模倣学習方策を実装しています。

拡散方策と関連用語の比較#

• アクションチャンク化： 1つの行動ではなく、将来の行動系列を予測します。拡散方策はノイズ除去によってチャンクを生成し、Action Chunking with Transformers (ACT)方策は条件付き変分オートエンコーダーで生成します。
• 拡散フォーシング： 系列の各トークンに独自のノイズレベルを与えることで、1つの因果モデルがステップごとに展開し、目標に向けて計画できるようにします。記憶が必要なロボットタスクでは、記憶を持たない拡散方策が失敗した一方、拡散フォーシングは80%の確率で成功しました。
• 拡散モデル： 段階的にノイズを加える過程を逆転させる方法を学ぶ生成モデルの一般的な系統で、画像、動画、音声の生成に使われます。拡散方策は、観測を条件とする行動上の拡散モデルです。