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

1. Il Problema

Il lavoro affronta il problema dell'ottimizzazione di traiettorie e politiche (TPO) in sistemi dinamici discreti con orizzonte finito. L'obiettivo è minimizzare il costo atteso di una traiettoria $J(\tau)$ sotto dinamiche differenziabili $f(x_t, u_t)$.
Il problema è formulato in due contesti principali:

• Ottimizzazione di Traiettoria (TO): Si fissa lo stato iniziale e si ottimizza una sequenza di controlli a ciclo aperto.
• Ottimizzazione di Politica (PO): Si ottimizza una politica di feedback $\pi_\theta$ (es. una rete neurale) per minimizzare il costo atteso su una distribuzione di stati iniziali.

Sfide principali:

• I metodi basati sul gradiente (come la programmazione non lineare o i gradienti di politica) sono efficienti ma tendono a rimanere intrappolati in minimi locali indesiderati a causa della non convessità del problema.
• I metodi basati sul campionamento (come MPPI o il metodo Cross-Entropy) esplorano meglio lo spazio degli stati ma spesso richiedono un numero elevato di campioni e non sfruttano le informazioni sui gradienti disponibili quando le dinamiche sono differenziabili.
• L'obiettivo è unificare i vantaggi dell'esplorazione stocastica e dell'ottimizzazione basata sui gradienti in un quadro coerente.

2. Metodologia: TSMC (Tempered Sequential Monte Carlo)

L'autore riformula il problema di ottimizzazione come un problema di inferenza. Invece di cercare un singolo parametro $\theta$, si cerca una distribuzione di probabilità $p(\theta)$ che minimizzi un obiettivo regolarizzato KL (Kullback-Leibler):
$$\min_p \mathbb{E}[J(\tau)] + \lambda D_{KL}(p \| p_0)$$
La soluzione ottima di questo problema è una distribuzione "Boltzmann-tilted" (o Gibbs posterior):
$$p^*(\theta) \propto p_0(\theta) \exp\left(-\frac{1}{\lambda} E(\theta)\right)$$
dove $E(\theta)$ è il costo atteso e $\lambda$ agisce come temperatura. Quando $\lambda \to 0$, la distribuzione si concentra sui minimi globali.

Campionare da questa distribuzione, specialmente a basse temperature (dove diventa multimodale e acuta), è difficile. Per questo viene proposto TSMC, che combina tre componenti chiave:

1. Temperatura Adattiva (Annealing): Si definisce un percorso di temperamento che deforma gradualmente una distribuzione a priori semplice $p_0$ nella distribuzione target $p^*$. Si utilizza una sequenza di temperature $\beta_k$ che aumenta da 0 a 1.
2. Importance Sampling e Resampling: A ogni livello di temperatura, le particelle (campioni di $\theta$) vengono pesate in base al rapporto di verosimiglianza incrementale. Viene utilizzato l'Effective Sample Size (ESS) per adattare dinamicamente il passo di temperatura, evitando il collasso dei pesi.
3. Rejuvenation con Hamiltonian Monte Carlo (HMC): Per mantenere la diversità delle particelle ed esplorare lo spazio dei parametri senza perdere la distribuzione target, viene applicato un passo di transizione MCMC.
   • Punto di forza: Sfrutta i gradienti esatti ottenuti tramite differenziazione automatica (backpropagation) attraverso le simulazioni delle traiettorie (rollout). Questo permette a HMC di fare "salti" distanti e direzionali nello spazio dei parametri, superando le barriere energetiche che intrappolano i metodi locali.

Estensioni per l'Ottimizzazione di Politica (PO):
Poiché il costo in PO è un'aspettativa su stati iniziali casuali (rendendo il calcolo esatto del gradiente difficile), l'autore propone due varianti:

• Approssimazione Deterministica: Sostituisce la distribuzione degli stati iniziali con un insieme fisso di stati (misura empirica), rendendo il costo e il gradiente deterministici.
• Spazio Esteso (Extended-Space): Tratta la casualità degli stati iniziali come variabili ausiliarie incluse nello stato della particella. Si esegue TSMC sullo spazio congiunto $(\theta, \text{stati iniziali})$, aggiornando alternativamente i parametri della politica e gli stati iniziali tramite HMC e Metropolis-Hastings.

3. Contributi Chiave

• Unificazione: Un quadro teorico unificato che tratta sia l'ottimizzazione di traiettoria che quella di politica come problemi di inferenza su distribuzioni Boltzmann.
• Algoritmo Ibrido: Introduzione di TSMC per il controllo, che integra l'esplorazione globale (tramite temperamento e resampling) con l'ottimizzazione locale efficiente (tramite HMC e gradienti esatti).
• Gestione della Multimodalità: Dimostrazione che TSMC è efficace nel trovare minimi globali in paesaggi energetici complessi e non convessi, dove i metodi puramente gradient-based falliscono.
• Implementazione Pratica: Sviluppo di varianti pratiche per gestire l'intrinseca stocasticità nell'ottimizzazione di politiche, con dimostrazione su benchmark reali.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su task di controllo classici e complessi, utilizzando dinamiche differenziabili (implementate in JAX).

• Ottimizzazione di Traiettoria (TO):

  • Pendolo Invertito: TSMC raggiunge costi vicini al limite inferiore certificato (SDP), paragonabile a metodi avanzati come IPOPT e NUTS parallelo, ma con maggiore robustezza.
  • Acrobot: TSMC supera significativamente tutti i baselines (NUTS, MALA, MPPI, IPOPT), essendo l'unico metodo a trovare una soluzione di "swing-up" riuscita in modo affidabile.
  • Pusher Slider (Dinamiche non lisce): In un ambiente con contatti e dinamiche non lisce, TSMC è l'unico metodo che riesce a raggiungere la posizione obiettivo, mentre gli altri falliscono o trovano soluzioni subottimali.

• Ottimizzazione di Politica (PO):

  • Confrontato con algoritmi RL standard come PPO e SAC.
  • Su task complessi come l'Acrobot e il "Double Pendulum on a Cart", TSMC trova politiche con costi di traiettoria significativamente inferiori rispetto a PPO e SAC, che spesso falliscono nel completare il task di swing-up entro il budget di addestramento.
  • TSMC non richiede l'apprendimento di una funzione valore (value function), semplificando la pipeline rispetto ai metodi Actor-Critic.

5. Significato e Implicazioni

Questo lavoro è significativo perché:

1. Supera i limiti dei minimi locali: Dimostra che combinare il campionamento sequenziale con i gradienti esatti permette di navigare efficacemente paesaggi di ottimizzazione altamente non convessi, tipici del controllo robotico.
2. Sfrutta le dinamiche differenziabili: Sfrutta appieno l'avvento di simulatori fisici differenziabili (come quelli basati su JAX/Brax) per guidare l'esplorazione stocastica, unendo il meglio dei mondi dell'ottimizzazione e dell'apprendimento per rinforzo.
3. Versatilità: Offre un approccio coerente sia per il controllo a ciclo aperto (TO) che per l'apprendimento di politiche (PO), suggerendo che la distinzione tra i due può essere sfumata quando si adotta una prospettiva di inferenza.
4. Futuro: Apre la strada all'uso di tecniche di inferenza bayesiana avanzata (come TSMC) per risolvere problemi di controllo complessi, specialmente in scenari con contatti e dinamiche non lisce, dove i gradienti tradizionali sono spesso inaffidabili o non informativi.

In sintesi, il paper propone un metodo robusto e scalabile che supera le limitazioni dei metodi attuali, offrendo soluzioni di alta qualità per problemi di controllo difficili attraverso una fusione intelligente di campionamento, temperamento e ottimizzazione basata sui gradienti.