Resumen Técnico: Escalado del Aprendizaje por Refuerzo de Modelos Mundiales mediante Optimización de Políticas de Difusión

1. Declaración del Problema

El aprendizaje por refuerzo (RL) basado en modelos que utiliza modelos mundiales ofrece un paradigma prometedor para la planificación eficiente en muestras, al permitir que los agentes "sueñen" y planifiquen dentro de entornos internos aprendidos. Sin embargo, la escalabilidad de estos enfoques sigue estando fundamentalmente limitada. Mientras que trabajos anteriores han atribuido las brechas de rendimiento al sesgo del modelo y a los errores de predicción acumulativos, este artículo identifica un cuello de botella más crítico y poco explorado: la desalineación estructural entre la búsqueda y el aprendizaje de valores.

En los enfoques existentes de modelos mundiales (por ejemplo, la serie TD-MPC), la mejora de la política suele depender de un planificador basado en muestreo (como el Integral de Trayectoria Predictiva de Modelo, MPPI) que opera sobre el modelo mundial aprendido. Por el contrario, la función de valor se entrena sobre trayectorias generadas por una política de búsqueda separada (a menudo la política de comportamiento o una política de búfer de repetición). Esto crea una discrepancia distribucional:

1. Brecha de Iteración de Valores: La función de valor se entrena con datos de la política $\beta$, pero es consultada por la política de búsqueda $\pi$. Cuando $\pi$ se desvía significativamente de $\beta$, las estimaciones de valor se vuelven poco fiables, lo que a menudo conduce a una sobreestimación espuria.
2. Brecha de Mejora de la Política: El procedimiento de búsqueda explota estos valores sobreestimados, reforzando los errores y provocando un ciclo vicioso de desplazamiento de distribución y amplificación de errores. En consecuencia, la búsqueda no constituye una verdadera optimización de políticas en el sentido convencional, fallando en garantizar una mejora monótona o la convergencia a la política óptima.

2. Metodología: Optimización de Políticas de Difusión Basada en Modelos (MBDPO)

Para abordar estas limitaciones, los autores proponen la Optimización de Políticas de Difusión Basada en Modelos (MBDPO), un marco que unifica la búsqueda y la optimización de políticas mediante representaciones de políticas de difusión dentro de modelos mundiales latentes.

Marco Central

MBDPO reformula la optimización de políticas como un proceso de difusión sobre trayectorias imaginadas en el espacio latente, en lugar de depender de planificadores explícitos como MPPI. El proceso implica:

• Modelo Mundial Latente: El agente aprende un codificador $E$, dinámicas latentes $F$ y un predictor de recompensas $R$ para operar en un espacio latente compacto $Z$.
• Función de Energía Implícita: Para anclar las actualizaciones de la política y prevenir el desplazamiento distribucional, MBDPO aprende una función de energía implícita $E_\theta$ a partir del conjunto de datos del búfer mediante un objetivo contrastivo. Esta función aproxima la densidad logarítmica de la política de comportamiento $\beta$, actuando efectivamente como una restricción KL.
• Ajuste de Puntuación mediante Trayectorias Imaginadas: En lugar de un proceso de difusión hacia adelante, la función de puntuación $\phi$ se estima directamente a partir de trayectorias imaginadas dentro del modelo mundial aprendido.
  • El algoritmo muestrea secuencias de acciones limpias candidatas desde una distribución de propuesta.
  • Estas secuencias se despliegan a través de las dinámicas latentes $F$ para calcular un retorno acumulado regularizado por energía $\tilde{G}$.
  • Las muestras se reponderan basándose en $\tilde{G}$ para estimar la función de puntuación, la cual guía el proceso de difusión inversa.

Unificación Teórica

El artículo establece una equivalencia teórica (Teorema 1) entre la condición de ajuste de puntuación en los modelos de difusión y la solución de un problema de optimización de políticas con restricción KL.

• El objetivo es maximizar el retorno esperado sujeto a una restricción de divergencia KL contra la política de comportamiento: $\max_{\pi_\phi} \mathbb{E}[G] \text{ s.t. } D_{KL}(\pi_\phi \| \beta) \leq \epsilon$.
• MBDPO demuestra que la política óptima para este problema restringido sigue una distribución de Gibbs.
• Al aprender la función de puntuación para igualar esta distribución, el proceso de difusión inversa realiza implícitamente la optimización restringida, unificando así la búsqueda (generación de trayectorias) y el aprendizaje de políticas (maximización de valores) en un único marco coherente.

Algoritmo Práctico

MBDPO emplea un paradigma de entrenamiento conjunto:

1. Modelo Mundial e Iteración de Valores: El codificador, las dinámicas, la recompensa, la función de valor ($\hat{Q}$) y la función de energía implícita se actualizan utilizando segmentos de trayectorias de un búfer de repetición.
2. Optimización de Políticas de Difusión: La red de puntuación se entrena para predecir la función de puntuación utilizando muestreo de Monte Carlo de trayectorias imaginadas ponderadas por sus retornos. Durante la inferencia, la red de puntuación sirve como una solución amortizada para la generación eficiente de secuencias de acciones.

3. Contribuciones Clave

1. Identificación de la Desalineación Estructural: El artículo proporciona un análisis teórico (Problemas 1 y 2) que demuestra que el desacoplamiento de la búsqueda y el aprendizaje de valores en los enfoques existentes de modelos mundiales conduce a la suboptimalidad y evita la mejora monótona de la política.
2. Marco MBDPO: La introducción de un marco unificado que formula la búsqueda y la optimización de políticas como un proceso de difusión. Este enfoque aprovecha una función de energía implícita para anclar la política a la distribución de comportamiento, mitigando la brecha de iteración de valores y la deriva distribucional.
3. Rendimiento Escalable: El desarrollo de un algoritmo de RL de difusión escalable que demuestra ganancias de rendimiento consistentes y monótonas a medida que aumenta la capacidad del modelo (de 1,7 M a 340 M de parámetros) en diversos entornos, incluido el preentrenamiento offline multi-tarea, el aprendizaje online y el ajuste fino de offline a online.
4. Validación Empírica: Evaluación exhaustiva en 121 tareas en cuatro puntos de referencia (DMControl, MetaWorld, ManiSkill2, MyoSuite) que muestra un rendimiento superior frente a las líneas base más avanzadas (SAC, DreamerV3, TD-MPC, TD-MPC2).

4. Resultados Experimentales

• Preentrenamiento Offline: MBDPO exhibe un comportamiento de escalado monótono claro. A medida que los parámetros del modelo aumentan de 1,7 M a 340 M, la puntuación normalizada en los conjuntos de datos de 80 y 30 tareas mejora consistentemente, superando significativamente a TD-MPC2, especialmente en escalas de modelo más pequeñas.
• Aprendizaje Online: En configuraciones online desde cero, MBDPO logra resultados superiores o competitivos en las 121 tareas. Demuestra una aceleración en la mejora de la política con menor varianza en las primeras etapas de entrenamiento en comparación con las líneas base.
• Mitigación de Cuellos de Botella: El análisis empírico muestra que MBDPO reduce significativamente los errores de diferencia temporal (TD) cruzada y la deriva de acciones en comparación con TD-MPC2. La política de difusión permanece estable y converge suavemente, mientras que los métodos de búsqueda heurística exhiben alta volatilidad y una deriva de política severa.
• Generalización: En experimentos de transferencia de offline a online (O2O), un agente generalista preentrenado en 70 tareas se ajustó finamente con éxito en 10 tareas no vistas con solo 40k interacciones por tarea, superando al entrenamiento desde cero.
• Eficiencia: MBDPO incurre en una sobrecarga insignificante en el entrenamiento online, con tiempos de entrenamiento por paso comparables o más rápidos que TD-MPC2 en la mayoría de los dominios.

5. Significado y Afirmaciones

El artículo afirma que MBDPO establece un nuevo paradigma para el aprendizaje de políticas dentro de modelos mundiales al resolver la desalineación estructural entre la búsqueda y el aprendizaje de valores. Las ideas clave incluyen:

• Escalabilidad de la Difusión: Los modelos de difusión actúan como una interfaz intrínseca que conecta los modelos mundiales y la optimización de políticas, permitiendo la iteración de valores continua y la mejora de políticas a través de trayectorias imaginadas.
• La Representación Importa: La calidad de la representación latente aprendida es crítica. La optimización conjunta de MBDPO ayuda a estructurar el espacio latente para alinearse con la intuición física, resultando en trayectorias de ciclo cerrado estructuradas para tareas cíclicas y caminos dirigidos a objetivos para la manipulación, que son más significativos físicamente que los producidos por métodos de búsqueda heurística.
• Fundamento Principiado: El trabajo sugiere que escalar el aprendizaje de políticas no requiere pipelines de búsqueda heurística más complejos, sino un fundamento principiado que vincule intrínsecamente los modelos mundiales con los paradigmas fundamentales del RL (específicamente a través de objetivos con restricción KL).

Los autores señalan limitaciones, incluido el enfoque actual en entornos simulados y la falta de escalado a tamaños de modelos fundacionales extremadamente grandes debido a restricciones computacionales. Las direcciones futuras incluyen desplegar MBDPO en robots del mundo real e integrarlo con representaciones visuales o multimodales preentrenadas más ricas.