Resumen Técnico: DWM: Separando los Efectos del Mundo de las Acciones en Modelos de Mundo Latentes

Declaración del Problema

Los modelos de mundo latentes son fundamentales para el control basado en modelos moderno, permitiendo la planificación y el aprendizaje de políticas directamente desde observaciones de píxeles. Sin embargo, las formulaciones actuales condicionadas a la acción sufren de una limitación a nivel de supervisión: predicen el siguiente estado latente utilizando un único objetivo indiferenciado. Esto obliga al modelo a absorber todas las fuentes de cambio —tanto aquellas inducidas por la acción del agente como las que surgen de la dinámica intrínseca del entorno— en una señal de aprendizaje monolítica.

En sistemas físicos del mundo real, las transiciones suelen provenir de dos fuentes heterogéneas:

1. Efectos impulsados por la acción: Cambios inducidos por el control del agente.
2. Efectos del mundo: Cambios invariantes a la acción que persisten incluso bajo una acción nula (por ejemplo, deslizamiento por gravedad, inercia, rebote por contacto o deriva persistente).

Fusionar estos elementos en un solo objetivo entrelaza ambos componentes dentro de la transición latente. En consecuencia, el modelo no puede atribuir explícitamente los cambios observados a sus causas subyacentes. Esta limitación suele estar oculta en entornos estáticos (por ejemplo, empuje sobre mesa plana), pero se vuelve crítica en entornos con dinámicas autónas persistentes, lo que conduce a rollouts de múltiples pasos inexactos y a una planificación subóptima.

Metodología: Modelo de Mundo Descompuesto (DWM)

Los autores proponen DWM (Decomposed World Model), un marco de entrenamiento que operativiza la separación de los efectos del mundo y de la acción sin alterar la arquitectura subyacente ni el flujo de inferencia.

Arquitectura Central

DWM conserva los componentes estándar de un modelo de mundo latente condicionado a la acción (codificador $\phi$, predictor condicionado a la acción $g$ y un cabezal de proyección):

• Cabezal de Predicción ($h_p$): Se mantiene del modelo base. Predice el estado latente completo $\hat{z}$ y es el único cabezal utilizado durante la inferencia.
• Cabezal del Mundo ($h_w$): Un MLP auxiliar ligero conectado en paralelo al estado del predictor $r_t$. Está entrenado para capturar el componente invariante a la acción de la transición ($\hat{z}_w \approx z_t + \Delta z^{world}_t$).
• Componente Impulsado por la Acción: Definido como el residuo $\hat{z}_a := \hat{z} - \hat{z}_w$.

El modelo impone la identidad $\hat{z} = \hat{z}_w + \hat{z}_a$, descomponiendo explícitamente la transición predicha en un componente del mundo invariante a la acción y un componente complementario impulsado por la acción.

Objetivo de Entrenamiento

La función de pérdida total combina la pérdida estándar del modelo de mundo con dos términos auxiliares:
$$L = L_{WM} + \lambda_{wc}L_{wc} + \lambda_{orth}L_{orth}$$

1. Pérdida de Contraste de Mundo ($L_{wc}$): Un objetivo InfoNCE aplicado al cabezal del mundo. Construye pares positivos alimentando el mismo historial de estados con diferentes acciones actuales (perturbaciones de acción) al cabezal del mundo. Esto obliga a $\hat{z}_w$ a ser invariante a la acción mientras permanece discriminativo entre diferentes estados.
2. Restricción de Ortogonalidad ($L_{orth}$): Definida como $|\cos(\hat{z}_w, \hat{z}_a)|$, este término fomenta que los componentes invariantes a la acción y los impulsados por la acción ocupen direcciones independientes en el espacio latente, evitando que el efecto de la acción sea absorbido en la representación del mundo.

Crucialmente, el cabezal del mundo auxiliar se descarta después del entrenamiento, asegurando que el flujo de inferencia sea idéntico al del modelo base.

Contribuciones Clave

1. Identificación de una Limitación de Supervisión: El artículo identifica que la supervisión de un solo objetivo en los modelos de mundo latentes no logra distinguir entre los cambios inducidos por el agente y la dinámica autónoma del entorno, lo que provoca representaciones entrelazadas.
2. Nuevo Conjunto de Benchmarks (W-Variants): Los autores introducen variantes controladas de tres benchmarks estándar —PushT-W, Reacher-W y TwoRoom-W— que inyectan dinámicas persistentes invariantes a la acción (deslizamiento por gravedad, gravedad en plano vertical y deriva constante) en tareas estándar. También evalúan en Ball-in-Cup, una tarea con dinámicas oscilatorias dependientes del estado.
3. Marco DWM: Un marco de entrenamiento que aumenta los modelos de mundo latentes existentes con un cabezal del mundo auxiliar y objetivos de separación complementarios, sin requerir cambios en la arquitectura o en la inferencia.
4. Validación Empírica: Demostración de que DWM mejora el éxito de la planificación y la precisión de la dinámica latente específicamente en entornos con efectos del mundo, manteniendo el rendimiento en tareas estándar.

Resultados Experimentales

Los autores evaluaron DWM frente al baseline LeWM utilizando el método de la media de coseno (CEM) para la planificación.

• Éxito de Planificación: En las tareas de variante W, DWM logró una mejora absoluta media del 13.1% en el éxito de la planificación CEM:
  • PushT-W: 32.0% $\to$ 44.0% (+12.0 pp)
  • Reacher-W: 60.0% $\to$ 70.7% (+10.7 pp)
  • TwoRoom-W: 60.0% $\to$ 76.7% (+16.7 pp)
  • Ball-in-Cup: 36.8% $\to$ 42.8% (+6.0 pp)
• Rendimiento Base: En sus contrapartes originales "planas" (sin efectos del mundo), DWM se mantuvo comparable al baseline LeWM de un solo cabezal, indicando que las pérdidas auxiliares no degradan el rendimiento cuando los efectos del mundo son insignificantes.
• Diagnóstico de Representación: El análisis mostró que el cabezal del mundo se volvió casi invariante a la acción (la varianza se redujo en órdenes de magnitud bajo perturbación de acción), mientras que el cabezal de predicción mantuvo su sensibilidad a las acciones.
• Calidad de Predicción: DWM redujo significativamente los errores de predicción autorregresiva de un paso y de horizonte-20 (por ejemplo, -89.8% de pérdida offline en PushT-W), demostrando que la descomposición conduce a una dinámica latente más precisa.
• Generalización Fuera de Distribución (OOD): En PushT-W con vectores de gravedad desplazados, DWM mantuvo un alto éxito de planificación (44.0%) comparado con la degradación del baseline (30.0%), y el componente impulsado por la acción aprendido se mantuvo estable ante cambios de gravedad (similitud de coseno $\approx$ 0.999).

Significado y Reivindicaciones

El artículo sostiene que desentrelazar los efectos del mundo invariantes a la acción de los cambios impulsados por la acción al nivel de la señal de entrenamiento sirve como un sesgo inductivo útil para los modelos de mundo latentes condicionados a la acción en entornos de control físicamente realistas.

Los autores enfatizan que su enfoque aborda una limitación fundamental en las estrategias de supervisión actuales sin requerir cambios estructurales profundos. Al separar explícitamente "lo que el mundo hace por sí mismo" de "lo que el agente cambia", DWM permite aprender dinámicas latentes más precisas, lo que se traduce directamente en una mejor planificación de múltiples pasos y robustez ante dinámicas ambientales que persisten independientemente de las acciones del agente. Los resultados sugieren que esta descomposición es particularmente vital para tareas donde el entorno evoluciona de forma autónoma, un escenario a menudo ignorado en los benchmarks estándar.