Persistent Robot World Models: Stabilizing Multi-Step Rollouts via Reinforcement Learning

1. El Problema: Sesgo de Exposición y Degradación en Despliegues Autoregresivos

Los modelos de mundo robótico basados en difusión de video (como Ctrl-World) han demostrado ser capaces de generar clips de video de alta fidelidad condicionados a secuencias de acciones del robot. Sin embargo, estos modelos enfrentan un desafío crítico cuando se despliegan para despliegues autoregresivos de largo horizonte (generar segundos de video coherente donde cada clip predicho se usa como contexto para el siguiente).

• Sesgo de Exposición (Exposure Bias): Durante el entrenamiento, los modelos se entrenan con "forzamiento de maestro" (teacher forcing), utilizando historias de frames reales (ground-truth). En la inferencia, deben condicionarse en sus propias salidas previas, que contienen imperfecciones.
• Compounding Error (Error Acumulativo): Cualquier error espacial o temporal en un paso se propaga al siguiente, corrompiendo el buffer de historia. Esto provoca una rápida degradación visual: los objetos pierden su identidad estructural (se disuelven en formas amorfas), los efectoros finales del robot se desvían de sus trayectorias y la escena pierde coherencia en cuestión de segundos.
• Limitación Actual: Los modelos existentes no están incentivados a ser robustos ante historias imperfectas, ya que nunca han visto tales historias durante el entrenamiento supervisado estándar.

2. Metodología: PersistWorld

Los autores proponen PersistWorld, un marco de post-entrenamiento mediante Aprendizaje por Refuerzo (RL) diseñado específicamente para cerrar la brecha entre el entrenamiento y la inferencia en modelos de difusión de video.

A. Post-entrenamiento con RL Contrastivo (Adaptado a Difusión)

En lugar de optimizar la verosimilitud en datos reales, el modelo se entrena directamente sobre sus propios rollouts autoregresivos.

• Adaptación de DiffusionNFT: Se adapta un método de RL contrastivo reciente (DiffusionNFT) para modelos de difusión que predicen el dato limpio ($x_0$) en lugar de un campo de velocidad.
• Mecanismo de Actualización:
  1. Se generan múltiples candidatos a partir de un estado compartido.
  2. Se asigna una recompensa basada en la calidad visual.
  3. Se construyen ramas "positivas" y "negativas" de predicción:
     • Rama Positiva: Extrapola la dirección en la que el modelo actual se ha desviado del modelo de referencia (si la recompensa es alta).
     • Rama Negativa: Construye una predicción contrafáctica que se aleja de la salida actual (si la recompensa es baja).
  4. Se minimiza una pérdida ponderada por la recompensa, reforzando las direcciones de alta calidad y penalizando las de baja calidad, sin necesidad de retropropagar a través de todo el proceso de denoising (lo cual sería prohibitivamente costoso).

B. Protocolo de Entrenamiento para Video Autoregresivo

Dado que la generación de video no tiene un "prompt" fijo como en la generación de imágenes, los autores diseñan un protocolo específico:

1. Prefijo Compartido (S1): Se genera un prefijo autoregresivo de longitud variable ($P$) a partir de una condición inicial real. Esto simula el estado de error acumulado del despliegue real.
2. Ramas de Candidatos (S2): Desde el estado congelado del prefijo, se generan $K$ (ej. 16) continuaciones independientes. Esto crea un grupo de candidatos que comparten el mismo contexto histórico (incluyendo sus errores acumulados).
3. Puntuación y Ranking (S3): Los candidatos se comparan con el ground-truth y se normalizan las recompensas dentro del grupo para obtener ventajas relativas.
4. Actualización (S4): Se actualiza el modelo utilizando las recompensas relativas para guiar el aprendizaje contrastivo.

C. Recompensas Visuales Multi-Vista

Para evaluar la calidad en entornos robóticos, se diseñan recompensas densas a nivel de clip que combinan métricas perceptuales de tres cámaras (dos externas y una en la muñeca):

• Métricas Combinadas: LPIPS (similitud perceptual profunda), SSIM (fidelidad estructural) y PSNR (desviación de píxeles).
• Normalización: Las recompensas se normalizan por grupo (z-score) para que el modelo aprenda a discriminar entre mejores y peores continuaciones, independientemente de la calidad absoluta del momento del rollout.
• Evaluación Centrada en la Tarea: Se utilizan máscaras para evaluar específicamente la fidelidad de los objetos manipulados y el robot, evitando que el modelo optimice solo el fondo estático.

3. Contribuciones Clave

1. Esquema de Post-entrenamiento RL: Introducen el primer método que optimiza un modelo de mundo robótico directamente sobre sus propios rollouts autoregresivos, adaptando garantías teóricas de convergencia de RL contrastivo a arquitecturas de predicción $x_0$.
2. Protocolo de Ramificación Variable: Desarrollan un método para generar candidatos comparables en video autoregresivo, exponiendo al modelo a regímenes de error desde etapas tempranas hasta tardías mediante la selección aleatoria de la longitud del prefijo.
3. Recompensas Visuales Eficientes: Crean señales de entrenamiento densas y de baja varianza combinando métricas perceptuales multi-vista, validando que las mejoras provienen de una mejor modelización de la dinámica de interacción y no solo de la preservación del fondo.
4. Nuevo Estado del Arte (SOTA): Demuestran resultados superiores en el conjunto de datos DROID, superando a los modelos base más fuertes en todas las métricas visuales.

4. Resultados Experimentales

Los experimentos se realizaron en el dataset DROID (manipulación robótica a gran escala) utilizando el modelo base Ctrl-World.

• Calidad Visual:
  • Cámaras Externas: Reducción del 14% en LPIPS y mejora del 9.1% en SSIM en la cámara de muñeca comparado con el mejor baseline.
  • Comparación Pareada: El modelo propuesto supera al baseline en aproximadamente el 98% de las muestras de validación en comparaciones uno a uno.
  • Estudio Humano: En un estudio ciego de preferencia humana, los evaluadores prefirieron los rollouts de PersistWorld en un 80% de los casos.
• Estabilidad Temporal: Mientras que los modelos base pierden coherencia estructural en segundos (los objetos se "deshacen"), PersistWorld mantiene la integridad estructural y la consistencia espacial durante despliegues de más de 11 segundos.
• Evaluación de Políticas: El modelo post-entrenado muestra una mayor correlación con el progreso real de tareas (como colocar objetos en cajas) al usarse como entorno de simulación para evaluar políticas, superando al modelo base en métricas de correlación de Pearson y violaciones de rango (MMRV).

5. Significado e Impacto

Este trabajo es fundamental para la viabilidad de los Modelos de Mundo en robótica práctica:

• Simuladores Fiables: Permite utilizar modelos de generación de video como entornos virtuales escalables y de alta fidelidad para la evaluación y mejora de políticas de visión-idioma-acción (VLA), algo que antes estaba limitado por la rápida degradación de la simulación.
• Superación del Sesgo de Exposición: Proporciona una solución generalizable al problema de la brecha entre entrenamiento y prueba en modelos generativos secuenciales, demostrando que el RL puede estabilizar la generación autoregresiva sin necesidad de cambiar la arquitectura base del modelo.
• Robustez en Interacciones: Al centrarse en la fidelidad de los objetos manipulados y la coordinación mano-ojo, el método aborda directamente las necesidades críticas de la manipulación robótica, donde la precisión estructural es más importante que la reconstrucción del fondo.

En resumen, PersistWorld transforma los generadores de video de "creadores de clips cortos" a "simuladores persistentes y estables", abriendo la puerta a la planificación y evaluación de robots en mundos virtuales generados por IA con una fiabilidad sin precedentes.