Resumen Técnico: SCALE: Incrustaciones Latentes Calibradas por el Estado para la Planificación JEPA en la Geometría Correcta

1. Declaración del Problema

Las Arquitecturas de Predicción de Incrustaciones Conjuntas (JEPA) permiten la planificación mediante la predicción de representaciones latentes futuras y la puntuación de secuencias de acciones candidatas basándose en la distancia entre una incrustación terminal predicha y una inclustación objetivo. Si bien modelos como LeWorldModel (LeWM) logran aprender representaciones no colapsadas y de extremo a extremo utilizando la regularización anti-colapso (SIGReg), existe una desconexión crítica entre la decodificabilidad y la influencia geométrica.

El artículo identifica que, si bien la información de estado relevante para la tarea suele ser decodificable a partir de las incrustaciones completas de LeWM, esto no significa necesariamente que modele la geometría latente utilizada por el planificador. Bajo la función de costo estándar de distancia euclidiana al cuadrado, las direcciones con alta varianza dominan la métrica de distancia. Si las variaciones relevantes para la tarea se codifican en direcciones de baja varianza, estas siguen siendo fáciles de decodificar mediante un cabezal de regresión, pero tienen una influencia insignificante en la selección de candidatos del planificador. Esto conduce a una situación en la que un modelo "conoce" el estado, pero no logra "usarlo" eficazmente para la planificación.

El problema central es que los objetivos de entrenamiento JEPA de extremo a extremo existentes no aseguran explícitamente que la palanca métrica (el peso que una dirección latente ejerce en el costo de planificación) se alinee con las variaciones de estado relevantes para la tarea.

2. Metodología: SCALE

Los autores proponen SCALE (State-CAlibrated Latent Embeddings - Incrustaciones Latentes Calibradas por el Estado), un regularizador durante el tiempo de entrenamiento diseñado para remodelar la geometría latente de modo que las distancias latentes por pares se correlacionen con las distancias en un espacio de estado estandarizado y relevante para la tarea.

Componentes Clave:

• Calibración de Estado: Durante el entrenamiento, las observaciones se emparejan con variables de estado del simulador ($q$). Estos estados se estandarizan (por ejemplo, usando seno-coseno para ángulos, excluyendo velocidades) para definir una geometría objetivo.
• Pérdida de Correlación ($L_{corr}$): SCALE toma muestras de pares de fotogramas (tanto dentro como entre episodios) y calcula la distancia euclidiana al cuadrado en el espacio latente ($x_k = \|z_i - z_j\|^2$) y en el espacio de estado estandarizado ($y_k = \|q_i - q_j\|^2$).
  • El objetivo minimiza la pérdida de correlación entre los perfiles estandarizados de estas distancias:
    $$L_{corr} = 1 - \frac{1}{K} \sum_{k=1}^K \tilde{x}_k \tilde{y}_k$$
  • Esto obliga al espacio latente a organizarse de tal manera que mayores diferencias de estado resulten en distancias latentes proporcionalmente mayores.
• Integración: El objetivo total combina la pérdida de predicción estándar de LeWM ($L_{pred}$), el término anti-colapso SIGReg y el nuevo término de correlación:
  $$L_{SCALE} = L_{pred} + \lambda_{sig} \text{SIGReg}(Z) + \lambda_{corr} L_{corr}$$
• Cero Sobrecarga de Inferencia: SCALE opera únicamente como un regularizador en el tiempo de entrenamiento. El codificador y el predictor permanecen basados solo en imágenes durante las pruebas; no se requieren cabezales auxiliares ni información de estado durante la planificación.

Control Mecanístico (Línea Base Auxiliar)

Para distinguir entre "codificar información" y "dar forma a la geometría", los autores introducen una línea base Aux. Esta línea base adjunta un cabezal de regresión solo para el entrenamiento para predecir el estado $q$ a partir del latente $z$. Esto asegura que el estado sea decodificable, pero deja la métrica latente sin restricciones. Los resultados muestran que una alta decodificabilidad (Aux) no garantiza una mejora en la planificación, confirmando que el alineamiento geométrico es el factor distintivo.

3. Contribuciones Clave

1. Identificación de la Palanca Métrica: El artículo define formalmente la palanca métrica como la contribución esperada de una dirección latente al costo de planificación de la distancia euclidiana al cuadrado, demostrando que es exactamente proporcional a la varianza de dicha dirección ($2\lambda_j$). Esto explica por qué la información relevante para la tarea en direcciones de baja varianza no logra influir en la planificación.
2. Introducción de SCALE: Un nuevo regularizador que correlaciona directamente las distidades latentes por pares con las distancias de estado relevantes para la tarea, "calibrando" efectivamente la geometría latente sin reemplazar el codificador aprendido.
3. Validación Empírica: Experimentos extensos en cinco entornos (Push-T, Reacher, Cube, Two-Room, PointMaze), tres solucionadores de planificación (CEM, iCEM, MPPI) y cinco presupuestos computacionales demuestran que SCALE supera consistentemente a la línea base LeWM.

4. Resultados

• Rendimiento de la Planificación: SCALE mejora la tasa de éxito sobre LeWM en todos los 15 promedios de tarea-solucionador. Las ganancias persisten a través de diferentes presupuestos computacionales. Por ejemplo, en la tarea Two-Room, SCALE logró una tasa de éxito media del 92.9% comparado con el 82.9% de LeWM bajo iCEM.
• Comparación con DINO-WM: Aunque DINO-WM (que utiliza un espacio de características preentrenado) muestra un fuerte alineamiento en algunas tareas, rinde por debajo en otras (por ejemplo, Push-T). SCALE, sin embargo, nunca cae por debajo de LeWM en ningún par tarea-solucionador, demostrando una mejora más consistente sin heredar los sesgos específicos de la tarea de un codificador preentrenado fijo.
• Análisis Geométrico:
  • Redistribución de la Varianza: SCALE desplaza la varianza hacia los componentes principales líderes (alta palanca métrica) en comparación con LeWM y la línea base Aux.
  • Predictibilidad del Estado: Los componentes principales líderes de SCALE son significativamente más predictivos del estado relevante para la tarea que los de LeWM o Aux.
  • Calidad del Rollout: SCALE mejora el acuerdo de clasificación de Kendall's $\tau$ entre las incrustaciones terminales predichas y los resultados reales del simulador, lo que indica que la geometría reorganizada sobrevive al proceso de predicción (rollout) y mejora la clasificación de candidatos.
• Decodificabilidad vs. Geometría: La línea base Aux iguala o supera a SCALE en la decodificabilidad del estado de la incrustación completa, pero no logra igualar las ganancias de planificación de SCALE. Esto confirma que simplemente hacer que el estado sea decodificable es insuficiente; el estado debe dar forma a la métrica utilizada por el planificador.

5. Significado y Afirmaciones

El artículo sostiene que el rendimiento de la planificación depende no solo de la presencia de información relevante para la tarea en la representación, sino de cómo esa información se organiza geométricamente.

• Influencia Geométrica: Los autores argumentan que, para los planificadores euclidianos, la "calidad" de una representación se define por si las variaciones relevantes para la tarea ocupan direcciones de alta varianza (alta palanca métrica).
• Flexibilidad de Extremo a Extremo: SCALE logra las propiedades geométricas favorables observadas en modelos de características preentrenadas (como DINO-WM) manteniendo la flexibilidad y adaptabilidad de una representación aprendida de extremo a extremo.
• Mínima Sobrecarga: El método añade solo un regularizador ligero durante el entrenamiento y no introduce sobrecarga computacional durante la fase de planificación.

El trabajo concluye que la planificación efectiva en el espacio latente requiere un alineamiento explícito entre la métrica latente y el espacio de estado relevante para la tarea, una propiedad que la regularización anti-colapso estándar por sí sola no garantiza.