Simulus: Combining Improvements in Sample-Efficient World Model Agents

Resumen Técnico: Simulus – Combinando Mejoras en Agentes de Modelos del Mundo Eficientes en Muestras

Enunciado del Problema
Los modelos del mundo (WM) representan una frontera en el aprendizaje por refuerzo (RL) eficiente en muestras, sin embargo, su potencial a menudo no se realiza debido a la complejidad de integrar mejoras prometedoras. Si bien existen avances individuales, combinarlos requiere una experiencia significativa y esfuerzo de ingeniería. Además, los modelos del mundo basados en tokens existentes (TBWM) están en gran medida restringidos a observaciones de imágenes y acciones discretas, limitando su aplicabilidad al control continuo y a configuraciones multimodales. Además, componentes como la motivación intrínseca (para la reducción de la incertidumbre epistémica) y la reproducción priorizada son frecuentemente pasados por alto en agentes eficientes en muestras, en parte debido al alto costo de integración y al riesgo de que la exploración intrínseca desperdicie presupuestos de interacción escasos en experiencias irrelevantes para la tarea. Este panorama refleja la era pre-Rainbow en el RL profundo, donde mejoras validadas individualmente no se habían combinado sistemáticamente.

Metodología
Los autores presentan Simulus, un agente modular de modelo del mundo basado en tokens diseñado para integrar y validar un conjunto de mejoras complementarias. El agente se construye sobre la arquitectura REM y comprende cuatro módulos distintos: un módulo de representación ($V$), un modelo del mundo ($M$), un controlador ($C$) y un búfer de reproducción.

1. Marco de Tokenización Flexible ($V$): Simulus emplea un sistema de tokenización modular que soporta combinaciones arbitrarias de modalidades de observación y acción. Maneja nativamente imágenes, vectores continuos, variables categóricas y "categóricas 2D" (grillas de variables categóricas similares a imágenes). Los codificadores producen secuencias de tokens de longitud fija, que se concatenan en una representación unificada. Esto permite que el agente opere en entornos diversos, incluyendo control continuo y tareas simbólicas multimodales.
2. Modelo del Mundo ($M$): El modelo de dinámicas utiliza una Red Retentiva (RetNet) aumentada con un mecanismo de Predicción de Observación Paralela (POP) para modelar las dinámicas del entorno como modelado de secuencias sobre tokens discretos.
   • Motivación Intrínseca: Para reducir la incertidumbre epistémica, Simulus emplea un conjunto de cabezas de predicción de la siguiente observación. El desacuerdo entre estas cabezas, medido mediante la Divergencia de Jensen-Shannon (JSD), genera recompensas intrínsecas ($r_{int}$). Estas se combinan con recompensas extrínsecas para guiar al controlador hacia regiones de alta incertidumbre durante la "imaginación" (planificación dentro del modelo).
   • Reproducción Priorizada: El búfer de reproducción mantiene los valores de pérdida del modelo del mundo para los ejemplos almacenados. El entrenamiento utiliza una mezcla de muestreo uniforme y priorizado (controlado por un único parámetro $\alpha$), donde las muestras se seleccionan proporcionalmente a sus pérdidas transformadas por softmax.
3. Controlador ($C$): El controlador es una arquitectura actor-crítico extendida que opera completamente sobre experiencia simulada generada por el modelo del mundo.
   • Regresión como Clasificación (RaC): Para manejar señales de recompensa y retorno ilimitadas de manera robusta, Simulus adopta una red de clasificación con bins exponenciales para la predicción de valor. Utiliza el método de pérdida HL-Gauss (basado en entropía cruzada) en lugar del error cuadrático medio tradicional, mejorando la estabilidad en la estimación de valor.
4. Ciclo de Entrenamiento: El agente sigue un ciclo de recolección de datos, aprendizaje de representación, aprendizaje de modelo del mundo y aprendizaje de control en la imaginación. El controlador se entrena utilizando retornos $\lambda$ y un objetivo REINFORCE con una línea base de valor.

Contribuciones Clave

• Integración de Componentes Pasados por Alto: El artículo identifica e integra la motivación intrínseca, la reproducción priorizada del modelo del mundo y un enfoque híbrido RaC. Demuestra que estos componentes no solo son efectivos individualmente, sino que producen ganancias sinérgicas cuando se combinan.
• Agentes Basados en Tokens Versátiles: Al adaptar la tokenización multimodal al régimen eficiente en muestras, Simulus establece a los TBWM como agentes versátiles capaces de manejar entradas visuales, continuas y multimodales, superando las limitaciones de los modelos de acción discreta solo con imágenes.
• Validación de la Exploración Inducida por el Modelo: El trabajo muestra que la exploración inducida por el modelo del mundo mejora significativamente el rendimiento incluso bajo presupuestos de interacción ajustados, desafiando la noción de que la motivación intrínseca es demasiado costosa para configuraciones eficientes en muestras.
• Fundamento de Código Abierto: Los autores proporcionan un punto de partida sólido para la investigación de TBWM poniendo a disposición el código y los pesos del modelo, reduciendo la barrera de entrada para desarrollar agentes de modelos del mundo competitivos.

Resultados
Simulus fue evaluado en tres benchmarks diversos con presupuestos estrictos de interacción:

• Atari Visual 100K: Simulus logró una eficiencia en muestras de vanguardia (SOTA) entre métodos sin planificación, superando a líneas base como DreamerV3, TWM, STORM y DIAMOND. Notablemente, es el primer modelo del mundo sin planificación en alcanzar puntuaciones medias intercuartílicas (IQM) y medianas a nivel humano, logrando un rendimiento sobrehumano en 13 de 26 juegos.
• Control Continuo (DMC Propriocepción 500K): En tareas de control continuo, Simulus coincidió consistentemente o superó ligeramente a DreamerV3, demostrando la efectividad de las arquitecturas basadas en tokens en dominios continuos con vocabularios compactos.
• Simbólico Multimodal (Craftax-1M): En un entorno de supervivencia 2D parcialmente observable que combina grillas simbólicas y características continuas, Simulus superó a los métodos concurrentes de modelos del mundo y a todas las líneas base sin modelo (incluyendo RND, PPO-RNN y E3B), manejando secuencias de tokens largas de manera eficiente.

Estudios de ablación confirmaron que cada componente contribuye individualmente, siendo las recompensas intrínsecas particularmente cruciales. La combinación de todos los componentes produjo el rendimiento más fuerte, indicando una sinergia significativa.

Significado
El artículo argumenta que el potencial no realizado de los agentes de modelos del mundo se debe a la dificultad de combinar mejoras modulares. Al inspirarse en el marco Rainbow en DQN, Simulus demuestra que una combinación sistemática de flexibilidad de tokenización, motivación intrínseca, reproducción priorizada y técnicas avanzadas de predicción de valor conduce a una eficiencia en muestras de vanguardia. El trabajo establece que los modelos del mundo basados en tokens son un enfoque versátil y escalable para diversos dominios de RL, siempre que los componentes se integren cuidadosamente para evitar la interferencia a menudo vista en arquitecturas monolíticas. Los autores sostienen que, aunque los agentes basados en tokens enfrentan actualmente limitaciones en la velocidad de entrenamiento en comparación con algunas líneas base, su diseño modular ofrece ventajas distintas para la inferencia de políticas y la escalabilidad de la investigación futura.