💬 NLP

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

Este artículo propone Temporal-Distance JEPA (TD-JEPA), un enfoque de modelo de mundo que extrae costos temporales dirigidos a partir de trayectorias fuera de línea libres de recompensas para cerrar la brecha entre el aprendizaje de representaciones y la planificación, superando así a los planificadores basados en JEPA existentes en múltiples entornos al permitir una clasificación del progreso hacia la meta más efectiva sin depender de la reconstrucción de píxeles o de recompensas explícitas.

Autores originales: Jiaxin Bai, Jiaxuan Xiong

Publicado 2026-07-29

📖 7 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Jiaxin Bai, Jiaxuan Xiong

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás enseñando a un robot a navegar por un mundo nuevo, pero no puedes darle un mapa ni una lista de reglas. Todo lo que tienes es un video de un humano realizando la tarea perfectamente. Este es el desafío del "aprendizaje offline" en la robótica: cómo enseñarle a una IA a planificar sus futuros movimientos simplemente observando grabaciones antiguas. Los científicos utilizan algo llamado "modelo de mundo" para resolver esto. Piensa en un modelo de mundo como una máquina de sueños dentro del cerebro del robot. En lugar de intentar redibujar cada píxel del video (lo que sería como intentar pintar una obra maestra de memoria), el robot aprende a predecir la esencia de lo que sucede después. Comprime el mundo en un "espacio latente" simplificado y abstracto: un mapa mental donde las situaciones similares están cerca unas de otras.

La gran pregunta es: una vez que el robot tiene este mapa mental, ¿cómo decide qué camino tomar? Normalmente, el robot solo mide la distancia en línea recta entre donde está y donde quiere estar en ese mapa mental. Pero aquí está el truco: en el mundo real, la línea más corta en un mapa no siempre es la forma más rápida de llegar allí. Podrías trazar una línea recta a través de una pared, pero el robot no puede atravesar paredes. Necesita entender el tiempo y la secuencia. Si el robot solo observa la distancia, podría confundirse sobre si se está acercando a la meta o si solo está dando vueltas en círculos. Este artículo aborda exactamente ese problema: cómo enseñar a un robot a entender el "flujo" del tiempo en sus sueños para que pueda planificar mejor, sin necesidad de recompensas adicionales o maestros humanos.

El problema del pensamiento de "línea recta"

Los investigadores partieron de un sistema potente ya existente llamado LeWM (Modelo de Mundo Latente). Este sistema es excelente prediciendo qué sucede después en su mapa mental abstracto. Sin embargo, cuando intenta planificar una ruta hacia un objetivo, se basa en una regla simple: "Cuanto más cerca esté la imagen mental del futuro de la imagen del objetivo, mejor". Esto es como intentar navegar por una ciudad mirando únicamente la distancia en línea recta entre dos puntos. Funciona aceptablemente en campos abiertos, pero si necesitas rodear un edificio o subir una escalera, esa distancia de línea recta te engaña.

Los autores observaron que, si bien el mapa mental del robot era bueno prediciendo el siguiente paso, no era muy bueno clasificando las posibilidades futuras según cuánto progreso real realizaban. En un video de un humano empujando un bloque, el robot podía predecir el siguiente fotograma, pero no sabía intrínsecamente que el fotograma 10 está "más cerca" de la meta que el fotograma 5 solo por el orden en que ocurrieron. Le faltaba un sentido de tiempo dirigido.

La solución: Minar el tiempo del pasado

Para solucionar esto, el equipo creó un nuevo sistema llamado TD-JEPA (Arquitectura de Predicción de Incrustación Conjunta de Distancia Temporal). En lugar de simplemente dejar que el robot adivine la distancia, le enseñaron a "minar" el concepto del tiempo directamente de los videos de demostración.

Así es como lo hicieron, usando una analogía sencilla: Imagina que estás viendo un programa de cocina. No necesitas un temporizador para saber que romper un huevo ocurre antes de batirlo, y batirlo ocurre antes de hornearlo. El orden de los eventos es el progreso. TD-JEPA observa los videos de entrenamiento del robot y dice: "Si el humano hizo el paso A, luego el paso B, luego el paso C, entonces el paso C está definitivamente más avanzado en el camino hacia la meta que el paso A".

El sistema aprende un "costo dirigido" especial. A diferencia de una distancia normal que es la misma si vas hacia adelante o hacia atrás (como la distancia entre tu casa y el parque, que es la misma en ambos sentidos), este nuevo costo entiende la dirección. Sabe que ir de "Inicio" a "Meta" requiere esfuerzo, pero ir de la "Meta" de regreso al "Inicio" es un tipo de viaje diferente. Al aprender este flujo dirigido de los videos, el robot construye un mapa mental donde la "distancia" al objetivo realmente refleja el número de pasos necesarios para llegar allí, no solo qué tan similares se ven las imágenes.

Dos formas diferentes de usar la nueva habilidad

Uno de los hallazgos más interesantes es que esta nueva habilidad "consciente del tiempo" funciona de manera diferente dependiendo de lo que el robot esté intentando hacer. Los autores descubrieron que la mejor manera de usar este nuevo conocimiento depende de la tarea:

Para navegación y alcance (Tareas topológicas): Cuando el robot necesita moverse a través del espacio, como navegar por un laberinto o alcanzar un objeto en el aire, el nuevo "costo de tiempo dirigido" es la herramienta perfecta. Actúa como una brújula que apunta directamente hacia el objetivo, ignorando callejones sin salida. En sus pruebas, el uso de este nuevo costo ayudó al robot a resolver una tarea de navegación de "Dos Habitaciones" el 100% de las veces, superando al sistema antiguo que solo tuvo éxito en un 97.4%.
Para manipulación rica en contacto (Tareas geométricas): Cuando el robot tiene que empujar, deslizar o apilar objetos (como la tarea "Push-T" donde empuja un bloque en forma de T), las cosas se complican. Aquí, la forma y el ángulo exactos del contacto importan más que la dirección general. En estos casos, el nuevo sistema consciente del tiempo funciona mejor cuando ayuda al método de distancia de "línea recta" anterior, en lugar de reemplazarlo. Refina el mapa mental del robot para que la distancia de línea recta sea más precisa. En la tarea "OGB-Cube", esta combinación mejoró las tasas de éxito en 14.2 puntos en comparación con el sistema anterior.

Lo que encontraron (y lo que no)

Los investigadores realizaron pruebas estrictas donde mantuvieron todo igual excepto el método de planificación. Encontraron que TD-JEPA consistentemente igualaba o superaba a los mejores métodos anteriores (LeWM y un método llamado RC-aux) en todos los entornos.

Sin embargo, también encontraron un límite. Cuando intentaron usar solo el nuevo costo basado en el tiempo para la compleja tarea "Push-T", el robot en realidad funcionó peor (69% de éxito) que cuando utilizó la antigua distancia geométrica (86% de éxito). ¿Por qué? Porque empujar un bloque requiere un control fino de los ángulos y los puntos de contacto. La señal de "tiempo" le decía al robot que se estaba acercando, pero no le daba suficiente detalle sobre cómo tocar el bloque correctamente. El robot necesitaba la información de la "forma" geométrica para tener éxito. Esto demuestra que, si bien entender el tiempo es crucial, no reemplaza la necesidad de entender la geometría física en cada situación.

La conclusión

El artículo muestra que puedes enseñar a un robot a entender el "flujo" del tiempo simplemente observando videos, sin necesidad de recompensas adicionales o retroalimentación humana. Al minar el orden de los eventos en los registros de demostración, crearon un sistema que sabe qué futuro está verdaderamente "más cerca" de la meta. Este nuevo sistema, TD-JEPA, cierra la brecha entre lo que el robot aprende y cómo planifica. Sugiere que, para algunas tareas, el robot debe seguir el flujo del tiempo como un río, mientras que para otras, debe usar ese flujo para agudizar su sentido de la forma física. El resultado es un robot más inteligente y adaptable que puede planificar mejor en el mundo real, ya sea navegando por una habitación o empujando un bloque a través de una mesa.

Resumen Técnico: TD-JEPA (Arquitectura de Predicción de Espacio Latente de Distancia Temporal)

Planteamiento del Problema

Las Arquitecturas de Predicción de Embebido Conjunto (JEPAs, por sus siglas en inglés) han surgido como modelos de mundo latentes eficaces, aprendiendo a predecir representaciones futuras en lugar de reconstruir píxeles. Sin embargo, existe una brecha estructural entre el objetivo de entrenamiento de las JEPAs y los requisitos del Control Predictivo de Modelo (MPC).

La Brecha Entrenamiento-Planificación: El entrenamiento de una JEPA típicamente optimiza la predicción latente de un horizonte corto. En contraste, la planificación requiere clasificar futuros imaginados de múltiples pasos basándose en el progreso hacia una meta.
La Limitación Geométrica: Los planificadores de JEPA existentes (por ejemplo, LeWM) dependen de la distancia euclidiana latente como un costo de planificación de facto. Si bien esta geometría puede correlacionarse con el progreso en algunos contextos, es un subproducto del aprendizaje de representaciones más que una señal extraída explícitamente para "pasos hacia la meta" o alcanzabilidad dirigida.
El Desafío: ¿Cómo pueden los registros de demostraciones offline sin recompensa proporcionar una estructura de progreso temporal que estreche esta brecha? Además, ¿cuándo debe desplegarse esta estructura extraída directamente como un costo de planificación frente a ser utilizada para dar forma a las representaciones para la planificación geométrica?

Metodología: TD-JEPA

Los autores proponen TD-JEPA (Temporal-Distance JEPA), que conserva la arquitectura de codificador-predictor de LeWM, pero la aumenta con un mecanismo de extracción de costo temporal dirigido. El método consta de tres componentes principales:

1. Backbone y Predicción

TD-JEPA utiliza la arquitectura estándar de LeWM: un codificador condicionado por la acción ( $Enc_\theta$ ) y un predictor ( $Pred_\phi$ ) que mapean las observaciones a estados latentes y pronostican futuros embebidos. Preserva la regularización SIGReg para evitar el colapso de la representación y la pérdida de predicción de un paso estándar ( $\mathcal{L}_{pred}$ ).

2. Extracción de Costo Temporal Dirigido

En lugar de depender únicamente de la distancia euclidiana, TD-JEPA aprende una función de energía no negativa y dirigida $d_\psi(z_s, z_g)$ que representa el costo de alcanzar un estado latente meta $z_g$ desde un estado latente origen $z_s$ .

Arquitectura (Cabezal MRN): El cabezal de costo utiliza una parametrización de Red de Residuo Métrico (MRN). Descompone el costo en un término de configuración simétrico y un residuo asimétrico:
$d_\psi(z_s, z_g) = \|\phi_{sym}(z_s) - \phi_{sym}(z_g)\|_2 + \max_k \text{ReLU}(\phi_{asym, k}(z_s) - \phi_{asym, k}(z_g))$
Este diseño asegura que el costo sea no negativo e inherentemente asimétrico ( $d_\psi(z_s, z_g) \neq d_\psi(z_g, z_s)$ ), reflejando la naturaleza dirigida del alcance de metas.
Supervisión (Etiquetas Temporales): El modelo se entrena sobre registros de demostraciones offline sin recompensas densas.
- Positivos: Para índices $i < j$ en la misma trayectoria, la brecha temporal $\tau(i, j) = j - i$ sirve como un proxy para los pasos hacia el subobjetivo. El modelo realiza una regresión de $d_\psi(z_i, z_j)$ hacia $\tau(i, j)$ .
- Negativos: Los pares de distintas trayectorias actúan como negativos heurísticos, penalizando costos sospechosamente bajos para pares de estados no relacionados mediante una pérdida de tipo hinge.
Consistencia de Rollout: Se añade una pérdida de forzado docente multietapa ( $\mathcal{L}_{roll}$ ) para alinear el costo aprendido con el horizonte de bucle abierto utilizado por el planificador, reduciendo los errores de acumulación durante el rollout de $H$ pasos.

3. Planificación y Despliegue

En tiempo de prueba, el planificador (usando el Método de la Cota de Entropía Cruzada o iCEM) busca secuencias de acciones que minimicen la energía terminal después de un rollout latente de $H$ pasos.

Estrategia de Despliegue Dual: El artículo identifica dos roles distintos para la señal extraída según la estructura de la tarea:
1. Tareas Dominadas por la Topología (ej. Navegación): El costo dirigido extraído $d_\psi$ se despliega directamente como el objetivo de planificación.
2. Tareas de Contacto Rico (ej. Manipulación): Se utiliza el mismo checkpoint entrenado temporalmente, pero el planificador depende de la distancia euclidiana latente ( $\ell_2$ ). Aquí, la supervisión temporal actúa como una señal de representación para dar forma al espacio latente, mejorando la eficacia de la planificación geométrica.

Contribuciones Clave

Diagnóstico de la Brecha Entrenamiento-Planificación de JEPA: Los autores demuestran que las JEPAs predictivas aprenden la dinámica de los registros, pero la planificación a menudo depende de la geometría del embebido no entrenada. Muestran que la extracción de un costo temporal dirigido de los registros reduce significativamente este desajuste de clasificación (la $\rho$ de Spearman aumenta de 0.65 para la distancia euclidiana a 0.95 para el costo extraído en Push-T).
Marco de Extracción de Distancia Temporal: TD-JEPA introduce un método para extraer la estructura de progreso temporal de registros sin recompensa utilizando supervisión temporal dirigida, negativos de trayectorias cruzadas y consistencia de rollout ajustada al horizonte.
Codiseño de Costo y Despliegue: El artículo proporciona evidencia de que la elección del costo de planificación (temporal extraído vs. geométrico) debe codiseñarse con la estructura de la tarea. Los costos temporales sobresalen en la navegación topológica, mientras que los costos geométricos siguen siendo superiores para la manipulación de contacto rico, incluso cuando la representación subyacente ha sido entrenada con señales temporales.

Resultados Experimentales

El método fue evaluado en cuatro entornos: Two-Room (navegación), Reacher (alcance), Push-T (manipulación) y OGB-Cube (manipulación de contacto rico). Las evaluaciones se realizaron bajo un protocolo "bloqueado" (semillas, checkpoints y manifiestos de episodios fijos) para asegurar una comparación justa contra LeWM y la línea base concurrente RC-aux.

Navegación y Alcance: El despliegue del costo extraído $d_\psi$ $d_{ψ}$ mejoró significativamente las tasas de éxito.
- Two-Room: Logró un 100.0% de éxito (vs. el 97.4% de LeWM).
- Reacher: Logró un 97.0% de éxito (vs. el 96.0% de LeWM).
Manipulación: En tareas de contacto rico, el entrenamiento temporal mejoró la representación latente, beneficiando la planificación geométrica.
- OGB-Cube: El uso del checkpoint entrenado temporalmente con planificación $\ell_2$ elevó el éxito en 14.2 puntos porcentuales sobre LeWM (82.2% vs. 68.0%).
- Push-T: Mejoró de 83.6% (LeWM) a 86.0%.
Estudios de Ablación: Eliminar el residuo dirigido (cabezal simétrico), los negativos de trayectorias cruzadas o la consistencia de rollout resultó en una degradación del rendimiento, confirmando la necesidad de cada componente.
Análisis de Fallos: En Push-T, el costo temporal puro ( $d_\psi$ ) tuvo un desempeño inferior al costo geométrico ( $\ell_2$ ) porque los fallos ocurrieron principalmente durante la alineación post-contacto, donde la geometría local es más crítica que el ranking de progreso temporal.

Significado y Afirmaciones

El artículo afirma que TD-JEPA estrecha la brecha entrenamiento-planificación para los planificadores de modelos de mundo JEPA mediante el descubrimiento de la estructura de progreso temporal en los registros offline y el codiseño de la forma del costo con el despliegue en tiempo de plan.

Afirmaciones Modestas: Los autores no afirman que los costos temporales reemplacen universalmente a los costos geométricos. En su lugar, abogan por una estrategia de despliegue condicional: usar el costo dirigido extraído cuando el progreso es topológico, y usarlo como una señal de representación para la planificación geométrica cuando la geometría de contacto domina.
Validación Empírica: Bajo evaluación bloqueada, TD-JEPA iguala o supera tanto a LeWM como a RC-aux en cada entorno, demostrando que la extracción temporal sin recompensa es una vía viable para mejorar el MPC latente sin requerir modelos de recompensa o interacción online.

El código está disponible en el repositorio de GitHub proporcionado.