Efficient Hypergradient Descent for Inverse Reinforcement Learning

Technische Samenvatting: Efficiënte Hypergradiëntafdaling voor Inverse Reinforcement Learning

Probleemstelling
Inverse Reinforcement Learning (IRL) streeft naar het terugwinnen van een beloningsfunctie uit expertdemonstraties, zodanig dat de optimale policy onder deze beloning het geobserveerde gedrag reproduceert. De auteurs richten zich op Maximum-Likelihood IRL (ML-IRL), wat geformuleerd kan worden als een bilevel optimalisatieprobleem. Het binnenste niveau optimaliseert een entropie-geregulariseerde policy onder een geleerde beloning $r_\phi$, terwijl het buitenste niveau de discrepantie (specifiek de KL-divergentie) tussen de geïnduceerde policy en de expertdata minimaliseert.

De primaire computationele bottleneck bij het oplossen van dit bilevel-probleem via impliciete differentiatie is de vereiste om de hypergradiënt te berekenen. Dit omvat een inverse-Hessiaan-vectorproduct van de binnenste doelfunctie. Het expliciet construeren van de Hessiaan vereist een kwadratisch geheugen ($O(d_\theta^2)$), terwijl iteratieve benaderingen (bijv. conjunctieve gradiënten) computationeel duur kunnen zijn en gevoelig voor conditionering. Hoewel Natural Hypergradient Descent (NHGD) is voorgesteld om de Hessiaan te vervangen door een Fisher-informatie-matrix-surrogaat, gaat standaard NHGD uit van een vaste datadistributie. In ML-IRL betreft de binnenste doelfunctie een reverse KL-divergentie over een door de policy geïnduceerde trajectdistributie, waardoor de directe toepassing van NHGD niet triviaal is.

Methodologie
Het artikel stelt een methode voor om de hypergradiënt voor ML-IRL efficiënt te berekenen door gebruik te maken van de geometrische structuur van het binnenste probleem en door matrix-sketching toe te passen.

1. Fisher-Hessiaan Equivalentie: De auteurs bewijzen dat onder exacte binnenste optimaliteit en realiseerbaarheid (waarbij de policy-klasse rijk genoeg is om de expertdistributie perfect te modelleren), de Hessiaan van de binnenste doelfunctie proportioneel is aan de gedisconteerde traject Fisher-informatie-matrix ($F_\theta$). Specifiek, $\nabla^2_\theta L_{inner} = \alpha F_\theta$. Dit maakt het mogelijk om het vereiste inverse-Hessiaan-vectorproduct te vervangen door een inverse-Fisher-vectorproduct.
2. Hypergradiënt Schatting: De methode deelt de hypergradiëntberekening op in drie componenten:
   • De buitenste gradiënt met betrekking tot de policy-parameters.
   • De oplossing van een gedempte lineair systeem met betrekking tot de Fisher-matrix en de buitenste gradiënt.
   • Een gemengde afgeleide-term met betrekking tot de beloningsgradiënt en de oplossing van het lineaire systeem.
3. Streaming Spectral Sketching (SCFD): Om de $O(d_\theta^2)$ geheugenkosten van het opslaan van de dichte Fisher-matrix te vermijden, benadert de auteur het vereiste inverse-Fisher-vectorproduct met behulp van een Streaming Spectral Compensation Frequent Directions (SCFD) sketch. In plaats van de Fisher-matrix te materialiseren, streamt het algoritme gewogen policy-score vectoren uit agent-trajecten naar een compacte sketch van grootte $m$. Dit reduceert de opslagcomplexiteit tot $O(md_\theta)$. De sketch wordt vervolgens gebruikt om het gedempte lineaire systeem benaderend op te lossen.

Belangrijkste Bijdragen

• Theoretische Afleiding: Het artikel leidt steekproefgebaseerde schatters af voor de ML-IRL impliciete hypergradiënt en bewijst dat de binnenste Hessiaan proportioneel is aan de traject Fisher-informatie-matrix bij een realiseerbare binnenste optimum. Dit wordt gepresenteerd als de eerste toepassing van deze identiteit voor impliciete hypergradiëntberekening in ML-IRL.
• Schaalbare Solver: De auteurs stellen een SCFD-gebaseerde solver voor die policy-score vectoren streamt. Deze aanpak vermijdt de constructie van de score-matrix of de dichte Fisher-matrix, wat de geheugenvereisten aanzienlijk vermindert.
• Empirische Validatie: De methode wordt geëvalueerd tegen een first-order stochastische bilevel baseline (single-loop ML-IRL) en een expliciete Fisher-solver in discrete (CartPole) en continue (LQR) controleomgevingen.

Resultaten
De experimenten werden uitgevoerd onder gelijke computationele budgetten (24-uurs trainingslimiet) met identieke architecturen en hyperparameter-tuning.

• Policy Prestaties: De Fisher-gebaseerde methoden bereikten competitieve policy-kwaliteit. Op CartPole bereikten alle methoden expert-niveau returns. Op LQR bereikte de single-loop ML-IRL de beste Policy Negative Log-Likelihood (NLL) en Environment Return, hoewel de marge ten opzichte van de "Fisher met Sketching"-methode klein was.
• Belonings Rangschikking: De Fisher-gebaseerde methoden vertoonden een sterke kwaliteit in de beloningsrangschikking. Op LQR behaalde de "Fisher met Sketching"-methode de hoogste RankCorrelation (0.972) tussen geleerde en grondwaarheid traject-returns, waarmee zij zowel de single-loop baseline als de expliciete Fisher-solver overtrof.
• Efficiëntie:
  • Geheugen: Sketching leverde het meest significante voordeel op in de hogere dimensie van de LQR-setting, waarbij het piekgeheugengebruik met wel 1.31$\times$ werd verminderd ten opzichte van de expliciete Fisher-solver.
  • Snelheid: Op CartPole leverde sketching een 1.29$\times$ versnelling op. Op LQR boden matige sketch-groottes (bijv. $m=32$) een lichte versnelling (1.04$\times$), terwijl grotere sketches ($m=256$) langzamer werden dan de expliciete solver vanwege de overhead van het sketching-proces, ondanks lagere uiteindelijke verlieswaarden.
  • Stabiliteit: De auteurs stelden vast dat de geschetste solver stabiel bleef onder aanzienlijk zwakkere dempingsparameters ($\lambda$) dan de expliciete formulering, wat suggereert dat het sketching-proces kan fungeren als een impliciete regularisator.

Betekenis en Claims
Het artikel claimt de schaalbaarheidsbottleneck van bilevel IRL aan te pakken door een directe link te leggen tussen de binnenste Hessiaan en de Fisher-informatie-matrix in de context van reverse-KL optimalisatie. Door dit theoretische inzicht te combineren met streaming matrix-sketching, demonstreren de auteurs dat het mogelijk is om impliciete hypergradiënten te berekenen zonder de kwadratische geheugenkosten van expliciete Hessiaan- of Fisher-matrixconstructie.

De auteurs concluderen bescheiden dat, hoewel de Fisher-gebaseerde methoden competitief zijn met single-loop ML-IRL wat betreft de uiteindelijke policy-kwaliteit, ze een duidelijk voordeel bieden in termen van beloningsrangschikkingsprestaties en computationele efficiëntie (met name geheugenreductie) voor problemen met hogere dimensies. Het werk suggereert dat matige sketch-groottes de optimale balans bieden tussen optimalisatiekwaliteit, geheugengebruik en runtime.