Recovering Hidden Reward in Diffusion-Based Policies

1. Enunciato del Problema

Le politiche basate sulla diffusione (diffusion-based policies) sono diventate un paradigma all'avanguardia per l'apprendimento di abilità di manipolazione robotica a partire da dimostrazioni esperte, grazie alla loro capacità di modellare distribuzioni di azioni complesse e multi-modali. Tuttavia, queste politiche si affidano tipicamente al Behavior Cloning (BC), che apprende a imitare le azioni senza modellare esplicitamente l'intento sottostante o la struttura della ricompensa.

• Limitazioni della Diffusione Standard: Mancano di una rappresentazione energetica esplicita, rendendole soggette a fallimenti quando le situazioni al momento del test si discostano dalla distribuzione di addestramento (Out-of-Distribution o OOD). Non possono fornire facilmente un segnale di ricompensa per il successivo Reinforcement Learning (RL) senza ricorrere a un addestramento avversario instabile o a un Inverse Reinforcement Learning (IRL) computazionalmente costoso.
• Il Divario: È necessario un framework che unifichi la modellazione generativa delle azioni con il recupero della ricompensa, consentendo sia un'imitazione di alta qualità sia l'estrazione di un segnale di ricompensa valido per il raffinamento della politica, garantendo al contempo teoremi di generalizzazione.

2. Metodologia: ENERGYFLOW

Gli autori propongono ENERGYFLOW, un framework che parametrizza una funzione energetica scalare $E_\theta(s, a)$ il cui gradiente funge da campo di denoising per la generazione delle azioni.

Architettura Principale

• Parametrizzazione Energetica Scalare: A differenza delle politiche di diffusione standard che regressano direttamente una previsione di rumore vettoriale $\epsilon_\theta(s, a)$, ENERGYFLOW apprende una funzione energetica scalare $E_\theta(a, s, t)$.
• Campo Vettoriale Conservativo: La funzione di punteggio (campo di denoising) è derivata tramite differenziazione automatica:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  Per costruzione, ciò garantisce che il campo vettoriale sia conservativo (senza rotore, $\nabla \times S = 0$), il che significa che è il gradiente di un potenziale scalare. Questo impone l'integrabilità, assicurando che le preferenze sulle azioni siano transitve e fisicamente realizzabili.
• Addestramento: Il modello è addestrato utilizzando il Denoising Score Matching (DSM). L'obiettivo minimizza la differenza tra il gradiente appreso e la direzione reale del rumore:
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• Inferenza: Le azioni sono generate integrando l'ODE del flusso di probabilità utilizzando il gradiente della funzione energetica appresa.

Estrazione della Ricompensa

Il framework estrae un segnale di ricompensa direttamente dal paesaggio energetico appreso. Per gestire la costante di integrazione dipendente dallo stato $c(s)$ (che rende i valori energetici grezzi ambigui tra stati diversi), gli autori propongono la Centered Shaping:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
Ciò sottrae una linea di base dipendente dallo stato (stimata tramite campionamento Monte Carlo), normalizzando la ricompensa per riflettere solo le preferenze relative tra le azioni all'interno di uno stato.

3. Contributi Teorici Chiave

Il paper fornisce dimostrazioni teoriche rigorose che stabiliscono la connessione tra il score matching della diffusione e il recupero della ricompensa:

1. Equivalenza Score-Ricompensa: Sotto l'ipotesi di Ottimalità a Massima Entropia, il paper dimostra che la funzione di punteggio appresa dai modelli di diffusione è proporzionale al gradiente della funzione Q morbida dell'esperto ($\nabla_a Q^* \propto S^*$). Pertanto, lo score matching recupera implicitamente il gradiente della ricompensa senza ottimizzazione avversaria min-max.
2. Riduzione della Complessità tramite Vincoli Conservativi: Gli autori dimostrano che vincolare lo spazio delle ipotesi a campi conservativi (gradienti di scalari) riduce strettamente la complessità di Rademacher rispetto ai campi vettoriali non vincolati. Nello specifico, il limite di complessità per i campi conservativi scala con la regolarità della rappresentazione ($L$), mentre i campi non vincolati scalano con la dimensione di uscita ($\sqrt{d}$).
3. Limiti di Generalizzazione OOD: Sfruttando la ridotta complessità, il paper deriva limiti di generalizzazione più stretti per scenari Out-of-Distribution (OOD). Il vincolo conservativo agisce come regolarizzatore geometrico, impedendo al modello di apprendere direzioni vettoriali arbitrarie e costringendolo a estrapolare la forma del paesaggio energetico, portando a una migliore robustezza sotto spostamenti di distribuzione.
4. Identificabilità e Propagazione dell'Errore: Il paper stabilisce che, sebbene il valore Q assoluto non possa essere univocamente identificato (a causa della costante sconosciuta dipendente dallo stato), il ranking delle azioni all'interno dello stato è esatto. Inoltre, delimita la propagazione degli errori di stima dello score sulle preferenze delle azioni, mostrando che gli errori scalano linearmente con la distanza tra le azioni (continuità di Lipschitz).

4. Risultati Sperimentali

Gli autori hanno valutato ENERGYFLOW sui benchmark RoboMimic e Meta-World, nonché su compiti robotici reali.

• Prestazioni di Imitazione (RQ1): ENERGYFLOW ha raggiunto prestazioni State-of-the-Art (SOTA).
  • RoboMimic: 93,8% di tasso di successo medio (vs. 91,2% per Diffusion Policy).
  • Meta-World: 92,5% di tasso di successo medio (vs. 90,7% per Diffusion Policy).
  • Ha superato significativamente altri metodi basati sull'energia (come Implicit BC) e baseline IRL.
• Deploy nel Mondo Reale (RQ2): Implementato su un robot AGIBOT G1 per i compiti "Bottle" e "Drawer". ENERGYFLOW ha raggiunto un tasso di successo del 100% con traiettorie fluide e esitazioni ridotte rispetto alle baseline.
• Qualità della Ricompensa per RL a valle (RQ3): Quando utilizzata come segnale di ricompensa per addestrare un agente Soft Actor-Critic (SAC):
  • La ricompensa "Centered Energy" ha permesso a SAC di raggiungere prestazioni vicine all'oracolo, superando significativamente le ricompense sparse e le ricompense energetiche grezze.
  • Ha fornito una guida densa, passo dopo passo, che ha prevenuto i plateau precoci nell'addestramento.
• Generalizzazione OOD (RQ4): Sotto crescenti perturbazioni della posizione iniziale, ENERGYFLOW ha mantenuto tassi di successo robusti laddove le politiche di diffusione e flow non vincolate si degradavano significativamente. Ciò conferma il beneficio teorico del vincolo conservativo.
• Efficienza (RQ6): ENERGYFLOW ha raggiunto velocità di inferenza comparabili alle Flow Policies non basate sull'energia (~10ms di latenza) fornendo al contempo i benefici di una funzione energetica esplicita, evitando i tempi di campionamento lenti dei tradizionali Energy-Based Models (EBM).

5. Significato e Impatto

• Unificazione di Imitazione e RL: ENERGYFLOW colma il divario tra Behavior Cloning e Reinforcement Learning. Permette a una politica di essere addestrata tramite imitazione esponendo simultaneamente un segnale di ricompensa valido per un ulteriore raffinamento, eliminando la necessità di progettazione manuale della ricompensa o di IRL avversario instabile.
• Garanzie Teoriche per la Diffusione: Il lavoro fornisce la prima giustificazione teorica del perché imporre un campo conservativo (integrabilità) nelle politiche di diffusione sia vantaggioso, collegandolo a una ridotta complessità dell'ipotesi e a una migliore generalizzazione OOD.
• Utilità Pratica: Il metodo è computazionalmente efficiente, scalabile a spazi di azioni ad alta dimensionalità ed efficace in ambienti robotici reali, offrendo una soluzione robusta per l'apprendimento di abilità di manipolazione complesse a partire da dimostrazioni limitate.

In sintesi, ENERGYFLOW dimostra che i vincoli strutturali richiesti per un'estrazione valida della ricompensa (campi conservativi) servono simultaneamente come forti bias induttivi che migliorano la generalizzazione e la robustezza delle politiche basate sulla diffusione.