TRAM: Test-Time Risk Adaptation with Mixture of Agents

Résumé technique : TRAM (Adaptation du risque au moment du test via un mélange d'agents)

1. Formulation du problème

L'article comble une lacune critique dans le déploiement de l'apprentissage par renforcement (RL) : les spécifications de sécurité changent souvent après l'entraînement. Un agent peut être entraîné sur un ensemble de tâches sources mais déployé dans un environnement comportant de nouvelles cartes de dangers, des seuils de risque révisés ou des contraintes d'alignement comportemental différentes (par exemple, un nouveau comportement de référence).

Les méthodes de RL sensibles au risque existantes fixent généralement le modèle de risque (par exemple, la variance, la CVaR ou une contrainte spécifique) pendant la phase d'entraînement. Cette approche est fragile car :

1. Sur-spécialisation : L'entraînement avec un proxy de risque spécifique (comme la variance des retours) peut effondrer la diversité des politiques sources, éliminant des comportements qui pourraient être sûrs sous une spécification de risque différente et future.
2. Représentation inadéquate du risque : Les risques basés sur la variance échouent à capturer les dangers spatiaux (par exemple, la proximité d'une barrière), la dérive comportementale par rapport à une référence, ou la volatilité locale, en particulier dans des environnements déterministes où la variance des retours de trajectoire est nulle.
3. Coûts de réentraînement : L'adaptation à une nouvelle spécification de risque nécessite généralement un réentraînement ou un ajustement fin, ce qui est coûteux en calcul et souvent impossible une fois le modèle déployé.

L'article formalise le problème comme une adaptation au moment du déploiement sans mise à jour. L'objectif est de réutiliser une bibliothèque fixe de politiques sources neutres au risque $\{\pi_j\}$ pour satisfaire une fonctionnelle de risque nouvellement spécifiée $\rho(\pi)$ au moment du test sans modifier aucun paramètre source. L'objectif est de maximiser le retour espéré $J(\pi)$ sous une contrainte de risque $\rho(\pi) \le \delta$, approximée via un substitut pénalisé :
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. Méthodologie : TRAM

La méthode proposée, TRAM, découple l'entraînement source de la spécification du risque. Elle fonctionne en deux phases :

Phase 1 : Entraînement source neutre au risque
Les politiques sources sont entraînées sur des tâches sources apparentées en utilisant des objectifs standards et neutres au risque. Cela préserve une « enveloppe comportementale » diverse, garantissant que la bibliothèque contient un large éventail de distributions d'occupation (modèles de visite d'état-action) qui peuvent être recombinées ultérieurement.

Phase 2 : Évaluation et composition des sources au moment du test
Au déploiement, TRAM ne réentraîne pas. Au lieu de cela, elle construit une politique déployée $\pi_T$ en sélectionnant des actions à partir de la bibliothèque source basée sur une règle de composition notée par source.

1. Évaluation du risque : Pour chaque politique source $\pi_j$, TRAM calcule un score de risque global mis en cache, basé sur l'occupation $\rho(d_{\pi_j})$. Cela permet des fonctionnelles de risque flexibles définies sur les occupations d'état-action décomptées $d_\pi$, telles que :
   • Exposition aux barrières : Temps passé près d'une région dangereuse.
   • Divergence comportementale : Divergence KL par rapport à un comportement de référence sûr.
   • Volatilité locale : Fluctuations par étape dans l'occupation.
2. Estimation de la valeur : La valeur de la tâche cible $Q^{\pi_j}_T(s, a)$ pour chaque source est estimée à l'état courant $s$. Cela peut être fait efficacement en utilisant des caractéristiques successeurs (SF) via un produit scalaire $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$, ou via des simulations/critiques.
3. Sélection d'action : L'action déployée est sélectionnée en maximisant le score de source ajusté au risque :
   $$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 source active peut changer le long de la trajectoire, créant une politique « localement assemblée » qui répond à l'état courant et à la spécification de risque spécifique.

Garanties théoriques :
L'article présente explicitement TRAM comme un substitut noté par source traitable, et non comme un optimiseur exact du problème complet de contrôle d'occupation de la politique assemblée.

• Bornes localisées : L'analyse théorique (Théorème 4.1) fournit des bornes sur la sous-optimalité de la politique TRAM par rapport à la politique optimale neutre au risque. La borne dépend de l'inadéquation de la récompense localisée entre la cible et la source la mieux adaptée, plus un terme proportionnel au poids de la pénalité de risque $c$ et à la constante de Lipschitz de la fonctionnelle de risque.
• Inadéquation de l'enveloppe source : L'article reconnaît que l'occupation réalisée de la politique assemblée peut différer de l'enveloppe convexe des occupations sources. L'Annexe C introduit un terme mesurable d'« inadéquation d'assemblage » ($\epsilon_{stitch}$) qui comble l'écart entre les risques sources mis en cache et le risque réalisé de la politique déployée.

3. Contributions clés

1. Formalisation de l'adaptation sans mise à jour : L'article formalise le cadre d'adaptation de politiques sources fixes à des risques basés sur l'occupation nouvellement spécifiés sans mises à jour de paramètres.
2. Algorithme TRAM : Introduction d'une règle de composition notée par source qui prend en charge divers types de risques (barrière, divergence, volatilité) au moment du test, contrairement aux méthodes liées à des substituts fixes comme la variance.
3. Fondements théoriques :
   • Preuve que l'entraînement source neutre au risque minimise le regret de risque de déploiement dans le pire des cas (Théorème 4.3), préservant l'enveloppe d'occupation la plus large réutilisable.
   • Garanties de transfert localisées qui séparent la qualité du transfert de récompense de la pénalité de risque de déploiement.
   • Analyse de l'« inadéquation d'assemblage » entre les risques sources mis en cache et les risques de politique réalisés.
4. Validation empirique : Démonstration à travers quatre domaines distincts :
   • Gridworlds : Validation de la modélisation expressive du risque (dangers spatiaux, volatilité par étape) là où les méthodes basées sur la variance échouent.
   • MuJoCo Reacher : Réduction des violations de barrière dans le contrôle continu tout en maintenant les performances d'atteinte de l'objectif.
   • Safety-Gymnasium : Atteinte de probabilités de collision avec les barrières plus faibles par rapport aux proxies de sécurité figés et aux bases de référence basées sur la variance.
   • Alignement des LLM : Adaptation de modèles de langage fixes de 7B à de nouvelles contraintes de risque comportemental (KL par rapport à une référence sûre) sans ajustement fin, montrant un meilleur alignement par rapport à la maximisation pure de la récompense.

4. Résultats

Les expériences montrent que TRAM améliore constamment la sécurité du déploiement (réduisant l'exposition aux dangers ou les collisions avec les barrières) tout en préservant les performances de récompense, sans nécessiter de réentraînement.

• Dans les Gridworlds, TRAM évite avec succès les dangers spatiaux et s'adapte aux conditions de risque changeantes, tandis que les bases de référence conscientes du risque (comme RaSF) restent uniformément conservatrices ou échouent à détecter les risques spatiaux.
• Dans Reacher, TRAM a réduit les taux d'échec de 85 % à 26 % lorsqu'une nouvelle barrière a été introduite après l'entraînement, avec un impact minimal sur la distance à l'objectif.
• Dans Safety-Gymnasium, TRAM a atteint la probabilité de collision avec les barrières la plus faible (0,209 contre 0,502 pour les bases de référence) et une CVaR normalisée compétitive.
• Dans l'alignement des LLM, TRAM a produit des réponses préférées par les comparaisons par paires de GPT-4 par rapport aux politiques de transfert sans risque, démontrant que la méthode peut prévenir la sur-optimisation du modèle de récompense dans les espaces textuels de haute dimension.

5. Signification et affirmations

L'article affirme que TRAM offre une couche de déploiement modulaire pour les systèmes RL. Sa signification principale réside dans le déplacement de la décision de risque de la phase d'entraînement vers la phase de déploiement.

• Modularité : Elle permet aux opérateurs de spécifier les exigences de sécurité (cartes de dangers, références comportementales) après que les comportements sources ont déjà été appris.
• Efficacité : Elle évite le coût computationnel du réentraînement ou de la résolution de nouveaux MDP contraints au moment du test, s'appuyant plutôt sur un score source rapide.
• Honnêteté sur la portée : Les auteurs sont modestes dans leurs affirmations, déclarant explicitement que TRAM est une méthode substitut. Elle ne prétend pas résoudre exactement le problème complet de contrôle d'occupation de la politique assemblée. Au lieu de cela, elle fournit une borne mesurable sur l'inadéquation entre le proxy de l'enveloppe source et le risque réalisé, en faisant une heuristique pratique pour les scénarios où les spécifications de risque sont dynamiques et où le réentraînement est interdit.

Ce travail suggère que pour de nombreux scénarios de déploiement (robotique, navigation, LLM), le maintien d'une bibliothèque diverse et neutre au risque de comportements et leur composition via un score au moment du test est une stratégie plus robuste que l'entraînement de politiques spécialisées et conscientes du risque pour chaque contrainte future potentielle.