Diffusion Policies

디퓨전 정책은 디퓨전 모델이 이미지를 생성하는 방식으로 행동을 생성하는 로봇 제어 정책입니다. 무작위 노이즈에서 시작해 로봇의 현재 관측값을 조건으로 노이즈를 단계적으로 제거하고, 모터 명령 시퀀스가 남을 때까지 반복합니다. 인간의 시연으로 학습하는 행동 복제의 한 형태이며, 여러 가지 뚜렷한 작업 수행 방식을 하나의 평균으로 합치지 않고 표현할 수 있습니다. 컬럼비아대학교, 도요타 연구소 및 MIT의 청 치, 슈란 쑹과 동료들은 RSS 2023에서 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%로 나뉘고 두 모드 사이의 행동은 약 1%뿐이므로, 각 샘플은 유효한 조작 하나를 선택합니다.

확산 정책 작동 방식#

학습 시 시연에서 행동 시퀀스 하나를 가져와 100개 단계 중 무작위로 선택한 단계에서 가우시안 노이즈를 섞고, 노이즈가 추가된 행동과 관측값으로부터 해당 노이즈를 예측하도록 신경망을 학습합니다. 실행 시 로봇은 순수 노이즈 시퀀스를 샘플링한 뒤 최신 관측값을 조건으로 노이즈를 제거합니다. 원 논문에서는 실제 로봇에 DDIM 노이즈 제거 단계 10개를 사용했으며, NVIDIA RTX 3080 GPU에서 추론 지연 시간은 0.1초였습니다.

정책은 청크 단위로 계획하고 이동 구간 반복 방식으로 재계획합니다. 논문에 제시된 대부분의 작업에서는 마지막 2개 타임스텝을 관측해 다음 16개 행동을 예측하고 처음 8개를 실행한 뒤, 새 관측값으로 새로운 시퀀스를 노이즈 제거합니다. 이렇게 하면 변화에 대응하면서도 움직임을 부드럽게 유지할 수 있습니다.

저자들은 노이즈 예측 네트워크 두 가지를 구축했습니다. 하나는 관측 특징을 FiLM 조건화로 주입하는 1D 시간 합성곱 신경망으로, 시작점으로 권장됩니다. 다른 하나는 행동 변화가 빠른 작업을 위한 트랜스포머입니다. 카메라 이미지는 처음부터 학습한 ResNet-18을 통과하며, 전역 평균 풀링 대신 공간 소프트맥스 풀링을, 배치 정규화 대신 그룹 정규화를 사용합니다. 속도 대신 목표 위치를 예측하면 일관되게 더 나은 성능을 보였고, 지연 시간에도 더 강건했습니다.

정책은 픽셀 대신 저차원 상태를 관측할 수도 있습니다. 로봇이 T자형 블록을 목표 지점으로 미는 Push-T 벤치마크에서 상태 기반 변형은 블록의 2D 키포인트 9개와 말단 장치의 위치를 관측합니다. 데이터셋 YAML에서 kpt_shape로 설정한 사용자 지정 키포인트 레이아웃을 사용해 학습한 Ultralytics YOLO26 자세 추정 모델은 result.keypoints.xy을 통해 카메라 프레임에서 이러한 키포인트를 반환합니다.

더 빠르고 규모가 큰 후속 모델#

노이즈 제거에는 행동 청크마다 네트워크를 여러 차례 실행해야 하므로, 후속 연구의 상당수는 실시간 추론을 목표로 합니다. 2024년 아디티야 프라사드와 동료들이 제안한 Consistency Policy는 디퓨전 정책을 한두 단계만 필요한 일관성 모델로 증류해 추론 속도를 10배 높입니다. 2024년 얀제 예와 동료들이 제안한 3D Diffusion Policy(DP3)는 이미지 대신 간결한 포인트 클라우드를 조건으로 사용합니다. 이를 통해 시뮬레이션 작업 72개 중 대부분을 시연 10회만으로 처리하고, 실제 로봇 작업 4개에서는 각각 시연 40회를 사용해 85%의 성공률을 달성했습니다.

이와 같은 방식은 이제 로봇 파운데이션 모델에도 적용됩니다. Physical Intelligence의 π0(2024)는 디퓨전과 밀접한 플로 매칭으로 학습한 행동 전문가를 30억 파라미터 비전-언어 모델에 연결하고, 최대 50Hz로 로봇을 제어하도록 행동 50개 청크를 생성합니다. Hugging Face의 오픈 소스 LeRobot 라이브러리는 체화된 AI 연구를 위해 다른 모방 학습 정책과 함께 Diffusion Policy를 구현합니다.

디퓨전 정책과 관련 용어 비교#

• 행동 청킹: 행동 하나 대신 미래 행동 시퀀스를 예측하는 방식입니다. 디퓨전 정책은 노이즈 제거로 청크를 생성하며, Action Chunking with Transformers(ACT) 정책은 조건부 변분 오토인코더로 청크를 생성합니다.
• 디퓨전 포싱: 시퀀스의 각 토큰에 고유한 노이즈 수준을 지정하므로, 인과 모델 하나로 단계별 롤아웃을 수행하고 목표를 향해 계획할 수 있습니다. 기억이 필요한 로봇 작업에서 이 방식은 80%의 성공률을 기록했지만, 기억을 사용하지 않는 Diffusion Policy는 실패했습니다.
• 디퓨전 모델: 점진적으로 노이즈를 추가하는 과정의 역과정을 학습하는 생성 모델의 일반적인 계열로, 이미지, 비디오 및 오디오 생성에 사용됩니다. 디퓨전 정책은 관측값을 조건으로 하는 행동 분포상의 디퓨전 모델입니다.