Résumé Technique : WAM-RL

Énoncé du Problème

Les modèles récents de type World-Action (WA) ont démontré une forte capacité de généralisation et d'efficacité des données dans l'apprentissage de politiques robotiques en modélisant conjointement les observations futures et les actions via des modèles génératifs basés sur la vidéo. Cependant, les modèles WA existants dépendent fortement de l'apprentissage supervisé à partir de trajectoires d'experts. Cette dépendance impose deux limitations fondamentales :

1. Contraintes Distributionnelles : La politique apprise est contrainte par le support des données d'entraînement, peinant à acquérir des compétences de manipulation fines qui se situent en dehors de la distribution des démonstrations.
2. Absence d'Amélioration Continue : Les modèles ne peuvent pas s'améliorer continuellement grâce à l'interaction avec le monde réel, ce qui entrave leur capacité à s'adapter à de nouveaux scénarios ou à corriger des erreurs lors de l'exécution.

Bien que l'apprentissage par renforcement (RL) soit une solution naturelle, l'application du RL aux modèles WA est non triviale. Contrairement aux modèles standard de type Vision-Language-Action (VLA) où la représentation visuelle est souvent fixe, les modèles WA se composent de deux composants étroitement couplés : un modèle de monde (générant des prédictions futures) et un modle d'action (mappant les prédictions en actions). L'application naïve du RL ou du fine-tuning en ligne peut provoquer des décalages de distribution dans l'espace latent du modèle de monde, entraînant une instabilité et une dégradation des performances car le modèle d'action est profondément dépendant de cet espace latent.

Méthodologie : WAM-RL

Les auteurs proposent WAM-RL, un cadre permettant l'optimisation conjointe du modèle de monde et du modèle d'action par l'interaction en ligne. L'idée centrale est que le modèle de monde fournit la structure prédictive et la capacité de planification, tandis que le modèle d'action agit principalement comme un traducteur convertissant les prédictions latentes en actions exécutables.

Le cadre utilise un schéma d'optimisation en deux parties :

1. SFT Vidéo en Ligne avec Régularisation KL (Modèle de Monde)

Pour affiner le modèle de monde, les auteurs utilisent des trajectoires collectées lors d'interactions réussies pour un apprentissage supervisé (SFT) auto-supervisé en ligne.

• Objectif : Le modèle de monde est entraîné pour prédire les observations futures en utilisant un objectif de modélisation vidéo standard (par exemple, le flow matching).
• Stabilisation : Pour empêcher la distribution des caractéristiques latentes de dériver drastiquement (ce qui briserait le modèle d'action), une terme de régularisation KL est introduit. Cela contraint les caractéristiques latentes du modèle de monde mis à jour à rester proches de celles d'une copie gelée du modèle pré-entraîné.
• Implémentation : Les auteurs approximent les distributions latentes comme des Gaussiennes, en utilisant une moyenne mobile exponentielle (EMA) pour estimer la covariance du modèle actuel tout en gardant la covariance de référence fixe. Cela permet une adaptation graduelle tout en préservant la géométrie attendue par l'acteur.

2. RL Basé sur la Reconstruction (Modèle d'Action)

Le modèle d'action est optimisé via le RL pour garantir que les trajectoires exécutées réalisent fidèlement les prédictions du modèle de monde.

• Signal de Récompense : Au lieu de récompenses de tâche éparses, les auteurs définissent une récompense dense basée sur la reconstruction. Cette récompense mesure la cohérence (similitude) entre les observations futures imaginées (prédites par le modèle de monde) et les observations réelles obtenues en exécutant l'action dans l'environnement.
• Métriques de Similitude : Les auteurs étudient diverses fonctions de similitude, incluant l'erreur quadratique moyenne (MSE) au niveau des pixels, la cohérence du flux optique et la similitude des caractéristiques sémantiques (DINOv2, V-JEPA2).
• Optimisation : L'acteur est mis à jour via des méthodes de gradient de politique en utilisant ces récompenses de reconstruction, comblant ainsi efficacement l'écart entre l'imagination du modèle de monde et l'exécution réelle.

Contributions Clés

1. Premier Cadre RL pour les Modèles WA : WAM-RL est le premier cadre intégrant l'apprentissage par renforcement dans le paradigme World-Action, permettant spécifiquement l'optimisation conjointe du modèle de monde et du modèle d'action.
2. Stratégie d'Entraînement Stable : Les auteurs proposent une nouvelle stratégie d'entraînement combinant le SFT vidéo en ligne avec la régularisation KL pour le modèle de monde et le RL basé sur la reconstruction pour l'acteur. Cette conception facilite la co-évolution des deux composants sans déstabiliser l'espace latent.
3. Perspectives Empiriques sur les Tâches à Long Horizon : Ce travail apporte des éclairages critiques sur les limites de l'optimisation de l'acteur seul. Les auteurs démontrent que l'optimisation de l'acteur seul produit des améliorations sur les tâches à court horizon mais échoue à fournir des gains significatifs sur les tâches à long horizon en raison de l'accumulation d'erreurs de prédiction. L'optimisation conjointe s'avère critique pour les performances complexes à long horizon.

Résultats Expérimentaux

La méthode a été évaluée sur les benchmarks LIBERO (manipulation centrée sur l'objet) et RLBench (compétences multi-étapes).

• Gains de Performance : WAM-RL surpasse systématiquement la ligne de base pré-entraînée (Base) et le RL de l'acteur seul (πRL).
  • Sur LIBERO-Object, les taux de succès passent de 68 % (Base) à 82 % (WAM-RL), surpassant nettement l'approche de l'acteur seul (78 %).
  • Sur RLBench (Water Plants), WAM-RL améliore le taux de succès de 19 % à 22 %, alors que la méthode de l'acteur seul (πRL) ne parvient pas à s'améliorer par rapport au modèle de base (18 %).
• Analyse Comportementale :
  • Comportements de Récupération : L'adaptation en ligne du modèle de monde a permis de générer des comportements de récupération plus réalistes (par exemple, tenter de nouveau une saisie ratée) plutôt que de prédire un échec immédiat.
  • Ablation sur les Récompenses : L'étude a révélé que si le flux optique fournissait le signal de récompense le plus discriminant entre le succès et l'échec, la MSE au niveau des pixels offrait les meilleures performances en aval. Les auteurs émettent l'hypothèse que cela est dû au fait que la MSE de pixels s'aligne mieux avec l'objectif d'entraînement du modèle de monde et impose une régularisation plus forte contre les actions hors distribution.
  • Ablation sur le SFT Vidéo : Les modèles entraînés sans SFT vidéo en ligne ont échoué à récupérer de leurs erreurs initiales, dérivant rapidement vers des comportements hors distribution. En revanche, les modèles avec SFT vidéo ont appris à anticiper l'échec et à générer des actions correctives au sein d'un seul bloc en boucle ouverte.

Signification et Revendications

L'article affirme que l'efficacité des modèles World-Action est fondamentalement limitée par la qualité du modèle de monde. Bien qu'une optimisation de l'acteur seul puisse améliorer les tâches à court horizon, elle ne peut pas traiter l'accumulation d'erreurs de prédiction dans les contextes à long horizon.

La portée de WAM-RL réside dans la démonstration que l'amélioration conjointe du modèle de monde et du modèle d'action est essentielle pour une prise de décision robuste à long horizon. En permettant au modèle de monde d'apprendre la dynamique de l'échec et les stratégies de récupération par l'interaction en ligne, le système peut anticiper et corriger les déviations, menant à une exécution plus robuste.

Les auteurs reconnaissent modestement les limites, notant que la régularisation KL contraint l'étendue de l'adaptation du modèle de monde, ce qui pourrait limiter l'expansion des capacités au-delà de la distribution pré-entraînée. De plus, la dépendance aux récompenses basées sur la reconstruction (souvent basées sur des représentations pré-entraînées) peut offrir un contraste limité entre le succès et l'échec, suggérant la nécessité de formulations de récompense plus expressives et sensibles aux tâches pour les travaux futurs.