De la Difusión al Flow

Diffusion Policy demostró que la generación basada en denoising es poderosa para acciones robóticas. Pero DDPM tiene una ineficiencia fundamental: los caminos que toma por el espacio de datos son curvos , requiriendo muchos pasos pequeños (50-100) para producir muestras de calidad.

Flow matching plantea: ¿qué pasaría si tomáramos caminos en línea recta? Aprender un campo de velocidades que transporte muestras de ruido a datos en líneas rectas requiere muchos menos pasos.

En robótica, donde las acciones deben generarse en tiempo real (a menudo a 10+ Hz), reducir los pasos de inferencia de 50-100 a 5-10 puede ser la diferencia entre un sistema viable y uno inviable.

💡 DDPM es como navegar por carreteras sinuosas — llegas, pero con muchas curvas. Flow matching es como tomar una autopista — línea recta del punto A al B.

Conditional Flow Matching

Conditional Flow Matching (CFM) (Lipman et al., 2023) (Tong et al., 2024) proporciona un framework libre de simulación para aprender flujos normalizadores continuos. La idea clave es elegante: definir un camino en línea recta entre ruido y datos, y luego aprender un campo de velocidades que lo siga.

Definimos una interpolación dependiente del tiempo entre una muestra de ruido $x_0 \sim \mathcal{N}(0, I)$ y una muestra de datos $x_1$ de nuestro dataset:

$$x_t = (1 - t)\, x_0 + t\, x_1, \quad t \in [0, 1]$$

Esto es simplemente interpolación lineal: en $t = 0$ tenemos ruido puro, en $t = 1$ tenemos datos limpios, y en el medio una mezcla ponderada. El camino de $x_0$ a $x_1$ es una línea recta en el espacio de datos.

La velocidad a lo largo de este camino — cuán rápido y en qué dirección debemos movernos — es constante y trivialmente computable:

$$v_{\text{target}} = \frac{dx_t}{dt} = x_1 - x_0$$

Esta es la dirección de ruido a datos — un solo vector, constante en el tiempo. Nuestra red neuronal $v_\theta(x_t, t)$ se entrena para predecir esta velocidad en cualquier punto $(x_t, t)$ a lo largo del camino.

La pérdida de entrenamiento es simplemente el error cuadrático medio entre las velocidades predicha y objetivo:

$$\mathcal{L}_{\text{CFM}} = \mathbb{E}_{t \sim U(0,1),\, x_0 \sim \mathcal{N}(0,I),\, x_1 \sim p_{\text{data}}} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]$$

Compara con la pérdida de DDPM $\|\epsilon - \epsilon_\theta(x_t, t)\|^2$. Son casi idénticas, pero la diferencia es crítica: DDPM predice ruido (qué sustraer), mientras flow matching predice velocidad (hacia dónde ir). La formulación de velocidad lleva a caminos más rectos.

En inferencia, resolvemos una ODE partiendo de ruido $x_0 \sim \mathcal{N}(0, I)$:

$$\frac{dx_t}{dt} = v_\theta(x_t, t)$$

Usando un solver de Euler simple con tamaño de paso $\Delta t$:

$$x_{t + \Delta t} = x_t + \Delta t \cdot v_\theta(x_t, t)$$

Con 5-10 pasos de Euler, flow matching produce calidad comparable a DDPM con 50-100 pasos.

📌 Nota técnica: CFM está estrechamente relacionado con los "rectified flows" y puede verse como un caso especial de flujos normalizadores continuos. La idea clave — usar caminos de interpolación en línea recta — simplifica tanto el entrenamiento como la inferencia en comparación con formulaciones generales de flow matching.

π₀ (pi-zero)

π₀ (Black et al., 2024) de Physical Intelligence es el primer VLA que combina un backbone VLM grande con una cabeza de flow matching. Los autores reportan rendimiento sólido en tareas diversas: doblar ropa, ensamblar cajas, limpiar mesas y recoger platos.

La arquitectura tiene tres componentes clave:

  • Backbone VLM: PaLI-Gemma (Beyer et al., 2024) (3B) para comprensión visual y lingüística.
  • Action expert: Parámetros transformer dedicados que procesan tokens de action chunk mediante cross-attention hacia las representaciones del VLM.
  • Cabeza de flow matching: Genera acciones continuas via CFM. Predice el campo de velocidades $v_\theta(A_t, t, O)$ que transporta chunks de acción ruidosos hacia limpios.

Una decisión de diseño crítica es la separación del action expert del backbone VLM . Durante el entrenamiento, los parámetros del backbone VLM se comparten entre dos objetivos:

  • Objetivo de lenguaje: Predicción estándar del siguiente token en texto, preservando las capacidades lingüísticas del VLM
  • Objetivo de acción: Pérdida de flow matching sobre chunks de acción, con gradientes que fluyen a través del action expert y hacia el backbone VLM

El action expert tiene su propio conjunto de capas transformer que se intercalan con las capas del backbone VLM mediante atención cruzada . El chunk de acción (como secuencia de tokens de acción ruidosos) atiende a las representaciones visuales y lingüísticas del VLM para extraer la información de condicionamiento que necesita.

El proceso de generación de acciones de π₀ en inferencia:

  • 1. Codificar la imagen y la instrucción de lenguaje a través del backbone VLM
  • 2. Muestrear un chunk de acción ruidoso $A^0 \sim \mathcal{N}(0, I) \in \mathbb{R}^{H \times d}$
  • 3. Ejecutar $N$ pasos de Euler (típicamente $N = 10$):
$$A^{n+1} = A^n + \frac{1}{N} \cdot v_\theta(A^n, \tfrac{n}{N}, O)$$
  • 4. El $A^N$ final es el chunk de acción predicho. Ejecutar las primeras $h$ acciones y repetir.

Flow vs Difusión en la Práctica

¿Cómo se comparan flow matching y difusión en el contexto de la generación de acciones robóticas?

  • Velocidad de inferencia: Flow matching necesita 5-10 pasos vs 20-100 para DDPM — ~5× más rápido.
  • Simplicidad de entrenamiento: Flow matching evita definir un schedule de ruido ($\beta_t$), que requiere ajuste cuidadoso en DDPM.
  • Calidad de muestras: Comparable o mejor, con caminos de interpolación más rectos.
  • Multi-modalidad: Ambos enfoques manejan distribuciones multi-modales bien. La diferencia es principalmente en eficiencia de inferencia, no en expresividad.

La conclusión práctica: flow matching tiende a lograr calidad comparable o mejor con significativamente menos pasos de inferencia . Para robótica, donde la inferencia en tiempo real es crítica, es una opción cada vez más preferida para nuevos sistemas.

Quiz

Pon a prueba tu comprensión de flow matching y π₀.

¿Cuál es la ineficiencia fundamental de DDPM que flow matching aborda?

En conditional flow matching, ¿cuál es el campo de velocidades objetivo para interpolación lineal de ruido x₀ a datos x₁?

¿Qué predice la red de flow matching comparado con DDPM?

¿Por qué π₀ separa su action expert del backbone VLM?

¿Cuántos pasos ODE necesita flow matching comparado con los pasos de denoising de DDPM?