Technische Samenvatting: Epistemische Monte Carlo Tree Search

Probleemstelling

De AlphaZero/MuZero (A/MZ) familie van algoritmen heeft aanzienlijk succes geboekt door Monte Carlo Tree Search (MCTS) te integreren met geleerde modellen van waarde en omgevingsdynamica. Er bestaat echter een kritieke beperking: terwijl geleerde modellen epistemische onzekerheid introduceren (onzekerheid die voortvloeit uit beperkte dekking van trainingsdata), houdt standaard MCTS geen rekening met de propagatie van deze onzekerheid tijdens het zoekproces. Bijgevolg kunnen A/MZ MCTS niet effectief benutten voor diepe exploratie in omgevingen met schaarse beloningen. Diepe exploratie vereist dat een agent zichzelf richt op nieuwe overgangen, ongeacht hun afstand tot de huidige staat; een vermogen dat essentieel is voor taken zoals algoritmeontwerp of programmeren, waar beloningen schaars zijn en de toestandsruimte enorm is. Zonder rekening te houden met epistemische onzekerheid kan de zoektocht convergeren naar suboptimale beleidslijnen gebaseerd op onnauwkeurige modelvoorspellingen, waardoor noodzakelijke gebieden van de toestandsruimte niet worden verkend.

Methodologie: Epistemische MCTS (EMCTS)

De auteurs stellen Epistemische MCTS (EMCTS) voor, een theoretisch onderbouwd kader dat epistemische onzekerheid integreert in het MCTS-proces om diepe exploratie te faciliteren. De methodologie omvat drie hoofdbestanddelen:

1. Het Formuleren van Zoeken met Onzekerheid

De auteurs modelleren het geleerde omgevingsmodel $\hat{M}$ als een stochastische variabele. Zij leiden een Upper Confidence Bound (UCB) af voor de optimale waardenfunctie $Q^*$ op basis van de variantie van de waardevoorspellingen binnen het geleerde model.

• Theoretische Basis: Stelling 1 stelt dat voor een geleerd model $\hat{M}$, de ware optimale waarde $Q^*(s, a)$ wordt begrensd door de maximale verwachte waarde in het model, plus een term evenredig met de standaardafwijking van die waarde, geschaald met een betrouwbaarheidsparameter $\delta$.
• Zoekbeleid: Het standaard PUCT (Predictor Upper Confidence Bound) selectiebeleid wordt gemodificeerd naar Epistemische P/UCT (EP/UCT). Het selectiecriteria wordt:
  $$a = \arg\max_{a} \left( q_{\hat{M}}(s, a) + \beta \sqrt{V[q_{\hat{M}}(s, a)]} + \text{Exploratieterm} \right)$$
  Hierbij vertegenwoordigt $q_{\hat{M}}$ de geschatte waarde, en $\sqrt{V[q_{\hat{M}}]}$ de epistemische onzekerheid. De hyperparameter $\beta$ regelt de afweging tussen exploitatie en exploratie.

2. Propagatie van Epistemische Onzekerheid

Een kernbijdrage is het mechanisme om onzekerheid door de zoekboom te propageren, in plaats van alleen de waarde.

• Onzekerheid in Backups: De onzekerheid van een backup-stap $\nu$ wordt berekend door de varianties van de directe beloning en de onzekerheid van de toekomstige afgepaste waarde op te tellen.
• Knopwaarde-onzekerheid: Omdat A/MZ hetzelfde model gebruikt tijdens de planning, zijn backup-retourwaarden gecorreleerd. Om het aannemen van onafhankelijkheid te vermijden, stellen de auteurs een bovengrens voor de variantie van de knopwaarde $V[q_{\hat{M}}(s, a)]$ voor, gebruikmakend van de som van de standaardafwijkingen van individuele backup-retourwaarden:
  $$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$$
• Schatters: De methode maakt gebruik van bestaande onzekerheidsschatters voor beloningen (bijv. Random Network Distillation (RND) of hash-gebaseerd tellen) en waarden (bijv. Uncertainty Bellman Equation (UBE)). Voor niet-geobserveerde overgangen wordt de variantie ingesteld op de maximaal mogelijke variantie voor een begrenste stochastische variabele.

3. Omgaan met Geleerde Overgangsmodellen

Hoewel de theoretische afleiding uitgaat van een bekend overgangsmodel, adresseren de auteurs de uitdagingen van geleerde overgangsdynamica (zoals in MuZero). Zij stellen een "maximaal optimistische" benadering voor waarbij, bij het tegenkomen van de eerste onzekere overgang in een traject, alle daaropvolgende voorspellingen in dat traject worden verondersteld maximale onzekerheid te hebben. Dit zorgt ervoor dat de UCB een geldige bovengrens blijft voor exploratiedoeleinden.

Belangrijkste Bijdragen

1. Epistemische MCTS (EMCTS): Een nieuw algoritme dat MCTS uitbreidt om epistemische onzekerheid te schatten en te propageren vanuit geleerde waarde- en/of beloningsmodellen, waardoor het zoekproces actief onzekere gebieden kan opzoeken.
2. Theoretisch Kader: Een afleiding van UCB-gebaseerde zoekbeleid (EP/UCT) die theoretisch is onderbouwd in de variantie van geleerde modellen, en die een formeel mechanisme biedt voor diepe exploratie.
3. Implementatie: Een geparalleliseerde JAX-implementatie van EMCTS, gekoppeld aan een AlphaZero-agent, toegepast op de Assembly-taal subleq-omgeving en de Deep Sea-benchmark.

Experimentele Resultaten

De auteurs evalueren EMCTS op twee uitdagende domeinen met schaarse beloningen:

1. Subleq Programmeertaak

• Taak: Het schrijven van code in de subleq assembly-taal om specifieke functies op te lossen (Negatie van Positieven en Identiteitsfunctie). Dit houdt in dat er wordt gezocht in een toestandsruimte van ongeveer $16^{10}$ toestanden.
• Resultaten: EMCTS gekoppeld aan AlphaZero (E-AZ) presteerde aanzienlijk beter dan de baseline AlphaZero. E-AZ loste de moeilijkere "Identiteitsfunctie"-taak op met veel minder samples dan de baseline. De methode toonde aan dat het gebruik van een passende onzekerheidsschatter (bijv. IO-hash versus volledige staat-hash) de sample-efficiëntie verder verbeterde.

2. Deep Sea Benchmark

• Taak: Een grid-wereldomgeving waarin de agent een unieke optimale traject moet vinden met schaarse beloningen. De kans om de oplossing te vinden via willekeurige exploratie neemt exponentieel af met de grid-grootte.
• Resultaten:
  • Diepe Exploratie: Baseline A/MZ-agenten slaagden er niet in om Deep Sea-varianten (zowel deterministische als stochastische beloningen) op te lossen binnen redelijke trainingsbudgetten. Daarentegen losten EMCTS-agenten (E-AZ en E-MZ) deze taken op, wat een sub-exponentiële schaling van de sample-complexiteit met de omgevingsgrootte aantoont.
  • Zoekvoordeel: EMCTS presteerde aanzienlijk beter dan een ablatie (A/MZ+UBE) die onzekerheid gebruikte voor actiekeuze, maar geen zoektocht gebruikte om die onzekerheid te schatten. Dit bevestigt dat de zoektocht zelf de kwaliteit van de onzekerheidsschatting verbetert, wat leidt tot efficiëntere exploratie.
  • Robuustheid: De methode bleef effectief, zelfs bij gebruik van MuZero's geleerde overgangsdynamica (waarde-equivalente abstractie) en in aanwezigheid van stochastische beloningen.

Betekenis en Claims

Het artikel claimt dat EMCTS een fundamentele kloof in modelgebaseerd versterkend leren aanpakt: het onvermogen van standaard MCTS om epistemische onzekerheid te benutten voor exploratie. Door propagatie van onzekerheid te integreren in de zoekboom, stelt de methode A/MZ-agenten in staat om:

• Aanzienlijk hogere sample-efficiëntie te bereiken in omgevingen met schaarse beloningen.
• Moeilijke exploratie-benchmarks (zoals Deep Sea) op te lossen die voor baseline A/MZ praktisch onoplosbaar zijn.
• Potentieel de betrouwbaarheid te verbeteren in offline RL en off-policy doelgeneratie door betere onzekerheidsschattingen voor waardevoorspellingen te bieden.

De auteurs positioneren EMCTS als een praktische en theoretisch gemotiveerde verbetering van de A/MZ-familie, waardoor deze algoritmen beter toegerust zijn voor real-world toepassingen die algoritmeontwerp en schaarse beloningen omvatten, waarbij diepe exploratie cruciaal is.