FlowMPC: Improving Flow Matching policies with World Models

Resumen Técnico: FlowMPC

Planteamiento del Problema

El Coincidencia de Flujo (Flow Matching, FM) ha emergido como un método poderoso para el clonado de comportamiento en espacios de acción multimodales, particularmente en la manipulación robótica donde múltiples trayectorias de acción distintas pueden resolver la misma tarea. Sin embargo, FM es fundamentalmente un método de aprendizaje por imitación entrenado para igualar el comportamiento experto en lugar de maximizar directamente el retorno esperado. En consecuencia, las políticas de FM están limitadas por la calidad y la cobertura de los datos de demostración. A menudo tienen dificultades para recuperarse de errores pequeños, errores acumulativos o estados que no están bien representados en el conjunto de datos experto.

Si bien trabajos recientes han intentado abordar esto refinando las políticas de FM mediante aprendizaje por refuerzo (RL), tales enfoques requieren modificar el propio objetivo de entrenamiento de FM. Este artículo investiga una pregunta diferente: ¿Puede un modelo de mundo aprendido mejorar una política de FM sin alterar el objetivo de entrenamiento de FM? El objetivo es determinar si la planificación basada en modelos puede complementar a las políticas de imitación basadas en flujo durante el tiempo de ejecución.

Metodología: Marco de Trabajo FlowMPC

El marco propuesto, FlowMPC, integra una política de FM aprendida por imitación con un modelo de mundo aprendido dentro de un bucle de planificación de Integral de Trayectoria de Camino de Modelo Predictivo (MPPI). La idea central es utilizar la política de FM como un mecanismo de propuesta estructurado para generar secuencias de acciones candidatas, las cuales son luego evaluadas y refinadas por el modelo de mundo.

1. Paradigma de Entrenamiento

El proceso de entrenamiento consta de cinco etapas:

1. Recolección de Datos Expertos: Un agente SAC con información de estado privilegiada genera demostraciones expertas.
2. Recolección de Trayectorias: Se recolectan trayectorias completas (RGB, propiocepción, acciones, recompensas).
3. Entrenamiento de la Política FM: Se entrena una política de FM sobre estas trayectorias expertas para aprender un prior de acción multimodal.
4. Preprocesamiento Visual: Las trayectorias expertas se recolectan nuevamente con observaciones RGB preprocesadas por el codificador visual de la política de FM para producir latentes visuales compactos ($z_{viz,t}$).
5. Entrenamiento del Modelo de Mundo: El modelo de mundo se entrena en línea utilizando un buffer de repetición precargado con datos expertos, utilizando los latentes visuales preprocesados combinados con el estado propioceptivo ($x_t = [s_t, z_{viz,t}]$).

2. Adaptaciones Arquitectónicas a TD-MPC2

El marco se basa en TD-MPC2, pero reemplaza su actor aprendido con la política de FM aprendida por imitación. Se realizaron varias modificaciones críticas para acomodar este híbrido:

• Predictor de Acción de Clonado de Comportamiento ($\pi_{BC}$): Una red dedicada mapea el estado actual a una acción. Esto se utiliza para el bootstrapping del valor terminal durante los rollouts de planificación para asegurar que la función de valor evalúe acciones dentro de la distribución, estabilizando el proceso de planificación.
• Interfaz de Representación Unificada: El modelo de mundo consume directamente el latente visual de FM concatenado con el estado propioceptivo, en lugar de aprender codificadores separados. Esto permite que el modelo de mundo aproveche las características visuales congeladas de la política de FM.
• Objetivos Simplificados: Dada la estructura de recompensa binaria de las tareas, las complejas recompensas discretizadas con transformación symlog del estándar TD-MPC2 fueron reemplazadas con entropía cruzada binaria para las recompensas y error cuadrático medio estándar para la regresión de valores Q.
• Dos Variantes del Modelo de Mundo:
  • Variante de Espacio de Estados: Predice directamente en el espacio de la concatenación de estado y latente visual. Se utiliza para la tarea PickCube ($H=1$).
  • Variante de Espacio Latente: Sigue el estándar de TD-MPC2 con una representación latente aprendida y rollouts de múltiples pasos. Se utiliza para la tarea más compleja PickSingleYCB ($H=3$).

3. Inferencia y Planificación

En el tiempo de ejecución, el planificador opera de la siguiente manera:

1. Propuesta: Dada la observación actual, el planificador muestrea $K$ trayectorias de acción candidatas. Un subconjunto ($K_\pi$) es propuesto por la política de FM, mientras que el resto se muestrea aleatoriamente.
2. Rollout y Calificación: Cada trayectoria se ejecuta en el modelo de mundo aprendido. Las trayectorias se califican sumando las recompensas predichas y una estimación de valor terminal ($\hat{Q}$) derivada del predictor de clonado de comportamiento.
3. Refinamiento: MPPI repondera y refina las trayectorias basándose en estas puntuaciones. Se ejecuta la primera acción de la secuencia optimizada.
4. Eficiencia: La codificación visual de FM se computa una vez por observación y se reutiliza a través de todas las muestras y pasos de integración.

Resultados Clave

El enfoque fue evaluado en dos tareas de manipulación de ManiSkill: PickCube-v1 y PickSingleYCB-v1.

• Éxito al Final del Episodio: Las ganancias más significativas se observaron en el mantenimiento del éxito hacia el final del episodio.
  • PickCube: Mejoró de 93.14% (solo FM) a 97.44% (FlowMPC).
  • PickSingleYCB: Mejoró de 56.81% (solo FM) a 66.41% (FlowMPC).
• Éxito en Cualquier Momento (Anytime Success): También se observaron mejoras en alcanzar un estado de éxito en cualquier punto, aunque las ganancias fueron menores.
  • PickCube: 95.78% $\to$ 98.68%.
  • PickSingleYCB: 68.77% $\to$ 69.78%.

Los resultados indican que el modelo de mundo es particularmente efectivo para corregir pequeños errores de acción y mantener un control estable cerca de la finalización de la tarea, en lugar de solo ayudar al agente a alcanzar un estado de éxito transitorio.

Significado y Reivindicaciones

El artículo postula que la planificación basada en modelos de mundo puede complementar eficazmente las políticas de imitación basadas en flujo sin requerir un nuevo objetivo de entrenamiento de FM.

• Enfoque Híbrido: El trabajo demuestra un híbrido viable entre el aprendizaje por imitación y el RL basado en modelos. Utiliza la imitación para aprender un prior de acción multimodal fuerte y la planificación para refinar las decisiones en el tiempo de ejecución.
• Sin Reentrenamiento Requerido: La mejora proviene de una mejor utilización de la política en el tiempo de ejecución, no del reentrenamiento de la política con un nuevo objetivo de RL.
• Alcance Modesto: Los autores declaran explícitamente que los resultados están limitados a dos tareas y que el diseño óptimo del modelo de mundo (espacio de estados vs. espacio latente, horizonte de planificación) varió entre tareas. No reclaman una configuración universalmente óptima, sino que sugieren que la combinación de políticas generativas expresivas y planificación con mirada hacia adelante (lookahead) es una dirección prometedora para la mejora de la robustez de las políticas.
• Línea Base de Rendimiento: El método mejoró un fuerte baseline de FM (que ya estaba cerca del rendimiento de TD-MPC2 en una tarea), lo que sugiere que el planificador "afina" una política fuerte en lugar de simplemente rescatar una débil.