SmolWAM: un modello latente del mondo e delle azioni per la manipolazione robotica su dispositivi edge

Abstract degli autori

I modelli del mondo e delle azioni (WAM) prevedono le azioni di un robot insieme alle osservazioni che queste produrranno. I modelli più potenti, però, si basano su transformer di diffusione video, troppo onerosi dal punto di vista computazionale per i dispositivi edge con risorse limitate. Presentiamo SmolWAM, un WAM latente che abbiamo ricavato da Cosmos Policy, un modello da 2 miliardi di parametri, mantenendone il tronco preaddestrato e ricostruendone l’interfaccia di inferenza: rappresentazioni visive latenti di un singolo fotogramma, caratteristiche testuali conservate nella cache che preservano gli attention sink del tronco, 12 dei 28 blocchi transformer e 1.008 token da cui si ricavano le azioni e le future rappresentazioni visive latenti in un solo passaggio di rectified flow. Addestrato in 3,6 ore GPU, SmolWAM ha successo nel 95,5% di 400 episodi LIBERO, con 57 ms e 2,3 TFLOP per gruppo di 16 azioni e 2,4 GiB di memoria GPU. È 11 volte più veloce del modello insegnante, richiede un numero di FLOP 18,5 volte inferiore e una quantità di memoria 3,2 volte inferiore; la sua testa dedicata al mondo continua inoltre a prevedere il movimento della scena.