Resumen Técnico: Búsqueda Árbol Monte Carlo Epistémica

Enunciado del Problema

La familia de algoritmos AlphaZero/MuZero (A/MZ) ha logrado un éxito significativo al integrar la Búsqueda Árbol Monte Carlo (MCTS) con modelos aprendidos de valor y dinámicas del entorno. Sin embargo, existe una limitación crítica: mientras que los modelos aprendidos introducen incertidumbre epistémica (incertidumbre que surge de la cobertura limitada de los datos de entrenamiento), el MCTS estándar no tiene en cuenta la propagación de esta incertidumbre durante el proceso de búsqueda. En consecuencia, A/MZ no puede aprovechar eficazmente el MCTS para la exploración profunda en entornos con recompensas dispersas. La exploración profunda requiere que un agente se dirija hacia transiciones novedosas independientemente de su distancia desde el estado actual, una capacidad esencial para tareas como el diseño de algoritmos o la programación, donde las recompensas son escasas y el espacio de estados es vasto. Sin tener en cuenta la incertidumbre epistémica, la búsqueda puede converger a políticas subóptimas basadas en predicciones de modelos inexactas, fallando en explorar las regiones necesarias del espacio de estados.

Metodología: MCTS Epistémico (EMCTS)

Los autores proponen MCTS Epistémico (EMCTS), un marco teóricamente motivado que integra la incertidumbre epistémica en el proceso de MCTS para facilitar la exploración profunda. La metodología implica tres componentes principales:

1. Formulación de la Búsqueda con Incertidumbre

Los autores modelan el modelo aprendido del entorno $\hat{M}$ como una variable aleatoria. Derivan una Cota Superior de Confianza (UCB) para la función de valor óptima $Q^*$ basada en la varianza de las predicciones de valor dentro del modelo aprendido.

• Base Teórica: El Teorema 1 establece que para un modelo aprendido $\hat{M}$, el valor óptimo verdadero $Q^*(s, a)$ está acotado por el valor esperado máximo en el modelo más un término proporcional a la desviación estándar de ese valor, escalado por un parámetro de confianza $\delta$.
• Política de Búsqueda: La política de selección estándar PUCT (Cota Superior de Confianza del Predictor) se modifica a EP/UCT (P/UCT Epistémico). El criterio de selección se convierte en:
  $$a = \arg\max_{a} \left( q_{\hat{M}}(s, a) + \beta \sqrt{V[q_{\hat{M}}(s, a)]} + \text{Término de Exploración} \right)$$
  Aquí, $q_{\hat{M}}$ representa el valor estimado y $\sqrt{V[q_{\hat{M}}]}$ representa la incertidumbre epistémica. El hiperparámetro $\beta$ controla la compensación entre explotación y exploración.

2. Propagación de la Incertidumbre Epistémica

Una contribución central es el mecanismo para propagar la incertidumbre a través del árbol de búsqueda, en lugar de solo el valor.

• Incertidumbre en las Actualizaciones (Backups): La incertidumbre de un paso de actualización $\nu$ se calcula sumando las varianzas de la recompensa inmediata y la incertidumbre del valor futuro descontado.
• Incertidumbre del Valor del Nodo: Dado que A/MZ utiliza el mismo modelo durante toda la planificación, los retornos de actualización están correlacionados. Para evitar asumir independencia, los autores proponen una cota superior para la varianza del valor del nodo $V[q_{\hat{M}}(s, a)]$ utilizando la suma de las desviaciones estándar de los retornos de actualización individuales:
  $$V[q_{\hat{M}}(s, a)] \leq \left( \frac{1}{N(s, a)} \sum_{i=1}^{N(s, a)} \sqrt{V[\nu_i(s, a)]} \right)^2$$
• Estimadores: El método utiliza estimadores de incertidumbre existentes para recompensas (por ejemplo, Distinción de Redes Aleatorias (RND) o conteo basado en hash) y valores (por ejemplo, Ecuación de Bellman de Incertidumbre (UBE)). Para transiciones no observadas, la varianza se establece en la varianza máxima posible para una variable aleatoria acotada.

3. Manejo de Modelos de Transición Aprendidos

Aunque la derivación teórica asume un modelo de transición conocido, los autores abordan los desafíos de las dinámicas de transición aprendidas (como en MuZero). Proponen una aproximación "máximamente optimista" donde, al encontrar la primera transición incierta en una trayectoria, se asume que todas las predicciones subsiguientes en esa trayectoria tienen incertidumbre máxima. Esto asegura que la UCB permanezca como una cota superior válida para fines de exploración.

Contribuciones Clave

1. MCTS Epistémico (EMCTS): Un algoritmo novedoso que extiende el MCTS para estimar y propagar la incertidumbre epistémica desde modelos aprendidos de valor y/o recompensa, permitiendo que el proceso de búsqueda busque activamente regiones inciertas.
2. Marco Teórico: Una derivación de políticas de búsqueda basadas en UCB (EP/UCT) que están fundamentadas teóricamente en la varianza de los modelos aprendidos, proporcionando un mecanismo formal para la exploración profunda.
3. Implementación: Una implementación paralelizada en JAX de EMCTS acoplada a un agente AlphaZero, aplicada al entorno de lenguaje ensamblador subleq y al punto de referencia Deep Sea.

Resultados Experimentales

Los autores evalúan EMCTS en dos dominios desafiantes con recompensas dispersas:

1. Tarea de Programación Subleq

• Tarea: Escribir código en el lenguaje ensamblador subleq para resolver funciones específicas (Negar Positivos y Función Identidad). Esto implica buscar en un espacio de estados de aproximadamente $16^{10}$ estados.
• Resultados: EMCTS acoplado con AlphaZero (E-AZ) superó significativamente a la línea base AlphaZero. E-AZ resolvió la tarea más difícil "Función Identidad" con muchas menos muestras que la línea base. El método demostró que el uso de un estimador de incertidumbre apropiado (por ejemplo, hash IO frente a hash de estado completo) mejoró aún más la eficiencia de las muestras.

2. Punto de Referencia Deep Sea

• Tarea: Un entorno de mundo en cuadrícula donde el agente debe encontrar una trayectoria óptima única con recompensas dispersas. La probabilidad de encontrar la solución mediante exploración aleatoria decae exponencialmente con el tamaño de la cuadrícula.
• Resultados:
  • Exploración Profunda: Los agentes de línea base A/MZ no lograron resolver las variaciones de Deep Sea (tanto recompensas deterministas como estocásticas) dentro de presupuestos de entrenamiento razonables. En contraste, los agentes EMCTS (E-AZ y E-MZ) resolvieron estas tareas, demostrando un escalado subexponencial de la complejidad de muestras con el tamaño del entorno.
  • Beneficio de la Búsqueda: EMCTS superó significativamente a una ablación (A/MZ+UBE) que utilizaba la incertidumbre para la selección de acciones pero no utilizaba la búsqueda para estimar dicha incertidumbre. Esto confirma que la búsqueda en sí misma mejora la calidad de la estimación de incertidumbre, lo que lleva a una exploración más eficiente.
  • Robustez: El método permaneció efectivo incluso al utilizar las dinámicas de transición aprendidas de MuZero (abstracción equivalente al valor) y en presencia de recompensas estocásticas.

Significado y Afirmaciones

El artículo afirma que EMCTS aborda una brecha fundamental en el aprendizaje por refuerzo basado en modelos: la incapacidad del MCTS estándar para utilizar la incertidumbre epistémica para la exploración. Al integrar la propagación de incertidumbre en el árbol de búsqueda, el método permite que los agentes A/MZ:

• Logren una eficiencia de muestras significativamente mayor en entornos con recompensas dispersas.
• Resuelvan puntos de referencia de exploración difícil (como Deep Sea) que son prácticamente irresolubles por A/MZ de línea base.
• Potencialmente mejoren la fiabilidad en RL offline y la generación de objetivos fuera de política al proporcionar mejores estimaciones de incertidumbre para las predicciones de valor.

Los autores posicionan a EMCTS como una mejora práctica y teóricamente motivada para la familia A/MZ, haciendo que estos algoritmos estén mejor equipados para aplicaciones del mundo real que involucran diseño de algoritmos y recompensas dispersas, donde la exploración profunda es crítica.