DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs

Sintesi Tecnica: DynaWM: Distillazione Consapevole della Dinamica con Modello del Mondo e Target di Momento per una Locomozione Fluida su Scale Continue

Problema
Sebbene i robot bipedi con ruote abbiano dimostrato successo nel superare pendenze e singoli gradini utilizzando il controllo predittivo del modello e l'apprendimento per rinforzo profondo, la navigazione di scale lunghe e continue rimane una sfida significativa. Le soluzioni attuali che si affidano a framework di distillazione teacher-student affrontano tre limitazioni primarie:

1. Rappresentazioni debolmente consapevoli della dinamica: Gli encoder del teacher, ottimizzati esclusivamente tramite gradienti di policy, tendono a codificare solo le caratteristiche del terreno direttamente rilevanti per le ricompense immediate, trascurando le informazioni geometriche più ampie essenziali per la modellazione della dinamica del suolo.
2. Collasso Dimensionale: Nella fase di addestramento concomitante teacher-student, gli aggiornamenti rapidi delle rappresentazioni del teacher possono causare il collasso dell'encoder dello student in rappresentazioni triviali (adattandosi solo alle predizioni medie) anziché catturare le caratteristiche essenziali del terreno.
3. Opacità: Lo spazio latente nei framework esistenti agisce come una "scatola nera", rendendo difficile verificare la reale capacità di modellazione dell'altezza del terreno o garantirne l'interpretabilità durante l'addestramento.

Questi problemi causano una mancanza di fluidità del movimento e di adattabilità al terreno quando i robot affrontano scale continue, ostacolando un affidabile deployment sim-to-real.

Metodologia: DynaWM
Gli autori propongono DynaWM (Distillazione consapevole della dinamica con Modello del Mondo e Target di Momento), un framework di apprendimento delle rappresentazioni progettato per migliorare la codifica del terreno e stabilizzare il trasferimento di conoscenza. Il framework è composto da tre moduli interconnessi:

1. Modello del Mondo come Regolarizzatore:
   Per imporre la consapevolezza della dinamica forward, viene introdotto un modello del mondo come regolarizzatore per l'encoder del teacher. L'encoder del teacher ($E_{\theta_t}$) genera variabili latenti ($z_t^t$) che, combinate con le osservazioni privilegiate, vengono alimentate in un modello del mondo ($W_\omega$) per predire lo stato privilegiato successivo ($\hat{s}_{t+1}$).

• Meccanismo: Una perdita di predizione ($L_{pred}$) viene minimizzata per addestrare il modello del mondo e propagare i gradienti all'encoder del teacher.
• Effetto: Ciò costringe l'encoder a conservare le informazioni essenziali per la dinamica forward (ad esempio, geometria del terreno e attrito) invece di scartarle come rumore irrilevante per le ricompense immediate. Il modello del mondo è implementato come una rete residua profonda con blocchi Dense-LayerNorm-Swish per catturare dipendenze temporali a lungo raggio.

2. Encoder Target di Momento per una Distillazione Stabile:
   Per affrontare la non-stazionarietà dell'encoder del teacher durante l'addestramento concomitante, il framework impiega un encoder target di momento ispirato a BYOL (Bootstrap Your Own Latent).

• Meccanismo: Invece di far regredire lo student direttamente verso il teacher online che cambia rapidamente, esso apprende da un encoder target ($E_{\hat{\theta}_s}$) che mantiene una media mobile esponenziale (EMA) dei propri pesi dello student. Viene aggiunta una testa predittiva allo student online per prevenire il collasso.
• Effetto: Questo fornisce target di distillazione coerenti e stabili, prevenendo il collasso dimensionale e permettendo allo student di apprendere rappresentazioni robuste della dinamica basate solo sulla storia propriocettiva.

3. Design della Ricompensa e Addestramento:
   La funzione di ricompensa è aumentata con vincoli specifici (intervallo dei giunti, distanza del passo, limiti della posizione della base e bilanciamento delle gambe) per mitigare problemi come l'eccessiva dipendenza dall'anca e l'oscillazione del corpo osservati nella salita delle scale. L'addestramento utilizza un paradigma Concurrent Teacher-Student (CTS) con PPO, incorporando la randomizzazione del dominio attraverso diversi tipi di terreno (piatto, pendenze, gradini, irregolare) e parametri fisici (massa, attrito, guadagni PD) per colmare il divario sim-to-real.

Contributi Chiave

• Regolarizzazione tramite Modello del Mondo: L'introduzione di un modello del mondo per imporre la consapevolezza della dinamica forward, garantendo che le rappresentazioni apprese catturino la geometria e la dinamica del terreno piuttosto che massimizzare solo le ricompense.
• Distillazione Stabile tramite Target di Momento: L'uso di un encoder target di momento per fornire target di distillazione stabili, prevenendo efficacemente il collasso dimensionale causato dagli aggiornamenti non stazionari del teacher.
• Valutazione Trasparente: L'applicazione della Principal Component Analysis (PCA) e di metriche quantitative (PC-Correlation, Linear Probing $R^2$, CCA) per verificare che gli encoder catturino gerarchicamente l'altezza del terreno, offrendo una visione trasparente della qualità della rappresentazione.

Risultati Sperimentali
Il metodo è stato valutato in simulazione (IsaacGym) e su un robot bipedale con ruote JiaRan nel mondo reale.

• Qualità della Rappresentazione: Le visualizzazioni PCA e le metriche quantitative hanno mostrato che il metodo proposto ha raggiunto una PC-Correlation di 0,96 per l'altezza del terreno nell'encoder del teacher e di 0,87 nell'encoder dello student, superando significativamente le ablazioni senza modello del mondo o target di momento (che soffrivano di rappresentazioni intrecciate o collasso).
• Prestazioni in Simulazione: Su scale continue diverse (larghezza e altezza variabile fino a 20 cm), DynaWM ha raggiunto tassi di successo fino al 100% per gradini larghi 26 cm e l'89,3% per gradini alti 20 cm e larghi, superando i baseline come CTS e RMA.
• Fluidità del Movimento: Il metodo ha dimostrato una fluidità di movimento superiore (minor jerk e accelerazione della base) e un consumo energetico comparabile rispetto a CTS.
• Deployment nel Mondo Reale: In esperimenti nel mondo reale su terreni non visti, inclusi gradini irregolari con protuberanze e gradini continui che superano le altezze di addestramento, il robot ha ottenuto un tasso di successo del 100% (5/5 prove) con un movimento significativamente più fluido rispetto al baseline CTS.

Significatività e Rivendicazioni
Il paper sostiene che DynaWM colma con successo il divario tra l'apprendimento privilegiato e il controllo deployable per terreni complessi. Garantendo che l'encoder catturi una geometria del terreno completa attraverso la regolarizzazione consapevole della dinamica e stabilizzando il processo di distillazione, il metodo consente ai robot bipedi con ruote di attraversare lunghe scale continue con alta adattabilità e fluidità. Gli autori affermano che questo approccio supera i limiti degli attuali framework teacher-student, che spesso falliscono nel mantenere la fluidità del movimento e l'adattabilità al terreno negli scenari di scale continue. Il lavoro dimosta che la modellazione esplicita della dinamica forward e la generazione di target stabili sono critiche per un robusto trasferimento sim-to-real nella locomozione legged-wheeled.