WoVR: World Models as Reliable Simulators for Post-Training VLA Policies with RL

Resumen Técnico: WoVR

1. El Problema: Alucinaciones en Modelos de Mundo para RL

Los modelos Visión-Lenguaje-Acción (VLA) han avanzado significativamente mediante aprendizaje por imitación, pero este paradigma está limitado por la calidad y cobertura de los datos de demostración. El Aprendizaje por Refuerzo (RL) ofrece un camino para superar estas limitaciones, pero su aplicación directa en robots físicos es impráctica debido al alto costo, la lentitud y los riesgos de la interacción real.

La solución propuesta anteriormente ha sido utilizar Modelos de Mundo (aprendidos a partir de datos) como simuladores para optimizar políticas en "imaginación" (rollouts cerrados). Sin embargo, el artículo identifica un problema fundamental: las alucinaciones.

• Definición: Una alucinación es una discrepancia sistemática entre los resultados imaginados y los reales en interacciones en bucle cerrado.
• Causas: Los errores de predicción se acumulan debido a la retroalimentación autorregresiva (el modelo se basa en sus propias salidas previas) y al desplazamiento de la distribución (la política evoluciona y se aleja de los datos de entrenamiento del modelo).
• Consecuencia: Si el RL confía ciegamente en estos rollouts alucinados, la política aprende a explotar los errores del modelo en lugar de mejorar en la tarea real, lo que lleva a un fracaso en la transferencia al mundo real.

2. Metodología: El Marco WoVR

Los autores proponen WoVR (World-model-based Reinforcement Learning), un marco diseñado específicamente para controlar la fiabilidad y mitigar las alucinaciones en tres niveles interconectados:

A. Nivel de Simulador: Modelo de Mundo Estabilizado y Condicionado a la Acción

• Arquitectura: Se basa en un modelo de difusión de video (Wan2.2-TI2V-5B) modificado para ser un generador condicionado a la acción.
• Inyección de Acción de Doble Canal: Utiliza dos vías para condicionar la generación: modulación tipo AdaLN-Zero en el nivel de características y reemplazo de embeddings de texto por embeddings de acción en la atención cruzada. Esto asegura que los cambios de estado respondan causalmente a las acciones del robot.
• Anclaje al Primer Cuadro (First-Frame Anchoring): Para evitar el colapso del fondo y la deriva espacial en rollouts largos, el modelo se condiciona constantemente en el primer cuadro de la escena junto con los marcos de memoria recientes. Esto mantiene la consistencia global de la escena.
• Entrenamiento: Se utiliza un objetivo de Rectified Flow y se inyecta ruido en los marcos de contexto durante el entrenamiento para reducir la brecha entre entrenamiento e inferencia.

B. Nivel de Interacción: Optimización de Políticas Consciente de Alucinaciones

• Rollouts Iniciados por Fotogramas Clave (Keyframe-Initialized Rollouts - KIR): En lugar de iniciar siempre desde el estado inicial (donde los errores se acumulan antes de llegar a puntos críticos), WoVR inicia los rollouts desde "fotogramas clave" cercanos a estados críticos o de fallo. Esto reduce la profundidad efectiva de error y permite que el modelo prediga correctamente las transiciones físicas difíciles.
• Optimización GRPO Enmascarada: Se utiliza Group Relative Policy Optimization (GRPO). Se enmascaran los pasos posteriores a la "éxito" en la imaginación para evitar que la política optimice sobre señales de éxito falsas generadas por el modelo. Además, se normaliza la longitud de la trayectoria para dar más peso a los segmentos cortos y críticos.

C. Nivel de Alineación: Co-evolución Política-Simulador (PACE)

• Estrategia PACE: A medida que la política mejora, su distribución de acciones se desvía de los datos originales del modelo de mundo. WoVR implementa una estrategia de co-evolución de baja frecuencia:
  1. Entrenar un modelo base (WMBase) con la política inicial.
  2. Optimizar la política en el simulador.
  3. Recopilar una pequeña cantidad de nuevos rollouts con la política evolucionada y refinar el modelo de mundo (WMEvo) una sola vez (o muy pocas veces).
• Esto alinea el simulador con la política actual sin requerir supervisión humana continua ni actualizaciones de alta frecuencia, manteniendo la fiabilidad del simulador.

3. Contribuciones Clave

1. Identificación del Problema de Fiabilidad: Demuestran que el uso de modelos de mundo para RL no es solo un problema de modelado, sino de fiabilidad, donde las alucinaciones corrompen las señales de optimización.
2. Marco WoVR: Presentan un sistema unificado que regula el diseño del simulador, el protocolo de interacción y la alineación política-simulador para permitir una optimización estable en imaginación.
3. Mecanismos Específicos: Introducen KIR para reducir la profundidad de error y PACE para mitigar el desplazamiento de distribución sin sobrecarga operativa.

4. Resultados Experimentales

Los experimentos se realizaron en el entorno de simulación LIBERO y en un robot físico real (Franka Emika Panda).

• Calidad del Modelo de Mundo:

  • WoVR superó a modelos de referencia (EVAC, Cosmos-Predict2, OpenSora) en todas las métricas de generación de video (LPIPS, FID, FVD, FloLPIPS), especialmente en horizontes largos (512 cuadros).
  • Logró una eficiencia de inferencia de 23 FPS, permitiendo rollouts en tiempo real.

• Rendimiento de la Tarea (LIBERO):

  • Mejora Promedio: WoVR aumentó la tasa de éxito promedio en LIBERO de 39.95% a 69.2% (+29.3 puntos) comparado con la política base.
  • Comparativa: Superó significativamente al RL online estándar (44.6%) y a otros métodos basados en modelos de mundo como WMPO (46.2%).
  • Tareas de Largo Horizonte: Fue particularmente efectivo en el conjunto LIBERO-Long, donde otros métodos fallaron debido a la acumulación de errores.

• Transferencia al Mundo Real:

  • En tareas de manipulación real (coger un plátano y ponerlo en un plato; coger pan), la política optimizada con WoVR logró una tasa de éxito del 91.7% (frente al 61.7% de la política base).
  • Esto demuestra una transferencia sim-to-real robusta sin necesidad de interacción online adicional durante la fase de optimización.

5. Significado e Impacto

El trabajo WoVR es significativo porque cambia el paradigma de asumir que los modelos de mundo son simuladores perfectos a tratarlos como simuladores imperfectos pero controlables.

• Viabilidad del RL en Imaginación: Demuestra que es posible realizar optimización de políticas estable y eficiente en robots físicos utilizando únicamente un simulador aprendido, eliminando la necesidad costosa de millones de interacciones reales.
• Control de Alucinaciones: Proporciona un conjunto de herramientas (anclaje, inicialización por fotogramas clave, co-evolución) para gestionar las limitaciones inherentes de los modelos generativos en entornos de control cerrado.
• Escalabilidad: Al reducir la dependencia de la interacción física, WoVR abre la puerta a la aplicación de RL avanzado en robótica de manipulación compleja, donde la recolección de datos es actualmente el cuello de botella principal.

En conclusión, WoVR establece que, mediante un diseño explícito para controlar la fiabilidad, los modelos de mundo aprendidos pueden servir como simuladores prácticos y efectivos para el entrenamiento de políticas de robots avanzados.