Résumé Technique : UniJEPA

Énoncé du Problème

Le développement de politiques robotiques généralistes capables de gérer des tâches diverses dans des environnements ouverts reste un défi central en intelligence artificielle incarnée. Les approches existantes reposent généralement soit sur des modèles vision-langage (VLM) pour la compréhension sémantique, soit sur des modèles génératifs pour la dynamique visuelle. Cependant, ces méthodes présentent des limitations significatives :

• Les approches basées sur les VLM souffrent souvent d'une dégradation des capacités fondamentales lors du fine-tuning sur des ensembles de données robotiques rares et hétérogènes. Elles négligent fréquemment les écarts fondamentaux entre les tâches vision-langage statiques et les tâches d'action robotique dynamiques.
• Les approches basées sur la génération (par exemple, la prédiction vidéo) facilitent l'apprentissage de représentations dynamiques mais échouent souvent à préserver l'alignement vision-langage crucial inhérent aux VLM pré-entraînés.
• Les modèles unifiés actuels tentent souvent d'unifier la compréhension et la génération dans un cadre de prédiction de tokens discrets, ce qui peut compromettre la robustesse de l'alignement vision-langage et la capacité à modéliser efficacement les dynamiques physiques continues.

L'article postule que l'apprentissage de politiques robotiques nécessite un paradigme intégrant la compréhension de tâches discrètes (compréhension) avec l'apprentissage de représentations d'états futurs continus (planification et dynamique).

Méthodologie : UniJEPA

Les auteurs proposent UniJEPA, un cadre Vision-Langage-Action (VLA) qui unifie l'apprentissage de représentations discrètes et continues grâce à une stratégie d'entraînement en deux étapes. L'architecture utilise une conception de Mélange de Transformers (MoT) avec des experts spécialisés par modalité.

1. Modélisation Unifiée des Embeddings Vision-Langage (Étape 1 : Pré-entraînement)

La première étape établit un paradigme de pré-entraînement cross-embodiment pour apprendre des représentations conjointes texte-image.

• Architecture : Le modèle est initialisé avec un VLM pré-entraîné (spécifiquement Paligemma) et introduit un Expert Générateur et un Expert Action (dans les étapes ultérieures) au sein du cadre MoT.
• Apprentissage de Représentations Discrètes : Le modèle est entraîné sur de grands ensembles de données vision-langage et des descriptions de tâches incarnées (annotées dans des formats de type VQA) pour apprendre des représentations linguistiques fines afin de comprendre les scènes et les instructions.
• Modélisation du Monde Continu : Contrairement aux travaux antérieurs qui prédisent des pixels d'image bruts, UniJEPA prédit des caractéristiques visuelles continues de haute dimension. Un encodeur visuel figé (par exemple, SigLIP) encode les observations futures, et le modèle apprend à prédire ces caractéristiques continues ($\hat{o}_{t+h}$) parallèlement aux tokens de texte discrets ($\hat{l}$).
• Objectif d'Entraînement : Le modèle emploie une fonction de perte conjointe combinant une perte d'entropie croisée pour la branche linguistique et une perte d'erreur quadratique moyenne (MSE) pour la branche de prédiction visuelle continue.

2. Modélisation Unifiée de l'Action (Étape 2 : Fine-tuning)

Dans la deuxième étape, le modèle est affiné sur des données spécifiques à l'embodiment pour mapper les représentations prédictives vers des tokens d'action.

• Expert Action : Un expert distinct est entraîné à partir de zéro pour gérer l'espace d'action. Il exploite le flow matching pour capturer la distribution continue et intrinsèquement multi-modale des actions.
• Optimisation Conjointe : Le modèle est entraîné pour prédire simultanément les états visuels futurs continus et exécuter des séquences d'actions. La fonction objectif combine la perte de prédiction visuelle continue avec une perte de flow matching pour la génération d'actions.
• Stratégie de Données : Le pré-entraînement utilise un mélange de vidéos robotiques (avec descriptions de sous-tâches), de vidéos de démonstrations humaines et de données génériques de QA vision-langage. Le fine-tuning est effectué exclusivement sur des données VLA provenant de simulations et de robots réels.

Contributions Clés

1. Architecture VLA Nouvelle : L'introduction d'un cadre VLA intégrant la prédiction de tokens discrets (pour la compréhension) avec la prédiction visuelle continue (pour la dynamique) au sein d'une architecture MoT unifiée.
2. Cadre d'Entraînement en Deux Étapes : Une stratégie qui aligne les représentations d'action tout en préservant les représentations intermédiaires apprises lors du pré-entraînement à grande échelle, transférant efficacement les connaissances des données à l'échelle d'Internet vers les tâches robotiques.
3. Prédiction de Caractéristiques Continues : La démonstration que la prédiction de caractéristiques visuelles continues (plutôt que de pixels bruts ou de tokens discrets) sert de signal plus efficace pour apprendre les informations dynamiques cruciales à la génération d'actions.

Résultats Expérimentaux

UniJEPA a été évalué sur deux benchmarks de simulation et deux plateformes robotiques réelles (un bras Franka Emika Panda à 7 DDL et une main dextre à 12 DDL).

• Benchmarks de Simulation :
  • Benchmarks Calvin : UniJEPA a atteint une longueur moyenne de séquence de sous-tâches de 4,11, surpassant le meilleur précédent (UP-VLA à 4,08) et dépassant significativement des bases comme GR-1 (3,06).
  • SimplerEnv : Sur les plateformes WindowX et Google Robot, UniJEPA a atteint des taux de réussite à la pointe de l'art de 71,0 % et 78,4 %, respectivement. Il a démontré des performances élevées et constantes sur toutes les sous-tâches, évitant les profils de performance "pointus" d'autres méthodes.
• Expériences Réelles :
  • Bras Franka Panda : UniJEPA a atteint les taux de réussite les plus élevés dans les quatre catégories de tâches (Prise et Dépose, Appuyer sur un Bouton, Routage de Câble, Tiroir), y compris sur des tâches non vues impliquant des objets nouveaux.
  • Main Dextre à 12 DDL : Le modèle a atteint le taux de réussite moyen le plus élevé sur neuf catégories de compétences, montrant une amélioration de 12 % sur les tâches non vues par rapport aux bases.
• Études d'Ablation :
  • Le retrait du pré-entraînement a entraîné une baisse d'environ 16 % du taux de réussite sur les tâches de prise et dépose réelles.
  • L'utilisation de caractéristiques visuelles continues pour la prédiction a généré un gain de performance d'environ 20 % par rapport aux modèles sans ce composant.
  • Parmi les méthodes d'encodage continu, les caractéristiques SigLIP ont fourni le meilleur alignement avec le socle VLM, surpassant les approches basées sur DINO et la distillation.

Signification et Revendications

L'article revendique que UniJEPA résout avec succès le compromis entre la compréhension sémantique et la prédiction dynamique dans l'apprentissage de politiques robotiques. En tirant parti d'un pré-entraînement à grande échelle sur des données de démonstrations robotiques et humaines, le modèle acquiert des capacités de généralisation robustes.

Les auteurs soulignent que leur approche permet aux robots de gérer des objets complètement nouveaux et des scénarios hors distribution (OOD) plus efficacement que les méthodes existantes. L'intégration de l'apprentissage de représentations futures continues permet à la politique de mieux anticiper les dynamiques physiques, tandis que l'alignement vision-langage préservé assure une interprétation précise d'instructions complexes. Les résultats suggèrent qu'une approche unifiée combinant compréhension, planification et apprentissage de représentations continues est une voie viable vers des politiques robotiques généralistes.