Fue en 2018, cuando surgió la idea de aprendizaje reforzado En el contexto de una red neuronal, se introdujo por primera vez un modelo mundial y pronto este principio fundamental se aplicó a los modelos mundiales. Algunos de los modelos destacados 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 estados recurrente. DreamerV2 demostró que el uso de latentes discretas podría dar como resultado una reducción de errores de composición, 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 mundiales, lo que indica que el progreso realizado en los modelos de visión generativa podría replicarse en beneficio de los modelos mundiales. Desde el uso de transformadores en procesamiento natural del lenguaje Los marcos ganaron popularidad, surgieron los marcos DALL-E y VQGAN. Los marcos implementaron codificadores automáticos discretos para convertir imágenes en tokens discretos y pudieron construir modelos generativos de texto a imagen altamente potentes y eficientes aprovechando las capacidades de modelado de secuencia de los transformadores autorregresivos. Al mismo tiempo, los modelos de difusión ganaron fuerza y, en la actualidad, 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 intentando 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 como política para el aumento de datos en Aprendizaje por refuerzo fuera de línea. 

Los modelos mundiales ofrecen un método prometedor para entrenar agentes de aprendizaje por refuerzo de forma 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 discretas. Inspirándonos en este cambio, en este artículo hablaremos de DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente de aprendizaje por refuerzo entrenado dentro de un modelo mundial de difusión. Exploraremos las opciones de diseño necesarias para que la difusión sea adecuada para el modelado mundial y mostraremos que los detalles visuales mejorados conducen a un mejor rendimiento de los agentes. DIAMOND establece un nuevo punto de referencia en la competitiva prueba Atari 100k, logrando una puntuación humana normalizada media de 1.46, la más alta para agentes entrenados íntegramente dentro de un modelo mundial. 

DIAMANTE: La difusión como modelo de sueños ambientales

Los modelos mundiales 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 cuanto a muestras, lo que limita significativamente sus aplicaciones en el mundo real. Por otro lado, los modelos mundiales han demostrado su capacidad para entrenar eficientemente agentes de aprendizaje por refuerzo 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 modelación mundial recientes generalmente modelan la dinámica del entorno como una secuencia de variables latentes discretas, y el modelo discretiza el espacio latente para evitar errores compuestos en horizontes temporales de varios pasos. Aunque el enfoque podría arrojar resultados sustanciales, también se asocia con una pérdida de información, lo que lleva a una pérdida de calidad de la reconstrucción y pérdida de generalidad. La pérdida de información podría convertirse en un obstáculo importante 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 información visual como el color del semáforo o el indicador de giro del vehículo de delante pueden cambiar la política de un agente. Aunque aumentar el número de latentes discretas puede ayudar a evitar la pérdida de información, dispara 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 ofrece una alternativa prometedora para eliminar la necesidad de discretización en el modelado mundial. Los modelos de difusión son conocidos por su capacidad para condicionarse fácilmente y modelar de manera flexible distribuciones multimodales complejas sin colapso modal. Estos atributos son cruciales para el modelado mundial, ya que el condicionamiento permite que un modelo mundial 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 multimodales ofrece una mayor diversidad de escenarios de entrenamiento para el agente, mejorando su desempeño general. 

Sobre la base de estas características, DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agente de aprendizaje por refuerzo entrenado dentro de un modelo mundial de difusión. El marco DIAMOND toma decisiones de diseño cuidadosas para garantizar que su modelo mundial de difusión siga siendo eficiente y estable durante largos horizontes de tiempo. El marco proporciona un análisis cualitativo para demostrar la importancia de estas opciones de diseño. DIAMOND establece un nuevo estado del arte con una puntuación humana normalizada media de 1.46 en el bien establecido punto de referencia Atari 100k, el más alto para agentes entrenados completamente dentro de un modelo mundial. Operar en el espacio de imágenes permite que el modelo mundial de difusión de DIAMOND sustituya sin problemas el entorno, ofreciendo mayores conocimientos sobre el modelo mundial y los comportamientos de los agentes. En particular, la mejora del rendimiento en ciertos juegos se atribuye a un mejor modelado 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 imágenes. Las funciones de transición describen la dinámica del entorno y la función de recompensa mapea las transiciones a recompensas escalares. La función de observación describe las probabilidades de observación y emite observaciones de imágenes, 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 las observaciones con las acciones con el intento de maximizar el rendimiento del descuento esperado con un factor de descuento. Los modelos mundiales son modelos generativos del entorno, y los modelos mundiales 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 mundial. La Figura 1 demuestra el desarrollo de la imaginación del marco DIAMOND a lo largo del tiempo. 

DIAMANTE : Metodología y Arquitectura

En esencia, los modelos de difusión son una clase de modelos generativos que generan una muestra invirtiendo el proceso de generación de ruido y se inspiran en gran medida en la termodinámica del desequilibrio. El marco DIAMOND considera un proceso de difusión indexado por una variable de tiempo continua con los marginales correspondientes y condiciones de contorno con una distribución previa no estructurada manejable. Además, para obtener un modelo generativo, que mapee el ruido a los datos, el marco DIAMOND debe invertir el proceso, siendo el proceso de reversión también un proceso de difusión, que retrocede en el tiempo. Además, en un momento dado, no es trivial estimar la función de puntuación, ya que el marco DIAMOND no accede a la función de puntuación verdadera, y el modelo supera este obstáculo implementando el objetivo de coincidencia de puntuación, un enfoque que facilita un marco para entrenar un modelo de puntuación 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 ambiental para que sirva como modelo mundial, y para lograr este propósito, el marco DIAMOND analiza el caso general del enfoque POMDP, en el que el marco puede hacer uso de observaciones y acciones pasadas para aproximarse al desconocido estado markoviano. Como se demuestra en la Figura 1, el marco DIAMOND hace uso de este historial para condicionar un modelo de difusión, estimar y generar la siguiente observación directamente. Aunque en teoría el marco DIAMOND puede recurrir a cualquier solucionador SDE u ODE, existe una compensación entre NFE o número de evaluaciones de funciones y la calidad de la muestra que afecta significativamente el costo de inferencia de los modelos de difusión. 

Sobre la base de los aprendizajes anteriores, veamos ahora la realización práctica del marco DIAMOND de un modelo mundial basado en difusión que incluye 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 del tiempo de difusión llamada programa de ruido. El marco selecciona los precondicionadores para mantener la variación de entrada y salida para cualquier nivel de voz. El entrenamiento de la red mezcla señal y ruido de forma adaptativa según el nivel de degradación, y cuando el ruido es bajo, el objetivo se convierte en la diferencia entre la señal limpia y la perturbada, es decir, el ruido gaussiano añadido. Intuitivamente, esto evita que el objetivo del entrenamiento se vuelva trivial en el régimen de bajo ruido. En la práctica, este objetivo es una alta variación en los extremos del programa de ruido, por lo que el modelo muestrea el nivel de ruido a partir de 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 2D U-Net estándar para el campo vectorial y mantiene un búfer de observaciones y acciones pasadas que el marco utiliza para condicionarse. Luego, el marco DIAMOND concatena estas observaciones pasadas con la siguiente observación ruidosa e ingresa acciones a través de capas de normalización de grupos adaptativos en los bloques residuales de U-Net. 

DIAMANTE: Experimentos y Resultados

Para una evaluación integral, el marco DIAMOND opta por el punto de referencia Atari 100k. El benchmark Atari 100k consta de 26 juegos diseñados para probar una amplia gama de capacidades de los agentes. En cada juego, un agente está limitado a 100 acciones en el entorno, lo que equivale aproximadamente a 2 horas de juego humano, para aprender el juego antes de la evaluación. A modo de comparación, los agentes de Atari sin restricciones suelen entrenarse para 50 millones de pasos, lo que representa un aumento de experiencia de 500 veces. Entrenamos a DIAMOND desde cero usando 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 de todos los juegos, la media y el IQM o media intercuartil de puntuaciones normalizadas por humanos. 

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

Los resultados muestran que DIAMOND tiene un rendimiento excepcional en el benchmark, superando a jugadores humanos en 11 partidas y alcanzando una media HNS sobrehumana de 1.46, lo que establece un nuevo récord para agentes entrenados completamente dentro de un modelo mundial. Además, el IQM de DIAMOND es comparable al de STORM y supera todas las demás líneas de base. DIAMOND destaca en entornos donde la captura de pequeños detalles es crucial, como Asterix, Breakout y RoadRunner. Además, como se mencionó anteriormente, el framework DIAMOND tiene la flexibilidad de implementar cualquier modelo de difusión en su pipeline. Si bien 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 de DDPM, el framework DIAMOND entrena ambas variantes con la misma arquitectura de red en el mismo conjunto de datos estáticos compartidos con más de 100 fotogramas recopilados con una política experta. El número de pasos de eliminación de ruido está directamente relacionado con el coste de inferencia del modelo del mundo, por lo que un menor número de pasos reducirá el coste de entrenamiento de un agente en trayectorias imaginarias. Para garantizar que nuestro modelo del mundo siga siendo computacionalmente comparable con otras líneas base, como IRIS, que requiere 16 NFE por paso de tiempo, nuestro objetivo es utilizar no más de unas decenas de pasos de eliminación de ruido, preferiblemente menos. Sin embargo, un número demasiado bajo de pasos de eliminación de ruido puede degradar la calidad visual, lo que genera errores de composición. Para evaluar la estabilidad de las diferentes variantes de difusión, mostramos trayectorias imaginarias generadas autorregresivamente hasta t = 1000 pasos de tiempo en la siguiente figura, utilizando diferentes números de pasos de eliminación de ruido n ≤ 10. 

Observamos que el uso de DDPM (a), en este régimen, produce graves errores de composición, lo que hace que el modelo mundial se salga rápidamente de la distribución. Por el contrario, el modelo mundial de difusión basado en EDM (b) permanece mucho más estable en horizontes temporales prolongados, incluso con un solo paso de eliminación de ruido. Trayectorias imaginadas con modelos del mundo de difusión Se muestran los basados ​​en DDPM (izquierda) y EDM (derecha). La observación inicial en t = 0 es la misma para ambos, y cada fila corresponde a un número decreciente de pasos de eliminación de ruido n. Observamos que la generación basada en DDPM sufre errores compuestos, con un número menor de pasos de eliminación de ruido que conducen a una acumulación de errores más rápida. Por el contrario, el modelo mundial basado en EDM de DIAMOND 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, que puede estar fuera de distribución si la distribución posterior es multimodal. . Mientras que algunos juegos, como Breakout, tienen transiciones deterministas que se pueden modelar con precisión con un solo paso de eliminación de ruido, otros juegos exhiben observabilidad parcial, lo que resulta en distribuciones de observación multimodales. En estos casos, es necesario 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 DIAMOND estableció n = 3 en todos nuestros experimentos.

La figura anterior compara el muestreo de un solo paso (fila superior) y de varios pasos (fila inferior) en el boxeo. Los movimientos del jugador negro son impredecibles, lo que hace que la eliminación de ruido de un solo paso se interpola entre posibles resultados, lo que da lugar a predicciones borrosas. Por el contrario, el muestreo de varios 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 mundial, lo que elimina la ambigüedad. Por tanto, tanto el muestreo de un solo paso como el de varios pasos 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 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 afectar significativamente el aprendizaje por refuerzo. Por ejemplo, un agente normalmente apunta a obtener recompensas y evitar 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 fotogramas consecutivos imaginados con IRIS (izquierda) y DIAMANTE (derecha). Los cuadros blancos resaltan inconsistencias entre cuadros, que solo aparecen en trayectorias generadas con IRIS. En Astérix (fila superior), un enemigo (naranja) se convierte en una recompensa (roja) en el segundo cuadro, luego vuelve a ser enemigo en el tercero y nuevamente en recompensa en el cuarto. En Breakout (fila del medio), los ladrillos y la puntuación son inconsistentes entre los cuadros. En Road Runner (fila inferior), las recompensas (pequeños puntos azules en el camino) se muestran de manera inconsistente entre cuadros. Estas inconsistencias no ocurren con DIAMOND. En Breakout, la puntuación se actualiza de forma fiable en +7 cuando se rompe un ladrillo rojo. 

Conclusión

En este artículo, hemos hablado de DIAMOND, un agente de aprendizaje por refuerzo entrenado dentro de un modelo de mundo de difusión. El marco DIAMOND toma decisiones de diseño cuidadosas para garantizar que su modelo mundial de difusión siga siendo eficiente y estable durante largos horizontes de tiempo. El marco proporciona un análisis cualitativo para demostrar la importancia de estas opciones de diseño. DIAMOND establece un nuevo estado del arte con una puntuación humana normalizada media de 1.46 en el bien establecido punto de referencia Atari 100k, el más alto para agentes entrenados completamente dentro de un modelo mundial. Operar en el espacio de imágenes permite que el modelo mundial de difusión de DIAMOND sustituya sin problemas el entorno, ofreciendo mayores conocimientos sobre el modelo mundial y los comportamientos de los agentes. En particular, la mejora del rendimiento en ciertos juegos se atribuye a un mejor modelado 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 imágenes. Las funciones de transición describen la dinámica del entorno y la función de recompensa mapea las transiciones a recompensas escalares.