Resumen Técnico: Modelos de Mundo Sensoriomotores (SMWM)

Declaración del Problema

El artículo aborda el desafío de aprender modelos de mundo latentes a partir de observaciones de alta dimensión (por ejemplo, píxeles) sin señales de recompensa o etiquetas de tareas. Aunque las Arquitecturas de Predicción de Incrustaciones Conjuntas (JEPA) ofrecen un marco prometedor al predecir estados futuros directamente en un espacio de incrustación (evitando la reconstrucción de píxeles), sufren de colapso de representación. Cuando un codificador y un modelo de dinámica hacia adelante se entrenan conjuntamente para minimizar el error de predicción, el codificador puede aprender a mapear todas las observaciones a una incrustación constante. Esta solución trivial minimiza la pérdida de propagación hacia adelante, pero deja al modelo inútil para el control o la planificación posteriores.

Las soluciones existentes para prevenir el colapso suelen depender de mecanismos complejos como:

• El congelamiento de codificadores preentrenados (ej. DINO-WM).
• La adición de regularizadores distribucionales (ej. el ajuste de las incrustaciones a una Gaussiana isotrópica mediante SIGReg).
• El uso de redes objetivo de Promedio Móvil Exponencial (EMA) con gradientes detenidos (stop-gradients).
• El empleo de funciones de pérdida multitérmino que involucran restricciones de varianza-covarianza.

Los autores postulan que estos enfoques pueden no abordar plenamente la cuestión central de qué debería retener una representación útil de un modelo de mundo. Basándose en el concepto de "percepción para la acción" y la percepción enactiva, el artículo argumenta que las representaciones deben ser moldeadas por su relevancia para las acciones, preservando las distinciones que importan para la dinámica controlable y descartando los distractores incontrolables.

Metodología: Modelo de Mundo Sensoriomotor (SMWM)

Los autores proponen SMWM, un modelo de mundo latente entrenado de extremo a extremo a partir de trayectorias sin recompensa $(o_t, a_t, o_{t+1})$. La innovación central es el uso de la regularización de la dinámica inversa como el único mecanismo para prevenir el colapso de la representación.

Arquitectura

El modelo consta de tres componentes entrenados conjuntamente:

1. Codificador ($f_\theta$): Mapea las observaciones $o_t$ a una incrustación latente compacta $z_t$.
2. Modelo de Dinámica hacia Adelante ($g_\phi$): Predice el siguiente estado latente $\hat{z}_{t+1}$ dado el estado actual $z_t$ y la acción $a_t$.
3. Modelo de Dinámica Inversa ($h_\psi$): Predice la acción $\hat{a}_t$ realizada entre dos incrustaciones consecutivas $z_t$ y $z_{t+1}$.

Objetivo de Entrenamiento

El modelo se optimiza minimizando una función de pérdida conjunta:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Pérdida hacia Adelante ($\mathcal{L}_{fwd}$): Error Cuadrático Medio (MSE) entre la siguiente incrustación predicha $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ y la siguiente observación codificada real $z_{t+1} = f_\theta(o_{t+1})$.
• Pérdida Inversa ($\mathcal{L}_{inv}$): MSE entre la acción predicha $\hat{a}_t = h_\psi(z_t, z_{t+1})$ y la acción real $a_t$.
• Peso de Regularización ($\lambda$): Un único hiperparámetro que controla la fuerza de la señal de la dinámica inversa.

Mecanismo de Acción

La pérdida inversa actúa como un mecanismo anti-colapso. Si el codificador colapsa hacia una incrustación constante, el modelo inverso recibe entradas idénticas para todas las transiciones y solo puede predecir una acción constante (la media de la distribución de acciones). Para reducir la pérdida inversa por debajo de este riesgo de predictor constante, el codificador debe preservar la información sobre la acción $a_t$ contenida en la transición $(o_t, o_{t+1})$. Esto obliga a la representación a retener información relevante para la acción (grados de libertad controlables) mientras filtra naturalmente los distractores incontrolables (ruido exógeno) que no influyen en la recuperación de la acción.

A diferencia de los regularizadores distribucionales que prescriben la geometría del espacio de incrustación (por ejemplo, forzándolo a ser una Gaussiana), la dinámica inversa ancla la representación a una cantidad fundamentada en la tarea (la acción) sin imponer supuestos geométricos específicos.

Contribuciones Clave

1. Arquitectura SMWM: Un modelo de mundo sensoriomotor simple y entrenable de extremo a extremo que utiliza la dinámica inversa como el único mecanismo anti-colapso, requiriendo solo un hiperparámetro adicional ($\lambda$).
2. Propiedades de la Representación: Evidencia empírica de que las representaciones aprendidas no son simplemente no colapsadas, sino que rastrean específicamente los grados de libertad controlables, preservan la geometría espacial (tanto cartesiana como periódica) y filtran los distractores incontrolables.
3. Desempeño en Tareas Posteriores: Demostración de un desempeño de planificación competitivo en tareas de control 2D y 3D, mostrando que las ventajas estructurales de la representación se traducen en una planificación condicionada a objetivos efectiva.

Resultados Experimentales

Bancos de Pruebas Controlados (Dot World & Sprite World)

• Recuperación de la Dimensión Intrínseca: En un entorno "dot world", el modelo recuperó con éxito la verdadera dimensión intrínseca 2D del espacio de estados a partir de los píxeles, incluso estando embebida en un espacio latente de 64 dimensiones.
• Filtrado de Distractores: Cuando se introdujeron objetos "distractores" incontrolables (moviéndose aleatoriamente), el codificador los ignoró. La dimensión latente efectiva coincidió con los grados de libertad controlables, no con la dimensión total de la observación.
• Percepción Dependiente del Control: En un experimento de "sprite world", el modelo aprendió a preservar las variables de pose (posición, orientación) solo cuando eran controlables. Cuando la orientación no era controlable, la representación la promediaba, resultando en una reconstrucción borrosa de la orientación.

Benchmarks de Planificación

Los autores evaluaron SMWM en cuatro tareas de control offline: TwoRoom (navegación 2D), Reacher (control continuo), Push-T (manipulación 2D) y OGBench-Cube (manipulación 3D).

• Baselines: Comparación contra SIGReg (un fuerte baseline de JEPA que usa regularizadores Gaussianos), Forward-only (sin pérdida inversa) y muestreo de acciones Random.
• Desempeño:
  • SMWM igualó el desempeño de SIGReg en las tres tareas 2D.
  • SMSM superó a SIGReg en la tarea 3D OGBench-Cube (84% de éxito frente a 59%).
  • Ambos métodos regularizados superaron significativamente al ablativo Forward-only, confirmando la necesidad de un mecanismo anti-colapso.
• Geometría Latente: Los experimentos de sondeo revelaron que las incrustaciones de SMWM forman subespacios de baja dimensión que reflejan la geometría del estado físico (ej. estructuras cilíndricas para ángulos de articulación periódicos en Reacher, mapas cartesianos para navegación).

Significado y Reivindicaciones

El artículo afirma que SMWM proporciona una instanciación sencilla y fundamentada de la "percepción para la acción" en el aprendizaje de modelos de mundo latentes. Al forzar a las incrustaciones latentes a preservar la información suficiente para recuperar las acciones del agente, el modelo sesga naturalmente la representación hacia los grados de libertad controlables del entorno.

Los puntos clave incluyen:

• Simplicidad: La dinámica inversa por sí sola es suficiente para estabilizar el entrenamiento de JEPA de extremo a extremo sin reconstrucción, señales de recompensa, codificadores congelados o supuestos distribucionales complejos.
• Sesgo Estructural: El método induce representaciones que son interpretables y están alineadas con la estructura causal del entorno (acciones y sus efectos), en lugar de solo correlaciones predictivas.
• Utilidad Práctica: Las representaciones aprendidas soportan una planificación latente efectiva, particularmente en entornos 3D complejos donde otros baselines regularizados tienen dificultades.

Los autores reconocen limitaciones, como la suposición de que las acciones son recuperables a partir de observaciones consecutivas y la incapacidad actual para capturar cantidades que no son identificables desde fotogramas únicos (como la velocidad) sin historiales de observación. Sin embargo, posicionan a SMWM como un enfoque robusto y sin recompensa para aprender modelos de mundo accionables a partir de datos offline.