Resumen Técnico: Modelos de Mundo Latentes de Longitud Variable para la Planificación de Largo Alcance

1. Planteamiento del Problema

Los modelos de mundo han emergido como un paradigma para construir agentes inteligentes mediante el aprendizaje de modelos predictivos de la dinámica del entorno. Si bien los modelos de mundo latentes recientes (por ejemplo, arquitecturas tipo JEWM) ofrecen alternativas eficientes a la predicción en el espacio de píxeles al operar en espacios de representación compactos, enfrentan limitaciones significativas en la planificación de largo alcance (long-horizon planning).

Los problemas principales identificados son:

• Errores Acumulativos (Compounding Errors): Los modelos existentes suelen depender de la predicción de un solo paso ($t \to t+1$). Para tareas de largo alcance, estos modelos deben realizar despliegues (rollouts) recursivos, lo que causa que los errores de predicción se acumulen rápidamente, provocando una deriva temporal y una degradación de la calidad de la trayectoria.
• Desajuste entre Entrenamiento y Planificación (Training-Planning Mismatch): La mayoría de los modelos se entrenan con un objetivo de predicción de paso fijo. Esto proporciona supervisión solo a una escala temporal local, fallando al no optimizar explícitamente la consistencia de largo alcance. En consecuencia, la señal de entrenamiento no se alinea con la naturaleza de los problemas de planificación de largo alcance.
• Rigidez Arquitectónica: Los modelos de mundo latentes actuales suelen incorporar acciones mediante mecanismos de condicionamiento dedicados (por ejemplo, capas FiLM o atención cruzada) que fusionan las acciones en los estados ocultos. Este diseño vincula estrechamente el modelo a un esquema de acción fijo (usualmente una acción por paso), lo que dificulta la capacidad de desentrelazar las contribuciones de múltiples acciones o manejar secuencias de acciones de longitud variable de forma natural.

2. Metodología: Modelos de Mundo Latentes de Longitud Variable (VLWM)

Los autores proponen VLWM, un marco diseñado para aprender la dinámica del entorno latente a través de múltiples horizontes de predicción dentro de un solo modelo.

2.1 Supervisión de Longitud Variable

En lugar de restringir la supervisión al siguiente paso inmediato, VLWM entrena un predictor $g_\phi$ para pronosticar el estado latente en un desfase futuro arbitrario $k \ge 1$. El objetivo de entrenamiento es:
$$\mathcal{L}_{VLWM} = \mathbb{E}_{t, k \sim p(k)} \left[ \| g_\phi(s_t, a_{t:t+k-1}) - \bar{g}(\bar{f}_{\bar{\theta}}(o_{t+k})) \|_2^2 \right]$$
donde $s_t$ es el estado latente actual, $a_{t:t+k-1}$ es un segmento de acción de longitud variable de longitud $k$, y el objetivo es el estado latente en el tiempo $t+k$. Esto unifica la dinámica de corto plazo y la consistencia semántica de largo alcance bajo un único objetivo.

2.2 Formulación de Acción-como-Token (Action-as-Token)

Para soportar nativamente secuencias de acciones de longitud variable, VLWM se aparta del condicionamiento interno de las capas. En su lugar, trata los embeddings de acción como tokens de primera clase en el Transformer.

• Secuencia de Entrada: El modelo recibe una secuencia $X$ que contiene un historial de tokens de estado seguidos por un número variable de tokens de acción ($e_t, \dots, e_{t+k-1}$).
• Mecanismo: Los tokens de estado y de acción se distinguen mediante embeddings de tipo aditivos. El transformer es causal, y el latente predicho $\hat{s}_{t+k}$ se lee directamente de la salida del último token de la secuencia.
• Beneficio: Esto elimina la necesidad de interfaces de condicionamiento de acción predefinidas, permitiendo que el modelo se condicione a un número arbitrario de acciones pasadas sin cambios arquitectónicos.

2.3 Estrategia de Entrenamiento por Currículo

Entrenar directamente con horizontes de predicción amplios desde el inicio es inestable debido al ruido en los gradientes y al fallo de optimización. VLWM emplea una estrategia de aprendizaje por currículo:

• El entrenamiento se divide en etapas. En la etapa $j$, el horizonte $k$ se muestrea uniformemente de $\{1, 2, \dots, j\}$.
• El modelo comienza con $k=1$ y expande progresivamente el rango de los horizontes muestreados hasta un máximo de $K_{max}$.
• Esto asegura que la semántica de largo alcance se construya sobre una dinámica de corto horizonte bien aprendida, estabilizando la optimización.

2.4 Planificación con Saltos de Longitud Variable

Durante la inferencia, el horizonte de planificación $H$ se particiona en fragmentos contiguos $H = k_1 + k_2 + \dots + k_M$. El planificador evalúa el despliegue como una secuencia de saltos latentes de longitud variable en lugar de predicciones recursivas de un solo paso. El artículo propone tres estrategias de planificación:

1. P1 (Fija pero Adaptativa a la Tarea): Utiliza un tamaño de fragmento constante $k^*$ seleccionado por tarea (pequeño para control fino, grande para objetivos de largo alcance).
2. P2 (Muestreo de Horizonte Aleatorio): Muestrea tamaños de fragmento uniformemente de $\{1, \dots, K_{max}\}$ en cada frontera.
3. P3 (De Largo a Corto): Comienza con fragmentos grandes para cubrir la estructura de largo alcance y reduce progresivamente los tamaños de los fragmentos hacia el final para refinar la dinámica de grano fino (una búsqueda de grueso a fino).

3. Contribuciones Clave

1. Marco Unificado: Introducción de VLWM, un marco que aprende la dinámica latente a través de múltiples horizontes y soporta naturalmente secuencias de acciones de longitud variable mediante una formulación de "acción-como-token".
2. Entrenamiento y Planificación Estables: Propuesta de una estrategia de aprendizaje por currículo escalable para un entrenamiento de longitud variable estable y un conjunto de algoritmos de planificación que aprovechan adaptativamente diferentes horizontes de predicción dentro de un mismo modelo.
3. Validación Empírica: Demostración de que VLWM supera consistentemente a los modelos de mundo de horizonte fijo (específicamente LeWM) en evaluaciones de planificación de largo alcance, con análisis que muestran una mejor predictibilidad de largo alcance y controlabilidad descendente.

4. Resultados Experimentales

Los autores evaluaron VLWM en tres entornos de planificación condicionados por objetivos: PushT (manipulación 2D), OGBench-Cube (manipulación robótica 3D) y TwoRoom (navegación 2D).

• Ganancias de Rendimiento: VLWM logró una mejora promedio del 13% sobre el estado del arte LeWM en los conjuntos de datos. Las ganancias fueron más significativas en regímenes de largo alcance (por ejemplo, en TwoRoom con un horizonte de 100 pasos, la tasa de éxito mejoró del 36% con LeWM al 68% con VLWM).
• Mitigación de Errores: Las visualizaciones mostraron que mientras los despliegues recursivos de LeWM derivaban y fallaban al cruzar fronteras ambientales (por ejemplo, puertas), los saltos de longitud variable por fragmentos de VLWM se mantenían cerca de las trayectorias reales.
• Calidad de la Representación: Las evaluaciones de sonda lineal en codificadores congelados revelaron que VLWM produjo latentes con una mayor correlación de Pearson para propiedades físicas (por ejemplo, el ángulo del bloque) comparado con LeWM, sugiriendo una mejor captura de los cambios de estado semántico lento.
• Necesidad del Currículo: Los estudios de ablación confirmaron que la estrategia de entrenamiento por currículo es esencial para un entrenamiento de longitud variable estable, particularmente en tareas que requieren una planificación extendida.

5. Significado y Reivindicaciones

El artículo sostiene que VLWM proporciona un paradigma simple pero efectivo para mejorar la predicción y la planificación de largo alcance en modelos de mundo latentes.

• Más allá de la Ingeniería: Los autores argumentan que la dificultad con la planificación de largo alcance no es meramente un problema de ingeniería, sino una consecuencia de los objetivos de entrenamiento. Al cambiar de una supervisión de paso fijo a una de longitud variable, VLWM alinea la señal de entrenamiento con los requisitos del razonamiento de largo alcance.
• Interfaz Unificada: El marco subsume una gama de estrategias de horizonte fijo y jerárquico dentro de un solo modelo entrenado, eliminando la necesidad de políticas de alto y bajo nivel separadas o módulos auxiliares.
• Dirección Futura: Los resultados sugieren que la supervisión de longitud variable es un paso fundacional hacia algoritmos de planificación más robustos, ofreciendo una interfaz flexible donde los planificadores pueden adaptar la granularidad de la predicción a las necesidades específicas de la tarea.