Résumé Technique : Recherche Arborescente de Monte Carlo Épistémique

Énoncé du Problème

La famille d'algorithmes AlphaZero/MuZero (A/MZ) a connu un succès significatif en intégrant la Recherche Arborescente de Monte Carlo (MCTS) avec des modèles appris de valeur et de dynamique de l'environnement. Cependant, une limitation critique subsiste : tandis que les modèles appris introduisent une incertitude épistémique (incertitude découlant d'une couverture limitée des données d'entraînement), la MCTS standard ne prend pas en compte la propagation de cette incertitude durant le processus de recherche. Par conséquent, A/MZ ne peut pas exploiter efficacement la MCTS pour une exploration profonde dans des environnements à récompenses rares. L'exploration profonde exige qu'un agent se dirige vers des transitions nouvelles, indépendamment de leur distance par rapport à l'état actuel, une capacité essentielle pour des tâches telles que la conception d'algorithmes ou la programmation, où les récompenses sont rares et l'espace d'états est vaste. Sans prise en compte de l'incertitude épistémique, la recherche peut converger vers des politiques sous-optimales basées sur des prédictions de modèle inexactes, échouant à explorer les régions nécessaires de l'espace d'états.

Méthodologie : MCTS Épistémique (EMCTS)

Les auteurs proposent la MCTS Épistémique (EMCTS), un cadre théoriquement motivé qui intègre l'incertitude épistémique dans le processus MCTS pour faciliter l'exploration profonde. La méthodologie comprend trois composantes principales :

1. Formulation de la Recherche avec Incertitude

Les auteurs modélisent le modèle appris de l'environnement $\hat{M}$ comme une variable aléatoire. Ils dérivent une borne supérieure de confiance (UCB) pour la fonction de valeur optimale $Q^*$ basée sur la variance des prédictions de valeur au sein du modèle appris.

• Fondement Théorique : Le théorème 1 établit que pour un modèle appris $\hat{M}$, la vraie valeur optimale $Q^*(s, a)$ est bornée par la valeur attendue maximale dans le modèle, plus un terme proportionnel à l'écart-type de cette valeur, mis à l'échelle par un paramètre de confiance $\delta$.
• Politique de Recherche : La politique de sélection standard PUCT (Predictor Upper Confidence Bound) est modifiée pour devenir P/UCT Épistémique (EP/UCT). Le critère de sélection devient :
  $$a = \arg\max_{a} \left( q_{\hat{M}}(s, a) + \beta \sqrt{V[q_{\hat{M}}(s, a)]} + \text{Terme d'Exploration} \right)$$
  Ici, $q_{\hat{M}}$ représente la valeur estimée, et $\sqrt{V[q_{\hat{M}}]}$ représente l'incertitude épistémique. L'hyperparamètre $\beta$ contrôle le compromis entre exploitation et exploration.

2. Propagation de l'Incertitude Épistémique

Une contribution centrale est le mécanisme permettant de propager l'incertitude à travers l'arbre de recherche, et non seulement la valeur.

• Incertitude dans les Sauvegardes : L'incertitude d'une étape de sauvegarde $\nu$ est calculée en sommant les variances de la récompense immédiate et de l'incertitude de la valeur future actualisée.
• Incertitude de la Valeur des Nœuds : Puisque A/MZ utilise le même modèle tout au long de la planification, les retours de sauvegarde sont corrélés. Pour éviter de supposer l'indépendance, les auteurs proposent une borne supérieure pour la variance de la valeur du nœud $V[q_{\hat{M}}(s, a)]$ en utilisant la somme des écarts-types des retours de sauvegarde individuels :
  $$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$$
• Estimateurs : La méthode utilise des estimateurs d'incertitude existants pour les récompenses (par exemple, Distillation de Réseau Aléatoire (RND) ou comptage basé sur hachage) et pour les valeurs (par exemple, Équation de Bellman d'Incertitude (UBE)). Pour les transitions non observées, la variance est fixée à la variance maximale possible pour une variable aléatoire bornée.

3. Gestion des Modèles de Transition Appris

Bien que la dérivation théorique suppose un modèle de transition connu, les auteurs abordent les défis des dynamiques de transition apprises (comme dans MuZero). Ils proposent une approximation « aussi optimiste que possible » où, lors de la rencontre de la première transition incertaine dans une trajectoire, toutes les prédictions subséquentes dans cette trajectoire sont supposées avoir une incertitude maximale. Cela garantit que l'UCB reste une borne supérieure valide à des fins d'exploration.

Contributions Clés

1. MCTS Épistémique (EMCTS) : Un algorithme novateur qui étend la MCTS pour estimer et propager l'incertitude épistémique à partir de modèles appris de valeur et/ou de récompense, permettant au processus de recherche de rechercher activement des régions incertaines.
2. Cadre Théorique : Une dérivation de politiques de recherche basées sur l'UCB (EP/UCT) ancrées théoriquement dans la variance des modèles appris, fournissant un mécanisme formel pour l'exploration profonde.
3. Implémentation : Une implémentation parallélisée en JAX de l'EMCTS couplée à un agent AlphaZero, appliquée à l'environnement de langage assembleur subleq et au benchmark Deep Sea.

Résultats Expérimentaux

Les auteurs évaluent l'EMCTS sur deux domaines à récompenses rares et exigeants :

1. Tâche de Programmation Subleq

• Tâche : Écrire du code dans le langage assembleur subleq pour résoudre des fonctions spécifiques (Négation des Positifs et Fonction Identité). Cela implique de rechercher un espace d'états d'environ $16^{10}$ états.
• Résultats : L'EMCTS couplé à AlphaZero (E-AZ) a nettement surpassé la ligne de base AlphaZero. E-AZ a résolu la tâche plus difficile « Fonction Identité » avec beaucoup moins d'échantillons que la ligne de base. La méthode a démontré que l'utilisation d'un estimateur d'incertitude approprié (par exemple, hachage IO vs hachage d'état complet) améliorait encore l'efficacité des échantillons.

2. Benchmark Deep Sea

• Tâche : Un environnement de grille où l'agent doit trouver une trajectoire optimale unique avec des récompenses rares. La probabilité de trouver la solution par exploration aléatoire décroît exponentiellement avec la taille de la grille.
• Résultats :
  • Exploration Profonde : Les agents de base A/MZ ont échoué à résoudre les variations de Deep Sea (récompenses déterministes et stochastiques) dans des budgets d'entraînement raisonnables. En revanche, les agents EMCTS (E-AZ et E-MZ) ont résolu ces tâches, démontrant une mise à l'échelle sous-exponentielle de la complexité des échantillons avec la taille de l'environnement.
  • Bénéfice de la Recherche : L'EMCTS a nettement surpassé une ablation (A/MZ+UBE) qui utilisait l'incertitude pour la sélection d'actions mais n'utilisait pas la recherche pour estimer cette incertitude. Cela confirme que la recherche elle-même améliore la qualité de l'estimation de l'incertitude, conduisant à une exploration plus efficace.
  • Robustesse : La méthode est restée efficace même lors de l'utilisation des dynamiques de transition apprises de MuZero (abstraction équivalente à la valeur) et en présence de récompenses stochastiques.

Importance et Revendications

L'article revendique que l'EMCTS comble une lacune fondamentale dans l'apprentissage par renforcement basé sur des modèles : l'incapacité de la MCTS standard à utiliser l'incertitude épistémique pour l'exploration. En intégrant la propagation de l'incertitude dans l'arbre de recherche, la méthode permet aux agents A/MZ de :

• Atteindre une efficacité d'échantillonnage nettement supérieure dans les environnements à récompenses rares.
• Résoudre des benchmarks d'exploration difficile (comme Deep Sea) qui sont pratiquement insolubles par A/MZ de base.
• Potentiellement améliorer la fiabilité dans l'apprentissage par renforcement hors ligne et la génération de cibles hors politique en fournissant de meilleures estimations d'incertitude pour les prédictions de valeur.

Les auteurs positionnent l'EMCTS comme une amélioration pratique et théoriquement motivée de la famille A/MZ, rendant ces algorithmes mieux équipés pour des applications réelles impliquant la conception d'algorithmes et des récompenses rares, où l'exploration profonde est critique.