Resumen Técnico: Volver a los Latentes Parsimoniosos: Aprendizaje de Modelos del Mundo Centrado en Tareas desde Fundamentos Visuales

Enunciado del Problema

Los modelos del mundo permiten que los agentes predigan la dinámica futura condicionada a las acciones, una capacidad central para la planificación y el control. Sin embargo, los enfoques existentes enfrentan una compensación fundamental en la representación latente:

1. Modelos basados en píxeles (p. ej., Dreamer, MuZero) aprenden la dinámica directamente desde píxeles o latentes compactos moldeados por señales de recompensa. Aunque son efectivos, a menudo incurren en altos costos computacionales o requieren interacción en línea/supervisión de recompensas para moldear el espacio latente.
2. Modelos de fundamentos de incrustación (p. ej., DINO-WM, V-JEPA) aprovechan codificadores visuales preentrenados congelados para predecir directamente en un espacio semántico de alta dimensión. Aunque ofrecen priores semánticos ricos y generalización de cero disparos, las incrustaciones contienen detalles excesivamente irrelevantes para la tarea (texturas, iluminación, fondo) que no afectan el control.

En configuraciones sin recompensa y fuera de línea, donde los agentes deben aprender de trayectorias fijas sin supervisión de recompensas, estos problemas se exacerban. Utilizar incrustaciones de fundamentos crudas como espacio de estados fuerza al modelo a desperdiciar capacidad en variaciones irrelevantes, lo que conduce a una baja eficiencia de modelado, una interpretabilidad de control reducida y un posible colapso de la representación, particularmente en espacios de acción de alta dimensión (p. ej., brazos robóticos de 7 grados de libertad). El desafío central es extraer una representación latente compacta y suficiente para la tarea de estas incrustaciones de fundamentos ricas pero ruidosas sin interacción en línea.

Metodología: TC-WM

Los autores proponen TC-WM (Modelo del Mundo Centrado en Tareas), un marco que trata las incrustaciones de fundamentos preentrenadas no como el espacio de estados final, sino como un "andamio semántico" del cual se destilan las dinámicas relevantes para la tarea.

Arquitectura Central

TC-WM opera sobre una cadena de incrustaciones visuales ($x_t$) y señales propioceptivas ($s^p_t$):

1. Incrustación Conjunta: Un codificador visual congelado (p. ej., DINOv2) mapea imágenes a características visuales, mientras que una capa lineal entrenable codifica la propiocepción. Estas se concatenan para formar una incrustación conjunta $x_t$.
2. Codificador Latente (Proyección): Un codificador lineal proyecta la incrustación conjunta de alta dimensión en un estado latente compacto $z_t$. Críticamente, $z_t$ se factoriza en dos subespacios:
   • $z^s_t$ (Relevante para el estado): Alineado con el estado físico.
   • $z^c_t$ (Complementario): Mantiene la estructura semántica residual (apariencia, fondo) para satisfacer las restricciones de reconstrucción.
3. Dinámicas Latentes: Un predictor basado en Transformer ($f_\theta$) modela las transiciones temporales en el espacio latente compacto ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. Alineación Centrada en Tareas: El modelo emplea un objetivo contrastivo (InfoNCE) para alinear el subespacio $z^s_t$ con el estado propioceptivo $s^p_t$. Esto fuerza activamente al codificador a separar los factores físicos relevantes para la tarea de las variaciones molestas.
5. Decodificador de Incrustación: Un decodificador lineal reconstruye la incrustación visual original desde $z_t$. Esto actúa como un regularizador, asegurando que el latente compacto permanezca anclado a la variedad de fundamentos y previniendo soluciones degeneradas.

Función de Pérdida de Entrenamiento

La pérdida total combina cuatro componentes:

• Pérdida de Dinámicas Latentes ($L^z_{dyn}$): Minimiza el error de predicción en el espacio latente.
• Pérdida de Dinámicas TC ($L^s_{dyn}$): Predice la propiocepción futura directamente desde el historial para imponer dinámicas centradas en la tarea.
• Pérdida de Alineación ($L_{align}$): Pérdida contrastiva que alinea $z^s_t$ con $s^p_t$.
• Pérdida de Reconstrucción ($L_{rec}$): Reconstruye la incrustación visual desde $z_t$ para preservar la estructura visual.

Planificación en Tiempo de Prueba

El espacio latente centrado en tareas aprendido soporta diversas estrategias de planificación:

• Control de baja dimensión (Navegación): Control Predictivo de Modelo (MPC) utilizando el Método de Entropía Cruzada (CEM) para optimizar secuencias de acciones en el espacio latente.
• Manipulación de alta dimensión: Planificación de Difusión Latente (LDP) combinada con un Modelo de Dinámica Inversa (IDM) para generar acciones ejecutables a partir de trayectorias latentes.

Contribuciones Clave

1. Descomposición Centrada en Tareas: El artículo introduce un método para descomponer incrustaciones de fundamentos de alta dimensión en un espacio latente compacto y suficiente para la tarea, separando explícitamente las dinámicas del estado físico de las variables molestas semánticas.
2. Identificabilidad Activa: A diferencia de trabajos anteriores que asumen pasivamente la diversidad de mecanismos para la identificabilidad, TC-WM utiliza la alineación contrastiva con la propiocepción para activamente inducir condicionales distintos para diferentes estados físicos, garantizando teóricamente la recuperación de factores centrados en la tarea hasta una transformación afín.
3. Aprendizaje Fuera de Línea sin Recompensa: El marco logra un modelado del mundo y una planificación de alta calidad en configuraciones fuera de línea sin señales de recompensa, confiando únicamente en observaciones visuales e información lateral propioceptiva.
4. Garantía Teórica: Los autores proporcionan un teorema que demuestra que, bajo suposiciones leves (dinámicas contextuales, variabilidad latente-observada y alineación suficiente), el latente compacto aprendido es identificable y recupera el estado físico verdadero hasta una reparametrización.

Resultados Experimentales

Los autores evalúan TC-WM en 9 benchmarks fuera de línea (Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper) contra líneas base que incluyen TD-MPC2, DreamerV3, MuZero y DINO-WM.

• Calidad del Modelo del Mundo: TC-WM logra el menor error de predicción de desarrollo latente (MSE) en casi todas las tareas, particularmente en entornos de manipulación donde las dinámicas de largo horizonte precisas son críticas. Aunque la reconstrucción en el espacio de imágenes es competitiva, la ganancia principal reside en la estabilidad y precisión de las dinámicas latentes.
• Rendimiento de Planificación:
  • En navegación (acciones de baja dimensión), TC-WM iguala el rendimiento de líneas base fuertes como TD-MPC2.
  • En manipulación (acciones de alta dimensión, ricas en contacto), TC-WM supera significativamente a todas las líneas base. Por ejemplo, en la tarea "Lift" de Robomimic, TC-WM logra una tasa de éxito del 54% en comparación con el 29% de DINO-WM y el 31% de TD-MPC2.
• Análisis de Representación:
  • Sondeo Lineal: Los regresores lineales entrenados en el espacio latente de TC-WM recuperan estados propioceptivos y de objetos con puntuaciones $R^2$ más altas que las líneas base, confirmando que el latente está fundamentado físicamente.
  • Anti-Colapso: TC-WM demuestra un mejor aprovechamiento del espacio latente (mayor rango efectivo) en comparación con DINO-WM, que sufre de colapso de representación al utilizar incrustaciones crudas directamente.
  • Robustez: TC-WM mantiene el rendimiento bajo perturbaciones visuales (ruido gaussiano, jitter de color) mejor que DINO-WM, con caídas menores en SSIM.

Significado y Afirmaciones

El artículo afirma que las incrustaciones de fundamentos se utilizan mejor como un andamio semántico en lugar de una representación de estado directa. Al proyectar estas incrustaciones en un subespacio centrado en la tarea alineado con señales físicas, los agentes pueden lograr:

• Latentes Parsimoniosos: Una reducción en la dimensionalidad que filtra detalles visuales irrelevantes mientras preserva las dinámicas necesarias para el control.
• Control Mejorado: Un rendimiento de planificación mejorado en entornos complejos de alta dimensión donde los modelos de fundamentos estándar fallan debido al ruido y variaciones irrelevantes.
• Solidez Teórica: Una justificación formal de por qué las proyecciones lineales y la alineación contrastiva son suficientes para recuperar factores relevantes para la tarea desde codificadores visuales congelados.

Los autores posicionan a TC-WM como un puente entre el poder de generalización de los modelos de fundamentos y la controlabilidad requerida para agentes físicos, abordando específicamente la brecha en escenarios de aprendizaje fuera de línea sin recompensas.