Sintesi Tecnica: Zero-Shot Off-Policy Learning (ZOL)

1. Definizione del Problema

Il paper affronta la sfida dell'Apprendimento Off-Policy Zero-Shot all'interno del paradigma dei Behavioral Foundation Models (BFM). I BFM sono agenti pre-addestrati su grandi dataset non etichettati e privi di ricompensa per apprendere prior comportamentali riutilizzabili, permettendo loro di adattarsi a nuovi compiti al tempo di test senza ulteriore addestramento (ad esempio, tramite prompt testuali, obiettivi o motion capture).

Tuttavia, questa adattazione basata sull'inferenza eredita un errore critico tipico dell'apprendimento reinforcement (RL) off-policy e offline: lo shift distributivo (distribuitional shift).

• Il Disallineamento: Una policy indotta da una nuova specifica di compito spesso interroga regioni stato-azione che sono scarsamente rappresentate o del tutto assenti nel dataset di pre-addestramento (Out-of-Distribution o OOD).
• La Conseguenza: L'inferenza zero-shot standard (ad esempio, la selezione di un vettore latente $z$ basata sulla proiezione della ricompensa) porta a errori sistematici nella stima del valore e nella sovrastima della funzione di valore. L'agente può selezionare latenti che appaiono ottimali sotto valori estrapolati, ma che inducono distribuzioni di visita non supportate dai dati, risultando in comportamenti subottimali o insicuri.
• Il Gap: Le soluzioni esistenti come DICE (DualDICE, OptiDICE) richiedono l'ottimizzazione specifica per il compito, il che contraddice il requisito di "zero-shot". Altri metodi di adattamento rapido (ReLA, LoLA) si affidano a interazioni con l'ambiente on-policy e rollout, che sono costose e violano la natura completamente offline del contesto BFM.

La domanda centrale posta è: È possibile stimare rapporti di occupazione stazionari ottimali per un'intera famiglia di policy indotte da BFM senza l'ottimizzazione on-policy per ogni compito, consentendo un adattamento rapido al tempo di test evitando al contempo le patologie off-policy?

2. Metodologia: Zero-Shot Off-Policy Learning (ZOL)

Gli autori propongono ZOL, una procedura di test-time priva di addestramento che colma il divario tra Successor Measures (SM) e Rapporti di Densità Stazionaria.

Intuizione Teorica

Il metodo si basa su una connessione teorica derivata dal framework Forward-Backward (FB) (Touati & Ollivier, 2021), utilizzato per stimare le Successor Measures.

1. Successor Measures (SM): La SM, $M^\pi$, caratterizza la visita futura scontata di una policy $\pi$. Nel framework FB, la densità successore $m_\pi$ è approssimata tramite una fattorizzazione a basso rango: $m_\pi(s_0, a_0, s, a) \approx F(s_0, a_0, z)^\top B(s, a)$, dove $F$ è un embedding forward e $B$ è un embedding backward.
2. Identità del Rapporto di Densità: Gli autori derivano un'identità esatta (Eq. 9) mostrando che il rapporto di densità stazionaria $w_{\pi/\beta} = d_\pi / d_\beta$ (il fattore di correzione necessario per allineare la visita della policy con il dataset) è proporzionale alla densità successore attesa sotto la distribuzione iniziale $\rho_0$:
   $$w_{\pi/\beta}(s, a) = (1-\gamma) \mathbb{E}_{(s_0, a_0) \sim \rho_0} [m_\pi(s_0, a_0, s, a)]$$
3. Stimatore a Basso Rango: Sostituendo l'approssimazione FB nell'identità, il rapporto di densità diventa una semplice funzione lineare dell'embedding backward:
   $$w_{\pi_z/\beta}(s, a) \approx (1-\gamma) W_{\pi_z}^\top B(s, a)$$
   Qui, $W_{\pi_z}$ è l'embedding forward atteso sulla distribuzione iniziale. Fondamentalmente, questo stimatore è basato solo sui dati; viene valutato esclusivamente sui campioni $(s, a)$ presenti nel dataset $D$, garantendo che la policy non assegni mai probabilità a regioni non supportate.

Algoritmo

ZOL opera in tre fasi senza modificare il pre-addestramento del BFM:

1. Pretraining: Addestrare gli embedding FB ($F$ e $B$) utilizzando obiettivi di differenza temporale standard sul dataset offline.
2. Inizializzazione Latente: Per un nuovo compito con ricompensa $r_T$, inferire un embedding iniziale del compito $c_T$ proiettando la ricompensa sugli embedding backward: $c_T = \mathbb{E}_{(s,a) \sim D}[B(s, a) r_T(s, a)]$.
3. Ottimizzazione Latente: Ottimizzare il vettore latente $z$ per massimizzare un ritorno pesato regolarizzato:
   $$J_{total}(z) = \mathbb{E}_{(s,a,r) \sim d_\beta} [w_{\pi_z/\beta}(s, a) (r(s, a) - \bar{r})] - \lambda_1 \chi^2(z) - \lambda_2 L_{trust}(z)$$
   • Ripesatura (Reweighting): Il termine $w_{\pi_z/\beta}$ ripesa la supervisione verso le regioni supportate dal dataset, correggendo lo shift distributivo.
   • Regolarizzazione:
     • Penalità $\chi^2$: Previene l'estrema varianza nei pesi di importanza.
     • Trust region ($L_{trust}$): Mantiene il latente ottimizzato vicino all'embedding iniziale del compito per prevenire il drift catastrofico.
     • Centratura della ricompensa: Previene soluzioni banali (es. politiche "non fare nulla") causate da offset costanti della ricompensa.

L'ottimizzazione viene eseguita tramite gradiente ascendente (Adam) su $z$ utilizzando il campionamento Monte Carlo degli stati iniziali per stimare $W_{\pi_z}$, senza richiedere interazioni con l'ambiente oltre al campionamento dello stato iniziale.

3. Contributi Chiave

1. Ponte Teorico: Il paper stabilisce un nuovo legame teorico mostrando che le rappresentazioni successore Forward-Backward contengono implicitamente i rapporti di densità stazionaria richiesti per la correzione off-policy. Ciò consente di calcolare i pesi di correzione ottimali direttamente dai componenti pre-addestrati.
2. Algoritmo ZOL: Un metodo di adattamento al tempo di test, training-free e interaction-free, che mitiga la sovrastima off-policy e gli errori di estrapolazione imponendo esplicitamente i vincoli di supporto tramite la ripesatura del rapporto di densità.
3. Validazione Empirica: ZOL è valutato su diversi benchmark (ExORL, OGBench, SMPL Humanoid) e supera costantemente i baselines zero-shot allo stato dell'arte (FB, HILP, TD-JEPA) e i metodi di adattamento on-line (LoLA, ReLA) in termini di robustezza e rendimento, specialmente in compiti a lungo orizzonte e ad alta dimensionalità.
4. Ablation e Analisi: Il paper fornisce un esempio toy minimale in 2D che dimostra come l'inferenza zero-shot standard fallisca a causa del bias di supporto, mentre ZOL riesce a recuperare la struttura del compito allineando l'inferenza con l'occupanza indotta.

4. Risultati Sperimentali

Gli autori hanno valutato ZOL su tre benchmark:

• ExORL (DeepMind Control Suite): ZOL ha ottenuto guadagni costanti rispetto ai baselines nei domini Walker, Cheetah, Quadruped e PointMass. In particolare, ha mostrato miglioramenti significativi nei compiti a lungo orizzonte e nei domini ad alta dimensionalità (es. Quadruped), superando i metodi on-line come LoLA e ReLA senza richiedere rollout nell'ambiente.
• OGBench (Long-Horizon Goal-Conditioned RL): Nei compiti che richiedono l'unione temporale (stitching) (es. AntMaze, manipolazione di cubi), ZOL ha raggiunto tassi di successo superiori rispetto a FB e altri baselines. Ha eguagliato o superato le prestazioni dei metodi on-line (LoLA/ReLA) operando interamente offline.
• SMPL Humanoid (Motion Tracking): Utilizzando un modello FB-CPR pre-addestrato, ZOL ha migliorato le prestazioni di motion tracking sul dataset AMASS, trovando policy sostanzialmente migliori rispetto all'inferenza baseline FB-CPR.

In tutti i benchmark, ZOL ha dimostrato robustezza alle scelte degli iperparametri, richiedendo una calibrazione minima.

5. Significato e Rivendicazioni

Il paper sostiene di unire i campi dell'apprendimento off-policy e dell'adattamento zero-shot. La sua importanza primaria risiede nel dimostrare che un adattamento zero-shot robusto non richiede necessariamente costose interazioni on-line o ottimizzazioni per ogni compito.

• Mitigazione del Gap di Inferenza: Gli autori sostengono che i gap di prestazione nei BFM derivino spesso non dalla capacità rappresentativa del modello, ma dalla subottimalità della procedura di inferenza stessa. ZOL affronta questo problema assicurando che il processo di adattamento rimanga ancorato alla distribuzione di pre-addestramento.
• Adattamento Training-Free: A differenza dei metodi che richiedono fine-tuning o rollout on-line, ZOL offre una soluzione "plug-and-play" che sfrutta la struttura FB esistente per correggere gli shift distributivi al tempo di test.
• Giustificazione Teorica: Il lavoro fornisce una giustificazione rigorosa del perché le rappresentazioni FB siano adatte alla correzione off-policy, andando oltre gli adattamenti euristici.

Limitazioni: Gli autori riconoscono che ZOL eredita il bias rappresentativo della fattorizzazione FB appresa. Se la stima del rapporto di densità è imprecisa o se un compito a valle richiede strettamente un comportamento prolungato fuori dal supporto, la natura conservativa del metodo (imposizione dei vincoli di supporto) potrebbe sacrificare le prestazioni in favore della sicurezza. Le direzioni future includono correzioni basate sull'incertezza ed estensioni a domini parzialmente osservabili.