Resumen Técnico: Modelos de Mundo Duelistas

Planteamiento del Problema

Los modelos de mundo latentes, particularmente aquellos que siguen la Arquitectura de Predicción de Embebido Conjunto (JEPA, por sus siglas en inglés), planifican prediciendo los embeddings de estados futuros a partir de estados actuales y acciones. Sin embargo, cuando una escena contiene movimiento que el agente no controla (distractores), estos modelos sufren un modo de fallo "silencioso": las predicciones para diferentes acciones se vuelven indistinguibles, incluso cuando el error de entrenamiento total (error de reconstrucción o de predicción) continúa mejorando.

Los remedios existentes intentan separar la dinámica controlable de los distractores exógenos mediante la adición de maquinaria externa, como pérdidas de reconstrucción, recompensas de tareas, objetivos auxiliares de dinámica inversa o ramas de red separadas. Estos enfoques suelen requerir supuestos específicos sobre el proceso del distractor o supervisión adicional. Los autores identifican que en los predictores monolíticos estándar $P(z, a)$, el canal de acción se entrelaza con la dinámica del distractor a menos que esté explícitamente fundamentado por etiquetas de acción reales o purgado mediante máscaras de segmentación. Crucialmente, la selección de modelos estándar basada en la pérdida de validación suele fallar al detectar esta "ceguera a la acción", seleccionando activamente puntos de control donde la señal de la acción ha colapsado.

Metodología: AD-JEPA

El artículo propone AD-JEPA (JEPA de Acción Descompuesta), una reestructuración mínima de los predictores condicionados por acción que aísla el canal de acción sin requerir recompensas, reconstrucción o pérdidas auxiliares específicas del distractor.

Mecanismo Central: Centrado de la Media de la Acción

El método descompone el predictor en dos flujos:

1. Flujo Pasivo ($B(z)$): Predice cómo evoluciona la escena con la acción marginalizada.
2. Flujo de Acción Centrado ($u(z, a)$): Representa el desplazamiento específico de la acción, calculado restando la media de la predicción sobre las acciones del desplazamiento de acción bruto.

La regla de actualización se define como:
$$\hat{z}_{t+1} = B(z_t) + \underbrace{[\Delta(z_t, a_t) - \bar{\Delta}(z_t)]}_{u(z_t, a_t)}$$
Donde $\bar{\Delta}(z) = \mathbb{E}_{a' \sim \nu}[\Delta(z, a')]$ es la media del desplazamiento sobre una distribución de propuesta $\nu$ de acciones.

Fundamento Teórico

Este enfoque adapta la descomposición de Dueling DQN ($Q = V + (A - \bar{A})$) de valores escalares a la dinámica de embebidos con valores vectoriales. Los autores demuestran que esta resta es una identidad algebraica en lugar de una tendencia aprendida:

• Cancelación de Modo Común: Cualquier variación que entre en los desplazos de acción de manera idéntica en todas las acciones (variación de modo común), incluyendo los distractores independientes de la acción, se cancela exactamente.
• Exactitud: La cancelación se mantiene tanto para conjuntos de acciones discretas (mediante enumeración exacta) como para conjuntos de acciones continuas (mediante muestreo de Monte Carlo), siempre que la media se estime sobre un conjunto finito de candidatos.
• Sin Supuesto de Exogeneidad: El método no requiere supuestos de exogeneidad a nivel de proceso (es decir, que los distractores sean independientes de la política del agente a nivel de proceso); solo requiere que la contribución del distractor a la transición sea de modo común a través de las acciones.

Objetivo de Entrenamiento

El modelo se entrena con una pérdida de predicción JEPA estándar contra un objetivo EMA con parada de gradiente, aumentada con:

• InfoNCE ($L_{act}$): Fomenta que la predicción de la acción ejecutada sea la más cercana al objetivo entre las acciones contrafácticas.
• Penalización de Compacidad ($L_{off}$): Desaconseja que el canal de acción absorba la persistencia del estado.
• Regularización VICReg: Evita el colapso de la representación.

Notablemente, los autores informan que añadir un mecanismo de compuerta (gating) aprendido para seleccionar dimensiones fue inerte, degenerando en un escalado uniforme hacia abajo, confirmando que el centrado es el mecanismo operativo.

Contribuciones Clave

1. Centrado de la Media de la Acción: Una reestructuración de una sola línea de los predictores condicionados por acción que vuelve el canal de acción invariante al modo común. Los autores proporcionan pruebas (Proposiciones 1 y 2) que muestran la cancelación exacta de los componentes de modo común para conjuntos de acciones discretos y muestreados.
2. Validación Empírica: El método fue validado en cuatro entornos distintos:
   • Gridworld (FourRooms): Con células distractoras binarias en movimiento.
   • Generadores Sintéticos: Con factores controlables y de molestia conocidos.
   • Control Continuo (DMC): Con oclusores móviles a nivel de píxel.
   • Atari (Freeway): Con distractores de tráfico naturales.
     En todos los casos, el canal centrado preservó el efecto del agente mientras reducía la filtración de la molestia hasta hacerla indistinguible de cero, mientras que los predictores entrelazados fallaron.
3. Aplicabilidad Solo de Lectura (Readout-Only): Debido a que el mecanismo es una propiedad de la lectura y no de la arquitectura, puede aplicarse post-hoc a modelos de mundo pre-entrenados y congelados (específicamente RePo y TIA). Los autores demostraron que restar la media de la acción de las lecturas brutas de estos modelos congelados recuperó con éxito un canal de acción que de otro modo era latente y no enrutado.
4. Límite Definido: El artículo identifica explícitamente la limitación: el método falla cuando la dinámica del distractor está correlacionada con la acción (por ejemplo, distractores reafferentes o agentes que reaccionan al usuario). En tales casos, el distractor no es de modo común y no se cancela.

Resultados

• Recuperación vs. Rechazo: En pruebas sintéticas, el canal centrado logró una alta recuperación del factor controlable ($R^2 \approx 0.87\text{--}0.93$) mientras rechazaba el factor de molestia ($R^2 \approx -0.00$). Las variantes no centradas recuperaron la señal pero filtraron pesadamente la molestia.
• Desempeño de Control: En el gridworld, un controlador de predicción de modelo usando el canal centrado mantuvo un alto éxito de alcance de meta ($\approx 0.92$) incluso con 30 celdas de distracción, mientras que el predictor estándar colapsó a $\approx 0.57$.
• Éxito Post-Hoc: Aplicar la resta a los modelos congelados RePo y TIA recuperó canales de acción con valores $R^2$ entre 0.19 y 0.52, mientras que las lecturas de los deltas de acción brutos eran cercanos a cero.
• Modos de Fallo: El método falló correctamente (colapsando al azar) en pruebas de estrés de acción correlacionada, alineándose con el límite teórico.

Significado y Reivindicaciones

El artículo afirma que el centrado es un complemento universal (plug-in) para los modelos de mundo condicionados por acción. Su importancia radica en:

• Simplicidad: Logra el rechazo de distractores sin la complejidad de la reconstrucción, la factorización basada en recompensas o las pérdidas auxiliares de dinámica inversa.
• Garantía Estructural: La separación es una propiedad algebraica de la parametrización, no una tendencia aprendida dependiente de señales de entrenamiento específicas.
• Capacidad de Retroajuste: Permite que los modelos de mundo existentes y congelados sean "des-distraídos" en el tiempo de inferencia sin reentrenamiento, haciendo emerger canales de acción que antes estaban ocultos por dinámicas entrelazadas.

Los autores son modestos respecto a la escalabilidad del control: aunque el canal permite el control en entornos de gridworld, un intento registrado para escalar esto al control continuo de DMC no logró superar a los baselines aleatorios. Atribuyen esto a cuellos de botella en la precisión del modelo en el planificador, más que a un defecto en el canal mismo. La contribución principal es, por tanto, a nivel de representación: proporcionar un canal de acción limpio e invariante al modo común donde los predictores entrelazados se vuelven "ciegos a la acción".