Sintesi Tecnica: WAM-RL

Definizione del Problema

I recenti modelli World-Action (WA) hanno dimostrato una forte capacità di generalizzazione ed efficienza nei dati nell'apprendimento di policy per la robotica, attraverso la modellazione congiunta di osservazioni e azioni future tramite modelli generativi basati su video. Tuttavia, gli esistenti modelli WA si affidano pesantemente all'apprendimento supervisionato da traiettorie esperte. Questa dipendenza impone due limitazioni fondamentali:

1. Vincoli Distribuzionali: La policy appresa è vincolata dal supporto dei dati di addestramento, faticando ad acquisire abilità di manipolazione fini che cadono al di fuori della distribuzione delle dimostrazioni.
2. Mancanza di Miglioramento Continuo: I modelli non possono migliorare continuamente attraverso l'interazione nel mondo reale, ostacolando la loro capacità di adattarsi a nuovi scenari o correggere errori durante l'esecuzione.

Sebbene il Reinforcement Learning (RL) sia una soluzione naturale, applicarlo ai modelli WA è non banale. A differenza dei classici modelli Vision-Language-Action (VLA), dove la rappresentazione visiva è spesso fissa, i modelli WA consistono in due componenti strettamente accoppiate: un modello del mondo (che genera predizioni future) e un modello dell'azione (che mappa le predizioni in azioni). L'applicazione ingenua di RL o del fine-tuning online può causare shift distribuzionali nello spazio latente del modello del mondo, portando a instabilità e degradazione delle prestazioni perché il modello dell'azione dipende profondamente da questo spazio latente.

Metodologia: WAM-RL

Gli autori propongono WAM-RL, un framework che consente l'ottimizzazione congiunta del modello del mondo e del modello dell'azione attraverso l'interazione online. L'intuizione centrale è che il modello del mondo fornisce la struttura predittiva e la capacità di pianificazione, mentre il modello dell'azione agisce principalmente come un traduttore che converte le predizioni latenti in azioni eseguibili.

Il framework impiega uno schema di ottimizzazione in due parti:

1. Video SFT Online con Regolarizzazione KL (Modello del Mondo)

Per perfezionare il modello del mondo, gli autori utilizzano traiettorie raccolte durante interazioni di successo per il fine-tuning auto-supervisionato (SFT) online.

• Obiettivo: Il modello del mondo viene addestrato per predire osservazioni future utilizzando un obiettivo standard di modellazione video (ad esempio, flow matching).
• Stabilizzazione: Per evitare che la distribuzione delle caratteristiche latenti si sposti drasticamente (il che romperebbe il modello dell'azione), viene introdotto un termine di regolarizzazione KL. Questo vincola le caratteristiche latenti del modello del mondo aggiornato a rimanere vicine a quelle di una copia congelata del modello preaddestrato.
• Implementazione: Gli autori approssimano le distribuzioni latenti come Gaussiane, utilizzando una media mobile esponenziale (EMA) per stimare la covarianza del modello corrente mantenendo fissa la covarianza di riferimento. Ciò consente un adattamento graduale preservando al contempo la geometria attesa dall'attore.

2. RL Basato sulla Ricostruzione (Modello dell'Azione)

Il modello dell'azione è ottimizzato tramite RL per garantire che le traiettorie eseguite realizzino fedelmente le predizioni del modello del mondo.

• Segnale di Reward: Invece di reward di task sparsi, gli autori definiscono un reward denso basato sulla ricostruzione. Questo reward misura la coerenza (similarità) tra le osservazioni future immaginate (predette dal modello del mondo) e le osservazioni effettive ottenute eseguendo l'azione nell'ambiente.
• Metriche di Similarità: Gli autori investigano varie funzioni di similarità, inclusi MSE a livello di pixel, coerenza del flusso ottico e similarità delle caratteristiche semantiche (DINOv2, V-JEPA2).
• Ottimizzazione: L'attore viene aggiornato tramite metodi di policy gradient utilizzando questi reward di ricostruzione, colmando efficacemente il divario tra l'immaginazione del modello del mondo e l'esecuzione nel mondo reale.

Contributi Chiave

1. Primo Framework RL per Modelli WA: WAM-RL è il primo framework che incorpora il reinforcement learning nel paradigma World-Action, abilitando specificamente l'ottimizzazione congiunta sia del modello del mondo che del modello dell'azione.
2. Strategia di Addestramento Stabile: Gli autori propongono una nuova strategia di addestramento che combina lo SFT video online con la regolarizzazione KL per il modello del mondo e il RL basato sulla ricostruzione per l'attore. Questo design facilita la co-evoluzione di entrambi i componenti senza destabilizzare lo spazio latente.
3. Intuizioni Empiriche su Task a Lungo Orizzonte: Il lavoro fornisce intuizioni critiche sui limiti dell'ottimizzazione del solo attore. Gli autori dimostrano che l'ottimizzazione del solo attore produce miglioramenti su task a breve orizzonte, ma non fornisce guadagni significativi su task a lungo orizzonte a causa dell'accumulo di errori di predizione. L'ottimizzazione congiunta si dimostra critica per le prestazioni complesse a lungo orizzonte.

Risultati Sperimentali

Il metodo è stato valutato sui benchmark LIBERO (manipolazione centrata sull'oggetto) e RLBench (skill multi-step).

• Guadagni di Performance: WAM-RL ha costantemente superato sia il baseline preaddestrato (Base) che il RL dell'attore-solo ($\pi$RL).
  • Su LIBERO-Object, il tasso di successo è migliorato dal 68% (Base) all'82% (WAM-RL), superando significativamente l'approccio del solo attore (78%).
  • Su RLBench (Water Plants), WAM-RL ha migliorato il tasso di successo dal 19% al 22%, mentre il metodo del solo attore ($\pi$RL) non è riuscito a migliorare rispetto al modello base (18%).
• Analisi del Comportamento:
  • Comportamenti di Recupero: L'adattamento online del modello del mondo ha permesso la generazione di comportamenti di recupero più realistici (ad esempio, riprovare una presa fallita) invece di predire il fallimento immediato.
  • Ablation sui Reward: Lo studio ha rilevato che, sebbene il flusso ottico fornisse il segnale di reward più discriminante tra successo e fallimento, il MSE a livello di pixel ha prodotto le migliori performance downstream. Gli autori ipotizzano che ciò sia dovuto al fatto che il MSE dei pixel si allinea meglio con l'obiettivo di addestramento del modello del mondo e impone una maggiore regolarizzazione contro le azioni out-of-distribution.
  • Ablation sullo SFT Video: I modelli addestrati senza SFT video online non sono riusciti a recuperare dagli errori iniziali, derivando rapidamente verso comportamenti out-of-distribution. Al contrario, i modelli con SFT video hanno imparato ad anticipare il fallimento e a generare azioni correttive entro un singolo chunk open-loop.

Significato e Rivendicazioni

L'articolo sostiene che l'efficacia dei modelli World-Action è fondamentalmente limitata dalla qualità del modello del mondo. Sebbene l'ottimizzazione del solo attore possa migliorare i task a breve orizzonte, non può affrontare l'accumulo di errori di predizione nei setting a lungo orizzonte.

La significatività di WAM-RL risiede nel dimostrare che migliorare congiuntamente il modello del mondo e il modello dell'azione è essenziale per un decision-making robusto a lungo orizzonte. Permettendo al modello del mondo di apprendere la dinamica del fallimento e le strategie di recupero attraverso l'interazione online, il sistema può anticipare e correggere le deviazioni, portando a un'esecuzione più robusta.

Gli autori riconoscono modestamente i limiti, notando che la regolarizzazione KL limita l'entità dell'adattamento del modello del mondo, potenzialmente limitando l'espansione delle capacità oltre la distribuzione preaddestrata. Inoltre, la dipendenza dai reward basati sulla ricostruzione (spesso basati su rappresentazioni preaddestrate) potrebbe offrire un contrasto limitato tra successo e fallimento, suggerendo la necessità di formulazioni di reward più espressive e consapevoli del task per il lavoro futuro.