Diffusion Policies

Диффузионная политика — это политика управления роботом, которая генерирует действия так же, как диффузионная модель генерирует изображения: она начинает со случайного шума и постепенно удаляет его, учитывая текущие наблюдения робота, пока не останется последовательность моторных команд. Обученная на человеческих демонстрациях, такая политика представляет собой разновидность клонирования поведения, способную представлять несколько разных способов выполнения задачи, а не усреднять их до одного. Чэн Чи, Шуран Сонг и их коллеги из Колумбийского университета, Исследовательского института Toyota и MIT представили её на RSS 2023 в статье Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. В среднем метод превзошёл более ранние методы обучения с подражанием на 46,9% по 12 задачам из 4 бенчмарков манипуляций в робототехнике.

Почему регрессия не справляется с мультимодальными демонстрациями#

В человеческих демонстрациях для одного и того же наблюдения часто встречается несколько допустимых действий. В одном примере человек может подойти к блоку слева, а в следующем — справа; водитель может объехать препятствие с любой стороны. Политика, обученная регрессией с функцией потерь в виде среднеквадратичной ошибки, предсказывает среднее этих вариантов, то есть действие, которого не совершал ни один демонстратор, например движение прямо на препятствие. В более ранних решениях действия разбивали на интервалы, подбирали гауссовские смеси или использовали модели на основе энергии, такие как неявное клонирование поведения, которому нужны отрицательные примеры во время обучения и которое часто было нестабильным. Диффузионная модель может представлять любое распределение действий, а её функция потерь на предсказание шума обеспечивает такое же стабильное обучение, как обычная регрессия.

В этом примере на 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% — поворот направо и лишь около 1% — между этими вариантами; таким образом, каждая выборка соответствует одному допустимому манёвру.

Как работают диффузионные политики#

При обучении берётся последовательность действий из демонстрации, к ней добавляется гауссовский шум на случайном шаге из 100, а нейронную сеть обучают предсказывать этот шум по зашумлённым действиям и наблюдению. Во время работы робот генерирует последовательность чистого шума и постепенно устраняет его, учитывая последнее наблюдение. В исходной статье на реальных роботах использовались 10 шагов устранения шума DDIM, обеспечивавших задержку инференса 0,1 секунды на GPU NVIDIA RTX 3080.

Политика планирует действия блоками и перепланирует их в цикле с сокращающимся горизонтом. В большинстве задач из статьи она наблюдает последние 2 временных шага, предсказывает следующие 16 действий, выполняет первые 8, а затем формирует новую последовательность по свежим наблюдениям, что обеспечивает плавность движения и позволяет реагировать на изменения:

Авторы создали две сети для предсказания шума: одномерную временную свёрточную нейронную сеть, которая вводит признаки наблюдений с помощью обусловливания FiLM и которую они рекомендуют использовать для начала, и трансформер для задач, где действия быстро меняются. Изображения с камеры проходят через ResNet-18, обученную с нуля; вместо глобального усредняющего пулинга используется пространственный softmax-пулинг, а вместо пакетной нормализации — групповая. Предсказание целевых положений вместо скоростей неизменно давало лучшие результаты и повышало устойчивость политики к задержкам.

Политики могут получать на вход не пиксели, а состояние с низкой размерностью. В бенчмарке Push-T, где робот перемещает блок в форме буквы T на целевое место, вариант, основанный на состоянии, использует 9 двумерных ключевых точек блока и положение рабочего органа робота. Модель Ultralytics YOLO26 для оценки позы, обученная с пользовательской схемой ключевых точек, заданной параметром kpt_shape в YAML-файле набора данных, извлекает такие точки из кадров камеры с помощью result.keypoints.xy.

Более быстрые и крупные преемники#

Для устранения шума требуется несколько проходов сети на каждый блок действий, поэтому многие последующие работы направлены на инференс в реальном времени. В 2024 году Адитья Прасад и его коллеги представили Consistency Policy — дистилляцию диффузионной политики в модель согласованности, которой нужен один или несколько шагов; это ускоряет инференс на порядок. В 2024 году Янцзе Зе и его коллеги представили 3D Diffusion Policy (DP3), которая использует компактные облака точек вместо изображений. Она справляется с большинством из 72 симулированных задач всего по 10 демонстрациям и достигает 85% успешных результатов в 4 задачах на реальном роботе, имея по 40 демонстраций для каждой.

Этот же подход теперь лежит в основе фундаментальных моделей для роботов. В модели π0 (2024) компании Physical Intelligence к модели «зрение — язык» с 3 миллиардами параметров добавлен эксперт действий, обученный с помощью сопоставления потоков — метода, близкого к диффузии. Модель генерирует блоки по 50 действий и управляет роботами с частотой до 50 Гц. Библиотека LeRobot с открытым исходным кодом от Hugging Face реализует Diffusion Policy и другие политики обучения с подражанием для исследований в области воплощённого ИИ.

Диффузионные политики и смежные понятия#

• Формирование блоков действий: предсказание последовательности будущих действий вместо одного. Диффузионные политики создают блоки, постепенно устраняя шум, а политика Action Chunking with Transformers (ACT) формирует их с помощью условного вариационного автокодировщика.
• Диффузионное принуждение: назначает каждому токену последовательности собственный уровень шума, чтобы одна причинно-следственная модель могла пошагово выполнять последовательность и планировать достижение целей. В робототехнической задаче, требовавшей памяти, метод достигал успеха в 80% случаев, тогда как Diffusion Policy без памяти не справлялась.
• Диффузионные модели: общее семейство генеративных моделей, обучающихся обращать постепенный процесс добавления шума; применяются для генерации изображений, видео и аудио. Диффузионная политика — это диффузионная модель действий, обусловленная наблюдениями.