Fue en 2018 cuando se introdujo la idea de aprendizaje por refuerzo en el contexto de un modelo de mundo neuronal, y pronto, este principio fundamental se aplicó en modelos de mundo. Algunos de los modelos más prominentes que implementan aprendizaje por refuerzo fueron el marco Dreamer, que introdujo aprendizaje por refuerzo desde el espacio latente de un modelo de espacio de estado recurrente. El DreamerV2 demostró que el uso de latentes discretos podría resultar en errores de compounding reducidos, y el marco DreamerV3 fue capaz de lograr un rendimiento similar al humano en una serie de tareas en diferentes dominios con hiperparámetros fijos.

Además, se pueden establecer paralelismos entre los modelos de generación de imágenes y los modelos de mundo, lo que indica que el progreso realizado en los modelos de visión generativa podría replicarse para beneficiar a los modelos de mundo. Desde que el uso de transformadores en marcos de procesamiento de lenguaje natural ganó popularidad, surgieron los marcos DALL-E y VQGAN. Estos marcos implementaron autoencoders discretos para convertir imágenes en tokens discretos y fueron capaces de construir modelos generativos de texto a imagen muy potentes y eficientes al aprovechar las capacidades de modelado de secuencias de los transformadores autoregresivos. Al mismo tiempo, los modelos de difusión ganaron tracción, y hoy en día, los modelos de difusión se han establecido como un paradigma dominante para la generación de imágenes de alta resolución. Debido a las capacidades ofrecidas por los modelos de difusión y el aprendizaje por refuerzo, se están haciendo intentos para combinar los dos enfoques, con el objetivo de aprovechar la flexibilidad de los modelos de difusión como modelos de trayectoria, modelos de recompensa, planificadores y políticas para la ampliación de datos en el aprendizaje por refuerzo fuera de línea.

Los modelos de mundo ofrecen un método prometedor para entrenar agentes de aprendizaje por refuerzo de manera segura y eficiente. Tradicionalmente, estos modelos utilizan secuencias de variables latentes discretas para simular la dinámica del entorno. Sin embargo, esta compresión puede pasar por alto detalles visuales cruciales para el aprendizaje por refuerzo. Al mismo tiempo, los modelos de difusión han ganado popularidad para la generación de imágenes, desafiando a los métodos tradicionales que utilizan latentes discretos. Inspirados por este cambio, en este artículo, hablaremos sobre DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente de aprendizaje por refuerzo entrenado dentro de un modelo de mundo de difusión. Exploraremos las decisiones de diseño necesarias para hacer que la difusión sea adecuada para el modelado de mundo y mostraremos que los detalles visuales mejorados conducen a un mejor rendimiento del agente. DIAMOND establece un nuevo estándar en la prueba competitiva Atari 100k, logrando una puntuación media normalizada humana de 1,46, la más alta para agentes entrenados completamente dentro de un modelo de mundo.

DIAMOND: Difusión como un Modelo de Sueños de Entorno

Los modelos de mundo o modelos generativos de entornos están surgiendo como uno de los componentes más importantes para que los agentes generativos planifiquen y razonen sobre sus entornos. Aunque el uso del aprendizaje por refuerzo ha logrado un gran éxito en los últimos años, los modelos que implementan el aprendizaje por refuerzo son conocidos por ser ineficientes en términos de muestras, lo que limita significativamente sus aplicaciones en el mundo real. Por otro lado, los modelos de mundo han demostrado su capacidad para entrenar agentes de aprendizaje por refuerzo de manera eficiente en diversos entornos con una eficiencia de muestras significativamente mejorada, lo que permite que el modelo aprenda de experiencias del mundo real. Los marcos de modelado de mundo recientes suelen modelar la dinámica del entorno como una secuencia de variables latentes discretas, con el modelo discretizando el espacio latente para evitar errores de compounding en horizontes de tiempo multi-paso. Aunque este enfoque puede producir resultados sustanciales, también se asocia con una pérdida de información, lo que conduce a una pérdida de calidad de reconstrucción y pérdida de generalidad. La pérdida de información puede convertirse en un obstáculo significativo para escenarios del mundo real que requieren que la información esté bien definida, como el entrenamiento de vehículos autónomos. En tales tareas, pequeños cambios o detalles en la entrada visual, como el color del semáforo o el indicador de giro del vehículo delantero, pueden cambiar la política de un agente. Aunque aumentar el número de latentes discretos puede ayudar a evitar la pérdida de información, esto aumenta significativamente los costos de cálculo.

Además, en los últimos años, los modelos de difusión han surgido como el enfoque dominante para la generación de imágenes de alta calidad, ya que los marcos construidos sobre modelos de difusión aprenden a revertir un proceso de ruido y compiten directamente con algunos de los enfoques más establecidos que modelan tokens discretos, ofreciendo así una alternativa prometedora para eliminar la necesidad de discretización en el modelado de mundo. Los modelos de difusión son conocidos por su capacidad para ser condicionados fácilmente y modelar distribuciones complejas y multi-modales sin colapso de modos. Estos atributos son cruciales para el modelado de mundo, ya que la condición permite que un modelo de mundo refleje con precisión las acciones de un agente, lo que conduce a una asignación de crédito más confiable. Además, modelar distribuciones multi-modales ofrece una mayor diversidad de escenarios de entrenamiento para el agente, mejorando su rendimiento general.

Basándonos en estas características, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), un agente de aprendizaje por refuerzo entrenado dentro de un modelo de mundo de difusión. El marco DIAMOND hace elecciones de diseño cuidadosas para asegurarse de que su modelo de mundo de difusión permanezca eficiente y estable en horizontes de tiempo largos. El marco proporciona un análisis cualitativo para demostrar la importancia de estas elecciones de diseño. DIAMOND establece un nuevo estado de la técnica con una puntuación media normalizada humana de 1,46 en el benchmark establecido de Atari 100k, la más alta para agentes entrenados completamente dentro de un modelo de mundo. Operar en el espacio de imágenes permite que el modelo de mundo de difusión de DIAMOND sustituya perfectamente el entorno, ofreciendo mayores conocimientos sobre el comportamiento del modelo de mundo y del agente. Notablemente, el mejor rendimiento en ciertos juegos se atribuye a una mejor modelización de detalles visuales críticos. El marco DIAMOND modela el entorno como un POMDP estándar o Proceso de Decisión de Markov Parcialmente Observable con un conjunto de estados, un conjunto de acciones discretas y un conjunto de observaciones de imagen. Las funciones de transición describen la dinámica del entorno, y la función de recompensa asigna las transiciones a recompensas escalares. La función de observación describe las probabilidades de observación y emite observaciones de imagen que luego son utilizadas por los agentes para ver los entornos, ya que no pueden acceder directamente a los estados. El objetivo principal del enfoque fue obtener una política que mapee observaciones a acciones con el intento de maximizar la recompensa esperada con un factor de descuento. Los modelos de mundo son modelos generativos del entorno, y los modelos de mundo se pueden utilizar para crear entornos simulados para entrenar agentes de aprendizaje por refuerzo en el entorno real y entrenar agentes de aprendizaje por refuerzo en el entorno del modelo de mundo. La Figura 1 muestra la imaginación desenrollada del marco DIAMOND con el tiempo.

DIAMOND: Metodología y Arquitectura

En su núcleo, los modelos de difusión son una clase de modelos generativos que generan una muestra invirtiendo el proceso de ruido, y se inspiran en la termodinámica no equilibrada. El marco DIAMOND considera un proceso de difusión indexado por una variable de tiempo continua con marginales y condiciones de frontera correspondientes con una distribución a priori no estructurada tratable. Además, para obtener un modelo generativo que mapee desde el ruido hasta los datos, el marco DIAMOND debe invertir el proceso, con el proceso de inversión también siendo un proceso de difusión que se ejecuta hacia atrás en el tiempo. Además, en cualquier punto dado en el tiempo, no es trivial estimar la función de puntuación, ya que el marco DIAMOND no tiene acceso a la función de puntuación real, y el modelo supera este obstáculo implementando un objetivo de emparejamiento de puntuación, un enfoque que facilita que un marco entrene un modelo de puntuación sin conocer la función de puntuación subyacente. El modelo de difusión basado en puntuación proporciona un modelo generativo incondicional. Sin embargo, se requiere un modelo generativo condicional de la dinámica del entorno para servir como un modelo de mundo, y para servir a este propósito, el marco DIAMOND considera el caso general del enfoque POMDP, en el que el marco puede utilizar observaciones y acciones pasadas para aproximar el estado markoviano desconocido. Como se demuestra en la Figura 1, el marco DIAMOND utiliza esta historia para condicionar un modelo de difusión para estimar y generar la siguiente observación directamente. Aunque el marco DIAMOND podría recurrir a cualquier solucionador de EDO o EDP en teoría, hay un trade-off entre el NFE o número de evaluaciones de función y la calidad de la muestra que impacta significativamente el costo de inferencia de los modelos de difusión.

Basándonos en estas lecciones, veamos ahora la realización práctica del marco DIAMOND de un modelo de mundo basado en difusión, incluyendo los coeficientes de deriva y difusión correspondientes a una elección particular de enfoque de difusión. En lugar de optar por DDPM, un candidato naturalmente adecuado para la tarea, el marco DIAMOND se basa en la formulación EDM y considera un núcleo de perturbación con una función de valor real de tiempo de difusión llamada programación de ruido. El marco selecciona los precondicionadores para mantener la varianza de entrada y salida para cualquier nivel de voz. La red de entrenamiento mezcla señal y ruido de manera adaptativa dependiendo del nivel de degradación, y cuando el ruido es bajo, el objetivo se convierte en la diferencia entre la señal limpia y la señal perturbada, es decir, el ruido gaussiano agregado. Intuitivamente, esto evita que el objetivo de entrenamiento se vuelva trivial en el régimen de bajo ruido. En la práctica, este objetivo tiene una alta varianza en los extremos de la programación de ruido, por lo que el modelo muestrea el nivel de ruido desde una distribución log-normal elegida empíricamente para concatenar el entrenamiento alrededor de las regiones de ruido medio. El marco DIAMOND utiliza un componente U-Net 2D estándar para el campo vectorial y mantiene un búfer de observaciones y acciones pasadas que el marco utiliza para condicionarse. El marco DIAMOND luego concatena estas observaciones pasadas con la siguiente observación ruidosa y acciones de entrada a través de capas de normalización de grupo adaptativo en los bloques residuales del U-Net.

DIAMOND: Experimentos y Resultados

Para una evaluación exhaustiva, el marco DIAMOND opta por el benchmark Atari 100k. El benchmark Atari 100k consiste en 26 juegos diseñados para probar una amplia gama de capacidades de los agentes. En cada juego, un agente está limitado a 100k acciones en el entorno, lo que equivale aproximadamente a 2 horas de juego humano, para aprender el juego antes de la evaluación. Para comparación, los agentes Atari no restringidos suelen entrenar durante 50 millones de pasos, lo que representa un aumento de 500 veces en la experiencia. Entrenamos a DIAMOND desde cero utilizando 5 semillas aleatorias para cada juego. Cada ejecución de entrenamiento requirió alrededor de 12 GB de VRAM y tomó aproximadamente 2,9 días en una sola Nvidia RTX 4090, lo que equivale a 1,03 años de GPU en total. La siguiente tabla proporciona la puntuación para todos los juegos, la media y la IQM o media intercuartil de puntuaciones normalizadas humanas.

Siguiendo las limitaciones de las estimaciones puntuales, el marco DIAMOND proporciona confianza de bootstrap estratificada en la media y la IQM o media intercuartil de puntuaciones normalizadas humanas, junto con perfiles de rendimiento y métricas adicionales, como se resume en la siguiente figura.

Los resultados muestran que DIAMOND se desempeña excepcionalmente bien en todo el benchmark, superando a los jugadores humanos en 11 juegos y logrando una puntuación media normalizada humana superhumana de 1,46, estableciendo un nuevo récord para agentes entrenados completamente dentro de un modelo de mundo. Además, la IQM de DIAMOND es comparable a STORM y supera a todas las demás líneas base. DIAMOND sobresale en entornos donde la captura de pequeños detalles es crucial, como Asterix, Breakout y RoadRunner. Además, como se discutió anteriormente, el marco DIAMOND tiene la flexibilidad de implementar cualquier modelo de difusión en su canal, aunque opta por el enfoque EDM, habría sido una elección natural optar por el modelo DDPM, ya que ya se está implementando en numerosas aplicaciones generativas de imágenes. Para comparar el enfoque EDM con la implementación DDPM, el marco DIAMOND entrena ambos variantes con la misma arquitectura de red en el mismo conjunto de datos estáticos compartido con más de 100k marcos recogidos con una política experta. El número de pasos de desruido está directamente relacionado con el costo de inferencia del modelo de mundo, y por lo tanto, menos pasos reducirán el costo de entrenar a un agente en trayectorias imaginadas. Para asegurarnos de que nuestro modelo de mundo permanezca computacionalmente comparable con otras líneas base, como IRIS, que requiere 16 NFE por tiempo, nos esforzamos por utilizar no más de decenas de pasos de desruido, preferiblemente menos. Sin embargo, establecer el número de pasos de desruido demasiado bajo puede degradar la calidad visual, lo que conduce a errores de compounding. Para evaluar la estabilidad de diferentes variantes de difusión, mostramos trayectorias imaginadas generadas autoregresivamente hasta t = 1000 pasos en la siguiente figura, utilizando diferentes números de pasos de desruido n ≤ 10.

Observamos que utilizar DDPM (a) en este régimen resulta en errores de compounding severos, lo que hace que el modelo de mundo se desvíe rápidamente de la distribución. En contraste, el modelo de mundo de difusión basado en EDM (b) permanece mucho más estable en horizontes de tiempo largos, incluso con un solo paso de desruido. Las trayectorias imaginadas con modelos de difusión basados en DDPM (izquierda) y EDM (derecha) se muestran. La observación inicial en t = 0 es la misma para ambos, y cada fila corresponde a un número decreciente de pasos de desruido n. Observamos que la generación basada en DDPM sufre de errores de compounding, con números más pequeños de pasos de desruido que conducen a una acumulación de errores más rápida. En contraste, el modelo de mundo de DIAMOND basado en EDM permanece mucho más estable, incluso para n = 1. La predicción óptima de un solo paso es la expectativa sobre posibles reconstrucciones para una entrada ruidosa dada, lo que puede estar fuera de la distribución si la distribución posterior es multi-modal. Mientras que algunos juegos, como Breakout, tienen transiciones deterministas que se pueden modelar con precisión con un solo paso de desruido, otros juegos exhiben observabilidad parcial, lo que resulta en distribuciones de observación multi-modales. En estos casos, un solucionador iterativo es necesario para guiar el procedimiento de muestreo hacia un modo específico, como se ilustra en el juego Boxing en la siguiente figura. En consecuencia, el marco DIAMOND establece n = 3 en todos nuestros experimentos.

La figura anterior compara el muestreo de un solo paso (fila superior) y el muestreo multi-paso (fila inferior) en Boxing. Los movimientos del jugador negro son impredecibles, lo que hace que el desruido de un solo paso interpole entre resultados posibles, lo que da como resultado predicciones borrosas. En contraste, el muestreo multi-paso produce una imagen clara al guiar la generación hacia un modo específico. Es interesante que, dado que la política controla al jugador blanco, sus acciones son conocidas por el modelo de mundo, lo que elimina la ambigüedad. Por lo tanto, tanto el muestreo de un solo paso como el muestreo multi-paso predicen correctamente la posición del jugador blanco.

En la figura anterior, las trayectorias imaginadas por DIAMOND generalmente exhiben una mayor calidad visual y son más fieles al entorno real en comparación con las imaginadas por IRIS. Las trayectorias generadas por IRIS contienen inconsistencias visuales entre los marcos (resaltadas por cajas blancas), como enemigos que se muestran como recompensas y viceversa. Aunque estas inconsistencias pueden afectar solo unos pocos píxeles, pueden impactar significativamente el aprendizaje por refuerzo. Por ejemplo, un agente generalmente apunta a objetivos y evita enemigos, por lo que estas pequeñas discrepancias visuales pueden hacer que sea más difícil aprender una política óptima. La figura muestra marcos consecutivos imaginados con IRIS (izquierda) y DIAMOND (derecha). Las cajas blancas resaltan las inconsistencias entre los marcos, que solo aparecen en las trayectorias generadas con IRIS. En Asterix (fila superior), un enemigo (naranja) se convierte en una recompensa (roja) en el segundo marco, luego se convierte en un enemigo en el tercer marco y nuevamente en una recompensa en el cuarto marco. En Breakout (fila media), los ladrillos y la puntuación son inconsistentes entre los marcos. En Road Runner (fila inferior), las recompensas (pequeños puntos azules en la carretera) se representan de manera inconsistente entre los marcos. Estas inconsistencias no ocurren con DIAMOND. En Breakout, la puntuación se actualiza de manera fiable en +7 cuando se rompe un ladrillo rojo.

Conclusión

En este artículo, hemos hablado sobre DIAMOND, un agente de aprendizaje por refuerzo entrenado dentro de un modelo de mundo de difusión. El marco DIAMOND hace elecciones de diseño cuidadosas para asegurarse de que su modelo de mundo de difusión permanezca eficiente y estable en horizontes de tiempo largos. El marco proporciona un análisis cualitativo para demostrar la importancia de estas elecciones de diseño. DIAMOND establece un nuevo estado de la técnica con una puntuación media normalizada humana de 1,46 en el benchmark establecido de Atari 100k, la más alta para agentes entrenados completamente dentro de un modelo de mundo. Operar en el espacio de imágenes permite que el modelo de mundo de difusión de DIAMOND sustituya perfectamente el entorno, ofreciendo mayores conocimientos sobre el comportamiento del modelo de mundo y del agente. Notablemente, el mejor rendimiento en ciertos juegos se atribuye a una mejor modelización de detalles visuales críticos. El marco DIAMOND modela el entorno como un POMDP estándar o Proceso de Decisión de Markov Parcialmente Observable con un conjunto de estados, un conjunto de acciones discretas y un conjunto de observaciones de imagen. Las funciones de transición describen la dinámica del entorno, y la función de recompensa asigna las transiciones a recompensas escalares.