Diffusion Policies

Policy khuếch tán là policy điều khiển robot tạo hành động theo cách model khuếch tán tạo hình ảnh: bắt đầu từ nhiễu ngẫu nhiên rồi loại bỏ nhiễu từng bước, dựa trên những gì robot đang quan sát, cho đến khi còn lại một chuỗi lệnh động cơ. Được huấn luyện trên các bản trình diễn của con người, đây là một dạng học bắt chước hành vi có thể biểu diễn nhiều cách thực hiện tác vụ riêng biệt thay vì gộp chúng thành một. Cheng Chi, Shuran Song và cộng sự tại Columbia University, Toyota Research Institute và MIT giới thiệu phương pháp này tại RSS 2023 trong bài Diffusion Policy: Visuomotor Policy Learning via Action Diffusion; phương pháp vượt các kỹ thuật học bắt chước trước đó trung bình 46,9% trên 12 tác vụ thuộc 4 benchmark thao tác robot.

Vì sao hồi quy thất bại với dữ liệu trình diễn đa phương thức#

Dữ liệu trình diễn của con người thường có nhiều hành động hợp lệ cho cùng một quan sát. Trong một bản trình diễn, người đẩy khối có thể tiếp cận từ bên trái, còn ở bản tiếp theo họ có thể tiếp cận từ bên phải; tài xế gặp chướng ngại vật cũng có thể tránh sang một trong hai phía. Policy được huấn luyện bằng hồi quy với hàm loss sai số bình phương trung bình sẽ dự đoán giá trị trung bình của những lựa chọn đó, có thể tạo ra hành động mà không người trình diễn nào thực hiện, chẳng hạn lái thẳng vào chướng ngại vật. Các giải pháp trước đây chia hành động thành bin, khớp hỗn hợp Gaussian hoặc dùng model dựa trên năng lượng như Implicit Behavioral Cloning, vốn cần mẫu âm trong lúc huấn luyện và thường không ổn định. Model khuếch tán có thể biểu diễn bất kỳ phân phối hành động nào, còn loss dự đoán nhiễu giúp huấn luyện ổn định như hồi quy thông thường.

Ví dụ PyTorch này huấn luyện một policy khuếch tán nhỏ trên các bản trình diễn có số lần điều khiển sang trái hoặc phải để tránh chướng ngại vật bằng nhau:

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

Policy hồi quy xuất 0.00, đi thẳng vào chướng ngại vật. Các mẫu từ policy khuếch tán chia thành 44% đánh lái sang trái và 55% sang phải, chỉ khoảng 1% nằm giữa hai mode, nên mỗi mẫu chọn hẳn một thao tác hợp lệ.

Cách policy khuếch tán hoạt động#

Trong quá trình huấn luyện, một chuỗi hành động từ bản trình diễn được lấy ra, trộn với nhiễu Gaussian tại một bước ngẫu nhiên trong số 100 bước, rồi một mạng nơ-ron được huấn luyện để dự đoán nhiễu đó từ các hành động đã nhiễu và quan sát. Khi chạy, robot lấy mẫu một chuỗi nhiễu thuần rồi khử nhiễu trong khi điều kiện hóa theo quan sát mới nhất. Bài báo gốc dùng 10 bước khử nhiễu DDIM trên robot thực, đạt độ trễ suy luận 0,1 giây trên GPU NVIDIA RTX 3080.

Policy lập kế hoạch theo từng đoạn và lập kế hoạch lại theo vòng chân trời trượt. Trong hầu hết tác vụ của bài báo, policy quan sát 2 bước thời gian gần nhất, dự đoán 16 hành động tiếp theo, thực hiện 8 hành động đầu tiên rồi khử nhiễu một chuỗi mới từ các quan sát mới; cách này giữ chuyển động mượt mà mà vẫn phản ứng với thay đổi:

Các tác giả xây dựng hai mạng dự đoán nhiễu: một mạng nơ-ron tích chập thời gian 1D truyền đặc trưng quan sát qua điều kiện hóa FiLM, được khuyến nghị làm điểm khởi đầu, và một Transformer dành cho các tác vụ có hành động thay đổi nhanh. Hình ảnh camera đi qua ResNet-18 được huấn luyện từ đầu, dùng pooling softmax không gian thay cho pooling trung bình toàn cục và chuẩn hóa nhóm thay cho batch normalization. Dự đoán vị trí mục tiêu thay vì vận tốc luôn cho kết quả tốt hơn và giúp policy vững vàng hơn trước độ trễ.

Policy cũng có thể quan sát trạng thái chiều thấp thay vì pixel. Trong benchmark Push-T, nơi robot đẩy khối hình chữ T lên mục tiêu, biến thể dựa trên trạng thái quan sát 9 keypoint 2D của khối cùng vị trí đầu công tác. Model ước tính tư thế Ultralytics YOLO26 được huấn luyện với bố cục keypoint tùy chỉnh, thiết lập bằng kpt_shape trong YAML của bộ dữ liệu, có thể trả về các keypoint như vậy từ khung hình camera thông qua result.keypoints.xy.

Các thế hệ kế nhiệm nhanh hơn và lớn hơn#

Khử nhiễu cần nhiều lượt chạy mạng cho mỗi đoạn hành động, vì vậy phần lớn công trình tiếp nối tập trung vào suy luận thời gian thực. Consistency Policy, do Aaditya Prasad và cộng sự giới thiệu năm 2024, chưng cất policy khuếch tán thành một model nhất quán chỉ cần một hoặc vài bước, tăng tốc suy luận lên một bậc độ lớn. 3D Diffusion Policy (DP3), do Yanjie Ze và cộng sự giới thiệu năm 2024, dùng point cloud nhỏ gọn thay cho hình ảnh để điều kiện hóa; phương pháp xử lý phần lớn trong số 72 tác vụ mô phỏng chỉ với 10 bản trình diễn và đạt tỷ lệ thành công 85% trên 4 tác vụ robot thực với 40 bản trình diễn mỗi tác vụ.

Cùng công thức này hiện đang vận hành các model nền tảng robot. π0 (2024) của Physical Intelligence gắn một chuyên gia hành động được huấn luyện bằng flow matching, một phương pháp gần với khuếch tán, vào model thị giác-ngôn ngữ 3 tỷ tham số và tạo các đoạn 50 hành động để điều khiển robot với tần suất lên đến 50 Hz. Thư viện mã nguồn mở LeRobot của Hugging Face triển khai Diffusion Policy cùng các policy học bắt chước khác phục vụ nghiên cứu AI hiện thân.

Policy khuếch tán và các thuật ngữ liên quan#

• Phân đoạn hành động: Dự đoán một chuỗi hành động tương lai thay vì một hành động đơn lẻ. Policy khuếch tán tạo các đoạn bằng cách khử nhiễu, còn policy Action Chunking with Transformers (ACT) tạo chúng bằng autoencoder biến phân có điều kiện.
• Diffusion forcing: Gán cho mỗi token trong chuỗi một mức nhiễu riêng để một model nhân quả có thể lần lượt tạo chuỗi từng bước và lập kế hoạch hướng tới mục tiêu. Trong một tác vụ robot cần ghi nhớ, phương pháp này thành công 80% số lần, trong khi Diffusion Policy không có bộ nhớ thất bại.
• Model khuếch tán: Họ model sinh tổng quát học cách đảo ngược quá trình thêm nhiễu dần dần, được dùng để tạo hình ảnh, video và âm thanh. Policy khuếch tán là model khuếch tán trên hành động, được điều kiện hóa theo các quan sát.