Diffusion Policies

Uma política de difusão é uma política de controle robótico que gera ações da mesma forma que um modelo de difusão gera imagens: começa com ruído aleatório e o remove passo a passo, condicionado às observações atuais do robô, até restar uma sequência de comandos motores. Treinada com demonstrações humanas, ela é uma forma de clonagem comportamental capaz de representar várias maneiras distintas de executar uma tarefa, em vez de reduzi-las a uma média. Cheng Chi, Shuran Song e colegas da Columbia University, do Toyota Research Institute e do MIT apresentaram-na no RSS 2023 em Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, superando métodos anteriores de aprendizado por imitação em uma média de 46,9% em 12 tarefas de 4 benchmarks de manipulação de robótica.

Por que a regressão falha com demonstrações multimodais#

As demonstrações humanas costumam incluir mais de uma ação válida para a mesma observação. Uma pessoa que empurra um bloco pode aproximar-se pela esquerda numa demonstração e pela direita na seguinte, e um motorista diante de um obstáculo pode contorná-lo por qualquer um dos lados. Uma política treinada por regressão com uma função de perda de erro quadrático médio prevê a média dessas escolhas, que pode ser uma ação que nenhum demonstrador realizou, como conduzir diretamente contra o obstáculo. Soluções anteriores discretizavam as ações em intervalos, ajustavam misturas gaussianas ou usavam modelos baseados em energia, como a Clonagem Comportamental Implícita, que exige amostras negativas durante o treinamento e era frequentemente instável. Um modelo de difusão pode representar qualquer distribuição de ações, e sua função de perda de previsão de ruído treina com a mesma estabilidade da regressão convencional.

Este exemplo de PyTorch treina uma pequena política de difusão com demonstrações que conduzem para a esquerda ou para a direita ao redor de um obstáculo em igual número:

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

A política de regressão produz 0.00, diretamente contra o obstáculo. As amostras da política de difusão se dividem em 44% de condução para a esquerda e 55% para a direita, com apenas cerca de 1% entre os dois modos, portanto cada amostra escolhe uma manobra válida.

Como funcionam as políticas de difusão#

O treinamento pega uma sequência de ações de uma demonstração, mistura nela ruído gaussiano numa etapa escolhida aleatoriamente entre 100 e treina uma rede neural para prever esse ruído a partir das ações ruidosas e da observação. Em tempo de execução, o robô amostra uma sequência de puro ruído e remove o ruído enquanto a condiciona à observação mais recente. O artigo original usou 10 etapas de remoção de ruído DDIM nos robôs reais, com uma latência de inferência de 0,1 segundo numa GPU NVIDIA RTX 3080.

A política planeja em blocos e replaneja num ciclo de horizonte móvel. Na maioria das tarefas do artigo, ela observa as 2 últimas etapas de tempo, prevê as próximas 16 ações, executa as 8 primeiras e, em seguida, remove o ruído de uma nova sequência com observações atualizadas, o que mantém o movimento suave e ainda permite reagir a mudanças:

Os autores criaram duas redes de previsão de ruído: uma rede neural convolucional temporal 1D que injeta características da observação por meio de condicionamento FiLM, que recomendam como ponto de partida, e um Transformer para tarefas cujas ações mudam rapidamente. As imagens da câmera passam por uma ResNet-18 treinada do zero, com agrupamento softmax espacial no lugar do agrupamento médio global e normalização de grupos no lugar da normalização em lote. Prever posições-alvo em vez de velocidades teve resultados consistentemente melhores e tornou a política mais robusta à latência.

As políticas também podem observar um estado de baixa dimensionalidade em vez de pixels. No benchmark Push-T, em que um robô empurra um bloco em forma de T até um alvo, a variante baseada em estado observa 9 pontos-chave 2D do bloco e a posição do efetuador final. Um modelo Ultralytics YOLO26 de estimativa de pose, treinado com uma disposição personalizada de pontos-chave definida por kpt_shape no YAML do conjunto de dados, retorna esses pontos-chave a partir de quadros de câmera por meio de result.keypoints.xy.

Sucessores mais rápidos e maiores#

A remoção de ruído exige várias passagens pela rede por bloco de ação, portanto muitos trabalhos posteriores visam a inferência em tempo real. A Consistency Policy, de Aaditya Prasad e colegas em 2024, destila uma política de difusão num modelo de consistência que exige uma ou poucas etapas, acelerando a inferência em uma ordem de grandeza. A 3D Diffusion Policy (DP3), de Yanjie Ze e colegas em 2024, condiciona-se a nuvens de pontos compactas em vez de imagens, resolve a maioria das 72 tarefas simuladas com apenas 10 demonstrações e alcança 85% de sucesso em 4 tarefas com robôs reais, com 40 demonstrações cada.

A mesma abordagem agora impulsiona modelos-base de robótica. O π0 (2024) da Physical Intelligence integra um especialista em ações treinado com correspondência de fluxo, um método intimamente relacionado à difusão, a um modelo de visão e linguagem de 3 bilhões de parâmetros e gera blocos de 50 ações para controlar robôs a até 50 Hz. A biblioteca de código aberto LeRobot da Hugging Face implementa a Diffusion Policy junto com outras políticas de aprendizado por imitação para pesquisa em IA incorporada.

Políticas de difusão e termos relacionados#

• Agrupamento de ações: prevê uma sequência de ações futuras em vez de uma só. As políticas de difusão geram seus blocos removendo ruído, e a política Action Chunking with Transformers (ACT) os gera com um autoencoder variacional condicional.
• Forçamento por difusão: atribui a cada token de uma sequência seu próprio nível de ruído para que um único modelo causal possa avançar passo a passo e planejar em direção a objetivos. Numa tarefa robótica que exigia memória, teve 80% de sucesso, enquanto a Diffusion Policy sem memória falhou.
• Modelos de difusão: a família geral de modelos generativos que aprendem a inverter um processo gradual de adição de ruído, usados para gerar imagens, vídeos e áudio. Uma política de difusão é um modelo de difusão aplicado a ações e condicionado a observações.