Résumé Technique : Lagrangian Perturbation Diffusion Steering (LP-DS)

1. Énoncé du Problème

Les politiques génératives à haute capacité, telles que celles basées sur les modèles de diffusion et le flow-matching, se sont imposées comme des paradigmes puissants pour le contrôle continu et la robotique. Ces modèles excellent dans la représentation de distributions d'actions multimodales et expressives apprises à partir de données hors ligne. Cependant, ils font face à deux limitations principales :

1. Contraintes de l'Apprentissage par Imitation (Behavior Cloning - BC) : Les politiques entraînées via BC héritent des limites de leurs jeux de données de démonstration, notamment une couverture incomplète de l'espace d'états, l'omission de comportements de récupération rares et une fragilité face au décalage de distribution (distribution shift).
2. Instabilité de l'Apprentissage par Renforcement (RL) : Bien que le RL offre un mécanisme d'amélioration en ligne, l'ajustement fin (fine-tuning) direct de décodeurs génératifs de grande taille est fréquemment instable et peu efficace en termes d'échantillonnage. L'interaction entre les objectifs d'optimisation et les dynamiques de débruitage ou d'intégration à long terme conduit souvent à une instabilité de l'entraînement.

Les alternatives existantes, telles que le Diffusion Steering via Reinforcement Learning (DSRL), tentent d'améliorer les politiques gelées en apprenant une politique dans l'espace du bruit latent. Cependant, DSRL remplace efficacement le prior de pré-entraînement par une politique latente apprise. Cette approche souffre de deux modes de défaillance :

• Dérive Hors-Manifold (Off-Manifold Drift) : Le bruit piloté peut dériver loin du support gaussien standard utilisé pour entraîner le décodeur, entraîissant un décodage instable.
• Effondrement de Mode (Mode Collapse) : La politique latente apprise peut réduire le comportement multimodal de la structure de base à un sous-ensemble étroit de modes, réduisant ainsi la diversité comportementale.

2. Méthodologie : Lagrangian Perturbation Diffusion Steering (LP-DS)

Les auteurs proposent LP-DS, un cadre d'adaptation léger qui améliore les politiques génératives gelées en apprenant un résidu conditionné par l'état dans l'espace du bruit latent, plutôt qu'en modifiant les poids du décodeur ou en remplaçant entièrement le prior.

Mécanisme Central

LP-DS traite le squelette génératif gelé $\Phi$ comme un décodeur boîte noire mappant le bruit latent $w$ et l'état $s$ vers des actions. Au lieu d'échantillonner le bruit directement à partir d'un prior $p_0$ (par exemple, $\mathcal{N}(0, I)$), LP-DS introduit un réseau de perturbation apprenable $\Delta_\theta(s)$ :
$$w = \epsilon + \Delta_\theta(s), \quad \text{où } \epsilon \sim \mathcal{N}(0, I)$$
Lorsque $\Delta_\theta(s) = 0$, la politique retrouve exactement la distribution de l'apprentissage par imitation originale.

Optimisation Contrainte via la Relaxation Lagrangienne

Pour empêcher la perturbation de pousser le bruit hors du manifold où le décodeur a été entraîné, LP-DS formule le problème d'apprentissage comme une optimisation contrainte. L'objectif est de maximiser la valeur Q attendue des actions générées tout en respectant une contrainte de région de confiance sur l'amplitude de la perturbation.

Les auteurs approximent la divergence de Kullback–Leibler (KL) attendue entre la distribution du bruit pilotée et le prior en utilisant l'amplitude au carré du décalage :
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
où $\delta$ est un hyperparamètre contrôlant l'écart autorisé.

Cela est résolu à l'aide d'une approche de relaxation lagrangienne. Une variable duale (multiplicateur de Lagrange) $\alpha$ est introduite pour imposer la contrainte de manière dynamique :
$$\mathcal{mathcal{L}}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
L'algorithme alterne entre :

1. Mise à jour de l'Acteur : Maximisation de $\mathcal{L}$ par rapport à $\theta$ (ascension de gradient).
2. Mise à jour du Dual : Maximisation de $\mathcal{L}$ par rapport à $\alpha$ (ascension de gradient projetée) pour imposer $\alpha \geq 0$. Si la perturbation viole la région de confiance, $\alpha$ augmente, pénalisant l'acteur et le forçant à adhérer plus étroitement au prior.

Composants de l'Algorithme

La procédure d'entraînement (Algorithme 1) comprend :

• Collecte de Données : Échantillonnage d'actions en utilisant la politique pilotée $w = \epsilon + \Delta_\theta(s)$.
• Mises à jour du Critique : Mise à jour d'un critique de la valeur de l'action $Q^A$ via l'apprentissage par différence temporelle (TD) et d'un critique de la valeur latente $Q^W$ via la distillation de $Q^A$ sur la distribution de bruit de base.
• Mises à jour Primal-Dual : Mise à jour du réseau de perturbation $\theta$ et du multiplicateur $\alpha$ pour équilibrer la maximisation de la récompense et la préservation du prior.

3. Contributions Clés

Le document présente quatre contributions principales :

1. Cadre LP-DS : Une méthode d'adaptation légère qui améliore les politiques génératives gelées via un contrôle de bruit résiduel conditionné par l'état, évitant l'instabilité et le coût d'un ajustement fin complet du décodeur.
2. Objectif de Région de Confiance Lagrangien : Un nouvel objectif qui contraint dynamiquement l'amplitude de la perturbation, fournissant un mécanisme contrôlable pour équilibrer la maximisation de la récompense et la préservation du prior latent pré-entraîné.
3. Atténuation de l'Effondrement de Mode : Preuves empiriques montrant que LP-DS préserve une plus grande diversité d'actions et atténue l'effondrement de mode par rapport aux méthodes de pilotage dans l'espace du bruit remplaçant le prior (comme DSRL).
4. Généralisabilité : Validation de la méthode au-delà des politiques de diffusion compactes, incluant les squelettes de flow-matching, un grand modèle vision-langage-action ($\pi_0$), et des expériences physiques avec le robot Franka.

4. Résultats Expérimentaux

Les auteurs ont évalué LP-DS à travers les benchmarks RoboMimic (manipulation), OpenAI Gym (locomotion) et Adroit (manipulation dextre), en comparant avec DSRL, DPPO (ajustement fin par gradient de politique) et des bases d'apprentissage Q-learning de diffusion offline-to-online.

• Performance : LP-DS a systématiquement amélioré les taux de succès et les retours épisodiques. Notamment, sur l'environnement WALKER2D, il a obtenu une amélioration d'environ 25 % du retour par rapport à la meilleure base de comparaison.
• Efficacité d'Échantillonnage : En mettant à jour uniquement un module de perturbation léger plutôt que le décodeur complet, LP-DS a démontré une efficacité d'échantillonnage supérieure et des dynamiques d'apprentissage plus stables.
• Préservation de la Diversité :
  • Dans une tâche de navigation multi-objectifs simplifiée, LP-DS a permis un contrôle explicite du compromis entre performance et diversité via la borne de la région de confiance $\delta$. Un $\delta$ plus petit préservait le comportement multimodal, tandis qu'un $\delta$ plus grand permettait la spécialisation. En revanche, DSRL s'est immédiatement effondré sur un seul mode.
  • Dans l'environnement AVOIDING, LP-DS a préservé un large ensemble de trajectoires réalisables pour l'évitement d'obstacles (multimodalité au niveau de la trajectoire) avec un petit $\delta$, alors que DSRL s'est effondré sur un chemin dominant unique.
  • Analyse de l'Entropie : LP-DS a systématiquement maintenu une entropie de l'espace d'action plus élevée que DSRL pendant l'adaptation en ligne, indiquant que la contrainte lagrangienne décourage efficacement l'effondrement prématuré tout en améliorant la performance.
• Études d'Ablation : Le retrait de la pénalité lagrangienne a conduit à un apprentissage instable et à une dégradation du succès. Le retrait de la pénalité lagrangienne et de la borne de bruit a entraîné un comportement hautement instable et de faibles taux de succès, confirmant que la contrainte de région de confiance est le mécanisme principal empêchant le pilotage hors distribution.
• Sensibilité : La méthode s'est révélée robuste à la valeur spécifique de la cible de la région de confiance $\delta$, des valeurs supérieures à 0,1 offrant généralement de fortes performances.

5. Signification et Revendications

Le document affirme que LP-DS offre une alternative stable et efficace à l'ajustement fin complet pour les politiques génératives. En traitant le squelette gelé comme un prior fixe et en apprenant uniquement un décalage résiduel borné, la méthode :

• Prévient la Dérive Hors-Manifold : Elle maintient les requêtes latentes dans la plage de fonctionnement valide du décodeur, évitant le comportement erratique observé avec le pilotage de bruit non contraint.
• Préserve la Multimodalité : Elle atténue l'effondrement de mode souvent observé lors du remplacement du prior de pré-entraînement, conservant ainsi la diversité comportementale inhérente au squelette gelé.
• Se Généralise à Diverses Architectures : L'approche n'est pas limitée aux modèles de diffusion mais s'applique à tout modèle de politique générative possédant une application latente-vers-action, y compris les modèles de flow-matching et les grands modèles vision-langage-action.
• Permet le Déploiement Physique : La méthode a été validée avec succès lors d'expériences physiques avec le robot Franka en utilisant l'adaptation latente basée sur la simulation, démontrant son applicabilité au-delà des simples benchmarks simulés.

Les auteurs positionnent LP-DS comme une solution pratique pour l'« amélioration en ligne » des politiques génératives à haute capacité, équilibrant le besoin d'optimisation spécifique à la tâche avec la nécessité de maintenir la structure multimodale robuste du modèle pré-entraîné.