Resumen Técnico: Cosmos Policy

Planteamiento del Problema

Los avances recientes en los modelos de generación de video han demostrado una capacidad impresionante para capturar interacciones físicas complejas, causalidad temporal y evolución de escenas. Sin embargo, adaptar estos modelos para el control robótico ha requerido tradicionalmente canales de entrenamiento complejos de múltiples etapas y la introducción de nuevos componentes arquitectónicos (por ejemplo, difusores de acción separados o modelos de dinámica inversa). Los modelos de video-acción unificados existentes a menudo no logran aprovechar los potentes sesgos espaciotemporales de los modelos de video preentrenados a gran escala debido a diseños personalizados entrenados desde cero. Además, aunque los modelos de Visión-Lenguaje-Acción (VLA) se han convertido en backbones populares, están entrenados principalmente en pares de imagen-texto estáticos, lo que potencialmente limita su comprensión de la dinámica temporal y la física implícita en comparación con los modelos de video.

El desafío central abordado por este trabajo es cómo adaptar eficazmente un modelo de video fundacional grande y preentrenado en una política robótica mediante un proceso simple de una sola etapa sin modificaciones arquitectónicas, permitiendo simultáneamente que la política realice la planificación basada en modelos utilizando modelos de mundo y funciones de valor aprendidas.

Metodología

Arquitectura de Cosmos Policy

Los autores presentan Cosmos Policy, una política robótica ajustada (fine-tuned) a partir del modelo de video fundacional NVIDIA Cosmos-Predict2-2B. El enfoque se distingue por la ausencia de cambios arquitectónicos; en su lugar, aprovecha el proceso de difusión latente nativo del modelo para manejar entradas y salidas multimodales.

1. Inyección de Fotogramas Latentes (Latent Frame Injection):
Para incorporar las modalidades no soportadas originalmente por el modelo de video (propiocepción robótica, acciones, valores de estado y cámaras multivista), los autores proponen la Inyección de Fotogramas Latentes.

• Mecanismo: Las modalidades adicionales se codifican como nuevos fotogramas latentes y se intercalan directamente en la secuencia de difusión latente del modelo de video.
• Estructura de la Secuencia: Para una configuración robótica con múltiples cámaras y propiocepción, la secuencia latente se estructura como $(s, a, s', V(s'))$, donde:
  • $s$: Estado actual (propiocepción del robot + imágenes de cámara actuales).
  • $a$: Chunk de acción (codificado como un fotograma latente).
  • $s'$: Estado futuro (propiocepción futura + imágenes de cámara futuras).
  • $V(s')$: Valor del estado futuro (recompensa acumulada esperada, codificada como un fotograma latente).
• Flexibilidad: Este método permite configuraciones robóticas arbitrarias (por ejemplo, variando el número de cámaras) simplemente añadiendo o eliminando fotogramas latentes en la secuencia sin alterar los pesos o la estructura del modelo.

2. Objetivos de Entrenamiento Conjunto:
El modelo se ajusta utilizando datos de demostración robótica mediante un objetivo de una sola etapa que aprende conjuntamente tres funciones:

• Política (Policy): Predice acciones, estados futuros y valores condicionados al estado actual ($p(a, s', V(s') | s)$).
• Modelo de Mundo (World Model): Predice estados futuros y valores condicionados al estado y la acción ($p(s', V(s') | s, a)$).
• Función de Valor (Value Function): Predice valores futuros condicionados al estado, la acción y el estado futuro ($p(V(s') | s, a, s')$).

Los datos de entrenamiento consisten en una mezcla de demostraciones de expertos (50% del lote) y ejecuciones de la política o rollouts (50% del lote). Los datos de los rollouts se utilizan para entrenar el modelo de mundo y la función de valor, permitiendo que el sistema aprenda tanto de trayectorias exitosas como fallidas.

3. Planificación con Búsqueda Basada en Modelos:
Cosmos Policy soporta dos modos de despliegue:

• Política Directa: Utiliza decodificación paralela para generar acciones directamente desde el estado actual.
• Planificación Basada en Modelos: Utiliza el modelo de mundo y la función de valor aprendidos para buscar acciones óptimas.
  • Proceso: El sistema muestrea múltiples propuestas de acción, utiliza el "modelo de planificación" (un checkpoint refinado entrenado intensamente con datos de rollouts) para predecir los estados futuros y valores resultantes, y selecciona la acción que conduce al valor predicho más alto.
  • Estrategia de Ensemble: Para manejar distribuciones multimodales y la varianza, el sistema realiza ensambles de predicciones consultando el modelo de mundo y la función de valor múltiples veces por cada acción, agregando los resultados mediante la "media de mayoría" para determinar la probabilidad de éxito.

Contribuciones Clave

1. Adaptación Simple de Modelos de Video: El artículo demuestra que un modelo de video preentrenado de gran tamaño puede adaptarse en una política robótica de vanguardia mediante una única etapa de post-entrenamiento con sin modificaciones arquitectónicas, utilizando la inyección de fotogramas latentes para codificar acciones y valores.
2. Política, Modelo de Mundo y Función de Valor Unificados: El enfoque unifica el aprendizaje de la política, el modelado del mundo y la estimación de valor dentro de un único marco de difusión latente, aprovechando los sesgos preentrenados del video para capturar distribuciones de acción complejas y la dinámica temporal.
3. Planificación mediante Refinamiento de Rollout: Los autores muestran que, al recolectar datos de ejecución de la política y ajustar los componentes de valor y modelo de mundo, el sistema puede realizar una planificación basada en modelos efectiva (muestreo best-of-N) para mejorar las tasas de éxito.
4. Rendimiento de Vanguardia (State-of-the-Art): El método alcanza nuevos hitos en simulación y manipulación en el mundo real, superando a políticas entrenadas desde cero, políticas basadas en video y VLAs ajustados.

Resultados Experimentales

Los autores evaluaron Cosmos Policy en tres entornos de prueba:

• LIBERO (Simulación): Logró una tasa de éxito promedio del 98.5%, superando a bases sólidas como Diffusion Policy, Dita, $\pi_0$, OpenVLA-OFT y CogVLA.
• RoboCasa (Simulación): Logró una tasa de éxito promedio del 67.1% utilizando solo 50 demostraciones humanas por tarea (significativamente menos de las 300+ utilizadas por muchos competidores), superando métodos como GR00T-N1.5 y Video Policy.
• Tareas ALOHA en el Mundo Real: En tareas de manipulación bimanual desafiantes (por ejemplo, doblar camisas, manipular bolsas ziploc), Cosmos Policy alcanzó la puntuación promedio más alta (93.6%).
  • Superó notablemente a los VLAs ajustados ($\pi_0.5$, OpenVLA-OFT+) en tareas que requieren agarres de alta precisión y manejo de alta multimodalidad de acciones.
  • Impacto de la Planificación: Al ser potenciado con la planificación basada en modelos usando datos de rollouts, la política experimentó un incremento promedio del 12.5% en las puntuaciones de finalización de tareas en las tareas más difíciles de ALOHA.

Estudios de Ablación:

• Eliminar los objetivos de aprendizaje auxiliares (predecir $s'$ y $V$ junto con $a$) resultó en una caída del 1.5% en la tasa de éxito.
• Entrenar el modelo desde cero (sin los pesos de video preentrenados) resultó en una caída del 3.9%, resaltando la importancia de los sesgos del modelo de video.
• La planificación basada en modelos utilizando una formulación $V(s')$ (que requiere un modelo de mundo) superó a una variante de modelo libre $Q(s, a)$, sugiriendo el valor de la dinámica del entorno aprendida.

Significancia y Reivindicaciones

El artículo afirma que Cosmos Policy representa un paso significativo hacia adelante en el aprendizaje robótico al demostrar que los sesgos espaciotemporales de los modelos de generación de video a gran escala son una base superior para el control de bajo nivel en comparación con los modelos VLA basados en imágenes estáticas.

Las afirmaciones clave incluyen:

• Simplicidad y Eficiencia: La capacidad de crear una política robótica poderosa sin cambios arquitectónicos ni canales de entrenamiento de múltiples etapas.
• Eficiencia de Datos: El método logra un alto rendimiento con menos demostraciones (50 frente a 300+) en comparación con otros métodos de vanguardia, probablemente debido a la fuerte inicialización desde el modelo de video fundacional.
• Capacidad de Planificación: La arquitectura unificada soporta naturalmente la planificación basada en modelos, permitiendo que el robot "imagine" estados futuros y valores para seleccionar mejores acciones, una capacidad que mejora el rendimiento en tareas de alta precisión y multimodalidad donde las políticas directas tienen dificultades.
• Generalización: El enfoque muestra una fuerte generalización tanto en escenarios in-distribution como out-of-distribution en entornos de simulación y del mundo real.

Los autores reconocen limitaciones, específicamente la velocidad de inferencia de la planificación basada en modelos (aproximadamente 5 segundos por chunk de acción) y la necesidad actual de una cantidad sustancial de datos de rollout para refinar el modelo de mundo, identificando estos como áreas para mejoras futuras.