Resumen Técnico: Predicción de Metas Latentes a partir de Lenguaje para la Planificación Basada en Modelos

Planteamiento del Problema

El aprendizaje por refuerzo (RL) basado en modelos que utiliza modelos de mundo (WM) ofrece un paradigma prometedor para el control con eficiencia de muestras al simular trayectorias futuras en un espacio latente compacto. Sin embargo, estos sistemas enfrentan dos cuellos de botella fundamentales:

1. Especificación de Metas: Definir metas optimizables es difícil. Los objetivos visuales (imágenes) proporcionan gradientes locales precisos pero ofrecen una guía deficiente para metas distantes y son impracticables de obtener para tareas complejas. El lenguaje ofrece flexibilidad pero sufre de un alineamiento cross-modal ruidoso en modelos contrastivos o de una ineficiencia computacional en modelos generativos grandes, los cuales no son adecuados para el muestreo de alta frecuencia requerido por la planificación basada en modelos.
2. Planificación de Largo Horizonte: La planificación sobre horizontes largos se ve obstaculizada por la acumulación de errores de predicción. Pequeñas imprecisiones en la dinámica latente se acumulan durante los rollouts, degradando la fiabilidad de las trayectorias imaginadas. Los métodos existentes suelen evitar esto aumentando la abstracción temporal (saltos de fotogramas) o restringiendo la evaluación a horizontes cortos, lo que limita su aplicabilidad en tareas complejas y temporalmente extendidas.

Los enfoques actuales luchan por combinar la precisión de los objetivos basados en imágenes con la flexibilidad de la especificación basada en lenguaje, manteniendo la robustez contra los errores de predicción en horizontes largos.

Metodología: LAGO

Los autores presentan LAGO (Latent Goal Prediction from Language), un marco diseñado para predecir secuencias de estados de metas intermedias a partir de instrucciones de lenguaje y rollouts condicionados por acciones dentro del mismo espacio latente.

Arquitectura

LAGO se construye sobre una Arquitectura de Predicción de Incrustación Conjunta (JEPA) que opera en un espacio latente compartido. Soporta dos modos de operación mediante un único predictor $\pi(\cdot)$:

1. Predicción de Transición de Siguiente Paso: Condicionado a acciones de bajo nivel para predecir el siguiente estado latente inmediato.
2. Predicción de Submetas: Condicionado a instrucciones de lenguaje y un escalar de completitud $\rho$ (que estima el progreso de la tarea) para predecir submetas latentes intermedias.

Componentes Clave:

• Codificador de Visión: Utiliza un DINOv2 ViT congelado para codificar observaciones RGB en tokens de parches.
• Condicionamiento Multimodal: Las instrucciones de lenguaje se incrustan y se proyectan en un espacio de condición compartido. Un escalar de completitud $\rho \in [0, 1]$ también se proyecta. Una incrustación específica del modo distingue entre el condicionamiento por acción y el condicionamiento por submeta.
• Entrenamiento Unificado: El modelo se entrena conjuntamente utilizando un objetivo de error cuadrático medio (MSE). Minimiza el error entre los estados predichos y reales (usando condicionamiento por acción) y entre los estados de submeta predichos y reales (usando condicionamiento por lenguaje y $\rho$).

Estrategia de Planificación

LAGO transforma la planificación de largo horizonte en una secuencia de objetivos localmente tratables:

1. Generación de Submetas: Dada una situación actual y una instrucción de lenguaje, LAGO genera $K$ submetas latentes en escalares de completitud distribuidos uniformemente ($\rho_k = k/K$). Crucialmente, estas submetas se predicen de forma independiente del estado actual (no autoregresivamente) para evitar la acumulación de errores.
2. Repredicción en Línea: Las submetas se repredicen en cada paso de Control Predictivo por Modelo (MPC), permitiendo que el agente se adapte a su estado evolutivo.
3. Costo de Trayectoria de Mínimo Suave (JLAGO): Para evaluar secuencias de acciones candidatas, LAGO emplea una función de costo de mínimo suave. En lugar de requerir que una trayectoria alcance una submeta específica en un tiempo específico, mide el mejor alineamiento logrado entre cada submeta y cualquier paso de tiempo a lo largo del rollout. Esto se controla mediante un parámetro de temperatura $\tau$, lo que permite una convergencia suave hacia las submetas en lugar de criterios de completitud discretos y frágiles.

Contribuciones Clave

1. Marco LAGO: Un modelo de mundo entrenado para mapear instrucciones de lenguaje natural directamente a estados objetivo dentro de su propio espacio latente, cerrando la breancia entre la precisión visual y la flexibilidad lingüística.
2. Descomposición Dinámica de Submetas: Una extensión de la predicción de metas para generar secuencias de submetas latentes intermedias. El enfoque utiliza la repredicción continua en línea y un costo de trayectoria de mínimo suave para permitir una planificación robusta de largo horizonte sin detección explícita de completitud de submetas.
3. Protocolo de Evaluación Basado en Distancia: Un método de evaluación sistemático que varía la dificultad de la tarea basándose en la distancia hacia la meta. Esto revela que los métodos previos se degradan bruscamente más allá de su horizonte de planificación efectivo, mientras que LAGO permanece robusto.

Resultados Experimentales

Los autores evaluaron LAGO en cuatro entornos: TwoRoom (navegación 2D), Kingdom (navegación 2D con diseños dinámicos y peligros irreversibles), OGBench Cube (manipulación) y OGBench Scene (manipulación compleja).

• Robustez de Largo Horizonte: En los entornos TwoRoom, Kingdom y OGBench Cube, las líneas base de metas de imagen (DINO-WM, Temporal Straightening, LeWM) fallaron completamente a medida que la dificultad de la tarea (distancia) aumentaba (por ejemplo, cayendo al 0% de tasa de éxito en dificultades altas). En contraste, LAGO mantuvo tasas de éxito altas (por ejemplo, 90% de éxito en la dificultad 115 en TwoRoom).
• Meta de Lenguaje vs. Meta Visual: LAGO condicionado a metas de texto superó significativamente a las líneas base de metas de lenguaje (LIV, PE-Core, Robo-Dopamine). Por ejemplo, en Kingdom, LAGO logró una tasa de éxito del 96.9% comparado con el 52.6% de la línea base LIV ajustada.
• Eficiencia: LAGO demostró una alta eficiencia de inferencia (aprox. 11ms por observación) en comparación con los grandes modelos de recompensa generativos como Robo-Dopamine (aprox. 451ms), lo que lo hace viable para MPC en tiempo real.
• Estudios de Ablación: El rendimiento se mantuvo estable para $K \ge 4$ submetas. Se encontró que la temperatura $\tau$ del mínimo suave es robusta dentro de un rango razonable.

Significado y Reivindicaciones

El artículo afirma que LAGO logra cerrar la brecha entre la precisión de las metas basadas en imágenes y la flexibilidad del control guiado por texto. Al fundamentar el lenguaje en el espacio latente del modelo de mundo y descomponer las tareas en submetas intermedias, el marco mitiga la acumulación de errores de predicción.

Los autores sostienen que su trabajo permite una planificación robusta y precisa de largo horizonte puramente desde el lenguaje, evitando la degradación aguda observada en métodos previos cuando los horizontes de planificación se extienden más allá del alcance inmediato. Enfatizan que esto se logra sin requerir supervisión de acciones o recompensas durante la fase de planificación, apoyándose en la guía estructurada de las submetas latentes predichas. El trabajo se presenta como un paso hacia el abordaje de desafíos de planificación cada vez más complejos al combinar las fortalezas de las modalidades visuales y lingüísticas dentro de un único y eficiente marco de planificación.