Résumé Technique : Un Cadre de Régularisation KL pour l'Apprentissage de la Planification avec des Priors Adaptatifs

1. Énoncé du Problème

Dans l'Apprentissage par Renforcement Basé sur un Modèle (MBRL), une exploration efficace demeure un défi critique, en particulier dans les tâches de contrôle continu à haute dimension où l'efficacité des échantillons est primordiale. Les approches récentes, telles que TD-MPC et BMPC, utilisent des politiques apprises comme distributions de proposition pour la planification par Intégrale de Chemin Prédictive par Modèle (MPPI). Cependant, une limitation fondamentale découle du décalage de politique : la politique d'échantillonnage ($\pi_{\theta_s}$) est souvent mise à jour indépendamment de la distribution d'action du planificateur ($\pi_P$).

Ce découplage crée un décalage de distribution où la fonction de valeur d'action bootstrap est entraînée sur des trajectoires générées par le planificateur MPPI (qui explore une distribution spécifique), tandis que la politique d'échantillonnage est mise à jour pour maximiser les valeurs d'action locales, divergeant souvent de la distribution du planificateur. Ce désalignement conduit à des estimations de valeur peu fiables, en particulier pour des horizons courts, et dégrade les performances à long terme. Les tentatives précédentes pour aligner ces distributions, telles que la minimisation de la divergence KL inverse (BMPC), souffrent d'une variance élevée car elles reposent sur des échantillons de planificateur « périmés » stockés dans le tampon de replay, qui représentent un mélange de politiques de planification obsolètes plutôt qu'une distribution cohérente actuelle.

2. Méthodologie : Cadre PO-MPC

Les auteurs proposent l'Optimisation de Politique – Contrôle Prédictif par Modèle (PO-MPC), un cadre unifié qui traite la distribution d'action du planificateur comme un prior dans l'optimisation de politique. PO-MPC formule l'étape d'apprentissage de la politique d'échantillonnage comme une instance d'Apprentissage par Renforcement Régularisé par KL.

Formulation Principale

L'objectif est d'apprendre une politique d'échantillonnage $\pi_{\theta_s}$ qui maximise le retour espéré tout en restant proche d'un prior adaptatif $\pi_p$ (dérivé du planificateur MPPI). La fonction objectif est :

$$J(\pi_{\theta_s}) = \mathbb{E} \left[ \sum_{t=0}^{T-1} \gamma^t \left( r(z_t, a_t) - \lambda \text{KL}[\pi_{\theta_s}(\cdot | z_t) \| \pi_p(\cdot | z_t)] \right) \right]$$

Où :

• $z_t$ est la représentation d'état latent.
• $\lambda$ est un hyperparamètre contrôlant la force de la régularisation KL.
• $\pi_p$ est le prior de politique.

Composants Clés

1. RL Régularisé par KL : Le cadre unifie les méthodes existantes comme des cas particuliers de cette formulation :

   • $\lambda = 0$ : Rétablit TD-MPC2, où la politique est mise à jour uniquement via la maximisation de la valeur d'action et la régularisation par entropie.
   • $\lambda \to \infty$ : Rétablit BMPC (avec KL inverse), où la politique est contrainte d'imiter la distribution du planificateur.
   • $\lambda$ intermédiaire : L'article explore des valeurs intermédiaires (par exemple, $\lambda \in [0.1, 9]$) pour équilibrer la maximisation du retour et l'alignement de distribution.

2. Prior Intermédiaire Adaptatif ($\pi_{\theta_p}$) :
   Pour résoudre la variance élevée causée par l'utilisation directe d'échantillons de planificateur périmés depuis le tampon de replay, PO-MPC introduit un prior intermédiaire appris. Au lieu de cloner la distribution brute des échantillons de planificateur stockés (qui peut être un mélange gaussien variant dans le temps), le cadre entraîne un réseau de neurones $\pi_{\theta_p}$ pour approximer la distribution actuelle induite par MPPI, $\pi_P$.

   • Ce prior est distillé à partir de la sortie du planificateur, protégeant les mises à jour de la politique d'échantillonnage de la variance des échantillons obsolètes.
   • Le prior peut être entraîné soit avec un KL Inverse (minimisant $\text{KL}[\pi_{\theta_p} \| \pi_P]$), soit avec un KL Direct (minimisant $\text{KL}[\pi_P \| \pi_{\theta_p}]$).
     • KL Inverse : Biais la politique vers l'ajustement à des modes spécifiques du planificateur (convergence plus rapide, potentiellement moins d'exploration).
     • KL Direct : Biais la politique vers la couverture du support de toutes les distributions échantillonnées (exploration améliorée, convergence potentiellement plus lente).

3. Boucle d'Algorithme :
   L'entraînement suit un cycle Planifier $\to$ Inférer $\to$ Régulariser :

   • Planifier : MPPI génère une distribution de trajectoire $\pi_P$ en utilisant la politique d'échantillonnage actuelle et la fonction de valeur.
   • Inférer : Le prior adaptatif $\pi_{\theta_p}$ est mis à jour pour approximer $\pi_P$.
   • Régulariser : La politique d'échantillonnage $\pi_{\theta_s}$ est mise à jour en utilisant l'objectif régularisé par KL, équilibrant la maximisation de la valeur Q et la pénalité KL vers $\pi_{\theta_p}$.

3. Contributions Clés

• Unification : L'article unifie des méthodes de RL basées sur MPPI disparates (TD-MPC2, BMPC) sous un seul cadre théorique (PO-MPC), clarifiant leurs choix de conception et leurs relations.
• Priors Adaptatifs : Il introduit un prior intermédiaire appris pour remplacer le clonage direct des échantillons du tampon de replay, démontrant théoriquement et empiriquement que cela réduit la variance des gradients et stabilise l'entraînement.
• Objectifs Flexibles : Le cadre permet l'exploration de variantes précédemment non étudiées, spécifiquement le réglage de la force KL $\lambda$ et le choix entre les pertes KL Direct et Inverse pour le prior.
• Performances de l'État de l'Art : La méthode atteint une efficacité d'échantillonnage et des performances finales supérieures dans des tâches de contrôle continu à haute dimension par rapport aux bases de référence existantes.

4. Résultats Expérimentaux

Les auteurs ont évalué PO-MPC sur 21 tâches de contrôle continu complexes issues de la DeepMind Control Suite (Humanoïde et Chien) et de HumanoidBench.

• Gains de Performance : PO-MPC avec des valeurs de $\lambda$ intermédiaires a constamment surpassé à la fois TD-MPC2 ($\lambda=0$) et BMPC ($\lambda \to \infty$).
  • Dans les tâches à haute dimension (par exemple, HumanoidBench), des valeurs de $\lambda$ intermédiaires (particulièrement $\lambda=1$ ou $\lambda=9$) ont produit des améliorations significatives du retour final et de l'efficacité des échantillons.
  • Dans certains environnements où les bases de référence ont échoué à apprendre (par exemple, HumanoidBench « Balance Simple », « Slide », « Stair »), PO-MPC a convergé avec succès.
• Efficacité du Prior Intermédiaire : L'utilisation du prior intermédiaire appris ($\pi_{\theta_p}$) a égalé ou dépassé les performances de l'utilisation d'échantillons bruts de planificateur depuis le tampon de replay, confirmant la réduction de la variance de mise à jour.
• Sensibilité de la Fonction de Perte : Le choix de la perte KL pour l'entraînement du prior importait. Le KL Direct était bénéfique dans les tâches axées sur l'exploration (par exemple, « Stair »), tandis que le KL Inverse a accéléré la convergence dans les tâches critiques pour la précision (par exemple, « Balance Simple »).
• Ablation : Les résultats ont montré que les performances se dégradent si $\lambda$ est trop faible (retour du décalage de politique) ou trop élevé (effondrement prématuré de mode), soulignant l'importance de régler ce compromis.

5. Signification et Revendications

L'article revendique que PO-MPC représente une unification principielle des approches basées sur MPPI. En alignant explicitement la politique d'échantillonnage avec le planificateur via un objectif régularisé par KL et un prior adaptatif, le cadre :

1. Ferme la boucle entre la planification et l'apprentissage de politique, assurant que la fonction de valeur est évaluée sur des distributions pertinentes pour le planificateur actuel.
2. Améliore la stabilité en atténuant la variance introduite par les échantillons de planification obsolètes dans le tampon de replay.
3. Fait progresser l'état de l'art en RL basé sur un modèle pour le contrôle continu, offrant un cadre flexible qui peut être réglé selon les exigences spécifiques de la tâche (exploration vs exploitation).

Les auteurs notent que, bien que le cadre soit puissant, il suppose des distributions gaussiennes pour les politiques et n'exploite pas pleinement toutes les trajectoires simulées pour l'apprentissage de la valeur. Des travaux futurs sont suggérés pour étendre les distributions de politique au-delà des gaussiennes et pour régler automatiquement l'hyperparamètre $\lambda$.