Sintesi Tecnica: PLUME

Definizione del Problema

La manipolazione destra multi-dito è altamente sensibile ai parametri fisici come la forma degli oggetti, le pose e i coefficienti di attrito. Questi parametri sono spesso incerti al momento della distribuzione, portando a dinamiche di sistema sconosciute. Sebbene la simulazione permetta la raccolta di grandi quantità di dati con parametri noti, le policy addestrate tramite la classica randomizzazione del dominio spesso non riescono a generalizzare su compiti precisi (ad esempio, girare un cacciavite) perché le strategie di manipolazione ottimali possono dover cambiare in base ai valori specifici dei parametri. Inoltre, la raccolta di dati di alta qualità per complessi compiti multi-dito è difficile, portando spesso a dataset eterogenei contenenti traiettorie sub-ottimali o fallimenti. I metodi esistenti di modellazione del mondo (world modeling) si affidano tipicamente a storie a orizzonte fisso per gestire implicitamente le variazioni delle dinamiche, il che è insufficiente per compiti precisi sensibili a tali variazioni. Inoltre, i metodi di apprendimento per rinforzo offline generalmente non si concentrano sul problema congiunto dell'identificazione del sistema e dell'apprendimento della policy, mentre i metodi di identificazione del sistema online richiedono spesso una complessa ingegneria delle ricompense e non supportano la pianificazione.

Metodologia

Gli autori propongono PLUME (Probabilistic Latent Unified world Modeling and parameter Estimation), un framework che apprende congiuntamente un modello del mondo e stima una credenza (belief) su parametri fisici latenti.

Architettura Core

PLUME consiste in tre componenti principali addestrate end-to-end:

1. Apprendimento della Rappresentazione dei Parametri: Un Autoencoder Wasserstein (WAE) con backbone Transformer apprende una rappresentazione latente unificata ($z_t$) di variabili eterogenee, inclusi parametri fisici (es. attrito, forma) e ricompense. L'encoder tokenizza questi input e utilizza il cross-attention pooling per produrre un embedding latente. Viene imposto un prior normale standard per facilitare l'inizializzazione ragionevole delle credenze sui parametri.
2. Modellazione della Storia: Un encoder Mamba-2 comprime la storia della traiettoria (azioni passate e osservazioni) in un vettore di contesto ($c_t$), mantenendo un contesto storico utile che le semplici assunzioni di Markov potrebbero perdere, specialmente con osservazioni rumorose o particelle limitate.
3. Joint Flow Matching: Un modello di flow matching funge da modello del mondo unificato. Esso apprende a stimare i parametri latenti e a modellare le dinamiche di sistema condizionate su tali parametri. Il modello utilizza un backbone transformer e una maschera ausiliaria per gestire diverse distribuzioni condizionali:
   • Stima dei Parametri: $p(Z_{t+1} | Z_t, y_t, u_{t+1}, y_{t+1}, c_t)$
   • Generazione di Rollout: $p(u_{t+1}, y_{t+1}, Z_{t+1} | Z_t, y_t, c_t)$
   • Dinamica Inversa: $p(u_{t+1} | Z_t, y_t, Z_{t+1}, y_{t+1}, c_t)$

Addestramento e Inferenza

• Addestramento: Il modello è addestrato su dataset offline contenenti traiettorie con diversi valori di parametri. La funzione di perdita combina una perdita di ricostruzione per i parametri, una regolarizzazione MMD (Maximum Mean Discrepancy) per allineare lo spazio latente a una normale standard, e la perdita di flow matching. Fondamentalmente, i gradienti dalla perdita di flow matching non vengono retropropagati all'encoder dei parametri per evitare che il modello ottimizzi il proprio target a scapito della capacità rappresentativa.
• Pianificazione (Deployment): Durante la distribuzione, PLUME mantiene una credenza sui parametri latenti rappresentata come un insieme di particelle ($Z_t$). Utilizza un approccio di campionamento della traiettoria a orizzonte scorrevole (receding-horizon):
  1. Rollout: Il modello campiona autoregressivamente molteplici traiettorie condizionate dall'osservazione corrente, dalla storia e dalla credenza delle particelle.
  2. Scoring: Le traiettorie vengono valutate in base a una combinazione di ricompense decodificate e probabilità di transizione stimata. La verosimiglianza è stimata utilizzando un modello di dinamica inversa; le basse verosimiglianze penalizzano le traiettorie ad alta ricompensa che mostrano dinamiche irrealistiche (rilevamento Out-of-Distribution).
  3. Esecuzione e Aggiornamento: Viene eseguita la prima azione della traiettoria con il punteggio più alto. La credenza latente viene poi aggiornata online utilizzando la nuova osservazione, permettendo al modello del mondo di adattarsi alle vere dinamiche senza ri-addestramento o fine-tuning.

Contributi Chiave

1. Modello del Mondo Unified Flow-Matching: Un modello che stima i parametri del sistema per condizionare i rollout del modello del mondo, consentendo a una singola policy di adattarsi a dinamiche variabili.
2. Rappresentazione Latente End-to-End: Un metodo per apprendere uno spazio latente unificato per parametri qualitativamente differenti (attrito, forma) e ricompense, ottimizzato congiuntamente con il modello del mondo per essere utile all'apprendimento delle dinamiche.
3. Pianificazione Pesata sulla Verosimiglianza: Una strategia di pianificazione che utilizza i rollout del modello del mondo e le stime della verosimiglianza delle transizioni per ottimizzare le azioni del robot, filtrando efficacemente i piani ad alta ricompensa ma irrealistici.

Risultati Sperimentali

Gli autori hanno valutato PLUME su quattro compiti simulati (girare un cacciavite, girare una valvola, scoccare un disco, sollevare un secchio) e un compito hardware (girare un cacciavite).

• Baseline: I confronti sono stati effettuati contro Unified World Model (UWM), Flow Q-Learning (FQL) e Decision Diffuser (DD).
• Performance:
  • Cacciavite Simulato: PLUME ha raggiunto una percentuale di validità dell'81,3%, superando significativamente UWM (70,7%), FQL (21,3%) e DD (19,3%). Ha inoltre ridotto la distanza dal target del 31% rispetto alla migliore baseline.
  • Cacciavite Hardware: PLUME ha raggiunto una percentuale di validità del 93,3% (28/30 prove) con trasferimento zero-shot dalla simulazione, superando UWM (86,7%).
  • Altri Compiti: PLUME ha superato significativamente tutte le baseline nei compiti di rotazione valvola e scatto del disco (misurati tramite distanza dal target) ed è stato l'unico metodo a riuscire costantemente nel compito di sollevamento del secchio DexArt.
• Ablazioni: Il paper nota che la capacità di PLUME di ragionare su variabili di stato parzialmente osservabili (posizione, velocità, orientamento) oltre ai parametri fisici contribuisce al suo successo, come dimostrato dalla stima accurata della ricompensa.

Significato e Claim

Il paper sostiene che, per compiti di manipolazione multi-dito precisi e sensibili, stimare esplicitamente e condizionarsi sulle rappresentazioni dei parametri di sistema porta a prestazioni superiori rispetto ai metodi che apprendono policy indipendenti dai parametri o che modellano la variazione implicitamente.

I claim principali includono:

• Trasferimento Zero-Shot: PLIME può distribuire una policy addestrata in simulazione a compiti hardware (girare un cacciavite) senza fine-tuning, adattandosi alle vere dinamiche attraverso l'inferenza dei parametri online.
• Robustezza alla Qualità dei Dati: Il metodo può apprendere modelli utili da dataset eterogenei contenenti traiettorie sub-ottimali e fallimenti, sfruttando le etichette di ricompensa per usare i dati di fallimento per l'apprendimento delle dinamiche pur evitando azioni problematiche tramite la pianificazione.
• Adattabilità: A differenza dei metodi che richiedono ri-addestramento o fine-tuning, PLUME allinea il suo modello del mondo alle vere dinamiche durante la distribuzione aggiornando la sua credenza sui parametri, permettendo a un singolo modello di gestire diverse condizioni fisiche.

Gli autori riconoscono dei limiti, specificamente la necessità di valori dei parametri "ground-truth" durante l'addestramento (disponibili in simulazione ma non sempre nel mondo reale) e la scalabilità quadratica della memoria dei transformer rispetto al numero di particelle. Suggeriscono che il lavoro futuro potrebbe coinvolgere il fine-tuning di modelli pre-addestrati su dati reali per allineare i prior senza richiedere etichette ground-truth.