Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility

Planteamiento del Problema

Los modelos de mundo-acción (WAM, por sus siglas en inglés) han surgido como una alternativa poderosa a los modelos de visión-lenguaje-acción (VLA) al predecir conjuntamente observaciones y acciones futuras, aprendiendo así representaciones espaciotemporales robustas y heredando prioris de la generación de video a gran escala. Sin embargo, los WAM generalistas actuales predicen casi exclusivamente latentes de imágenes RGB futuras. Estos latentes, entrenados puramente para la reconstrucción de píxeles, capturan detalles de apariencia pero a menudo fallan en alinearse con la geometría 3D explícita y la semántica a nivel de objeto requerida para la manipulación robótica precisa.

Supervisar los WAM directamente en geometría y semántica suele incurrir en tres costos significativos:

1. Sensores Adicionales: Requiere nuevas modalidades de hardware (por ejemplo, sensores de profundidad).
2. Nuevo Pre-entrenamiento: Necesita el entrenamiento de nuevos prioris visuales desde cero.
3. Latencia de Inferencia: Ralentiza el despliegue debido al procesamiento multimodal complejo.

El artículo plantea la siguiente pregunta: ¿Cómo podemos amplificar las fortalezas de los WAM (predicción conjunta y fuertes prioris visuales) para incluir la geometría 3D y la semántica sin estos sacrificios?

Metodología: FLEX-π

Los autores presentan FLEX-π, un WAM de 6 mil millones de parámetros que predice simultáneamente observaciones RGB futuras, mapas de puntos 3D y semántica centrada en objetos de DINO. La innovación central es lograr este "almuerzo gratis" de fundamentación multimodal sin nuevos sensores ni pre-entrenamiento.

1. Espacio Latente Compartido mediante un VAE Congelado

FLEX-π utiliza un Autoencoder Variacional (VAE) único y congelado de un modelo de generación de video pre-entrenado (Wan-2.2).

• Codificación RGB: El VAE codifica imágenes RGB como de costumbre.
• Codificación de Mapas de Puntos 3D: Sorprendentemente, los autores encuentran que este mismo VAE congelado, entrenado solo en píxeles RGB, puede codificar y decodificar mapas de puntos 3D de forma casi sin pérdidas. Esto permite que la geometría 3D se proyecte en el mismo espacio latente que el RGB sin entrenar un nuevo codificador.
• Codificación Semántica: La semántica centrada en objetos se deriva utilizando un codificador DINOv3 congelado.

2. Backbone de Mezcla de Transformers (MoT) de Múltiples Flujos

El modelo incrusta tres flujos visuales (RGB, Mapa de Puntos, DINO) en un espacio latente compartido de flujos de tokens.

• Arquitectura: Un backbone de Mezcla de Transformers (5 mil millones de parámetros) procesa estos flujos. Utiliza un backbone de transformer visual compartido pero mantiene redes de proyección feedforward independientes para cada modalidad.
• Experto de Acción: Un segundo y más estrecho Transformer Experto de Acción (1 mil millones de parámetros) realiza atención cruzada con los flujos visuales para generar fragmentos de acción (action chunks).
• Generación Conjunta Causal: Los tokens de acción atienden a las observaciones actuales y realizan atención cruzada con los tokens visuales futuros activos. Esto permite que la generación de acciones se beneficie de la representación evolutiva del futuro del modelo.

3. Entrenamiento Flexible mediante Dropout de Flujo y Forzado de Cross-Modalidad

Para permitir que un único checkpoint soporte cualquier combinación de entradas y salidas en la inferencia, FLEX-π emplea dos estrategias clave de entrenamiento:

• Dropout de Flujo Visual: Durante el entrenamiento, los flujos de entrada se eliminan aleatoriamente (con una probabilidad de 0.5), asegurando que el modelo pueda operar con observaciones parciales (por ejemplo, solo RGB).
• Forzado de Cross-Modalidad: Crucialmente, incluso si un flujo es eliminado del input, el modelo sigue siendo forzado a denotar y predecir ese flujo en el output. Por ejemplo, si el mapa de puntos 3D falta en la entrada, el modelo debe sintetizar el mapa de puntos futuro basándose en las entradas de RGB y DINO. Esto obliga al modelo a interiorizar una representación unificada donde la apariencia, la geometría y la semántica son mutuamente predictivas.

4. Flexibilidad de Inferencia

En el despliegue, los usuarios pueden intercambiar latencia y rendimiento seleccionando qué flujos generar:

• Modo Solo-Acción: El modelo predice solo acciones (inferencia más rápida).
• Generación Conjunta Completa: El modelo predice acciones más RGB futuros, mapas de puntos y características DINO (máximo rendimiento).
• Generación Parcial: Se puede generar cualquier subconjunto de flujos.

Contribuciones Clave

1. Fundamentación Geométrica y Semántica Gratuita: Demostrar que un VAE de video pre-entrenado en RGB puede codificar mapas de puntos 3D sin pérdidas, permitiendo que los WAM sean supervisados en geometría 3D y semántica DINO sin nuevos sensores o pre-entrenamiento.
2. Flexibilidad de Cómputo: Un único checkpoint entrenado soporta cualquier subconjunto de flujos de entrada/salida, permitiendo a los profesionales elegir su punto de operación en la frontera de velocidad-rendimiento en el momento del despliegue.
3. Forzado de Cross-Modalidad: Un régimen de entrenamiento que obliga al modelo a predecir las modalidades faltantes a partir de las observadas, mejorando significativamente la representación compartida y el rendimiento de la política.

Resultados Experimentales

Los autores evalúan FLEX-π en entornos de simulación (RoboTwin, LIBERO, LIBERO-Plus) y tareas de manipulación bimanual en el mundo real con un robot YAM.

• Rendimiento en el Mundo Real: En cinco tareas manuales precisas (incluyendo "Self-Repair Gripper" y "Soft-Bag Zipping"), FLEX-π supera a los fuertes baselines (π0.5, ManiFlow, Fast-WAM) por un factor de 2 a 7× en tasas de éxito.
  • Notablemente, la variante de solo acción de FLEX-π (que no genera futuros visuales en la inferencia) ya supera a todos los baselines y es más rápida que π0.5.
  • La generación conjunta completa aumenta aún más el rendimiento, logrando hasta una mejora absoluta del 42.7% en la completitud de tareas en tareas difíciles comparado con el baseline más fuerte.
• Eficiencia de Datos: FLEX-π es significativamente más eficiente en datos. En RoboTwin, supera a los baselines por 1.9–4.5× en regímenes de bajos datos (50 demostraciones). Incluso cuando se entrena con solo el 50% de los datos reales, FLEX-π supera a los baselines entrenados con el conjunto de datos completo.
• Generalización (OOD): FLEX-π demuestra robustez ante cambios fuera de la distribución (objetos no vistos, desorden, distractores). Mientras que los baselines como ManiFlow y π0.5 sufren caídas significativas de rendimiento (25–26%) bajo cambios de distribución, FLEX-π cae menos del 5%.
• Ablaciones:
  • Eliminar las entradas de DINO o el Mapa de Puntos durante el entrenamiento reduce las tasas de éxito significativamente (en un 6.8% y 20% respectivamente).
  • Eliminar el Forzado de Cross-Modalidad perjudica las tasas de éxito en un 21%, confirmando que forzar al modelo a sintetizar las modalidades faltantes es crítico para aprender una representación compartida sólida.

Significancia y Reivindicaciones

El artículo afirma que FLEX-π demuestra que un WAM puede obtener fundamentación semántica y geométrica esencialmente gratis. Al aprovechar el espacio latente de un modelo de video pre-entrenado y emplear el forzado de cross-modalidad, el modelo logra:

• Mejor Rendimiento de la Política: Tasas de éxito significativamente más altas en tareas complejas, ricas en contacto y de largo horizonte.
• Sin Sobrecarga de Hardware: No requiere sensores adicionales más allá de RGB.
• Sin Sobrecarga de Pre-entrenamiento: No necesita entrenar nuevos prioris visuales.
• Flexibilidad de Despliegue: La capacidad de ejecutar un modo rápido de solo acción o un modo de generación conjunta de alto rendimiento usando el mismo checkpoint.

Los autores concluyen que, si bien FLEX-π es altamente efectivo, todavía requiere más épocas de ajuste fino para converger que los modelos más simples, y su modo de generación conjunta completa es más lento que los VLAs con parámetros comparables. Sugieren que las mejoras futuras podrían provenir de combinar FLEX-π con backbones de VLM más fuertes y acceso a conjuntos de datos robóticos aún más grandes.