Regularized Latent Dynamics Prediction is a Strong Baseline For Behavioral Foundation Models

1. Problema e Contesto

Il lavoro si concentra sul Reinforcement Learning (RL) a zero-shot, un paradigma in cui un agente apprende una politica capace di risolvere qualsiasi compito in un ambiente senza ulteriore addestramento o pianificazione dopo una fase iniziale di pre-addestramento. Questo approccio è fondamentale per lo sviluppo di Modelli Fondamentali Comportamentali (BFM), agenti generalisti che possono adattarsi a ricompense sconosciute.

Attualmente, i BFM si basano sull'apprendimento di rappresentazioni degli stati ($\phi$) che catturano le misure di successione (successor measures) o le feature di successione (successor features). La performance di questi modelli dipende criticamente dalla capacità della rappresentazione di coprire (span) un ampio spazio di funzioni di ricompensa.
Tuttavia, i metodi allo stato dell'arte (come FB, PSM, HILP) presentano diverse limitazioni:

• Richiedono obiettivi di apprendimento complessi basati su backup di Bellman (valutazione o ottimalità), che introducono instabilità, bias e difficoltà di convergenza.
• Sono sensibili alla copertura del dataset: in scenari con bassa copertura (low-coverage), dove le azioni proposte dalle politiche di esplorazione non sono presenti nei dati, i backup di Bellman portano a generalizzazioni errate e rappresentazioni degeneri.
• Spesso richiedono una definizione a priori di una classe di politiche o funzioni di ricompensa.

Il paper si chiede: Sono necessari obiettivi complessi di apprendimento delle rappresentazioni per il RL a zero-shot? È possibile ottenere performance superiori con un approccio più semplice?

2. Metodologia: RLDP (Regularized Latent Dynamics Prediction)

Gli autori propongono RLDP, un metodo semplice ed efficace per apprendere rappresentazioni degli stati per i BFM. Il metodo si articola in due fasi principali:

A. Apprendimento delle Rappresentazioni (Fase Offline)

Invece di stimare direttamente le misure di successione tramite backup di Bellman, RLDP utilizza l'apprendimento della dinamica latente (latent dynamics learning).

1. Obiettivo Principale: Prevedere lo stato futuro latente dato lo stato corrente e l'azione.
   • Si apprende un encoder $\phi: S \to \mathbb{R}^d$ e un modello di dinamica latente $g(\phi(s), a)$.
   • L'obiettivo è minimizzare l'errore di previsione: $\phi(s') \approx g(\phi(s), a)^\top w$.
2. Il Problema del Collasso: Gli autori osservano che l'obiettivo di previsione della dinamica, se usato da solo, porta a una forma di collasso delle feature. Le rappresentazioni di stati diversi tendono a diventare sempre più simili (alta similarità coseno) durante l'addestramento, riducendo lo spazio di ricompense che il modello può rappresentare.
3. Soluzione: Regularizzazione Ortogonale: Per mitigare il collasso, viene introdotta una semplice regolarizzazione ortogonale.
   • Si proiettano le rappresentazioni degli stati su una ipersfera.
   • Si aggiunge un termine di perdita che minimizza la similarità coseno tra coppie di stati distinti: $L_r(\phi) = \mathbb{E}[\phi(s)^\top \phi(s')]$.
   • La perdita totale è: $L_{RLDP} = L_{dynamics} + \lambda L_{orthogonal}$.

B. Inferenza Zero-Shot (Fase di Test)

Una volta appresa la rappresentazione $\phi$, si addestra un BFM per prevedere le misure di successione usando $\phi$ come base fissa.

• Per una nuova funzione di ricompensa $r_{test}(s)$, si risolve un problema di regressione lineare per trovare il vettore di embedding del task $z$ tale che $r_{test}(s) \approx \phi(s)^\top z$.
• La politica ottimale viene quindi derivata direttamente dalle feature di successione apprese, senza bisogno di ulteriore addestramento della politica per la nuova ricompensa.

3. Contributi Chiave

1. Proposta di RLDP: Un nuovo obiettivo di apprendimento delle rappresentazioni basato sulla previsione della dinamica latente regolarizzata, che evita la complessità dei backup di Bellman.
2. Identificazione e Mitigazione del Collasso: Dimostrazione empirica e teorica che la previsione della dinamica da sola causa un aumento della similarità tra le feature (collasso) e che una semplice regolarizzazione ortogonale risolve questo problema, mantenendo la diversità delle rappresentazioni.
3. Robustezza in Scenari a Bassa Copertura: Dimostrazione che RLDP supera i metodi basati su Bellman (come FB e PSM) in dataset con copertura limitata, dove gli errori di generalizzazione fuori distribuzione (OOD) tipicamente degradano le performance degli altri approcci.
4. Performance Competitiva: RLDP raggiunge o supera lo stato dell'arte in una vasta gamma di domini (controllo continuo, robotica umanoide) sia in setting offline che online.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi benchmark:

• Controllo Continuo (DeepMind Control Suite - ExoRL):

  • Su task come Walker, Cheetah, Quadruped e Pointmass, RLDP ha mostrato performance paragonabili o superiori a FB, PSM e HILP.
  • In molti casi, RLDP ha ottenuto i migliori risultati medi, confermando che la semplicità dell'obiettivo non compromette la qualità della rappresentazione.

• Robotica Umanoide (SMPL Humanoid):

  • In un ambiente complesso con spazi di stato e azione ad alta dimensionalità (358 osservazioni, 69 azioni), RLDP ha competuto efficacemente con metodi basati su regolarizzazione della politica (CPR) e rappresentazioni di Hilbert (HILP), superando i baseline in diversi task specifici (es. raise arms, lie on ground).

• Scenari a Bassa Copertura (D4RL - MuJoCo):

  • Questo è il risultato più significativo. Su dataset medium e medium-expert (dove le azioni OOD sono frequenti), i metodi basati su Bellman (FB) hanno fallito o mostrato performance scadenti.
  • RLDP ha superato tutti i baseline in 5 su 6 task, dimostrando che l'indipendenza dalla politica (policy-independence) dell'obiettivo di dinamica latente è cruciale per evitare bias di generalizzazione.

• Ablation Studies:

  • La rimozione della regolarizzazione ortogonale ($\lambda=0$) ha portato a un crollo delle performance e a un aumento della similarità delle feature.
  • L'uso della normalizzazione sferica (spherical normalization) è stato identificato come un componente critico per la stabilità.
  • L'aggiunta di backup di Bellman all'obiettivo RLDP ha spesso peggiorato le performance, confermando che la complessità aggiuntiva non è necessaria e può essere dannosa.

5. Significato e Implicazioni

Il lavoro di Jajoo et al. sfida la convinzione comune secondo cui gli obiettivi complessi di apprendimento delle rappresentazioni (basati su misure di successione esplicite e backup di Bellman) sono indispensabili per il RL a zero-shot.

• Semplificazione: RLDP offre un'alternativa più semplice, stabile e scalabile. Non richiede la definizione a priori di classi di politiche o ricompense, né l'uso di backup di Bellman che sono noti per essere instabili.
• Generalizzazione: La capacità di funzionare robustamente in scenari a bassa copertura rende RLDP una scelta pratica per applicazioni reali (es. robotica) dove la raccolta di dati esaustivi è costosa o impossibile.
• Fondamento per i BFM: Il metodo stabilisce un nuovo baseline forte per i Modelli Fondamentali Comportamentali, suggerendo che la diversità delle rappresentazioni (garantita dalla regolarizzazione ortogonale) è più importante della complessità dell'obiettivo di apprendimento.

In sintesi, il paper dimostra che la previsione della dinamica latente regolarizzata è un approccio potente, spesso superiore, per apprendere rappresentazioni universali necessarie per l'adattamento a zero-shot, offrendo una via più diretta verso agenti RL generalisti.