TRAM: Test-Time Risk Adaptation with Mixture of Agents

Riepilogo Tecnico: TRAM (Adattamento del Rischio in Fase di Test tramite Mixture of Agents)

1. Formulazione del Problema

Il documento affronta un divario critico nel deployment dell'Apprendimento per Rinforzo (RL): le specifiche di sicurezza spesso cambiano dopo l'addestramento. Un agente può essere addestrato su un insieme di task sorgente ma deployato in un ambiente con nuove mappe di pericolo, soglie di rischio riviste o vincoli diversi di allineamento comportamentale (ad esempio, un nuovo comportamento di riferimento).

I metodi RL sensibili al rischio esistenti tendono a fissare il modello di rischio (ad esempio, varianza, CVaR o un vincolo specifico) durante la fase di addestramento. Questo approccio è fragile perché:

1. Sovra-specializzazione: L'addestramento con un proxy di rischio specifico (come la varianza del ritorno) può collassare la diversità delle policy sorgente, eliminando comportamenti che potrebbero essere sicuri sotto una specifica di rischio futura e diversa.
2. Rappresentazione del Rischio Inadeguata: I rischi basati sulla varianza non riescono a catturare pericoli spaziali (ad esempio, la vicinanza a una barriera), la deriva comportamentale da un riferimento o la volatilità locale, specialmente in contesti deterministici dove la varianza del ritorno della traiettoria è zero.
3. Costi di Riaddestramento: L'adattamento a una nuova specifica di rischio richiede solitamente il riaddestramento o il fine-tuning, che è computazionalmente costoso e spesso non fattibile una volta che un modello è deployato.

Il documento formalizza il problema come adattamento al momento del deployment senza aggiornamenti. L'obiettivo è riutilizzare una libreria fissa di policy sorgenti neutre al rischio $\{\pi_j\}$ per soddisfare una funzionale di rischio $\rho(\pi)$ specificata di recente al momento del test, senza modificare alcun parametro sorgente. L'obiettivo è massimizzare il ritorno atteso $J(\pi)$ soggetto a un vincolo di rischio $\rho(\pi) \le \delta$, approssimato tramite un surrogato penalizzato:
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. Metodologia: TRAM

Il metodo proposto, TRAM, disaccoppia l'addestramento sorgente dalla specifica del rischio. Opera in due fasi:

Fase 1: Addestramento Sorgente Neutrale al Rischio
Le policy sorgente vengono addestrate su task sorgente correlati utilizzando obiettivi standard e neutri al rischio. Questo preserva un "involucro comportamentale" (behavioral hull) diversificato, garantendo che la libreria contenga un'ampia gamma di distribuzioni di occupazione (pattern di visita stato-azione) che possono essere ricombinate successivamente.

Fase 2: Valutazione e Composizione delle Sorgenti in Fase di Test
Al momento del deployment, TRAM non riaddestra. Invece, costruisce una policy deployata $\pi_T$ selezionando azioni dalla libreria sorgente basandosi su una regola di composizione con punteggio sorgente.

1. Valutazione del Rischio: Per ogni policy sorgente $\pi_j$, TRAM calcola un punteggio di rischio globale basato sull'occupazione, memorizzato nella cache, $\rho(d_{\pi_j})$. Ciò consente funzionali di rischio flessibili definiti sulle occupazioni stato-azione scontate $d_\pi$, come:
   • Esposizione alla Barriera: Tempo trascorso vicino a una regione di pericolo.
   • Divergenza Comportamentale: Divergenza KL da un comportamento di riferimento sicuro.
   • Volatilità Locale: Fluttuazioni passo-passo nell'occupazione.
2. Stima del Valore: Il valore del task target $Q^{\pi_j}_T(s, a)$ per ogni sorgente viene stimato allo stato corrente $s$. Ciò può essere fatto efficientemente utilizzando le Caratteristiche Successore (SF) tramite un prodotto scalare $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$, o tramite rollout/critici.
3. Selezione dell'Azione: L'azione deployata viene selezionata massimizzando il punteggio sorgente aggiustato per il rischio:
   $$S_j(s, a) = Q^{\pi_j}_T(s, a) - c \rho(d_{\pi_j})$$
   $$\pi_T(a|s) \in \arg\max_{b \in A} \max_{j \in [n]} S_j(s, b)$$
   La sorgente attiva può cambiare lungo la traiettoria, creando una policy "cucita localmente" che risponde allo stato corrente e alla specifica di rischio particolare.

Garanzie Teoriche:
Il documento inquadra esplicitamente TRAM come un surrogato con punteggio sorgente trattabile, non come un ottimizzatore esatto del problema completo di controllo dell'occupazione della policy cucita.

• Limiti Localizzati: L'analisi teorica (Teorema 4.1) fornisce limiti sulla subottimalità della policy TRAM rispetto alla policy neutrale al rischio ottimale. Il limite dipende dal disallineamento del reward localizzato tra il target e la sorgente migliore, più un termine proporzionale al peso della penalità di rischio $c$ e alla costante di Lipschitz del funzionale di rischio.
• Disallineamento Involucro-Sorgente: Il documento riconosce che l'occupazione realizzata della policy cucita può differire dall'involucro convesso delle occupazioni sorgente. L'Appendice C introduce un termine misurabile di "disallineamento di cucitura" ($\epsilon_{stitch}$) che colma il divario tra i rischi sorgente memorizzati nella cache e il rischio realizzato della policy deployata.

3. Contributi Chiave

1. Formalizzazione dell'Adattamento Senza Aggiornamenti: Il documento formalizza il contesto di adattamento di policy sorgenti fisse a rischi basati sull'occupazione specificati di recente senza aggiornamenti dei parametri.
2. Algoritmo TRAM: Introduzione di una regola di composizione con punteggio sorgente che supporta tipi di rischio diversi (barriera, divergenza, volatilità) al momento del test, in contrasto con metodi legati a surrogati fissi come la varianza.
3. Fondamenti Teorici:
   • Dimostrazione che l'addestramento sorgente neutrale al rischio minimizza il rimpianto di rischio di deployment nel caso peggiore (Teorema 4.3), preservando l'involucro di occupazione riutilizzabile più ampio.
   • Garanzie di trasferimento localizzato che separano la qualità del trasferimento del reward dalla penalità di rischio di deployment.
   • Analisi del "disallineamento di cucitura" tra i rischi sorgente memorizzati nella cache e i rischi della policy realizzata.
4. Validazione Empirica: Dimostrazione attraverso quattro domini distinti:
   • Gridworlds: Validazione della modellazione espressiva del rischio (pericoli spaziali, volatilità passo-passo) dove i metodi basati sulla varianza falliscono.
   • MuJoCo Reacher: Riduzione delle violazioni di barriera nel controllo continuo mantenendo le prestazioni di raggiungimento dell'obiettivo.
   • Safety-Gymnasium: Raggiungimento di probabilità di impatto con barriera inferiori rispetto a proxy di sicurezza congelati e baselines basati sulla varianza.
   • Allineamento LLM: Adattamento di modelli linguistici fissi da 7B a nuovi vincoli di rischio comportamentale (KL rispetto a un riferimento sicuro) senza fine-tuning, mostrando un migliore allineamento rispetto alla massimizzazione pura del reward.

4. Risultati

Gli esperimenti mostrano che TRAM migliora costantemente la sicurezza di deployment (riducendo l'esposizione ai pericoli o gli impatti con barriera) preservando le prestazioni di reward, senza richiedere riaddestramento.

• Nei Gridworlds, TRAM evita con successo i pericoli spaziali e si adatta a condizioni di rischio variabili, mentre le baselines consapevoli del rischio (come RaSF) rimangono uniformemente conservative o non riescono a rilevare i rischi spaziali.
• Nel Reacher, TRAM ha ridotto i tassi di fallimento dall'85% al 26% quando è stata introdotta una nuova barriera dopo l'addestramento, con un impatto minimo sulla distanza dall'obiettivo.
• In Safety-Gymnasium, TRAM ha raggiunto la probabilità di impatto con barriera più bassa (0,209 contro 0,502 per le baselines) e un CVaR normalizzato competitivo.
• Nell'allineamento LLM, TRAM ha prodotto risposte preferite dalle comparazioni a coppie di GPT-4 rispetto alle policy di trasferimento prive di rischio, dimostrando che il metodo può prevenire l'over-ottimizzazione del modello di reward negli spazi testuali ad alta dimensionalità.

5. Significato e Affermazioni

Il documento afferma che TRAM offre un layer di deployment modulare per sistemi RL. Il suo significato principale risiede nello spostare la decisione sul rischio dalla fase di addestramento alla fase di deployment.

• Modularità: Consente agli operatori di specificare requisiti di sicurezza (mappe di pericolo, riferimenti comportamentali) dopo che i comportamenti sorgente sono già stati appresi.
• Efficienza: Evita il costo computazionale del riaddestramento o della risoluzione di nuovi MDP vincolati al momento del test, affidandosi invece a una rapida valutazione delle sorgenti.
• Onestà sul Campo di Applicazione: Gli autori sono modesti nelle loro affermazioni, dichiarando esplicitamente che TRAM è un metodo surrogato. Non afferma di risolvere esattamente il problema completo di controllo dell'occupazione della policy cucita. Invece, fornisce un limite misurabile sul disallineamento tra il proxy dell'involucro sorgente e il rischio realizzato, rendendolo un euristica pratica per scenari in cui le specifiche di rischio sono dinamiche e il riaddestramento è proibito.

Il lavoro suggerisce che per molti scenari di deployment (robotica, navigazione, LLM), mantenere una libreria diversificata e neutrale al rischio di comportamenti e comporli tramite valutazione al momento del test è una strategia più robusta rispetto all'addestramento di policy specializzate e consapevoli del rischio per ogni possibile vincolo futuro.