Gated Memory Policy

1. Problématique

Les tâches de manipulation robotique présentent des besoins en mémoire très variés :

• Tâches Markoviennes : Ne nécessitent aucune mémoire (ex: saisir et placer un objet simple).
• Mémoire intra-essai (In-trial) : Nécessite de se souvenir des actions et observations passées au sein d'une même exécution (ex: attendre un signal visuel avant d'agir).
• Mémoire inter-essais (Cross-trial) : Nécessite d'agréger des informations sur plusieurs tentatives pour inférer des propriétés physiques non observables (ex: friction, masse) et adapter le comportement.

Le défi principal : Une approche naïve consistant à étendre simplement l'historique des observations (long-context) pour gérer ces tâches non markoviennes entraîne souvent une baisse significative des performances sur les tâches markoviennes. Cela est dû à :

1. Le surapprentissage (Overfitting) : L'augmentation de la dimensionnalité des entrées avec des données fixes.
2. Le décalage de distribution (Distribution Shift) : La présence de quelques frames bruyantes ou inconnues dans un long historique peut sortir l'ensemble de la séquence de la distribution d'entraînement.
3. Coût computationnel : L'attention auto-attentionnelle (Self-Attention) dans les Transformers a une complexité quadratique $O(H^2)$ par rapport à la longueur de l'historique $H$, rendant l'inférence lente et coûteuse.

Il existe un compromis difficile entre la robustesse (nécessitant de l'historique) et la réactivité/efficacité (nécessitant de l'ignorer quand il est inutile).

2. Méthodologie : Gated Memory Policy (GMP)

Les auteurs proposent GMP, une politique visuomotrice basée sur un Transformer de diffusion (Diffusion Transformer) qui apprend quand rappeler la mémoire et quoi rappeler de manière efficace.

A. Architecture de base

• Backbone : Utilisation d'une politique de diffusion basée sur un Transformer (DiT).
• Encodage de l'historique :
  • Les actions historiques sont encodées via un MLP.
  • Les observations visuelles sont extraites par un encodeur ViT pré-entraîné, puis agrégées en un seul token par image via un Multi-head Attention Pooling (MAP) pour réduire la charge computationnelle.
  • Un module de Cross-Attention est utilisé spécifiquement pour conditionner la politique actuelle sur un cache d'historique (fenêtre glissante), évitant ainsi le coût quadratique de l'auto-attention sur tout l'historique.

B. Trois innovations clés

1. Mécanisme de Porte Mémoire (Memory Gate) - "Quand rappeler ?"

   • Un module apprenant (un MLP avec activation sigmoïde) produit une valeur binaire $\mu_t \in \{0, 1\}$ à chaque pas de temps.
   • Si $\mu_t = 0$, l'attention croisée sur l'historique est ignorée (la politique se base uniquement sur l'observation actuelle).
   • Si $\mu_t = 1$, l'historique est activé.
   • Calibration : Pour éviter que la porte ne s'active inutilement (surapprentissage), les auteurs utilisent une procédure de calibration self-supervised. Ils entraînent deux politiques (avec et sans mémoire) sur un jeu de validation, comparent leurs erreurs de prédiction d'actions, et étiquettent les pas de temps où la politique sans mémoire échoue significativement comme "nécessitant de la mémoire". Cette porte est ensuite figée et la politique est réentraînée sur l'ensemble complet.

2. Augmentation par Bruit de Diffusion - "Robustesse"

   • Pour éviter que la politique ne dépende trop d'historiques "parfaits" (ce qui n'existe pas en déploiement réel), du bruit de diffusion est injecté dans les actions historiques durant l'entraînement (et l'inférence).
   • Le niveau de bruit diminue au fur et à mesure que le processus de débruitage avance. Cela force le modèle à être robuste aux erreurs accumulées et aux imprécisions de l'historique.

3. Efficacité Computationnelle

   • L'utilisation de la Cross-Attention (au lieu de l'Auto-Attention sur tout l'historique) rend la complexité linéaire par rapport à la longueur de l'historique.
   • La porte mémoire permet de sauter complètement le calcul de l'attention historique lorsque celle-ci n'est pas nécessaire, réduisant le temps d'inférence à une constante minimale.

3. Contributions Clés

1. Gated Memory Policy (GMP) : Un cadre pratique qui sélectionne dynamiquement l'historique pertinent, maintenant la robustesse et la réactivité tout en gérant les tâches non markoviennes complexes.
2. MemMimic : Un nouveau benchmark (simulation et réel) spécifiquement conçu pour évaluer les capacités de mémoire intra-essai et inter-essais, comblant le vide laissé par les benchmarks existants (comme RoboMimic) qui se concentrent sur des tâches markoviennes.
3. Étude Empirique Complète : Une analyse approfondie des choix de conception (porte binaire vs continue, injection de bruit, calibration) et de leur impact sur la performance, la robustesse et l'efficacité.

4. Résultats Expérimentaux

Les expériences ont été menées sur des tâches simulées (MuJoCo) et réelles (robots UR5 et ARX X5).

• Benchmark MemMimic (Tâches non markoviennes) :
  • GMP surpasse les méthodes de base à historique long (Long-hist DP, Long-hist PTP) avec une amélioration moyenne de 30,1% du taux de réussite.
  • Sur des tâches complexes comme le "Casting itératif" (inférence de friction) ou le "Poussage itératif", les méthodes naïves échouent car elles ne peuvent pas filtrer l'historique ou s'adapter aux dynamiques physiques changeantes. GMP réussit à apprendre par essais-erreurs.
• Benchmark RoboMimic (Tâches markoviennes) :
  • Contrairement aux politiques à long historique qui voient leurs performances chuter drastiquement, GMP maintient des performances compétitives, prouvant que le mécanisme de porte n'interfère pas avec les tâches simples.
• Efficacité et Robustesse :
  • Temps d'inférence : GMP est significativement plus rapide que les baselines à long historique lorsque la porte est fermée (pas d'attention historique).
  • Robustesse au bruit : L'injection de bruit dans l'historique permet à GMP de fonctionner correctement même lorsque les données passées sont imparfaites, là où d'autres méthodes échouent.
  • Visualisation : Les poids d'attention montrent que le modèle se concentre automatiquement sur les moments clés (ex: le moment où les couleurs ont été vues dans une tâche de correspondance de couleur) sans supervision explicite.

5. Signification et Impact

Ce travail apporte une solution élégante au dilemme "mémoire vs réactivité" en robotique. Au lieu de forcer le modèle à traiter un contexte temporel infini (ce qui est coûteux et instable), GMP apprend à activer la mémoire à la demande.

• Généralisation : La méthode est agnostique à la tâche et fonctionne aussi bien sur des simulations que dans le monde réel avec des objets et environnements non vus.
• Efficacité : Elle rend viable l'utilisation de longs contextes historiques sans pénalité computationnelle excessive, ouvrant la voie à des robots capables d'apprentissage continu et d'adaptation contextuelle complexe.
• Reproductibilité : Le code, les données et les instructions de déploiement sont publics, facilitant l'adoption par la communauté.

En résumé, GMP représente une avancée majeure vers des politiques robotiques capables de gérer la complexité temporelle des tâches réelles tout en restant rapides et robustes.