Résumé technique : Échantillonnage de tâches adaptatif robuste distributionnellement (DRATS)

Énoncé du problème

L'apprentissage par renforcement multi-tâches (MTRL) vise à entraîner un agent unique à optimiser les performances simultanément sur plusieurs tâches, en exploitant le transfert de connaissances pour améliorer l'efficacité des données. Cependant, les approches MTRL standard souffrent souvent d'un apprentissage déséquilibré : les agents résolvent rapidement les tâches faciles tout en faisant des progrès lents sur les tâches plus difficiles.

Les travaux antérieurs ont largement attribué ce déséquilibre à des gradients de tâches conflictuels, proposant des solutions impliquant la manipulation des gradients (par exemple, la chirurgie des gradients) ou des architectures de réseaux spécialisées. Cet article identifie une cause distincte et sous-explorée de déséquilibre : l'allocation déséquilibrée des données. Le MTRL standard alloue un budget d'interaction égal à toutes les tâches. Étant donné que les tâches difficiles nécessitent nettement plus d'expérience pour être résolues que les tâches faciles, une allocation uniforme entraîne une sur-allocation de données aux tâches faciles (qui ont besoin de peu de données supplémentaires) et une sous-allocation de données aux tâches difficiles. Bien que les méthodes d'apprentissage par curriculum (CL) adaptent les distributions d'échantillonnage, elles privilégient généralement les tâches faciles en premier pour faciliter le transfert, ce qui peut exacerber le déséquilibre des données dans des contextes où le transfert positif est faible ou où les tâches sont en conflit.

Méthodologie : Échantillonnage de tâches adaptatif robuste distributionnellement (DRATS)

Les auteurs proposent DRATS, un algorithme qui priorise de manière adaptative l'échantillonnage des tâches les plus éloignées d'une résolution. La méthode est dérivée selon les étapes suivantes :

1. Formalisation comme problème de faisabilité

Au lieu de maximiser le retour moyen sur l'ensemble des tâches, le MTRL est formalisé comme un problème de faisabilité : trouver une politique $\theta$ telle que le retour $J_i(\theta)$ dans chaque tâche $i$ atteigne un objectif de référence $J^{ref}_i$.
L'objectif est transformé en la minimisation de l'écart de retour $g_i(\theta) = \max(J^{ref}_i - J_i(\theta), 0)$. Pour gérer les cas où les contraintes ne peuvent être satisfaites simultanément (en raison de limites de capacité ou de conflits), l'objectif devient la minimisation de l'écart de retour maximal sur l'ensemble des tâches (un objectif minimax) :
$$\min_{\theta} \max_{i \in [k]} g_i(\theta)$$

2. Optimisation minimax et robustesse distributionnelle

La maximisation interne est reformulée comme une maximisation sur une distribution de tâches $q \in \Delta_k$ (le simplexe de probabilité). Résoudre cela directement conduit à une solution non lisse et gourmande qui concentre toute la probabilité sur la seule tâche aux performances les plus faibles, provoquant une instabilité et une inefficacité des données.

Pour remédier à cela, DRATS emploie l'Optimisation Robuste Distributionnellement (DRO) en restreignant la distribution de tâches $q$ à se situer dans une boule de Kullback-Leibler (KL) de rayon $\epsilon$ centrée sur une distribution de base $p_0$ (uniforme). Cela donne l'optimisation contrainte suivante :
$$\min_{\theta} \max_{q \in \mathcal{Q}} \mathbb{E}_{i \sim q}[g_i(\theta)]$$
où $\mathcal{Q} = \{q \in \Delta_k \mid \text{KL}(q \| p_0) \leq \epsilon\}$.

3. Solution sous forme fermée et algorithme

En utilisant la relaxation lagrangienne, la distribution de tâches optimale $q^*$ est dérivée sous forme fermée comme un softmax des écarts de retour :
$$q^*_i = \frac{\exp(\eta g_i(\theta))}{\sum_j \exp(\eta g_j(\theta))}$$
Ici, $\eta$ agit comme un paramètre de température inverse contrôlant la netteté de la priorisation.

Algorithme DRATS (Algorithme 1) :

1. Collecte de données : Échantillonnage des tâches à partir de la distribution actuelle $q_t$ et collecte des trajectoires.
2. Estimation de l'écart : Calcul des estimations Monte Carlo des retours $\hat{J}_i$ et mise à jour des retours de référence $J^{ref}_i$ (estimés en ligne comme le retour maximal observé).
3. Normalisation de l'écart : Calcul des écarts de retour normalisés $\hat{g}_i$.
4. Mise à jour de la distribution : Mise à jour de $q_t$ par ascension miroir (gradient exponentiel) sur l'objectif $L(\theta, q) = \mathbb{E}_{i \sim q}[g_i(\theta)] - \frac{1}{\eta}\text{KL}(q \| p_0)$.
5. Projection : Imposition d'une probabilité d'échantillonnage minimale $\epsilon$ pour garantir que toutes les tâches sont échantillonnées suffisamment pour une estimation stable de l'écart.
6. Mise à jour de la politique : Mise à jour de la politique $\theta$ en utilisant un algorithme on-policy (par exemple, PPO) pour maximiser le retour pondéré.

4. Garanties de convergence

Les auteurs prouvent que, sous des hypothèses standard (incluant une garantie de faible regret pour la mise à jour de la politique), DRATS converge vers l'écart de retour au pire cas optimal à un taux de $O(1/\sqrt{T})$, où $T$ est le nombre d'itérations.

Contributions clés

1. Algorithme DRATS : Introduction d'un algorithme de RL multi-tâches qui priorise de manière adaptative les tâches en fonction de l'écart de retour, dérivé d'une formulation minimax du MTRL.
2. Convergence théorique : Preuve que DRATS converge vers une erreur $\epsilon$ de l'écart de retour au pire cas optimal à un taux standard de $O(1/\sqrt{T})$.
3. Validation empirique : Démonstration que DRATS atteint une efficacité des données supérieure, un retour agrégé plus élevé et un retour au pire cas significativement plus élevé par rapport aux références d'apprentissage par curriculum existantes (par exemple, Uniforme, Progrès d'apprentissage, Potentiel d'apprentissage, Difficile en premier).
4. Orthogonalité : Démonstration que DRATS est compatible avec les architectures de réseaux multi-tâches existantes (par exemple, Soft Modularization, MOORE) et améliore les performances indépendamment du fait que les tâches présentent un transfert positif ou un conflit.

Résultats expérimentaux

Les auteurs ont évalué DRATS sur quatre benchmarks : un Gridworld à 4 tâches, MuJoCo6 (6 tâches de contrôle continu), MetaWorld-MT10 et MetaWorld-MT50.

• Performance agrégée : DRATS a obtenu le retour agrégé le plus élevé sur tous les benchmarks.
• Performance au pire cas : L'amélioration a été principalement pilotée par des performances supérieures sur les tâches les plus difficiles. Par exemple, dans MuJoCo6, DRATS était la seule méthode à résoudre avec succès les tâches difficiles Ant et Humanoid, tandis que les références échouaient à faire des progrès sur Ant.
• Efficacité des données : DRATS a atteint les niveaux de performance des références avec nettement moins d'interactions. Dans MT10, il a dépassé le retour final de la référence Uniforme sur les tâches difficiles (par exemple, pick-place, push) après seulement 6 millions d'étapes.
• Comparaison avec les références :
  • Uniforme : Sur-allocation aux tâches faciles ; échec à résoudre les tâches difficiles.
  • Progrès d'apprentissage/Potentiel : Ont tendance à privilégier les tâches avec des taux d'apprentissage élevés ou une erreur d'estimation de valeur élevée, ce qui correspond souvent aux tâches faciles, négligeant les tâches difficiles qui sont bloquées.
  • Difficile en premier (SMT) : Souffre d'un effondrement des performances sur les tâches difficiles ; une fois qu'une tâche est déclarée "résolue" sur la base d'un seuil, la probabilité d'échantillonnage chute brusquement, obligeant l'agent à oublier la tâche et nécessitant un réapprentissage.
• Ablations :
  • Échantillonnage adaptatif vs Pondération : L'échantillonnage adaptatif (DRATS) est nettement plus efficace en termes de données que la pondération de l'objectif de données échantillonnées uniformément. La pondération conduit à des estimations de gradient à forte variance pour les tâches difficiles en raison d'une collecte de données insuffisante.
  • Hyperparamètres : DRATS est robuste au choix de la température $\eta$ et du pas $\alpha$.

Importance et affirmations

L'article affirme que DRATS adresse une inefficacité fondamentale dans le MTRL : la mauvaise allocation des ressources de données. En déplaçant l'accent du conflit de gradients vers l'allocation des données, DRATS fournit un mécanisme efficace en données pour équilibrer les performances entre les tâches.

Les auteurs soulignent que :

• DRATS est la première application de l'Optimisation Robuste Distributionnellement au RL multi-tâches dans des contextes de contrôle continu.
• Contrairement à de nombreuses méthodes d'apprentissage par curriculum qui supposent un transfert positif des tâches faciles vers les tâches difficiles, DRATS cible explicitement l'écart de retour au pire cas, ce qui le rend efficace même lorsque les tâches sont en conflit ou que le transfert est faible.
• La méthode est modulaire ; elle ne modifie que le processus de collecte de données et peut être combinée avec des architectures spécialisées ou des techniques de manipulation des gradients.

L'article conclut que DRATS offre une approche fondée sur des principes et théoriquement solide pour réaliser un apprentissage multi-tâches équilibré, améliorant à la fois l'efficacité de l'apprentissage et la robustesse de la politique finale sur l'ensemble des tâches.