Diffusion Policies

Una política de difusión es una política de control robótico que genera acciones del mismo modo que un modelo de difusión genera imágenes: parte de ruido aleatorio y lo elimina paso a paso, condicionada por lo que observa el robot en ese momento, hasta obtener una secuencia de órdenes motoras. Entrenada con demostraciones humanas, es una forma de clonación del comportamiento que puede representar varias maneras distintas de realizar una tarea en lugar de promediarlas en una sola. Cheng Chi, Shuran Song y sus colaboradores de Columbia University, Toyota Research Institute y MIT la presentaron en RSS 2023 en Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, donde superó a métodos anteriores de aprendizaje por imitación en un 46,9 % de media en 12 tareas de 4 benchmarks de manipulación de robótica.

Por qué falla la regresión con demostraciones multimodales#

Las demostraciones humanas suelen contener más de una acción válida para una misma observación. Una persona puede acercarse a un bloque por la izquierda en una demostración y por la derecha en la siguiente, y un conductor puede pasar un obstáculo por cualquiera de los dos lados. Una política entrenada mediante regresión con una función de pérdida de error cuadrático medio predice el promedio de esas opciones, que puede ser una acción que ningún demostrador realizó, como conducir directamente contra el obstáculo. Las soluciones anteriores discretizaban las acciones en intervalos, ajustaban mezclas gaussianas o utilizaban modelos basados en energía como Implicit Behavioral Cloning, que necesita muestras negativas durante el entrenamiento y a menudo era inestable. Un modelo de difusión puede representar cualquier distribución de acciones, y su función de pérdida de predicción del ruido se entrena con la misma estabilidad que una regresión normal.

Este ejemplo de PyTorch entrena una pequeña política de difusión con demostraciones que indican girar a la izquierda o a la derecha alrededor de un obstáculo en igual proporción:

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

La política de regresión devuelve 0.00, es decir, sigue recto hacia el obstáculo. Las muestras de la política de difusión se reparten entre un 44 % que gira a la izquierda y un 55 % que gira a la derecha, con solo alrededor de un 1 % entre ambos modos, así que cada muestra opta por una maniobra válida.

Cómo funcionan las políticas de difusión#

Durante el entrenamiento se toma una secuencia de acciones de una demostración, se le añade ruido gaussiano en un paso aleatorio de un total de 100 y se entrena una red neuronal para predecir ese ruido a partir de las acciones ruidosas y la observación. En tiempo de ejecución, el robot muestrea una secuencia de ruido puro y la elimina mientras la condiciona a su observación más reciente. El artículo original utilizó 10 pasos de eliminación de ruido DDIM en sus robots reales, lo que proporcionó una latencia de inferencia de 0,1 segundos en una GPU NVIDIA RTX 3080.

La política planifica en bloques y vuelve a planificar en un bucle de horizonte móvil. En la mayoría de las tareas del artículo, observa los últimos 2 pasos temporales, predice las siguientes 16 acciones, ejecuta las primeras 8 y luego elimina el ruido de una nueva secuencia a partir de observaciones recientes, lo que mantiene la fluidez del movimiento y permite reaccionar a los cambios:

Los autores crearon dos redes de predicción del ruido: una red neuronal convolucional temporal 1D que inyecta características de observación mediante condicionamiento FiLM, recomendada como punto de partida, y un transformer para tareas cuyas acciones cambian rápidamente. Las imágenes de cámara pasan por una ResNet-18 entrenada desde cero, con agrupación softmax espacial en lugar de agrupación de promedio global y normalización por grupos en lugar de normalización por lotes. La predicción de posiciones objetivo en vez de velocidades dio mejores resultados de forma sistemática y aumentó la robustez de la política frente a la latencia.

Las políticas también pueden observar un estado de baja dimensionalidad en lugar de píxeles. En el benchmark Push-T, donde un robot empuja un bloque con forma de T hasta un objetivo, la variante basada en estados observa 9 puntos clave 2D del bloque, además de la posición del efector final. Un modelo de estimación de pose Ultralytics YOLO26 entrenado con una disposición personalizada de puntos clave, definida mediante kpt_shape en el YAML del conjunto de datos, devuelve esos puntos clave a partir de los fotogramas de cámara mediante result.keypoints.xy.

Sucesores más rápidos y de mayor escala#

La eliminación del ruido requiere varias pasadas por la red para cada bloque de acciones, por lo que muchos trabajos posteriores se centran en la inferencia en tiempo real. Consistency Policy, de Aaditya Prasad y sus colaboradores en 2024, destila una política de difusión en un modelo de consistencia que necesita uno o unos pocos pasos y acelera la inferencia en un orden de magnitud. 3D Diffusion Policy (DP3), de Yanjie Ze y sus colaboradores en 2024, se condiciona a nubes de puntos compactas en lugar de imágenes, resuelve la mayoría de sus 72 tareas simuladas con solo 10 demostraciones y alcanza un 85 % de éxito en 4 tareas robóticas reales con 40 demostraciones cada una.

La misma receta impulsa ahora modelos fundacionales de robótica. π0 de Physical Intelligence (2024) añade un experto en acciones entrenado con flow matching, un método estrechamente relacionado con la difusión, a un modelo de visión y lenguaje de 3000 millones de parámetros y genera bloques de 50 acciones para controlar robots a hasta 50 Hz. La biblioteca de código abierto LeRobot de Hugging Face implementa Diffusion Policy junto con otras políticas de aprendizaje por imitación para la investigación en IA incorporada.

Políticas de difusión y términos relacionados#

• Segmentación de acciones: Predice una secuencia de acciones futuras en lugar de una sola. Las políticas de difusión generan sus bloques eliminando ruido, y la política Action Chunking with Transformers (ACT) los genera con un autocodificador variacional condicional.
• Forzado de difusión: Asigna a cada token de una secuencia su propio nivel de ruido para que un único modelo causal pueda desplegarse paso a paso y planificar hacia objetivos. En una tarea robótica que requería memoria, tuvo éxito el 80 % de las veces, mientras que Diffusion Policy, que no tiene memoria, fracasó.
• Modelos de difusión: Familia general de modelos generativos que aprenden a invertir un proceso gradual de adición de ruido y se utilizan para generar imágenes, vídeo y audio. Una política de difusión es un modelo de difusión aplicado a acciones y condicionado a observaciones.