Flexible Multitask Learning with Factorized Diffusion Policy

A continuación se presenta un resumen técnico detallado del artículo "Aprendizaje por Imitación Multitarea Flexible con Política de Difusión Factorizada" (FDP).

1. Planteamiento del Problema

El artículo aborda los desafíos del aprendizaje por imitación multitarea en robótica, específicamente la dificultad de modelar distribuciones de acción altamente multimodales y diversas.

• Limitaciones de los Modelos Monolíticos: Las políticas tradicionales de red única (por ejemplo, la Política de Difusión estándar) luchan por generalizar a través de tareas diversas, a menudo subajustando distribuciones complejas o fallando en representar modos de comportamiento distintos (por ejemplo, diferentes estrategias de agarre para objetos distintos).
• Limitaciones de los Modelos Modulares Existentes: Aunque las arquitecturas de Mezcla de Expertos (MoE) intentan descomponer las tareas, sufren de:
  • Inestabilidad en el Entrenamiento: Dificultad para equilibrar la utilización de los expertos.
  • Problemas de Enrutamiento: La selección discreta de expertos puede conducir a un "desequilibrio de enrutamiento" (algunos expertos nunca se utilizan, otros se sobreutilizan).
  • Mala Interpretabilidad: Los expertos a menudo tienen roles superpuestos o poco claros, y las habilidades pueden distribuirse a través de capas en lugar de estar desacopladas.
  • Olvido Catastrófico: Adaptarse a nuevas tareas a menudo requiere reentrenar todo el modelo, arriesgando la pérdida de habilidades previamente aprendidas.

2. Metodología: Política de Difusión Factorizada (FDP)

Los autores proponen FDP, una arquitectura de política modular que factoriza la distribución de acción en una composición de modelos de difusión especializados.

A. Arquitectura Central

• Factorización: En lugar de una única política $p(a_t|o_t)$, FDP modela la distribución de acción como un producto ponderado de $N$ distribuciones componentes:
  $$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$$
  Aquí, cada $p_i$ es un modelo de difusión especializado (experto) que captura un modo de comportamiento distinto (sub-habilidad), y $w_{t,i}$ son pesos dependientes de la observación.
• Muestreo Composicional:
  • Enrutador: Un MLP ligero toma la observación $o_t$ y genera pesos continuos $\{w_{t,i}\}$ para todos los componentes.
  • Agregación de Puntuaciones: A diferencia de MoE discreto (que selecciona un experto), FDP agrega las estimaciones de puntuación (gradientes del log-probabilidad) de todos los componentes:
    $$\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$$
  • Proceso de Eliminación de Ruido: El paso iterativo de eliminación de ruido utiliza esta suma ponderada de puntuaciones para guiar la generación de la trayectoria de acción. Esto corresponde a muestrear desde el producto de distribuciones, enmarcando efectivamente la generación de acción como un problema de satisfacción de restricciones.

B. Entrenamiento y Adaptación

• Entrenamiento Conjunto: Todos los componentes de difusión y el enrutador se entrenan de extremo a extremo utilizando una pérdida de Error Cuadrático Medio (MSE) sobre la predicción de puntuación agregada.
• Adaptación Eficiente (Reciclaje): Para adaptarse a una nueva tarea sin reentrenar todo el modelo:
  1. Se introduce un nuevo componente de difusión.
  2. Se inicializa mediante reciclaje (copiando pesos) desde un componente existente.
  3. Solo el nuevo componente y el enrutador se ajustan finamente; todos los componentes anteriores se congelan.
  4. Esto previene el olvido catastrófico y reduce significativamente el número de parámetros entrenables.

3. Contribuciones Clave

1. Arquitectura Modular Novel: Se introduce FDP, que utiliza agregación continua de puntuaciones en lugar de selección discreta de expertos. Esto asegura un entrenamiento estable, previene el desequilibrio de enrutamiento y promueve una especialización clara entre los componentes.
2. Interpretación Probabilística: Fundamentado en Modelos Basados en Energía (EBM) y ajuste de puntuaciones, proporciona una formulación principista donde la política compuesta representa la intersección de restricciones de comportamiento.
3. Estrategia de Adaptación Escalable: Se propone un mecanismo de "ajuste selectivo" utilizando reciclaje y congelamiento, permitiendo una expansión eficiente de habilidades con datos mínimos y sin olvido catastrófico.
4. Validación Empírica: Demostró un rendimiento superior frente a líneas base sólidas (DP Monolítico, Política de Difusión Dispersa y MoDE) tanto en entornos de simulación como del mundo real.

4. Resultados Experimentales

Los autores evaluaron FDP en MetaWorld, RLBench, LIBERO y en manipulación robótica del mundo real (brazo UR5e).

• Rendimiento Multitarea:

  • FDP superó consistentemente a la Política de Difusión (DP) monolítica y a las líneas base modulares (SDP, MoDE) en más de 30 tareas.
  • MetaWorld: FDP logró un 74.8% de éxito promedio frente al 70.8% (DP) y 69.8% (SDP).
  • RLBench: FDP logró un 63.9% de éxito promedio frente al 45.6% (DP) y 42.1% (SDP).
  • Mundo Real: FDP mostró tasas de éxito más altas en tareas complejas (por ejemplo, colgar tazas, recoger cubos de colores) donde las líneas base fallaron debido al sobreajuste o la incapacidad de capturar distribuciones multimodales.

• Adaptación y Transferencia de Tareas:

  • Eficiencia de Datos: Al adaptarse a nuevas tareas con datos limitados (10–25 demostraciones), FDP utilizando la estrategia "Añadir Nuevo Módulo" superó al ajuste fino de parámetros completos y al ajuste fino parcial.
  • Aprendizaje Continuo: En una prueba de aprendizaje continuo de 12 tareas (LIBERO), FDP mantuvo un alto rendimiento a medida que se añadían nuevos expertos, mientras que otros métodos se degradaron.
  • Retención de Conocimiento: Congelar los componentes antiguos preservó efectivamente las habilidades previas. Usar un búfer de replay mínimo (5 demostraciones/tarea) permitió que el modelo retuviera casi el 100% del rendimiento original después de la adaptación.

• Análisis y Escalado:

  • Escalado de Componentes: El rendimiento mejoró a medida que el número de componentes aumentaba de 2 a 4, luego se estabilizó.
  • Escalado de Tareas: La brecha de rendimiento entre FDP y los modelos monolíticos se amplió a medida que aumentaba el número de tareas, destacando la escalabilidad de FDP.
  • Interpretabilidad: La visualización mostró que los componentes individuales se especializaban en sub-habilidades distintas (por ejemplo, alineación, aproximación, agarre) en lugar de niveles de ruido o capas arbitrarias.
  • Velocidad de Inferencia: FDP soporta "reconstrucción parcial" (usando solo los componentes top-$k$), logrando una aceleración de 2x en la inferencia con solo una caída de rendimiento de ~19%.

5. Significado

• Aprendizaje Robótico: FDP ofrece una solución práctica para desplegar robots en entornos dinámicos donde deben aprender nuevas habilidades rápidamente sin olvidar las antiguas.
• Estabilidad: Al reemplazar el enrutamiento discreto con la composición continua de puntuaciones, FDP resuelve los problemas de inestabilidad en el entrenamiento que han afectado a las políticas de difusión basadas en MoE.
• Interpretabilidad: La factorización permite a los investigadores inspeccionar y comprender habilidades subespecíficas (por ejemplo, "agarrar" frente a "aproximarse") dentro de la política, ayudando en la depuración y la verificación de seguridad.
• Direcciones Futuras: El marco abre vías para diseños de módulos heterogéneos (mezcla de arquitecturas) y sistemas de aprendizaje de por vida donde los robots acumulan continuamente una biblioteca de habilidades especializadas.

En resumen, FDP representa un avance significativo en el aprendizaje por imitación robótica al combinar el poder generativo de los modelos de difusión con una arquitectura modular robusta, interpretable y adaptable.