EXPO: Stable Reinforcement Learning with Expressive Policies

1. Enunciato del Problema

Il paper affronta la sfida del fine-tuning di classi di policy espressive (come le Diffusion Policies e le Flow-Matching policies) utilizzando Reinforcement Learning (RL) Online dato un dataset offline.

• Il Conflitto Centrale: Sebbene le policy espressive eccellano nel catturare distribuzioni di comportamento complesse e multi-modali tramite apprendimento per imitazione (IL), sono difficili da ottimizzare direttamente con RL. A differenza delle semplici policy Gaussiane, le policy espressive sono parametriche tramite lunghe catene di denoising. La retropropagazione dei gradienti del valore attraverso queste catene per aggiornare i parametri della policy è computazionalmente costosa e instabile, portando spesso al collasso dell'addestramento o a una scarsa convergenza.
• Il Divario: I metodi RL online esistenti sono progettati per semplici policy Gaussiane e non riescono a sfruttare la capacità dei modelli espressivi pre-addestrati. Al contrario, i metodi che tentano di eseguire il fine-tuning di policy espressive spesso lottano con la stabilità o richiedono modifiche complesse al processo di denoising.
• Obiettivo: Sviluppare un algoritmo RL online efficiente dal punto di vista del campione e stabile che possa eseguire il fine-tuning di policy espressive pre-addestrate utilizzando dati offline senza ottimizzare direttamente i parametri della policy espressiva rispetto alla funzione di valore.

2. Metodologia: EXPO (Expressive Policy Optimization)

Gli autori propongono EXPO, un framework che disaccoppia la policy espressiva dalla massimizzazione diretta del valore. Invece di retropropagare i gradienti attraverso la policy espressiva, EXPO costruisce una policy "on-the-fly" (OTF) che massimizza i valori Q utilizzando due componenti distinte:

A. Architettura della Policy a Due Componenti

1. Policy Espressiva di Base ($\pi_{base}$):
   • Ruolo: Genera candidati di azione iniziali basati sulla distribuzione del comportamento.
   • Addestramento: Addestrata utilizzando un obiettivo stabile di Apprendimento per Imitazione (IL) (ad es. Denoising Diffusion Probabilistic Models). Non viene mai addestrata esplicitamente per massimizzare il valore Q, preservando la sua stabilità e la capacità di catturare modalità complesse.
2. Policy di Modifica Leggera ($\pi_{edit}$):
   • Ruolo: Una piccola policy Gaussiana che prende un'azione $a$ campionata da $\pi_{base}$ e produce una modifica residua $\hat{a}$ per generare un'azione raffinata $\tilde{a} = a + \hat{a}$.
   • Addestramento: Addestrata utilizzando i gradienti di policy RL standard (massimizzando $Q(s, a + \hat{a})$) con regolarizzazione dell'entropia.
   • Vincolo: La modifica è vincolata (scalata da un iperparametro $\beta$) per rimanere vicina all'azione originale, garantendo che la policy di modifica risolva un problema di ottimizzazione locale piuttosto che globale. Ciò consente un addestramento efficiente e stabile.

B. Estrazione della Policy On-the-Fly (OTF)

La policy effettiva utilizzata per il campionamento e il backup Temporal-Difference (TD) non è una singola rete neurale, ma un processo di selezione dinamico:

1. Campionamento: Per un dato stato $s$, il sistema campiona $N$ azioni da $\pi_{base}$.
2. Modifica: Ogni azione campionata viene raffinata da $\pi_{edit}$.
3. Selezione: Il sistema valuta il valore Q sia per le azioni originali ($a_i$) che per le azioni modificate ($\tilde{a}_i$).
4. Massimizzazione: L'azione con il valore Q più alto tra tutti i candidati viene selezionata per l'esecuzione e utilizzata come target per il backup TD.

C. Backup dell'Entropia per Regimi con Dati Limitati

In scenari con piccoli dataset offline, gli autori introducono un backup dell'entropia. Poiché le policy espressive (come la diffusione) non hanno un'entropia in forma chiusa, costruiscono una distribuzione di campionamento soft sui $N$ candidati di azione base+modificati, ponderati in base ai loro valori Q. Ciò consente l'inclusione di un bonus di entropia nella perdita RL per incoraggiare l'esplorazione senza richiedere calcoli analitici dell'entropia.

3. Contributi Chiave

1. Strategia di Massimizzazione del Valore Stabile: Il paper introduce un approccio innovativo per evitare la propagazione diretta dei gradienti attraverso le catene di denoising espressive. Spostando la massimizzazione del valore su una policy di modifica leggera e un meccanismo di selezione, EXPO raggiunge la stabilità laddove l'ottimizzazione diretta fallisce.
2. Efficienza del Campione: EXPO ottiene un miglioramento di 2-3 volte nell'efficienza del campione rispetto ai metodi precedenti all'avanguardia sia nelle impostazioni RL online che nel fine-tuning da offline a online.
3. Agnosticismo della Policy: Il framework è agnostico rispetto alla specifica parametrizzazione della policy di base (funziona con Diffusion, Flow-Matching, ecc.) e non richiede una funzione di valore pre-addestrata, rendendolo applicabile a qualsiasi policy robotica pre-addestrata.
4. Backup TD On-the-Fly: Gli autori dimostrano che l'uso dell'azione che massimizza il valore per il target TD (non solo per il campionamento) è critico per la velocità di apprendimento, colmando efficacemente il divario tra aggiornamenti simili a SARSA e simili a Q-learning nel contesto delle policy espressive.

4. Risultati Sperimentali

Il metodo è stato valutato su 12 compiti di controllo continuo sfidanti in 4 domini (AntMaze, Adroit, Robomimic, MimicGen), tutti caratterizzati da ricompense sparse.

• RL Online (Nessun Pre-addestramento): EXPO ha superato significativamente le baseline (inclusi RLPD, IDQL e QSM) nell'efficienza del campione, convergendo spesso più velocemente anche senza sfruttare il dataset offline per il pre-addestramento.
• Fine-tuning da Offline a Online:
  • EXPO ha eseguito con successo il fine-tuning di policy pre-addestrate senza il collasso delle prestazioni spesso osservato in altri metodi (ad es. DAC è collassato, IDQL è stagnato).
  • Ha mantenuto prestazioni elevate in compiti di manipolazione fine (ad es. infilare un ago, impilare blocchi) dove altri metodi faticavano a migliorare oltre le prestazioni iniziali dell'apprendimento per imitazione.
• Studi di Ablazione:
  • Backup TD OTF: Rimuovere la selezione che massimizza il valore per il target TD (utilizzando solo il campionamento) ha ridotto drasticamente le prestazioni, confermando l'importanza dell'allineamento del valore Q nel backup.
  • Modifiche delle Azioni: Rimuovere la policy di modifica ha portato a scarse prestazioni, specialmente in compiti che richiedevano esplorazione oltre la stretta distribuzione del comportamento.
  • Qualità dei Dati Offline: Le prestazioni erano fortemente correlate alla qualità del dataset offline. Tuttavia, la variante con backup dell'entropia ha permesso a EXPO di avere successo anche con dati offline molto scadenti (dove l'apprendimento per imitazione da solo ha raggiunto <10% di successo).

5. Significato

EXPO rappresenta un passo significativo avanti nella Robotica e RL risolvendo il collo di bottiglia del "fine-tuning di policy espressive".

• Impatto Pratico: Consente il dispiegamento di grandi modelli generativi pre-addestrati (che eccellono nell'imitare dati umani) in ambienti reali dove devono adattarsi e migliorare tramite interazione online.
• Insight Teorico: Mette in discussione la necessità di retropropagare attraverso catene generative complesse per il RL, mostrando che una strategia di "modifica locale + selezione globale" è più stabile ed efficiente.
• Direzione Futura: Il lavoro apre la porta al fine-tuning di qualsiasi policy generativa pre-addestrata (diffusione, flusso, ecc.) in ambienti online, potenzialmente accelerando l'adozione di modelli fondazionali nella robotica.

Limitazioni: Il metodo richiede il campionamento di più azioni per ogni aggiornamento del batch, il che è computazionalmente costoso. Inoltre, assume l'esistenza di un prior ragionevole (dataset offline o policy pre-addestrata); non è progettato per l'apprendimento da zero senza conoscenze pregresse.