Fue en 2018 cuando se introdujo por primera vez la idea de aprendizaje por refuerzo en el contexto de un modelo de mundo de red neuronal, y pronto, este principio fundamental se aplicó en los modelos de mundo. Algunos de los modelos más prominentes que implementan el aprendizaje por refuerzo fueron el marco Dreamer, que introdujo el 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 pudo 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 procesamiento de lenguaje natural ganó popularidad, surgieron los marcos DALL-E y VQGAN. Los marcos implementaron autoencoders discretos para convertir imágenes en tokens discretos y pudieron construir modelos generativos de texto a imagen muy potentes y eficientes aprovechando 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 que ofrecen los modelos de difusión y el aprendizaje por refuerzo, se están haciendo intentos para combinar las dos aproximaciones, con el objetivo de aprovechar la flexibilidad de los modelos de difusión como modelos de trayectoria, modelos de recompensa, planificadores y como política 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 a 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 los métodos tradicionales que utilizan latentes discretos. Inspirados por este cambio, en este artículo, hablaremos sobre DIAMONDO (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 elecciones de diseño necesarias para hacer que la difusión sea adecuada para la modelización del mundo y mostraremos que los detalles visuales mejorados conducen a un mejor rendimiento del agente. DIAMONDO 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 enteramente dentro de un modelo de mundo.

DIAMONDO: DIffusion As a Model Of eNvironment Dreams

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 éxito considerable 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 a agentes de aprendizaje por refuerzo de manera eficiente en diversos entornos con una eficiencia de muestra significativamente mejorada, lo que permite que el modelo aprenda de experiencias del mundo real. Los marcos de modelización del 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 lleva 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 de adelante, 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 los marcos de 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, y por lo tanto ofrecen una alternativa prometedora para eliminar la necesidad de discretización en la modelización del mundo. Los modelos de difusión son conocidos por su capacidad para ser condicionados fácilmente y para modelar distribuciones complejas y multi-modales sin colapso de modos. Estos atributos son cruciales para la modelización del mundo, ya que el condicionamiento permite que un modelo de mundo refleje con precisión las acciones de un agente, lo que lleva 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, lo que mejora su rendimiento general.

Partiendo de estas características, DIAMONDO, (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 DIAMONDO 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. DIAMONDO 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 enteramente dentro de un modelo de mundo. Operar en el espacio de imágenes permite que el modelo de mundo de difusión de DIAMONDO 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 DIAMONDO modela el entorno como un POMDP o Proceso de Decisión Markoviano Parcialmente Observable estándar 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 era 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 a agentes de aprendizaje por refuerzo en el entorno real y entrenar a agentes de aprendizaje por refuerzo en el entorno del modelo de mundo. La Figura 1 muestra la imaginación desenrollada del marco DIAMONDO con el tiempo.

DIAMONDO: 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 DIAMONDO considera un proceso de difusión indexado por una variable de tiempo continua con marginales y condiciones de límite correspondientes con una distribución de prioridad no estructurada trazable. Además, para obtener un modelo generativo que mapee desde el ruido hasta los datos, el marco DIAMONDO 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 DIAMONDO no tiene acceso a la función de puntuación real, y el modelo supera este obstáculo implementando el objetivo de emparejamiento de puntuaciones, un enfoque que facilita que un marco de puntuación se entrene sin conocer la función de puntuación subyacente. El modelo de difusión basado en puntuaciones 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 DIAMONDO examina 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 muestra en la Figura 1, el marco DIAMONDO utiliza esta historia para condicionar un modelo de difusión, para estimar y generar la próxima observación directamente. Aunque el marco DIAMONDO podría recurrir a cualquier solucionador de EDO o EDE en teoría, hay un compromiso 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.

Partiendo de estas lecciones, veamos ahora la realización práctica del marco DIAMONDO de un modelo de mundo de 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 el DDPM, un candidato naturalmente adecuado para la tarea, el marco DIAMONDO 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 la 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 es de alta varianza en los extremos de la programación de ruido, así 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 DIAMONDO 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 DIAMONDO luego concatena estas observaciones pasadas a la próxima 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.

DIAMONDO: Experimentos y Resultados

Para una evaluación exhaustiva, el marco DIAMONDO 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 de Atari no restringidos suelen entrenar durante 50 millones de pasos, lo que representa un aumento de 500 veces en la experiencia. Entrenamos a DIAMONDO desde cero utilizando 5 semillas aleatorias para cada juego. Cada ejecución de entrenamiento requirió alrededor de 12GB 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 DIAMONDO proporciona una 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 DIAMONDO 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 enteramente dentro de un modelo de mundo. Además, la IQM de DIAMONDO es comparable a la de STORM y supera a todas las demás líneas base. DIAMONDO 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 DIAMONDO 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 DIAMONDO entrena ambas variantes con la misma arquitectura de red en el mismo conjunto de datos estáticos compartido con más de 100k frames recopilados 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 de paso, 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 lleva a errores de compounding. Para evaluar la estabilidad de las diferentes variantes de difusión, mostramos trayectorias imaginadas generadas autoregresivamente hasta t = 1000 pasos de tiempo 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 mundo 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 llevan a una acumulación de errores más rápida. En contraste, el modelo de mundo de DIAMONDO 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 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, se requiere un solucionador iterativo 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 DIAMONDO establece n = 3 en todos nuestros experimentos.

La figura anterior compara el muestreo de un solo paso (fila superior) y el muestreo de múltiples pasos (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 de múltiples pasos produce una imagen clara al guiar la generación hacia un modo específico. Curiosamente, 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 de múltiples pasos predicen con precisión la posición del jugador blanco.

En la figura anterior, las trayectorias imaginadas por DIAMONDO 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 cuadros (resaltadas por cuadros blancos), como enemigos que se muestran como recompensas y viceversa. Aunque estas inconsistencias pueden afectar solo a unos pocos píxeles, pueden impactar significativamente el aprendizaje por refuerzo. Por ejemplo, un agente generalmente apunta a targeting recompensas y evitar enemigos, así que estas pequeñas discrepancias visuales pueden hacer que sea más difícil aprender una política óptima. La figura muestra cuadros consecutivos imaginados con IRIS (izquierda) y DIAMONDO (derecha). Los cuadros blancos resaltan inconsistencias entre los cuadros, 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 cuadro, luego se convierte en un enemigo en el tercer cuadro y nuevamente en una recompensa en el cuarto cuadro. En Breakout (fila media), los ladrillos y la puntuación son inconsistentes entre los cuadros. En Road Runner (fila inferior), las recompensas (pequeños puntos azules en la carretera) se representan de manera inconsistente entre los cuadros. Estas inconsistencias no ocurren con DIAMONDO. 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 DIAMONDO, un agente de aprendizaje por refuerzo entrenado dentro de un modelo de mundo de difusión. El marco DIAMONDO 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. DIAMONDO 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 enteramente dentro de un modelo de mundo. Operar en el espacio de imágenes permite que el modelo de mundo de difusión de DIAMONDO 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 DIAMONDO modela el entorno como un POMDP o Proceso de Decisión Markoviano Parcialmente Observable estándar 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.