DRILL: addestrare i modelli del mondo per migliorare le policy, non per prevedere i pixel

Abstract degli autori

I modelli del mondo sono sempre più utilizzati per ridurre le interazioni con l’ambiente nella fase di post-addestramento, tramite apprendimento per rinforzo, delle policy visione-linguaggio-azione (VLA). Di norma, però, vengono addestrati a prevedere i pixel del fotogramma successivo, anche se in questo contesto il loro ruolo non è soltanto generare video visivamente fedeli, ma produrre, a partire da rollout immaginati, segnali di apprendimento che migliorino la policy. Mostriamo che questa discrepanza è sostanziale: un MSE sui pixel più basso può produrre gradienti della policy peggiori, e un modello del mondo addestrato per essere utile alla policy può portare a policy finali migliori nonostante un MSE sui pixel del fotogramma successivo superiore del 17%. Proponiamo il Downstream-Return Imagination Learning Loop (DRILL), un framework bilevel che aggiorna il modello del mondo per massimizzare il ritorno della policy dopo un aggiornamento GRPO interno basato su rollout immaginati. DRILL ha due varianti. DRILL-VWI è una funzione surrogata in forma chiusa che non richiede ulteriori rollout nel simulatore e ripesa gli errori di previsione in base alla frequenza di visita della policy, al vantaggio GRPO relativo al gruppo e al grado di confidenza del modello. DRILL-IMG è un metodo basato sul meta-gradiente completo che differenzia attraverso l’aggiornamento interno usando prodotti hessiana-vettore troncati. Mostriamo che, sotto le ipotesi di compatibilità con lo score e di residui piccoli, DRILL-VWI recupera la componente locale a un passo del meta-gradiente di DRILL-IMG. In cinque simulatori di manipolazione e con due backbone VLA, DRILL-IMG aumenta il tasso medio di successo di 13,1 punti rispetto a WMPO e di 7,9 punti rispetto a RLVR-World, e integra il pre-addestramento RLVR anziché competere con esso. Più in generale, riteniamo che i modelli del mondo basati su video per l’apprendimento per rinforzo delle VLA possano essere addestrati efficacemente in base ai gradienti della policy che producono, non soltanto in base alla previsione dei pixel.