Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

Probleemstelling

Het artikel richt zich op het optimaliseren van trajecten en beleidsstrategieën (policy) in discrete tijdsystemen met differentieerbare dynamica. Het doel is om een controleer $\theta$ te vinden die de verwachte kosten van een traject minimaliseert. Dit omvat twee fundamentele scenario's:

1. Trajectoptimalisatie (TO): Het vinden van een open-loop controlesequentie voor een specifieke beginstaat.
2. Beleidsoptimalisatie (PO): Het vinden van een state-feedback beleidsfunctie (bijv. een neurale netwerken) die de verwachte kosten minimaliseert over een verdeling van beginstaten.

De uitdaging:

• Lokale minima: Gradient-based methoden (zoals backpropagation door simulaties) zijn efficiënt maar kunnen vastlopen in ongewenste lokale minima, vooral bij niet-convexe problemen.
• Sampling-moeilijkheden: Sampling-based methoden (zoals MPPI) verkennen de ruimte beter maar vereisen vaak veel samples en missen de efficiëntie van gradienten.
• Multimodaliteit: De optimale verdeling van parameters is vaak scherp en multimodaal (vele lokale minima), wat het moeilijk maakt om te sampleen, vooral bij lage "temperatuur" (waar de oplossing zich concentreert op de beste optima).

Het artikel stelt de centrale vraag: Kunnen we een unificerend algoritme ontwerpen dat de voordelen van sampling (globale exploratie) en gradient-based optimalisatie (lokale verfijning) combineert onder differentieerbare dynamica?

Methodologie

De auteur introduceert een raamwerk dat optimalisatie herschrijft als inference (sluiten). In plaats van één parameter $\theta$ te optimaliseren, wordt een verdeling $p(\theta)$ over de parameters geoptimaliseerd.

1. Van Optimalisatie naar Inference (KL-geregulariseerd)

Het probleem wordt omgezet in het minimaliseren van een KL-geregulariseerde verwachte kostenfunctie:
$$\min_p \mathbb{E}_{x_0, \theta} [J(\tau)] + \lambda D_{KL}(p \| p_0)$$
waarbij $p_0$ een prior is en $\lambda$ een temperatuurparameter. De optimale oplossing $p^\star$ heeft een gesloten vorm: een Boltzmann-tilted verdeling:
$$p^\star(\theta) \propto p_0(\theta) \exp\left(-\frac{1}{\lambda} E(\theta)\right)$$
Hierbij is $E(\theta)$ de verwachte kosten (energie). Bij lage $\lambda$ concentreert deze verdeling zich rond de globale minima. Het probleem is echter om te sampleen uit deze scherpe, multimodale verdeling.

2. Tempered Sequential Monte Carlo (TSMC)

Om uit deze verdeling te sampleen, wordt Tempered Sequential Monte Carlo (TSMC) voorgesteld. Dit is een annealing-scheme dat een pad volgt van een makkelijke prior naar de doelprior. Het proces bestaat uit drie stappen per temperingsstap:

1. Tempering Path: Er wordt een reeks tussenliggende distributies gedefinieerd met toenemende temperatuur $\beta_k$ (van 0 naar 1).
   $$p_k(\theta) \propto p_0(\theta) \exp\left(-\frac{\beta_k}{\lambda} E(\theta)\right)$$
2. Reweighting & Resampling: Deeltjes (particles) worden opnieuw gewogen op basis van de incrementele waarschijnlijkheidsverhouding. Als de gewichten te sterk uiteenlopen (gemeten via Effective Sample Size of ESS), wordt er geresampleerd om deeltjes met lage kosten te behouden en degeneratie te voorkomen.
3. MCMC Rejuvenation (Hamiltonian Monte Carlo): Om deeltjesdiversiteit te behouden en lokale exploratie mogelijk te maken, wordt Hamiltonian Monte Carlo (HMC) toegepast.
   • Omdat de dynamica differentieerbaar is, kunnen exacte gradienten $\nabla_\theta E(\theta)$ worden berekend via backpropagation door de rollouts (adjoint methode).
   • HMC gebruikt deze gradienten om ver uit de buurt liggende voorstellen te doen in de parameterruimte, wat deeltjes helpt om uit lokale minima te ontsnappen en de ruimte efficiënter te verkennen.

3. Toepassing op Policy Optimization (PO)

Voor PO is de verwachte kosten $E(\theta)$ een verwachting over beginstaten, wat exact berekenen vaak onmogelijk is. De auteur biedt twee praktische oplossingen:

• Deterministische Benadering: De verdeling van beginstaten wordt benaderd door een vaste set van $B$ beginstaten (empirische maat). Hierdoor wordt de energie deterministisch en gradienten berekenbaar.
• Extended-Space Constructie: Deeltjes worden uitgebreid met een batch van beginstaten als "auxiliary variables". TSMC wordt uitgevoerd in deze uitgebreide ruimte, waarbij zowel de beleidsparameters als de beginstaten worden bijgewerkt via HMC en Metropolis-Hastings stappen.

Kernbijdragen

1. Unificatie: Een unificerend raamwerk dat TO en PO behandelt als het sampleen uit een Boltzmann-tilted verdeling.
2. TSMC voor Controle: De introductie van TSMC, gecombineerd met HMC, specifiek voor optimalisatieproblemen met differentieerbare dynamica. Dit combineert de globale exploratie van sampling met de efficiëntie van gradient-based methoden.
3. Exacte Gradienten: Het gebruik van exacte gradienten (via autodifferentiatie) binnen de HMC-stap, wat essentieel is voor het efficiënt verkennen van complexe, niet-convexe landschappen.
4. Praktische Variaties: Effectieve methoden om TSMC toe te passen op beleidsoptimalisatie met onzekerheid in beginstaten (via deterministische benadering of extended-space).

Resultaten

De methode is getest op diverse benchmarks, zowel voor TO als PO, en vergeleken met state-of-the-art baselines (Parallel-NUTS, Parallel-MALA, Parallel-MPPI, Parallel-IPOPT, PPO, SAC).

• Trajectoptimalisatie (TO):

  • Inverted Pendulum: TSMC bereikt prestaties vergelijkbaar met de beste baselines (nabij de theoretische ondergrens).
  • Acrobot & Pusher Slider: TSMC overtreft alle baselines significant. Bij de Acrobot is TSMC de enige methode die succesvol een swing-up bereikt (kosten ~2389 vs ~2967 voor baselines). Bij de Pusher Slider (met niet-gladde contactdynamica) bereikt alleen TSMC het doel.
  • Conclusie: De combinatie van tempering en HMC is superieur bij het ontsnappen uit lokale minima in complexe landschappen.

• Beleidsoptimalisatie (PO):

  • Sparse Rewards & Acrobot: TSMC leert beleidsstrategieën die beter presteren dan PPO en SAC, zelfs bij moeilijke taken waar RL-methoden vaak falen of instabiel zijn.
  • Double Pendulum on a Cart: TSMC bereikt lagere kosten dan PPO en SAC.
  • Efficiëntie: Hoewel TSMC soms langzamer is in runtime dan simpele baselines (door de complexiteit van tempering en HMC), levert het kwalitatief superieure oplossingen op die andere methoden niet vinden.

Betekenis en Toekomstperspectief

• Wetenschappelijke Impact: Het artikel versterkt het "control-as-inference" paradigma en toont aan dat geavanceerde Bayesian inference technieken (TSMC) zeer effectief kunnen zijn voor optimalisatieproblemen in de robotica en besturingstechniek.
• Praktische Toepassing: Het biedt een robuust alternatief voor traditionele gradient-based methoden die gevoelig zijn voor lokale minima, en voor pure sampling-methoden die te veel samples vereisen.
• Beperkingen: De methode is rekenintensief en vraagt veel GPU-geheugen, vooral bij het differentiëren door lange rollouts en het onderhouden van veel deeltjes.
• Toekomst: De auteur ziet kansen in het integreren met differentieerbare contactoplossers (voor contact-rijke taken) en het versnellen van TSMC via aangepaste GPU-kernen.

Kortom, dit artikel presenteert een krachtige, principieel onderbouwde methode die de kloof tussen sampling-based en gradient-based optimalisatie overbrugt, met name voor complexe, niet-convexe problemen in dynamische systemen.