Resumen Técnico: WAM-RL

Planteamiento del Problema

Los modelos de Acción-Mundo (WA) recientes han demostrado una fuerte generalización y eficiencia de datos en el aprendizaje de políticas robóticas mediante el modelado conjunto de observaciones futuras y acciones a través de modelos generativos basados en video. Sin embargo, los modelos WA existentes dependen fuertemente del aprendizaje supervisado de trayectorias expertas. Esta dependencia impone dos limitaciones fundamentales:

1. Restricciones Distribucionales: La política aprendida está limitada por el soporte de los datos de entrenamiento, lo que dificulta la adquisición de habilidades de manipulación finas que se encuentran fuera de la distribución de las demostraciones.
2. Falta de Mejora Continua: Los modelos no pueden mejorar continuamente mediante la interacción con el mundo real, lo que obstaculiza su capacidad para adaptarse a nuevos escenarios o corregir errores durante la ejecución.

Si bien el Aprendizaje por Refuerzo (RL) es una solución natural, aplicar RL a los modelos WA no es trivial. A diferencia de los modelos estándar de Visión-Lenguaje-Acción (VLA) donde la representación visual suele ser fija, los modelos WA constan de dos componentes estrechamente acoplados: un modelo de mundo (que genera predicciones futuras) y un modelo de acción (que mapea las predicciones a acciones). Aplicar RL o ajuste fino (fine-tuning) en línea de forma ingenua puede causar cambios de distribución en el espacio latente del modelo de mundo, provocando inestabilidad y degradación del rendimiento debido a que el modelo de acción depende profundamente de dicho espacio latente.

Metodología: WAM-RL

Los autores proponen WAM-RL, un marco que permite la optimización conjunta del modelo de mundo y el modelo de acción mediante la interacción en línea. La idea central es que el modelo de mundo proporciona la estructura predictiva y la capacidad de planificación, mientras que el modelo de acción actúa principalmente como un traductor que convierte las predicciones latentes en acciones ejecutables.

El marco emplea un esquema de optimización de dos partes:

1. SFT de Video en Línea con Regularización KL (Modelo de Mundo)

Para refinar el modelo de mundo, los autores utilizan trayectorias recolectadas durante interacciones exitosas para el ajuste fino supervisado (SFT) auto-supervisado en línea.

• Objetivo: El modelo de mundo se entrena para predecir observaciones futuras utilizando un objetivo de modelado de video estándar (por ejemplo, flow matching).
• Estabilización: Para evitar que la distribución de las características latentes cambie drásticamente (lo que rompería el modelo de acción), se introduce un término de regularización KL. Esto restringe las características latentes del modelo de mundo actualizado para que permanezcan cercanas a las de una copia congelada del modelo preentrenado.
• Implementación: Los autores aproximan las distribuciones latentes como Gaussianas, utilizando un promedio móvil exponencial (EMA) para estimar la covarianza del modelo actual mientras mantienen fija la covarianza de referencia. Esto permite una adaptación gradual preservando la geometría esperada por el actor.

2. RL Basado en la Reconstrucción (Modelo de Acción)

El modelo de acción se optimiza mediante RL para asegurar que las trayectorias ejecutadas realicen fielmente las predicciones del modelo de mundo.

• Señal de Recompensa: En lugar de recompensas de tarea dispersas, los autores definen una recompensa densa basada en la reconstrucción. Esta recompensa mide la consistencia (similitud) entre las observaciones futuras imaginadas (predichas por el modelo de mundo) y las observaciones reales obtenidas al ejecutar la acción en el entorno.
• Métricas de Similitud: Los autores investigan diversas funciones de similitud, incluyendo MSE a nivel de píxel, consistencia de flujo óptico y similitud de características semánticas (DINOv2, V-JEPA2).
• Optimización: El actor se actualiza mediante métodos de gradiente de política utilizando estas recompensas de reconstrucción, cerrando efectivamente la brecha entre la imaginación del modelo de mundo y la ejecución en el mundo real.

Contribuciones Clave

1. Primer Marco de RL para Modelos WA: WAM-RL es el primer marco que incorpora el aprendizaje por refuerzo al paradigma de Acción-Mundo, permitiendo específicamente la optimización conjunta tanto del modelo de mundo como del modelo de acción.
2. Estrategia de Entrenamiento Estable: Los autores proponen una estrategia de entrenamiento novedosa que combina SFT de video en línea con regularización KL para el modelo de mundo y RL basado en la reconstrucción para el actor. Este diseño facilita la co-evolución de ambos componentes sin desestabilizar el espacio latente.
3. Perspectivas Empíricas sobre Tareas de Largo Horizonte: El trabajo proporciona información crítica sobre las limitaciones de la optimización de solo el actor. Los autores demuestran que optimizar únicamente el actor produce mejoras en tareas de corto horizonte, pero no logra ganancias significativas en tareas de largo horizonte debido a la acumulación de errores de predicción. Se demuestra que la optimización conjunta es crítica para el rendimiento complejo de largo horizonte.

Resultados Experimentales

El método fue evaluado en los entornos de referencia LIBERO (manipulación centrada en objetos) y RLBench (habilidades de múltiples pasos).

• Ganancias de Rendimiento: WAM-RL superó consistentemente tanto al modelo base preentrenado (Base) como al RL de solo actor (πRL).
  • En LIBERO-Object, las tasas de éxito mejoraron del 68% (Base) al 82% (WAM-RL), superando significativamente el enfoque de solo actor (78%).
  • En RLBench (Water Plants), WAM-RL mejoró las tasas de éxito del 19% al 22%, mientras que el método de solo actor (πRL) no logró mejorar respecto al modelo base (18%).
• Análisis de Comportamiento:
  • Comportamientos de Recuperación: La adaptación en línea del modelo de mundo permitió la generación de comportamientos de recuperación más realistas (por ejemplo, reintentar agarres fallidos) en lugar de predecir un fallo inmediato.
  • Ablación sobre Recompensas: El estudio encontró que, si bien el flujo óptico proporcionó la señal de recompensa más discriminativa entre el éxito y el fallo, el MSE a nivel de píxel produjo el mejor rendimiento descendente. Los autores hipotetizan que esto se debe a que el MSE de píxeles se alinea mejor con el objetivo de entrenamiento del modelo de mundo e impone una regulación más fuerte contra acciones fuera de la distribución.
  • Ablación sobre Video SFT: Los modelos entrenados sin SFT de video en línea fallaron al recuperarse de errores iniciales, derivando rápidamente hacia comportamientos fuera de la distribución. En contraste, los modelos con video SFT aprendieron a anticipar el fallo y generar acciones correctivas dentro de un único bloque de lazo abierto (open-loop chunk).

Significado y Afirmaciones

El artículo afirma que la efectividad de los modelos de Acción-Mundo está fundamentalmente limitada por la calidad del modelo de mundo. Si bien la optimización de solo el actor puede mejorar las tareas de corto horizonte, no puede abordar los errores de predicción acumulados en entornos de largo horizonte.

La importancia de WAM-RL radica en demostrar que mejorar conjuntamente el modelo de mundo y el modelo de acción es esencial para una toma de decisiones robusta de largo horizonte. Al permitir que el modelo de mundo aprenda la dinámica de fallos y las estrategias de recuperación mediante la interacción en línea, el sistema puede anticipar y corregir desviaciones, lo que conduce a una ejecución más robusta.

Los autores reconocen modestamente las limitaciones, señalando que la regularización KL restringe el alcance de la adaptación del modelo de mundo, lo que podría limitar la expansión de la capacidad más allá de la distribución preentrenada. Además, la dependencia de recompensas basadas en la reconstrucción (a menudo basadas en representaciones preentrenadas) puede ofrecer un contraste limitado entre el éxito y el fallo, lo que sugiere la necesidad de formulaciones de recompensa más expresivas y conscientes de la tarea en trabajos futuros.