Amortizing Trajectory Diffusion with Keyed Drift Fields

1. El Problema

Los planificadores de trayectorias basados en difusión han demostrado ser fundamentales para la toma de decisiones robótica y el aprendizaje por refuerzo (RL) fuera de línea (offline RL), ya que pueden modelar distribuciones complejas y multimodales de comportamientos futuros directamente desde los datos. Sin embargo, presentan un desafío crítico en entornos de control en tiempo real:

• Costo de Inferencia: Los modelos de difusión requieren múltiples pasos secuenciales de "desruido" (denoising) para generar una trayectoria. En un bucle de control de horizonte deslizante (receding-horizon), esto implica evaluar la red neuronal muchas veces por paso de control, lo que genera una latencia inaceptable bajo presupuestos de cómputo estrictos.
• Fallo en la Condicionamiento Naïve: Intentar entrenar un generador de un solo paso (one-step) simplemente emparejando la distribución de datos en el espacio completo de la ventana de trayectoria falla. En espacios de alta dimensión, las métricas de distancia están dominadas por componentes futuros no restringidos. Esto provoca que el modelo se incline hacia "trayectorias promedio", colapsando la diversidad de acciones y generando comportamientos estáticos o inertes.

La pregunta central: ¿Cómo se puede lograr un comportamiento de planificación similar al de la difusión con inferencia de un solo paso, manteniendo la capacidad de muestrear planes diversos y condicionarse al estado actual?

2. Metodología: Keyed Drifting Policies (KDP)

Los autores proponen KDP (Keyed Drifting Policies), un generador condicional de trayectorias de un solo paso que amortiza el refinamiento iterativo durante el entrenamiento en lugar de realizarlo en la inferencia.

Conceptos Clave:

1. Vecindarios Conscientes del Condicionamiento (Conditioning-Aware Neighborhoods):

   • La observación clave es que la generación condicional requiere una noción de "vecindad" que respete la condición (el estado actual).
   • En lugar de medir la similitud entre ventanas completas de trayectorias (donde el futuro no restringido domina la distancia), KDP define la vecindad en un espacio de claves (key space) compacto alineado con la condición (ej. el estado inicial $s_0$).
   • Esto asegura que el modelo aprenda a atraer trayectorias hacia ventanas de datos que coinciden con la condición actual, evitando el colapso hacia promedios globales.

2. Campo de Deriva Clave (Keyed Drift Field):

   • El entrenamiento utiliza un objetivo de "campo de deriva" que tiene dos roles:
     • Atracción: Atrae las muestras generadas hacia ventanas de datos que coinciden con la condición (calculada en el espacio de claves).
     • Repulsión: Repulsa las muestras generadas de otras muestras generadas cercanas en el mismo lote para mantener la diversidad multimodal.
   • Aunque la similitud se calcula en el espacio de claves, las actualizaciones (deriva) se aplican en el espacio completo de la trayectoria ($\mathbb{R}^{H \times D}$).

3. Amortización con Objetivo Desviado (Stop-Gradient Drifted Target):

   • En lugar de aplicar iterativamente el campo de deriva durante la inferencia, KDP "amortiza" este refinamiento en el entrenamiento.
   • Se calcula un objetivo de trayectoria refinado ($\tilde{x}$) sumando el vector de deriva a la salida del generador.
   • Se utiliza un stop-gradient en este objetivo para que el modelo aprenda a predecir directamente la trayectoria refinada en un solo paso de forward pass.
   • Se aplican máscaras de restricción (clamping) para asegurar que las dimensiones condicionadas (como el estado inicial) no se modifiquen durante el refinamiento.

4. Planificación en Tiempo Real:

   • En la inferencia, el modelo genera ventanas completas de trayectorias en una sola pasada hacia adelante.
   • Se pueden generar múltiples candidatos en paralelo ($K$) y seleccionar el mejor mediante un evaluador (scorer), aprovechando el paralelismo masivo de hardware moderno en lugar de la profundidad secuencial.

3. Contribuciones Clave

1. Formulación de Campo de Deriva Clave: Una formulación práctica para la generación condicional de trayectorias que evita el colapso modal bajo restricciones de condicionamiento estrictas, utilizando un espacio de claves alineado con la condición.
2. Receta de Planificación de Un Paso: Un método que preserva la interfaz de los planificadores de difusión (muestreo de ventanas de horizonte) pero elimina el bucle secuencial de $T$ pasos, reduciendo drásticamente la latencia.
3. Evaluación End-to-End: Validación exhaustiva que abarca desde benchmarks de simulación estándar (D4RL) hasta despliegues en hardware de tiempo real (drones y brazos robóticos), demostrando viabilidad bajo presupuestos de latencia estrictos.

4. Resultados Experimentales

Los autores evaluaron KDP en varios dominios:

• Control de Locomoción (D4RL): KDP es competitivo con los métodos más avanzados de RL offline y supera a los planificadores de difusión tradicionales (como Diffuser) en tareas de locomoción, logrando un rendimiento alto con solo 1 paso de inferencia (NFE=1) en comparación con los 20 pasos de Diffuser.
• Planificación de Largo Horizonte (Maze2D, AntMaze): En tareas con recompensas dispersas y horizontes largos, KDP mantiene una ventaja de latencia masiva (milisegundos vs. segundos) mientras logra un rendimiento superior o comparable. La ablación confirma que sin el "espacio de claves", el rendimiento colapsa.
• Manipulación Dexterosa (Adroit): En tareas de alta complejidad (pen, door, hammer, relocate), KDP supera significativamente a Diffuser tanto en puntuación como en tiempo. La capacidad de generar acciones precisas en el primer paso es crucial para evitar errores de contacto.
• Despliegue en Hardware (Crazyflie y SO-100):
  • En navegación de drones y manipulación de brazos, KDP permitió tasas de replanificación mucho más altas (38 Hz vs 4 Hz para Diffuser).
  • Esto resultó en una mayor robustez ante perturbaciones y fallos, validando que la planificación reactiva de tiempo real es posible con modelos generativos.

Métricas de Eficiencia:

• Latencia de Planificación: KDP reduce la latencia de ~200-300 ms (Diffuser) a ~8-10 ms en hardware.
• Costo Computacional: Reduce las evaluaciones secuenciales de la red neuronal de $T$ (ej. 20) a 1, permitiendo escalar el número de candidatos ($K$) mediante paralelismo en lugar de secuencialidad.

5. Significado e Impacto

Este trabajo es significativo porque cierra la brecha entre la capacidad de modelado de los métodos generativos (difusión) y los requisitos de tiempo real de la robótica.

• Cambio de Paradigma: Demuestra que no es necesario sacrificar la calidad del modelo o la diversidad multimodal para ganar velocidad. La "refinación" puede aprenderse durante el entrenamiento mediante objetivos de deriva, eliminando la necesidad de iteraciones costosas en tiempo de ejecución.
• Viabilidad Práctica: Permite el uso de planificadores basados en difusión en sistemas de control de bucle cerrado con frecuencias altas, lo cual era previamente impráctico debido a la latencia.
• Robustez: Al permitir una replanificación frecuente, el sistema se vuelve más reactivo a errores de modelado y perturbaciones del entorno, un factor crítico para el éxito en tareas de manipulación y navegación en el mundo real.

En resumen, KDP ofrece una solución elegante para el cuello de botella de inferencia de la difusión, transformando un proceso iterativo lento en un generador condicional rápido y robusto, sin perder las ventajas de la modelación de distribuciones multimodales.