Policy-Guided World Model Planning for Language-Conditioned Visual Navigation

Resumen Técnico: PiJEPA

1. Planteamiento del Problema

La navegación de agentes autónomos hacia un objetivo visual especificado mediante instrucciones de lenguaje natural sigue siendo un desafío fundamental en la IA encarnada (Embodied AI). El problema se define por la necesidad de generar una secuencia de acciones que lleve al robot desde una observación actual ($o_t$) a una imagen de objetivo ($o_g$), respetando la semántica de una instrucción ($\ell$).

Existen dos enfoques principales con limitaciones inherentes:

• Políticas Reactivas (ej. Octo, NoMaD): Modelos de acción-idioma-visión (VLA) que predicen acciones en un solo paso. Aunque son rápidos y siguen bien las instrucciones a corto plazo, carecen de capacidad para razonar sobre las consecuencias a largo plazo de sus decisiones, lo que lleva a errores acumulativos en tareas de horizonte largo.
• Modelos del Mundo (World Models): Permiten planificar simulando futuros estados. Sin embargo, los planificadores basados en muestreo (como MPPI o CEM) que operan sobre estos modelos suelen partir de una distribución a priori no informada (Gaussiana) en espacios de acción de alta dimensión. Esto resulta en una convergencia lenta y una alta susceptibilidad a quedar atrapados en mínimos locales.

El objetivo es combinar la eficiencia semántica de las políticas reactivas con la capacidad de razonamiento a largo plazo de los modelos del mundo.

2. Metodología: PiJEPA

Los autores proponen PiJEPA, un marco de dos etapas que utiliza una política aprendida para "calentar" (warm-start) la planificación de un modelo del mundo. El sistema se basa en tres componentes principales:

1. Codificador de Visión Congelado: Se utiliza un encoder preentrenado y congelado (DINOv2 o V-JEPA-2) para mapear las observaciones visuales a un espacio latente ($z_t$). Tanto la política como el modelo del mundo comparten este espacio de representación.
2. Política Prior (Etapa 1):
   • Se utiliza un modelo Octo (una política generalista basada en transformadores) afinado (fine-tuned) en el dataset CAST.
   • Dada la observación actual y la instrucción, la cabeza de difusión de Octo genera $N_\pi$ muestras de bloques de acción.
   • Estas muestras se transforman del marco global al marco local del cuerpo del robot y se calculan sus estadísticas empíricas: media ($\mu_\pi$) y desviación estándar ($\sigma_\pi$). Esta distribución actúa como un prior informado.
3. Planificación MPPI con Modelo del Mundo JEPA (Etapa 2):
   • Se emplea un Modelo del Mundo JEPA (Joint-Embedding Predictive Architecture) entrenado con el mismo encoder congelado. Este modelo predice futuros estados latentes ($\hat{z}_{t+1}$) dados los estados actuales y las acciones, sin necesidad de reconstrucción a nivel de píxel.
   • Se utiliza el algoritmo MPPI (Model Predictive Path Integral) para optimizar la secuencia de acciones.
   • Innovación clave: A diferencia de los métodos tradicionales que inician el muestreo MPPI desde una distribución Gaussiana no informada, PiJEPA inicializa la distribución de muestreo con $(\mu_\pi, \sigma_\pi)$ derivada de la política.
   • El planificador refina iterativamente estas trayectorias prometedoras minimizando la distancia en el espacio latente entre el estado final predicho y el objetivo codificado ($z_g$).

Flujo de ejecución:

1. La política genera un prior de acción.
2. El planificador MPPI usa este prior para iniciar la búsqueda.
3. El modelo JEPA evalúa las trayectorias candidatas proyectando su futuro latente.
4. Se ejecuta la primera acción de la secuencia óptima y el proceso se repite en el siguiente paso (replanificación).

3. Contribuciones Clave

1. Marco Unificado PiJEPA: Propone la primera integración efectiva de una política VLA afinada con planificación basada en modelos del mundo JEPA, utilizando la salida de la política para inicializar el planificador.
2. Análisis Sistemático de Codificadores: Evalúa y compara dos arquitecturas de visión modernas: DINOv2 (encoder de imágenes auto-supervisado, fuerte en características espaciales) y V-JEPA-2 (encoder de video auto-supervisado, captura dinámicas temporales). Se prueban todas las combinaciones posibles (política y modelo del mundo con el mismo o diferente encoder).
3. Validación Empírica: Demuestran que la planificación guiada por políticas supera significativamente tanto a la ejecución pura de políticas reactivas como a la planificación de modelos del mundo sin información previa, mejorando la precisión de llegada al objetivo y la fidelidad al seguir instrucciones.

4. Resultados Experimentales

Los experimentos se realizaron en el dataset CAST (navegación visual a gran escala con instrucciones) utilizando robots en entornos reales simulados.

• Rendimiento General: PiJEPA logra la mejor precisión posicional (menor Error de Trayectoria Absoluta - ATE) en ambas configuraciones de codificadores.
  • Con V-JEPA-2, PiJEPA alcanza un RMSE de posición XY de 1.65 m, superando a la política Octo pura (1.72 m) y a la planificación MPPI sin prior (1.87 m).
• Análisis de Componentes:
  • Política Reactiva (Octo): Excelente control local y predicción de orientación (heading), pero sufre en trayectorias largas debido a la acumulación de errores.
  • MPPI No Informado: Logra una buena alineación de orientación gracias al objetivo en el espacio latente, pero falla estrepitosamente en la precisión posicional debido a la ineficiencia de la búsqueda en espacios de alta dimensión.
  • Scoring de Modelo del Mundo (Octo-WM): Filtrar las propuestas de la política usando el modelo del mundo ya aporta grandes beneficios, pero la optimización iterativa de MPPI (PiJEPA) añade precisión adicional.
• Latencia: La sobrecarga computacional de la etapa de planificación MPPI es mínima (0.35s) comparada con el tiempo de muestreo de difusión de la política (2.13s), haciendo el enfoque práctico para implementación en tiempo real.
• Análisis de Fallos: Se identificó que en instrucciones ambiguas (ej. "Sigue el edificio" cuando hay varios), el modelo del mundo puede estancarse prediciendo estados latentes idénticos. El prior de la política ayuda a mitigar esto anclando la búsqueda en una región productiva, aunque no lo resuelve completamente.

5. Significado e Impacto

El trabajo de PiJEPA es significativo porque resuelve el cuello de botella de la inicialización en la planificación basada en modelos del mundo. Al demostrar que una política reactiva, aunque insuficiente para la planificación a largo plazo por sí sola, proporciona un prior de acción altamente informativo, el método permite que el planificador se centre en refinar trayectorias prometedoras en lugar de explorar ciegamente todo el espacio de acciones.

Esto establece un nuevo paradigma de "división de trabajo" en la navegación robótica:

• La política se encarga de la comprensión semántica rápida y el control local.
• El modelo del mundo se encarga de la simulación a largo plazo y la corrección de trayectorias.

La capacidad de combinar la eficiencia de los modelos VLA con la robustez de la planificación en espacio latente (JEPA) sin reconstrucción de píxeles sugiere un camino viable hacia agentes robóticos más autónomos y capaces de navegar en entornos complejos siguiendo instrucciones naturales.