Technische Samenvatting: StaQ: Een Finite Memory Benadering van Discrete Action Policy Mirror Descent

Probleemstelling

In Reinforcement Learning (RL) regulariseren Policy Mirror Descent (PMD) algoritmen beleidsupdates met behulp van de Kullback-Leibler (KL) divergentie om grote afwijkingen tussen opeenvolgende beleidsfuncties te bestraffen. Een belangrijke theoretische eigenschap van PMD is het vermogen om fouten in de beleidsevaluatie uit te middelen, wat onvermijdelijk is bij het gebruik van functiebenaderingen zoals neurale netwerken. De gesloten vorm van de exacte PMD omvat een oneindige gewogen som van alle voorgaande Q-functies. Deze oneindige som is echter in de praktijk meestal onhandelbaar, vooral bij niet-lineaire functiebenaderingen die een gesloten vorm voor beleidsuitdrukkingen verhinderen.

Bestaande praktische benaderingen, zoals Natural Policy Gradient of Actor-Critic (AC) methoden (bijv. TRPO, MDPO, ECPO), proberen deze update te benaderen. Deze methoden introduceren echter vaak beleidsupdatefouten omdat ze vertrouwen op het optimaliseren van een surrogaatdoelstelling of het uitvoeren van gradiëntstappen, wat mogelijk niet perfect overeenkomt met de theoretische PMD-update. Dit creëert een kloof tussen de theoretische voordelen van PMD (foutuitmiddeling) en de praktische implementatie, wat kan leiden tot suboptimale convergentie of hogere foutvloeren.

Methodologie: StaQ

De auteurs stellen StaQ (Stacked Q-functions) voor, een PMD-variant ontworpen voor discrete actieruimtes die de onhandelbaarheid van de oneindige som aanpakt en tegelijkertijd beleidsupdatefouten vermijdt.

Kernalgoritme

In plaats van een oneindige geschiedenis bij te houden of het beleid via gradiëntafdaling te benaderen, houdt StaQ een eindig geheugen bij van de laatste $M$ Q-functie schattingen. Het beleid op iteratie $k$ wordt gedefinieerd als een Boltzmann-distributie over een gewogen gemiddelde van deze $M$ Q-functies:

$$\pi_k \propto \exp\left( \frac{\alpha}{1 - \beta^M} \sum_{i=0}^{M-1} \beta^i q_{k-i} \right)$$

waarbij $q_{k-i}$ de Q-functie schattingen zijn en $\alpha$ en $\beta$ temperatuur- en wegingparameters zijn afgeleid van de KL-regularisatieweg $\eta$ en de entropieweg $\tau$.

Implementatiedetails

1. Finite Memory Truncation: Het algoritme implementeert een "First-In-First-Out" (FIFO) mechanisme. Na elke beleidsevaluatie worden de gewichten van de nieuwe Q-functie op een stapel van neurale netwerken geplaatst. Als de stapel de grootte $M$ overschrijdt, wordt de oudste Q-functie verwijderd.
2. Optimization-Free Update: In tegenstelling tot Actor-Critic methoden die gradiëntstappen vereisen om het beleidsnetwerk bij te werken, is de beleidsupdate van StaQ optimization-free. Het houdt simpelweg in dat de bevroren gewichten van de Q-functies worden gestapeld. De beleidslogits $\xi_k$ worden direct berekend uit de gewogen som van de opgeslagen Q-functies.
3. Efficiënte GPU-parallellisatie: De auteurs maken gebruik van gebatchte operaties op GPU's om meerdere Q-waarden parallel te berekenen. Ze berekenen de logits vooraf voor de replay buffer om de computationele overhead van een grote $M$ tijdens de beleidsevaluatiefase te minimaliseren.

Theoretische Analyse

Het artikel biedt een convergentieanalyse voor StaQ onder het kader van Entropy-Regularized Policy Mirror Descent (EPMD).

• Convergentie: De auteurs bewijzen dat StaQ, voor een voldoende grote $M$, convergeert naar het optimale beleid in de afwezigheid van evaluatiefouten.
• Foutuitmiddeling: De analyse toont aan dat de foutvloer in StaQ afhankelijk is van het gewogen gemiddelde van de laatste $M$ evaluatiefouten. De afwijking van de exacte PMD foutvloer bestaat uit extra termen die exponentieel snel afnemen met betrekking tot $M$.
• Geen Beleidsupdatefout: Cruciaal is dat omdat de beleidsupdate direct is afgeleid van de gesloten vorm (afgekapt tot $M$), StaQ geen additionele fouten introduceert door de beleidsupdate stap, in tegenstelling tot AC-benaderingen die lijden onder optimalisatiefouten.

Belangrijkste Bijdragen

1. Uitbreiding van PMD-analyse: De auteurs breiden eerdere PMD-analyses uit naar Policy Iteration (PI) en demonstreren dat het uitmiddelingseffect van beleidsevaluatiefouten kan worden bereikt met slechts een DKL-penalisatie op de beleidsupdate, zonder dat een DKL-regularisatie in de evaluatiestap vereist is.
2. StaQ Algoritme: Ze stellen StaQ voor en analyseren dit als een finite-memory PMD-variant. Ze tonen theoretisch aan dat StaQ, voor een voldoende grote $M$, convergeert zonder beleidsupdatefouten en de eigenschappen van foutuitmiddeling van exacte PMD benadert met exponentieel afnemende fouttermen.
3. Empirische Validatie: Ze bieden een efficiënte GPU-implementatie die aantoont dat het vergroten van $M$ de prestaties verbetert tot een drempelwaarde, waarna StaQ ononderscheidbaar wordt van Exact PMD (dat zij evalueren tot 5 miljoen tijdstappen).
4. Runtime Efficiëntie: De studie toont aan dat voor middelgrote omgevingen (bijv. MinAtar) het vergroten van $M$ een marginale impact heeft op de runtime, waardoor StaQ sneller is dan Actor-Critic baselines vanwege het ontbreken van beleidsoptimalisatiestappen.

Experimentele Resultaten

De auteurs evalueerden StaQ op MinAtar (5 omgevingen) en Classic Control taken (4 omgevingen), waarbij het werd vergeleken met:

• Baselines: Natural Policy Gradient (TRPO), SAC (discrete versie), ECPO (M-Projection), MDPO (Direct Optimization), DQN, en M-DQN.
• Exact PMD: Geëvalueerd met $M=1000$ (effectief een oneindige som binnen de 5M tijdstap window).

Bevindingen:

• Prestaties versus Geheugengrootte ($M$): De prestaties verbeteren bijna monotoon met $M$. Op bijna elke taak levert $M \ge 300$ prestaties die vrijwel ononderscheidbaar zijn van Exact PMD.
• Vergelijking met AC Methoden: StaQ presteert over het algemeen beter dan de Actor-Critic baselines (TRPO, ECPO, MDPO) over de taken heen. De AC-methoden, die beleidsupdatefouten introduceren, blijven achter bij StaQ met een gematigde $M$ (bijv. $M=50$ tot $M=300$).
• Runtime: StaQ is sneller dan de AC baselines. Zo duurde StaQ ($M=100$) op MinAtar/SpaceInvaders ongeveer 10,4 uur vergeleken met 13,9 uur voor SAC, ondanks de extra geheugenoverhead, omdat het de gradiëntstappen voor het beleid vermijdt.
• Exploratie: De auteurs merken op dat hoewel entropieregularisatie helpt, StaQ (net als andere entropieregulariserende methoden) kan moeite hebben met exploratie in omgevingen met een ijle beloning (sparse reward) zoals MountainCar zonder aanvullende mechanismen (bijv. "sticky" acties of $\epsilon$-greedy exploratie).

Betekenis en Claims

Het artikel claimt dat StaQ een praktische en theoretisch onderbouwde implementatie van PMD biedt voor discrete actieruimtes. De betekenis ligt in:

• Brug tussen Theorie en Praktijk: Het demonstreert dat de gewenste eigenschappen van exacte PMD (specifiek het uitmiddelen van evaluatiefouten) kunnen worden teruggewonnen zonder de onhandelbare oneindige som of de benaderingsfouten van Actor-Critic methoden.
• Elimineren van Beleidsupdatefout: Door een eindig geheugen van gestapelde Q-functies te gebruiken, vermijdt StaQ de "foutvloer" die wordt geïntroduceerd door inexacte beleidsupdates in standaard AC-algoritmen.
• Schaalbaarheid: De efficiënte GPU-implementatie suggereert dat het bijhouden van een eindige geschiedenis van Q-functies computationeel haalbaar en voordelig is, wat het idee uitdaagt dat exacte PMD puur theoretisch is.

De auteurs concluderen dat StaQ een veelbelovende setting biedt voor het testen van andere RL-componenten (zoals beleidsevaluatietechnieken), omdat het de beleidsoptimalisatiemechanisme isoleert, waardoor wordt gegarandeerd dat geobserveerde prestatiewinsten of -tekortkomingen niet worden verward met beleidsoptimalisatiefouten. Ze merken bescheiden op dat terwijl het geheugengebruik aanzienlijk kan worden voor grootschalige taken (zoals volledige Atari), toekomstig werk de mogelijkheid kan verkennen om oude Q-functies te comprimeren om de efficiëntie te behouden.