Coupled Local and Global World Models for Efficient First Order RL

Resumen Técnico: Modelos de Mundo Acoplados Local y Global para un Aprendizaje por Refuerzo de Primer Orden Eficiente

Planteamiento del Problema

El Aprendizaje por Refuerzo (RL) ha logrado un éxito significativo en tareas de locomoción mediante simuladores en GPU paralelizados, pero enfrenta barreras sustanciales en el control perceptual y la manipulación. Dos problemas primarios obstaculizan el progreso:

1. Ineficiencia de Muestreo: El aprendizaje en línea práctico directamente desde interacciones con robots reales suele ser inviable debido al alto costo de la recolección de datos.
2. Fidelidad de la Simulación: Simular tareas de manipulación es inherentemente más complejo que la locomoción. El estado del entorno cambia dinámicamente con las acciones del agente (por ejemplo, objetos deformables, dinámica granular), introduciendo una variabilidad infinita que los simuladores de física actuales luchan por modelar con la fidelidad suficiente para la transferencia de simulación a realidad (sim-to-real).

Si bien el aprendizaje de agentes dentro de modelos de mundo diferenciables utilizando métodos de Gradiente de Primer Orden (FoG) ofrece una vía para resolver la ineficiencia de muestreo, los enfoques existentes enfrentan un cuello de botella de escalabilidad. Los modelos de mundo a gran escala (por ejemplo, modelos de difusión) requeridos para una dinámica de alta fidelidad en el mundo real son computacionalmente costosos de diferenciar, lo que hace que la retropropagación estándar sea intratable para entrenar políticas FoG. Por el contrario, los modelos ligeros adecuados para la diferenciación carecen de la expresividad necesaria para capturar la compleja dinámica del mundo real.

Metodología

El artículo propone un marco novedoso que desacopla la simulación hacia adelante de la diferenciación hacia atrás para permitir un RL basado en modelos de primer orden (FoG-MBRL) eficiente utilizando modelos de mundo a gran escala. La arquitectura central consta de dos componentes acoplados:

1. Modelo de Mundo Global (Paso hacia Adelante - Forward Pass):

   • Un modelo a gran escala, entrenado fuera de línea (offline), actúa como un simulador "global".
   • Este modelo opera en el espacio de píxeles (o en un espacio latente de alta dimensión) para generar trayectorias hacia adelante ($s^{\rightarrow}_{t+1}$) de alta fidelidad que capturen con precisión la compleja dinámica visual y de contacto.
   • Se entrena con datos de juego (play data) no estructurados y de alta entropía recolectados de robots reales, eliminando la necesidad de simuladores de física diseñados manualmente.

2. Modelo de Mundo Local (Paso hacia Atrás - Backward Pass):

   • Un Modelo de Espacio de Estado Recurrente (RSSM) ligero actúa como un sustituto "local" en un espacio latente de baja dimensión.
   • Este modelo aproxima la dinámica local y proporciona los gradientes necesarios para la optimía de la política.
   • Crucialmente, los gradientes se calculan evaluando los Jacobianos del modelo local en los puntos de estado generados por el modelo global ($s^{\rightarrow}_{t+1}$), en lugar de las propias predicciones del modelo local. Este desacoplamiento permite que el sistema aproveche la precisión del modelo global para la generación de trayectorias, mientras utiliza el modelo local para una estimación de gradientes computacionalmente tratable.

3. Protocolo de Entrenamiento:

   • Recolección de Datos: Se recolectan datos de juego no estructurados en robots reales sin una segmentación estricta de tareas.
   • Modelado de Recompensa: Dependiendo de la tarea, las señales de recompensa se derivan de funciones de energía contrastiva (entrenadas en datos de juego/demostración) o cabezales de clasificación ligeros que identifican comportías relevantes para la tarea.
   • Optimización: La política se optimiza enteramente dentro del modelo de mundo utilizando el algoritmo de Optimización de Política Basada en Modelos con Gradiente Desacoplado (DMO). El modelo local se ajusta continuamente durante el entrenamiento de RL para mantener la precisión relativa a la distribución de acciones de la política actual.

Contribuciones Clave

• Marco FoG-MBRL Desacoplado: Los autores introducen un método que separa la simulación hacia adelante de alta fidelidad (vía grandes modelos de difusión/flujo) de la diferenciación hacia atrás eficiente (vía RSSM ligeros). Esto resuelve el compromiso entre la complejidad del modelado y el costo de computación del gradiente.
• Aprendizaje Libre de Simuladores a partir de Datos Reales: El enfoque entrena agentes completamente desde cero utilizando datos de robots reales, evitando la necesidad de simuladores de física y abordando directamente la brecha de simulación a realidad.
• Despliegue Zero-Shot: El marco demuestra la capacidad de entrenar políticas en el modelo de mundo y desplegarlas en robots reales sin necesidad de ajuste fino adicional o guía de clonación de comportamiento.
• Robustez ante Dinámicas Complejas: El método maneja con éxito tareas que involucran observabilidad parcial, dinámicas de contacto complejas y altos grados de libertad (por ejemplo, manipulación humanoide).

Resultados Experimentales

Los autores evaluaron su enfoque en tres tareas robóticas del mundo real de complejidad creciente:

1. Push-T: Una tarea de manipulación en mesa para empujar un objeto en forma de T.
2. Ego-Centric Push Cube: Un robot cuadrúpedo empujando un cubo hacia una meta, involucrando dinámicas de marcha y oclusión parcial.
3. Humanoid Grasp and Lift: Un humanoide de alto DoF (grados de libertad) agarrando y levantando una caja usando una mano diestra.

Hallazgos Clave:

• Eficiencia de Muestreo y Tiempo: El método DMO propuesto superó significativamente a PPO (Optimización de Política Próxima) tanto en eficiencia de muestreo como en tiempo de ejecución (wall-clock time) en todas las tareas. PPO tuvo dificultades para converger o requirió muestras excesivas, cayendo a menudo en soluciones locales subóptimas.
• Necesidad de Modelos Globales: Los estudios de ablación ("No Diffusion") que reemplazan el modelo de difusión global con un RSSM estándar para los pasos hacia adelante resultaron en un fallo total de la tarea (0% de tasa de éxito), confirmando que los modelos hacia adelante de alta fidelidad son esenciales para la manipulación compleja.
• Superioridad sobre la Clonación de Comportamiento: Las políticas aprendidas superaron a las líneas base de Clonación de Comportamiento (Action Chunking Transformer), demostrando la capacidad de descubrir estrategias a largo plazo y manejar regiones del espacio de estados no bien cubiertas por las demostraciones.
• Éxito Zero-Shot: Las políticas entrenadas enteramente en el modelo de mundo lograron altas tasas de éxito (por ejemplo, 9/10 en Push-T, 8/10 en Humanoid Grasp) cuando se desplegaron en robots reales sin entrenamiento adicional en el mundo real.

Significado y Reivindicaciones

El artículo sostiene que el aprendizaje de agentes enteramente dentro de modelos de mundo basados en datos es una vía factible y prometedora para resolver tareas difíciles de modelar en el espacio de la imagen. Al desacoplar los requisitos de precisión hacia adelante y la tratabilidad hacia atrás, los autores demuestran que los modelos de mundo a gran escala y computacionalmente pesados pueden utilizarse eficazmente para la optimización de políticas eficientes.

Este trabajo se posiciona como una solución a la dependencia de los simuladores de física diseñados a mano, ofreciendo un método que escala a entornos complejos y no estructurados mediante la adquisición de datos del mundo real. Los autores señalan que, si bien su implementación actual requiere entrenar modelos globales desde cero y depende de recompensas etiquetadas manualmente para ciertas tareas, el marco proporciona un proxy seguro para el entrenamiento de robots reales, permitiendo la optimización de políticas escalable sin el riesgo de daños al hardware. El trabajo futuro se identifica en el aprovechamiento de modelos de video pre-entrenados a escala de internet para reducir aún más los requisitos de datos y automatizar el etiquetado de recompensas.