Resumen Técnico: Modelo de Mundo Prismático (PRISM-WM)

1. Declaración del Problema

El Aprendizaje por Refuerzo Basado en Modelos (MBRL) y la planificación en dominios robóticos están fundamentalmente limitados por la naturaleza híbrida de la dinámica física. Los sistemas del mundo real, particularmente en la locomoción bípeda y la manipulación, involucran movimiento continuo interrumpido por eventos discretos (por ejemplo, contactos, impactos, adherencia frente a deslizamiento). Estas interacciones crean dinámicas no suaves y altamente no lineales con discontinuidades agudas en el campo vectorial subyacente.

Los modelos de mundo latente convencionales suelen emplear redes neuronales monolíticas que imponen continuidad global. Aunque teóricamente capaces de aproximación universal, estos modelos aprenden prácticamente una representación promediada que suaviza en exceso modos dinámicos distintos. Para los planificadores de optimización de trayectorias, este suavizado excesivo induce errores acumulativos durante predicciones a largo plazo, particularmente en los límites físicos. Esto resulta en:

• Difuminado espectral a través de los límites de contacto, causando un desajuste significativo de gradientes.
• Deriva acumulativa en espacios de configuración de alta dimensión, haciendo que la planificación a largo plazo sea poco fiable.
• Colapso de la política en entornos ricos en contactos donde los modelos estándar no logran distinguir entre regímenes estables e inestables.

2. Metodología: El Modelo de Mundo Prismático (PRISM-WM)

Para abordar estas limitaciones, los autores introducen PRISM-WM, una arquitectura estructurada diseñada para descomponer dinámicas híbridas complejas en primitivas composables. El modelo se basa en un marco Mezcla de Expertos (MoE) consciente del contexto con un objetivo novedoso de ortogonalización latente.

Arquitectura Central

• Gestión Consciente del Contexto: Una red de gestión ($G$) recibe un vector de contexto $c_t$ (compuesto por el estado $z_t$, la acción $a_t$ y opcionalmente la incrustación de tarea $e_{task}$) y produce pesos normalizados mediante softmax $w_t$. Este mecanismo identifica implícitamente el régimen físico activo (por ejemplo, vuelo frente a apoyo) sin supervisión explícita.
• Expertos Especializados: El estado se enruta a un conjunto de $K$ redes expertas $\{E_k\}$. Cada experto se especializa en aprender el flujo continuo dentro de un régimen dinámico local específico, evitando la interferencia de modos conflictivos.
• Predicción Residual: El siguiente estado latente se predice como una combinación ponderada de las salidas de los expertos sumada al estado actual:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  donde $v_k$ representa las dinámicas residuales predichas por el $k$-ésimo experto.

Innovación Clave: Ortogonalización Latente

Para prevenir el colapso de modos (donde los expertos aprenden características redundantes) y asegurar la diversidad, PRISM-WM introduce una capa de ortogonalización Gram-Schmidt diferenciable.

• Después de que cada experto emite un vector de la penúltima capa $u_k$, el modelo aplica un proceso Gram-Schmidt estabilizado para generar vectores de base ortogonales $v_k$.
• Esto obliga a las representaciones de los expertos a residir en una variedad de Stiefel, asegurando que abarquen un subespacio de rango máximo y no redundante para las dinámicas residuales.
• Aunque los modos de contacto físico no están ordenados, esta ortogonalización asegura que los expertos capturen componentes distintos de la variedad residual, permitiendo que la red de gestión reconstruya modos físicos no ordenados mediante una suma ponderada.

Integración con la Planificación

PRISM-WM se integra en el marco de planificación TD-MPC. Al factorizar dinámicas complejas en expertos no redundantes, el modelo reduce el residuo del modelo de un paso de $\epsilon$ a aproximadamente $\epsilon/K\alpha$. Esta descomposición estructural mitiga la deriva acumulativa, permitando rodajes de alta fidelidad a largo plazo y una propagación de gradientes estable para el aprendizaje de políticas.

3. Contribuciones Clave

1. Descomposición Arquitectónica: Introducción de PRISM-WM, que descompone dinámicas híbridas utilizando una arquitectura MoE consciente del contexto acoplada con ortogonalización latente, alejándose de aproximaciones globales monolíticas.
2. Objetivo de Ortogonalización: Un objetivo de entrenamiento novedoso que utiliza ortogonalización Gram-Schmidt diferenciable para imponer la diversidad de expertos y prevenir el colapso de modos, alineando estructuralmente el espacio latente con transiciones físicas discretas.
3. Validación Teórica y Empírica: Demostración de que esta descomposición reduce los errores acumulativos en espacios de alta dimensión y mejora la geometría del paisaje de pérdida tanto para la planificación (TD-MPC) como para el aprendizaje directo de políticas (PWM).
4. Capacidad de Transferencia Zero-Shot: Evidencia de que la base ortogonal permite que el modelo generalice entre simuladores con físicas de contacto diferentes (por ejemplo, IsaacGym de contacto blando a MuJoCo de contacto rígido) aislando modos dinámicos específicos (por ejemplo, fricción frente a soporte vertical).

4. Resultados Experimentales

Los autores evaluaron PRISM-WM en benchmarks de control continuo, incluyendo humanoides de alta dimensión, configuraciones de múltiples tareas y transferencia sim-a-sim zero-shot.

• Locomoción de Alta Dimensión (DiffRL & Humanoid-Bench): PRISM-WM logró una mayor eficiencia de muestras y rendimiento asintótico en comparación con las líneas base monolíticas (TD-MPC2, DreamerV3, PWM, SAC). Notablemente, mientras que los modelos monolíticos experimentaron colapso de la política en tareas de locomoción bípeda (Humanoid, SNU-Humanoid) debido al promediado de eventos de contacto no suaves, PRISM-WM mantuvo una optimización estable.
• Generalización Multi-Tarea: En un benchmark de 30 tareas del DeepMind Control Suite, PRISM-WM alcanzó una puntuación normalizada media de 0.531, una mejora del 23.5% sobre la línea base monolítica TD-MPC2 (0.430). Superó a la línea base en 26 de 30 tareas, mitigando efectivamente la transferencia negativa entre dinámicas heterogéneas.
• Transferencia Sim-a-Sim Zero-Shot: Al transferir de IsaacGym a MuJoCo, la línea base MLP monolítica sufrió una caída del 47.8% en el rendimiento con una tasa de éxito del 12%. PRISM-WM redujo esta caída al 11.8% con una tasa de éxito del 88%. La ortogonalización permitió que el modelo ajustara los pesos para expertos específicos (por ejemplo, fricción horizontal) sin comprometer la variedad de dinámicas globales.
• Estabilidad y Eficiencia: Los estudios de ablación confirmaron que la ortogonalización previene la degradación de la estimación de recompensas observada en modelos MoE sin restricciones. Las pruebas de eficiencia de inferencia en una NVIDIA RTX 4090 mostraron que PRISM-WM (K=4) alcanzó 7,712 FPS, superando ligeramente a la línea base monolítica (7,499 FPS) con una sobrecarga de memoria mínima.

5. Significado y Afirmaciones

El artículo afirma que PRISM-WM proporciona un sustrato de alta fidelidad para algoritmos de optimización de trayectorias, actuando como un modelo fundamental para agentes basados en modelos. Al descomponer la física compleja en primitivas ortogonales, el modelo:

• Mitiga la deriva a largo plazo en horizontes de planificación, permitiendo un razonamiento fiable a largo plazo en entornos ricos en contactos.
• Estabiliza la propagación de gradientes para el aprendizaje de políticas, transformando paisajes de pérdida fracturados en mezclas convexas funcionalmente suaves.
• Permite una toma de decisiones robusta en sistemas híbridos donde los modelos de mundo estándar fallan debido al suavizado excesivo.

Los autores concluyen que descomponer la física compleja en primitivas ortogonales es "fundamentalmente esencial para dominar el control rico en contactos". Aunque el marco actual requiere un ajuste manual de la cantidad de expertos $K$, el trabajo establece una base estructural para construir modelos de mundo de propósito general y fundamentados físicamente para futuros agentes encarnados.