Résumé Technique : Retour aux Latents Économes : Apprentissage de Modèles du Monde Centrés sur la Tâche à partir de Fondations Visuelles

Énoncé du Problème

Les modèles du monde permettent aux agents de prédire les dynamiques futures conditionnées par les actions, une capacité centrale pour la planification et le contrôle. Cependant, les approches existantes font face à un compromis fondamental dans la représentation latente :

1. Les modèles basés sur les pixels (par exemple, Dreamer, MuZero) apprennent les dynamiques directement à partir des pixels ou de latents compacts façonnés par des signaux de récompense. Bien qu'efficaces, ils entraînent souvent des coûts computationnels élevés ou nécessitent une interaction en ligne/une supervision par récompense pour façonner l'espace latent.
2. Les modèles d'embedding de fondation (par exemple, DINO-WM, V-JEPA) exploitent des encodeurs visuels pré-entraînés figés pour prédire directement dans un espace sémantique de haute dimension. Bien qu'ils offrent des priors sémantiques riches et une généralisation zero-shot, les embeddings contiennent des détails excessifs sans rapport avec la tâche (textures, éclairage, arrière-plan) qui n'affectent pas le contrôle.

Dans les contextes hors ligne sans récompense, où les agents doivent apprendre à partir de trajectoires fixes sans supervision par récompense, ces problèmes sont exacerbés. L'utilisation d'embeddings de fondation bruts comme espace d'état force le modèle à gaspiller sa capacité sur des variations non pertinentes, conduisant à une faible efficacité de modélisation, une interprétabilité réduite du contrôle et un effondrement potentiel de la représentation, en particulier dans les espaces d'action de haute dimension (par exemple, bras robotiques à 7 degrés de liberté). Le défi central consiste à extraire une représentation latente compacte et suffisante pour la tâche de ces embeddings de fondation riches mais bruyants, sans interaction en ligne.

Méthodologie : TC-WM

Les auteurs proposent TC-WM (Task-Centric World Model), un cadre qui traite les embeddings de fondation pré-entraînés non pas comme l'espace d'état final, mais comme un « échafaudage sémantique » à partir duquel les dynamiques pertinentes pour la tâche sont distillées.

Architecture Principale

TC-WM opère sur un pipeline d'embeddings visuels ($x_t$) et de signaux proprioceptifs ($s^p_t$) :

1. Embedding Joint : Un encodeur visuel figé (par exemple, DINOv2) mappe les images vers des caractéristiques visuelles, tandis qu'une couche linéaire entraînables encode la proprioception. Ceux-ci sont concaténés pour former un embedding joint $x_t$.
2. Encodeur Latent (Projection) : Un encodeur linéaire projette l'embedding joint de haute dimension vers un état latent compact $z_t$. Crucialement, $z_t$ est factorisé en deux sous-espaces :
   • $z^s_t$ (Pertinent pour l'état) : Aligné avec l'état physique.
   • $z^c_t$ (Complémentaire) : Conserve la structure sémantique résiduelle (apparence, arrière-plan) pour satisfaire les contraintes de reconstruction.
3. Dynamiques Latentes : Un prédicteur basé sur Transformer ($f_\theta$) modélise les transitions temporelles dans l'espace latent compact ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. Alignement Centrée sur la Tâche : Le modèle emploie un objectif contrastif (InfoNCE) pour aligner le sous-espace $z^s_t$ avec l'état proprioceptif $s^p_t$. Cela force activement l'encodeur à séparer les facteurs physiques pertinents pour la tâche des variations parasites.
5. Décodeur d'Embedding : Un décodeur linéaire reconstruit l'embedding visuel original à partir de $z_t$. Cela agit comme un régularisateur, assurant que le latent compact reste ancré sur la variété de fondation et empêchant les solutions dégénérées.

Fonction de Perte d'Entraînement

La perte totale combine quatre composantes :

• Perte de Dynamique Latente ($L^z_{dyn}$) : Minimise l'erreur de prédiction dans l'espace latent.
• Perte de Dynamique Centrée sur la Tâche ($L^s_{dyn}$) : Prédit la proprioception future directement à partir de l'historique pour imposer des dynamiques centrées sur la tâche.
• Perte d'Alignement ($L_{align}$) : Perte contrastive alignant $z^s_t$ avec $s^p_t$.
• Perte de Reconstruction ($L_{rec}$) : Reconstruit l'embedding visuel à partir de $z_t$ pour préserver la structure visuelle.

Planification au Moment du Test

L'espace latent centré sur la tâche appris supporte diverses stratégies de planification :

• Contrôle de basse dimension (Navigation) : Contrôle Prédictif de Modèle (MPC) utilisant la Méthode de l'Entropie Croisée (CEM) pour optimiser les séquences d'actions dans l'espace latent.
• Manipulation de haute dimension : Planification par Diffusion Latente (LDP) combinée à un Modèle de Dynamique Inverse (IDM) pour générer des actions exécutables à partir de trajectoires latentes.

Contributions Clés

1. Décomposition Centrée sur la Tâche : L'article introduit une méthode pour décomposer les embeddings de fondation de haute dimension en un espace latent compact et suffisant pour la tâche en séparant explicitement les dynamiques d'état physique des variables parasites sémantiques.
2. Identifiabilité Active : Contrairement aux travaux antérieurs qui supposent passivement la diversité des mécanismes pour l'identifiabilité, TC-WM utilise un alignement contrastif avec la proprioception pour activer l'induction de conditionnelles distinctes pour différents états physiques, garantissant théoriquement la récupération des facteurs centrés sur la tâche jusqu'à une transformation affine.
3. Apprentissage Hors Ligne Sans Récompense : Le cadre atteint une modélisation du monde et une planification de haute qualité dans des contextes hors ligne sans signaux de récompense, reposant uniquement sur des observations visuelles et des informations latérales proprioceptives.
4. Garantie Théorique : Les auteurs fournissent un théorème prouvant que sous des hypothèses légères (dynamiques contextuelles, variabilité latente-observée et alignement suffisant), le latent compact appris est identifiable et récupère l'état physique réel jusqu'à une reparamétrisation.

Résultats Expérimentaux

Les auteurs évaluent TC-WM sur 9 benchmarks hors ligne (Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper) par rapport à des bases de comparaison incluant TD-MPC2, DreamerV3, MuZero et DINO-WM.

• Qualité du Modèle du Monde : TC-WM atteint la plus faible erreur de prédiction de déroulement latent (MSE) sur presque toutes les tâches, en particulier dans les environnements de manipulation où des dynamiques à long horizon précises sont critiques. Bien que la reconstruction dans l'espace image soit compétitive, le gain principal réside dans la stabilité et la précision des dynamiques latentes.
• Performance de Planification :
  • Dans la navigation (actions de basse dimension), TC-WM égale la performance de bases de comparaison solides comme TD-MPC2.
  • Dans la manipulation (actions de haute dimension, riches en contacts), TC-WM surpasse significativement toutes les bases de comparaison. Par exemple, sur la tâche « Lift » de Robomimic, TC-WM atteint un taux de réussite de 54 % contre 29 % pour DINO-WM et 31 % pour TD-MPC2.
• Analyse de la Représentation :
  • Sondage Linéaire : Des régresseurs linéaires entraînés sur l'espace latent de TC-WM récupèrent les états proprioceptifs et des objets avec des scores $R^2$ supérieurs aux bases de comparaison, confirmant que le latent est ancré physiquement.
  • Anti-Effondrement : TC-WM démontre une meilleure utilisation de l'espace latent (rang effectif plus élevé) par rapport à DINO-WM, qui souffre d'un effondrement de la représentation lors de l'utilisation directe d'embeddings bruts.
  • Robustesse : TC-WM maintient sa performance sous des perturbations visuelles (bruit gaussien, jitter de couleur) mieux que DINO-WM, avec des baisses plus faibles du SSIM.

Signification et Revendications

L'article revendique que les embeddings de fondation sont mieux utilisés comme un échafaudage sémantique plutôt que comme une représentation d'état directe. En projetant ces embeddings dans un sous-espace centré sur la tâche aligné avec des signaux physiques, les agents peuvent atteindre :

• Latents Économes : Une réduction de dimensionnalité qui filtre les détails visuels non pertinents tout en préservant les dynamiques nécessaires au contrôle.
• Contrôle Amélioré : Une performance de planification accrue dans des environnements complexes et de haute dimension où les modèles de fondation standards échouent en raison du bruit et des variations non pertinentes.
• Solidité Théorique : Une justification formelle expliquant pourquoi les projections linéaires et l'alignement contrastif suffisent à récupérer les facteurs pertinents pour la tâche à partir d'encodeurs visuels figés.

Les auteurs positionnent TC-WM comme un pont entre le pouvoir de généralisation des modèles de fondation et la contrôlabilité requise pour les agents physiques, adressant spécifiquement le fossé dans les scénarios d'apprentissage hors ligne sans récompense.