Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

Resumen Técnico de Traj-LeWM: Planificación de Modelos de Mundo Consciente de la Trayectoria mediante el Costo de Trayectoria Latente

1. Planteamiento del Problema

El artículo aborda dos limitaciones específicas en LeWM (Lightweight Visual World Model), una arquitectura de predicción de incrustación conjunta (JEPA) que aprende la dinámica latente a partir de píxeles para realizar una planificación condicionada por objetivos. Aunque LeWM logra un rendimiento sólido en tareas simples, su tasa de éxito disminuye significativamente en entornos más complejos (por ejemplo, de un 96% en Push-T a un 74% en OGBench-Cube). Los autores identifican dos causas raíz para esta degradación:

1. Falta de Supervisión a Nivel de Trayectoria: El objetivo de entrenamiento de LeWM depende de la predicción del siguiente paso (transiciones locales) y de un término de regularización (SIGReg) para prevenir el colapso de la representación. Ninguno de estos componentes requiere explícitamente que el codificador compartido preserve información sobre la trayectoria completa en relación con el objetivo de la tarea. En consecuencia, el espacio latente no está entrenado directamente para evaluar la calidad de una trayectoria candidata completa.
2. Insensibilidad de la Calificación Basada Solo en el Punto Final: Durante la planificación, LeWM clasifica las secuencias de acciones candidatas basándose únicamente en la distancia euclidiana entre la representación latente del punto final predicho y la representación del objetivo. Este enfoque ignora la información del camino intermedio. Dos trayectorias con puntos finales idénticos pero evoluciones intermedias diferentes (una exitosa y otra fallida) reciben la misma puntuación, lo que conduce a una selección de candidatos subóptima cuando las predicciones del modelo divergen de los resultados reales de ejecución.

2. Metodología: Traj-LeWM

Los autores proponen Traj-LeWM, un marco que conserva el objetivo original de dinámica local de LeWM y la calificación de punto final, introduciendo un Costo de Trayectoria Latente (LTC) Condicionado por el Objetivo. El método opera en tres etapas:

A. Costo de Trayectoria Latente (LTC) Condicionado por el Objetivo
La innovación central es una función de costo escalar aprendible, $S_\psi(\tau | g)$, que evalúa una trayectoria latente completa $\tau$ dado un objetivo $g$. A diferencia de la distancia al punto final, el LTC agrega información a lo largo de todo el camino. Se define como:
$$S_\psi(\tau | g) = \frac{1}{T} \sum_{t=0}^{T-1} \ell_\psi(z_t, \Delta z_t, d^g_t, \eta_t)$$
Donde:

• $z_t$: Estado latente actual.
• $\Delta z_t$: Evolución local (diferencia entre estados adyacentes).
• $d^g_t = z_g - z_t$: Cantidad relativa al objetivo (condicionamiento).
• $\eta_t$: Fase temporal normalizada.
• $\ell_\psi$: Una función de contribución por paso no negativa y aprendible.

Esta estructura permite que el costo varíe según cómo se alcanza un estado, no solo dónde termina.

B. Supervisión de Preferencia de Trayectoria
Para entrenar el LTC y dar forma al codificador compartido, los autores introducen un mecanismo de aprendizaje basado en preferencias utilizando pérdida logística por pares. El sistema aprende a asignar costos más bajos a las trayectorias "preferidas" (que coinciden con el experto) y costos más altos a las trayectorias "negativas". Se utilizan dos tipos de ejemplos negativos:

1. Preferencias Sintéticas:
   • Desajuste de Objetivo: Una trayectoria experta emparejada con un objetivo desplazado y no relacionado.
   • Perturbación de Preservación de Punto Final: Una trayectoria experta con ruido añadido a los estados intermedios (manteniendo fijos el inicio y el final). Esto obliga al modelo a distinguir entre diferentes caminos intermedios incluso si los puntos finales son idénticos.
2. Preferencias de Fallo de Bucle Cerrado:
   • El modelo actual realiza la planificación utilizando solo la calificación de punto final (sin LTC) y ejecuta el plan en el entorno.
   • Se minan las ejecuciones fallidas y se emparejan con trayectorias expertas que tienen el mismo inicio y objetivo.
   • Estos fallos del mundo real proporcionan ejemplos negativos que la métrica de solo punto final no logró identificar, calibrando el LTC contra los modos de fallo específicos del modelo.

C. Entrenamiento Conjunto y Planificación

• Entrenamiento: La pérdida total combina la pérdida de predicción original de LeWM ($L_{LeWM}$), la pérdida de preferencia sintética ($L_{path}$) y la pérdida de fallos minados ($L_{mined}$). La rama sintética actualiza tanto el LTC como el codificador compartido, mientras que la rama minada actualiza solo el LTC (mediante gradientes desvinculados/detached) para evitar confundir el aprendizaje de la dinámica del predictor.
• Planificación: Los candidatos se clasifican utilizando una Puntuación Conjunta:
  $$C_\lambda(\hat{\tau}, g) = C_0(\hat{\tau}, g) + \lambda \tilde{S}_\psi(\hat{\tau} | g)$$
  Donde $C_0$ es la distancia original al punto final y $\tilde{S}_\psi$ es una puntuación LTC calibrada. Esto permite que el planificador utilice la información del camino intermedio cuando las predicciones del punto final son ambiguas o inexactas.

3. Contribuciones Clave

1. Costo de Trayectoria Latente (LTC): Introducción de un costo aprendido y condicionado por el objetivo que evalúa trayectorias predichas completas, agregando la evolución latente relativa al objetivo para proporcionar una señal de planificación sensible al camino.
2. Aprendizaje de Preferencia de Trayectoria: Desarrollo de un mecanismo que utiliza negativos sintéticos (desajustes de objetivo y perturbaciones intermedias) y fallos de bucle cerrado minados para entrenar el LTC y dar forma a la representación compartida sin alterar el objetivo original del predictor.
3. Evaluación Exhaustiva: Pruebas extensas en cuatro tareas simuladas (Push-T, OGBench-Cube, Reacher, Two-Room) y un robot físico Franka FR3, demostrando mejoras sobre la línea base.

4. Resultados Experimentales

Resultados de Simulación:
Traj-LeWM supera a LeWM en las cuatro tareas, con mejoras en la tasa de éxito de:

• Push-T: +3% (96% $\to$ 99%)
• OGBench-Cube: +14% (74% $\to$ 88%)
• Reacher: +7% (86% $\to$ 93%)
• Two-Room: +7% (87% $\to$ 94%)

Evaluación con Robot Físico:
En un Franka FR3 físico con 20 tareas de inicio-objetivo, Traj-LeWM logró una tasa de éxito del 70% (14/20) comparado con el 50% (10/20) de LeWM. Los fallos debidos a colisiones disminuyeron del 20% al 10%.

Experimentos Controlados y Ablaciones:

• Moldeado de Representación (Brecha 1): Incluso cuando se elimina el LTC en el tiempo de prueba, el codificador compartido de Traj-LeWM (entrenado con preferencias de trayectoria) produce puntuaciones de punto final con un AUC más alto para distinguir éxito vs. fallo en comparación con LeWM.
• Conciencia del Camino (Brecha 2): Cuando las predicciones del punto final son inexactas, añadir LTC mejora el orden correcto de los candidatos. En pruebas de "control de punto final" (donde los estados inicial y final son idénticos), el LTC distingue con éxito trayectorias exitosas de fallidas en >94% de los casos en Reacher, probando que utiliza la información del camino intermedio.
• Complementariedad: Los estudios de ablación muestran que la supervisión de trayectoria (entrenamiento) y la calificación explícita de LTC (planificación) son complementarias. La puntuación conjunta ofrece los mejores resultados, indicando que no son intercambiables.

5. Significado y Reivindicaciones

El artículo afirma que Traj-LeWM aborda con éxito las limitaciones de los modelos de mundo de solo dinámica local mediante la integración de información a nivel de trayectoria tanto en el aprendizaje de la representación como en la clasificación de candidatos.

• Reivindicaciones Modestas: Los autores declaran explícitamente que los resultados del robot real son "evidencia preliminar de la viabilidad en robots reales más que una ventaja de rendimiento concluyente" debido al número limitado de tareas emparejadas (20).
• Perspectiva Central: El trabajo demuestra que la distancia al punto final por sí sola es insuficiente para una planificación fiable en entornos complejos. Al aprender un costo sobre rollouts latentes completos, el modelo puede aprovechar la evidencia del camino intermedio para clasificar candidatos de manera más efectiva, particularmente cuando las predicciones del punto final son inciertas.
• Eficiencia: El método añade menos de 1M de parámetros al modelo, manteniendo una arquitectura ligera comparable al LeWM original mientras aumenta significativamente el rendimiento en tareas desafiantes de manipulación y navegación.