Resumen Técnico: Correspondencia de Tokens Identificables para Modelos del Mundo

1. Declaración del Problema

Los modelos del mundo basados en transformadores han demostrado un rendimiento sólido en el aprendizaje por refuerzo visual (RL) al tratar secuencias de estados y acciones pasados como flujos de tokens para predecir estados futuros. Sin embargo, estos modelos sufren frecuentemente de inconsistencia temporal durante despliegues de largo horizonte. Los modos de fallo comunes incluyen la duplicación de objetos, su desaparición y su transmutación en objetos diferentes.

Los autores identifican una causa principal de estos errores: los enfoques existentes tratan la predicción del siguiente cuadro puramente como un problema de generación de tokens. No modelan explícitamente la correspondencia entre tokens a través del tiempo. En entornos visuales realistas, los tokens en cuadros sucesivos a menudo corresponden a las mismas entidades subyacentes que persisten y se mueven. Cuando se obliga a un modelo a regenerar cada token desde cero sin rastrear la identidad, le resulta difícil preservar estructuras persistentes, lo que conduce a una preservación de identidad frágil y a errores acumulativos a lo largo del tiempo.

2. Metodología: Correspondencia de Tokens Identificables (ITC)

El artículo propone la Correspondencia de Tokens Identificables (ITC), un modelo del mundo basado en transformadores que formula la predicción del siguiente cuadro como un problema de inferencia probabilística estructurada con variables latentes de correspondencia de tokens.

Mecanismo Central

En lugar de generar cada token del siguiente cuadro de novo, ITC determina para cada token del siguiente cuadro si debe ser:

1. Copiado desde un token del cuadro anterior (preservando la identidad).
2. Generado como un nuevo token (manejando cambios o contenido nuevo).

Esta asignación se modela como un problema de Transporte Óptimo (OT).

Arquitectura Técnica

1. Tokenización y Embebido:

   • Las observaciones visuales se convierten en tokens discretos utilizando un tokenizador de búsqueda de parches por vecino más cercano.
   • Los embebidos de tokens se aumentan con Embebidos de Posición Rotatoria 3D (RoPE). A diferencia del RoPE 1D estándar, esta formulación codifica las coordenadas temporal, vertical y horizontal de forma independiente, permitiendo al modelo capturar estructuras relacionales localizadas tanto en el espacio como en el tiempo.

2. Solucionador de Transporte Óptimo:

   • El modelo construye un grafo donde los vértices fuente representan los tokens del estado anterior ($u$) y los tokens candidatos del estado siguiente predichos por el transformador ($\tilde{p}$). Los vértices de destino representan los tokens finales del estado siguiente ($\hat{u}$).
   • Se calcula una matriz de afinidad para puntuar la relación entre los tokens anteriores y los tokens predichos. Esta matriz incluye:
     • $A^{(prev)}$: Puntuación de la afinidad entre tokens anteriores y predicciones, penalizada por la distancia espacial para imponer plausibilidad física (por ejemplo, los objetos no pueden teletransportarse arbitrariamente lejos).
     • $A^{(gen)}$: Puntuación de la "bonificación" de admitir un token generado recientemente ( comodín) en lugar de reutilizar uno anterior, sujeto a una penalización constante.
   • El solucionador (utilizando el algoritmo de Sinkhorn) calcula un plan de transporte que minimiza el costo de mover masa desde las fuentes hacia los destinos.

3. Binarización y Decodificación:

   • El plan de transporte continuo se convierte en una matriz de asignación binaria utilizando un procedimiento de binarización codiciosa (Algoritmo 2). Esto asegura una mapeo estricto uno a uno.
   • Regla de Decodificación: Para cada posición en el siguiente cuadro, si el plan de transporte asigna un token anterior, ese token se copia. Si el plan asigna un token de "generar", el modelo muestrea desde la distribución predicha por el transformador.

Bucle de Entrenamiento

El modelo del mundo y la política se entrenan concurrentemente. El modelo del mundo genera "imaginaciones" (trayectorias sintéticas) que se utilizan junto con datos reales del entorno para actualizar la política mediante Optimización de Política Proximal (PPO).

3. Contribuciones Clave

• Formulación de la Correspondencia de Tokens: El artículo reformula la predicción del siguiente cuadro modelando explícitamente la correspondencia de tokens como un problema de transporte óptimo, separando las tareas de "reutilizar" entidades persistentes y "generar" contenido nuevo.
• Integración de Transporte Óptimo: Introduce un mecanismo diferenciable (vía Sinkhorn) pero discreto (vía binarización) para seleccionar entre copiar tokens anteriores y generar nuevos, reduciendo efectivamente las alucinaciones.
• Codificación Posicional 3D: La implementación de RoPE 3D permite al modelo distinguir entre diferencias temporales y espaciales, un componente necesario para rastrear objetos en movimiento en secuencias de video.

4. Resultados Experimentales

Los autores evaluaron ITC en cuatro benchmarks desafiantes: Craftax-clásico, Craftax, MinAtar y Atari 100K.

• Craftax-clásico: ITC logró un retorno del 72.5% y una puntuación del 35.6%, superando significativamente al estado del arte anterior (Dedieu et al., 2025) que logró un retorno del 67.4% y una puntuación del 27.9%. También demostró una convergencia más rápida.
• Craftax: En este entorno más complejo, ITC logró un retorno del 7.09% y una puntuación del 2.40%, superando a las líneas base que incluyen Simulus y Dedieu et al.
• MinAtar: ITC superó al estado del arte anterior para RL basado en modelos en los cuatro juegos (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Aplicado a un conjunto de 26 juegos, ITC logró un nuevo rendimiento del estado del arte en la Media Intercuartílica (IQM) de 1.092 y una Brecha de Optimalidad de 0.376, superando al modelo del mundo basado en tokens Simulus.
• Precisión de Predicción: En un conjunto de datos de 10,000 transiciones, ITC mejoró la precisión general de predicción en un 3.39% en comparación con la línea base, con una mejora notable del 3.45% en transiciones que involucran criaturas en movimiento.
• Mejoras Cualitativas: El análisis visual confirmó que ITC elimina artefactos comunes como la duplicación y desaparición de objetos vistos en modelos base como DreamerV3 y Dedieu et al.
• Eficiencia: El método introdujo una sobrecarga insignificante, aumentando el tiempo total de entrenamiento de extremo a extremo solo en un 2.8%.

5. Significado y Afirmaciones

El artículo afirma que al modelar explícitamente la correspondencia entre tokens a través del tiempo, ITC resuelve un problema generalizado en los modelos del mundo donde los agentes fallan en mantener la identidad de los objetos a lo largo de horizontes largos. Los autores declaran que este enfoque permite al modelo "reutilizar selectivamente tokens que tienen una fuerte correspondencia con el cuadro anterior", aprovechando así las similitudes entre cuadros en lugar de depender únicamente del transformador para regenerar cada token.

El significado radica en lograr un rendimiento del estado del arte en diversos entornos (desde mundos abiertos basados en cuadrículas hasta juegos estándar de Atari) manteniendo la eficiencia computacional. Los autores concluyen que ITC previene efectivamente las inconsistencias temporales como la duplicación y la desaparición, lo que lleva a imaginaciones de mayor calidad y a un mejor rendimiento de la política.