MATE: Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddings

Technische Samenvatting: MATE – Oplossen van Contextuele Markov Beslissingsprocessen met Geheugen van Geaccumuleerde Overgangsembeddings

1. Probleemstelling

Het artikel adresseert de uitdaging om Contextuele Markov Beslissingsprocessen (CMDP's) op te lossen, een klasse van MDP's geparametriseerd door een niet-geobserveerde latente context $c$. In deze settings hangen de overgangskern $P_c$ en de beloningsfunctie af van $c$, wat voor de agent niet direct waarneembaar is. Om optimaal te handelen, moet een agent de posterior-credo over contexten, $p(c|x_{1:t})$, afleiden op basis van zijn interactiegeschiedenis $x_{1:t} = (s_0, a_0, r_0, \dots, s_t)$.

Standaard op toestand afhankelijke beleidsregels falen in CMDP's omdat ze niet kunnen adapteren aan de latente context. Hoewel geheugen-gebaseerd Versterkend Leren (RL) een oplossing biedt door geschiedenis te coderen in een latente geheugentoestand, staan bestaande architecturen voor aanzienlijke afwegingen:

• Recurrente Neuronale Netwerken (RNN's): lijden onder gradiëntinstabiliteiten over lange horizonnen en sluiten temporale parallelisatie tijdens trainingsupdates uit.
• Transformers: blinken uit in langetermijngeheugen maar brengen kwadratische rekentkosten ($O(T^2)$) met zich mee ten opzichte van de sequentielengte tijdens zowel trainingsupdates als trajectrollouts, wat vaak beperkte contextvensters noodzakelijk maakt.

Bovendien zijn standaard sequentiemodellen (zoals causale Transformers) orde-gevoelig, terwijl de optimale posterior-credo in CMDP's theoretisch permutatie-invariant is ten opzichte van de volgorde van overgangen.

2. Methodologie: MATE Architectuur

De auteurs stellen MATE (Geheugen van Geaccumuleerde Overgangsembeddings) voor, een geheugenarchitectuur ontworpen om de permutatie-invariantie van de context-posterior te benutten, terwijl de rekentechnische en optimalisatienadelen van RNN's en Transformers worden vermeden.

Kernmechanisme:
MATE representeert de geheugentoestand $m_t$ als een eenvoudige som-aggregatie van overgangsembeddings:
$$m_t = \sum_{i=1}^{t} E_\psi(x_i)$$
waarbij $E_\psi$ een overgangscoder is (bijvoorbeeld een neuronaal netwerk) die elke overgang $x_i$ afbeeldt op een embeddingsvector.

Belangrijke Ontwerpeigenschappen:

• Permutatie-invariantie: Per ontwerp is de som-operatie commutatief. Dit sluit aan bij het theoretische kenmerk dat de posterior $p(c|x_{1:t})$ invariant is ten opzichte van de ordening van overgangen (afgeleid via de regel van Bayes en de conditionele onafhankelijkheid van overgangen gegeven $c$).
• Rekentechnische Efficiëntie:
  • Rollout: MATE behoudt een constante inferentiekost per stap ($O(1)$) via de update-regel $m_{t+1} = m_t + E_\psi(x_{t+1})$, wat overeenkomt met de efficiëntie van RNN's.
  • Training: In tegenstelling tot RNN's ondersteunt MATE temporale parallelisatie over de sequentie-dimensie tijdens de update-stap, vergelijkbaar met Transformers, maar met lineaire rekentechnische complexiteit ($O(T)$) in plaats van kwadratisch.
• Hypersferische Normalisatie: Om te voorkomen dat de norm van het geaccumuleerde geheugen onbeperkt groeit (wat interne covariaatverschuiving veroorzaakt), introduceren de auteurs een normalisatietechniek. Het geheugen wordt geprojecteerd op een hypersfeer: $\hat{m}_t = (m_t + \Psi) / \|m_t + \Psi\|$, waarbij $\Psi$ een trainbaar offset-vector is die informatieverlies tijdens normalisatie voorkomt.

3. Belangrijkste Bijdragen

1. Theoretische Expressiviteit: Het artikel levert een theoretisch bewijs (Propositie 3.1) dat een geheugenarchitectuur die som-aggregatie van overgangsembeddings gebruikt, bewezen toereikend is om het optimale beleid in CMDP's te representeren. Specifiek wordt aangetoond dat met een voldoende grote embeddendimensie ($m = 2nT + 1$), een eenvoudige één-laags MLP-coder gecombineerd met sommatie de multiset van overgangen injectief kan afbeelden op een geheugentoestand, waardoor de informatie die vereist is voor het optimale beleid behouden blijft.
2. Architecturale Innovatie: MATE vervangt complexe attentiemechanismen en recurrente verbindingen door eenvoudige som-aggregatie. Het is de eerste black-box geheugen-gebaseerde RL-agent die volledig gebruikmaakt van de permutatie-invariantie van de CMDP-posterior zonder bijsturende inferentie-Doelstellingen.
3. Empirische Validatie: Uitgebreide experimenten tonen aan dat MATE prestaties behaalt die vergelijkbaar zijn met of superieur zijn aan standaard RNN (LSTM) en Transformer (GPT-2) baselines over diverse benchmarks, terwijl het superieure rekentechnische efficiëntie biedt.

4. Experimentele Resultaten

De auteurs evalueerden MATE op drie benchmarks: MuJoCo, Meta-World en T-Maze.

• MuJoCo (Adaptatie aan Dynamica/Beloningen): MATE presteerde op gelijke hoogte met diverse geheugen-gebaseerde baselines (LSTM, Transformer-varianten) in taken die variaties in beloningen omvatten (Ant-Dir, Cheetah-Vel) en systeemdynamica (Walker-Param, Hopper-Param).
• Meta-World (Multi-task Adaptatie): MATE presteerde consequent beter dan RNN- en Transformer-baselines, met hogere returns en verbeterde sample-efficiëntie ondanks zijn lichtgewicht architectuur.
• T-Maze (Langetermijngeheugen & Credit Assignment):
  • In de Passieve setting behaalden MATE en GPT-2 optimale returns over alle contextlengtes, terwijl de prestaties van LSTM degradeerden naarmate de sequentielengte toenam.
  • In de Actieve setting (waarbij exploratie vereist is om de context te vinden), behield MATE consistente prestaties over variërende ganglengtes (tot 600 stappen), terwijl GPT-2 lagere returns liet zien en de prestaties van LSTM afnamen bij langere lengtes.
• Rekentechnische Kosten: Zoals samengevat in de complexiteitsanalyse van het artikel, bereikt MATE $O(T)$ rollout- en updatekosten met paralleliseerbare updates, in contrast met de $O(T^2)$ updatekosten van Transformers en de niet-paralleliseerbare updates van RNN's.

5. Betekenis en Aanspraken

Het artikel stelt dat MATE een rekentechnisch efficiënt en theoretisch onderbouwd alternatief biedt voor bestaande geheugenarchitecturen voor CMDP's. De betekenis hiervan ligt in:

• Brug slaan tussen Theorie en Praktijk: Het valideert dat de inductieve bias van permutatie-invariantie, vaak genegeerd ten gunste van orde-gevoelige sequentiemodellen, niet alleen toereikend is voor optimaliteit maar ook rekentechnisch voordelig.
• Efficiëntie: Het lost de afweging op tussen de gradiëntstabiliteit/parallelisatie van Transformers en de constante-stap-efficiëntie van RNN's.
• Eenvoud: Het toont aan dat complexe attentiemechanismen niet strikt noodzakelijk zijn voor effectieve contextinferentie in CMDP's als de architectuur is afgestemd op de inherente symmetrieën van het probleem.

De auteurs concluderen dat MATE effectief relevante contextuele informatie identificeert tijdens rollout en suggereert toekomstig werk in het integreren van deze architectuur met diepe, zelftoezicht-gebaseerde RL-kaders om de schaalbaarheid verder te verbeteren.