Résumé Technique : Modèles du Monde Basés sur des Caractéristiques Visuelles via une Action Latente Résiduelle

Énoncé du Problème

Les modèles du monde visent à prédire les transitions d'états futures à partir d'observations et d'actions, servant de fondement à l'apprentissage de politiques et au raisonnement. Le paradigme dominant actuel repose sur la génération vidéo dans des espaces latents alignés sur les pixels (par exemple, les VAE). Bien que visuellement convaincantes, ces approches souffrent d'une surcharge computationnelle élevée et d'une propension à l'hallucination (génération d'états physiquement implausibles).

À l'inverse, les modèles du monde basés sur des caractéristiques visuelles (par exemple, DINO-WM) prédisent des caractéristiques visuelles futures (telles que les tokens DINO) plutôt que des pixels bruts. Cette approche est plus efficace et moins sujette aux hallucinations. Cependant, les méthodes basées sur les caractéristiques existantes reposent sur une régression directe pour prédire les tokens futurs. Dans les interactions 3D complexes, cela conduit à une « régression vers la moyenne », entraînant des prédictions floues ou effondrées. Une alternative apparemment viable — l'utilisation de modèles génératifs (comme la diffusion) directement dans l'espace des caractéristiques — est entravée par le fléau de la dimensionnalité ; les tokens DINO ont une dimensionalité significativement plus élevée (par exemple, 1 million de dimensions pour une image de 512×512) par rapport aux latents VAE standard (16 000), rendant la modélisation générative computationnellement prohibitive et contrecarrant les objectifs d'efficacité des modèles basés sur les caractéristiques.

L'article aborde deux questions clés :

1. Comment développer un modèle du monde efficace et précis dans un espace de caractéristiques visuelles qui s'adapte à la manipulation 3D complexe ?
2. Comment exploiter de tels modèles du monde pour améliorer les politiques robotiques en aval ?

Méthodologie

1. Action Latente Résiduelle (RLA)

Les auteurs proposent une nouvelle représentation d'action latente appelée Action Latente Résiduelle (RLA). Au lieu de prédire l'état futur absolu $s_{t+h}$ ou de le générer via une diffusion lourde, la RLA encode le résidu entre les tokens DINO actuels ($s_t$) et les tokens DINO futurs ($s_{t+h}$) dans un vecteur latent compact $z$.

• Architecture : Un autoencodeur est entraîné où un encodeur $f_{enc}$ traite le résidu $(s_{t+h} - s_t)$ ainsi que des requêtes apprenables pour produire un latent compact $z$. Un décodeur $f_{dec}$ reconstruit $s_{t+h}$ en utilisant $s_t$ et $z$.
• Propriétés Clés :
  • Suffisance Prédictive : Contrairement aux actions latentes antérieures utilisées comme conditionnement faible pour une génération itérative, la RLA permet la reconstruction précise des tokens futurs en un seul passage feedforward.
  • Généralisabilité : L'encodeur RLA se généralise à de nouvelles scènes et motifs de mouvement, même lorsqu'il est entraîné sur des données limitées et agnostiques de la tâche.
  • Topologie Temporelle : L'espace latent encode naturellement la progression temporelle. L'interpolation linéaire entre un bruit gaussien et un vecteur RLA produit des représentations correspondant à des états temporels intermédiaires.

2. Modèle du Monde RLA (RLA-WM)

S'appuyant sur la RLA, les auteurs proposent RLA-WM, un modèle du monde qui prédit le vecteur RLA compact $z$ plutôt que le $s_{t+h}$ de haute dimension directement.

• Appariement de Flux (Flow Matching) : Le modèle utilise l'appariement de flux pour prédire la vitesse nécessaire pour transformer un bruit gaussien en la RLA cible $z$.
  • Entrée : État actuel $s_t$ et un bloc d'actions $a_{t:t+h}$.
  • Processus : Un réseau de conditionnement traite $s_t$ et les actions pour générer des tokens de condition. Ceux-ci sont concaténés avec un latent bruité $z_\tau$. Un réseau de flux prédit la vitesse $\hat{v}$ pour transformer itérativement le bruit en la RLA finale $z_1$.
  • Sortie : La RLA prédite $z_1$ est décodée avec $s_t$ pour reconstruire l'état futur $s_{t+h}$.
• Efficacité : En effectuant l'appariement de flux dans l'espace RLA compact (par exemple, 2048 dimensions) plutôt que dans l'espace complet des tokens DINO, le modèle atteint une haute fidélité avec un coût computationnel (FLOPs) significativement inférieur par rapport aux modèles de diffusion vidéo.

3. Applications à l'Apprentissage Robotique

L'article démontre deux applications de la RLA et du RLA-WM :

• Modèle d'Action du Monde Minimaliste (WAM) pour Vidéos sans Action :

  • Étend une politique standard d'Imitation par Clonage de Comportement (BC) (ResNet-18) en ajoutant une seule couche linéaire pour prédire la RLA $z$ à partir des observations.
  • Entraîné sur un mélange de vidéos avec et sans étiquettes d'actions. Pour les vidéos sans action, la perte d'action est masquée, et le réseau apprend à prédire la cible RLA (dérivée de la paire vidéo) pour guider la politique.
  • Cela permet d'apprendre à partir de vidéos à grande échelle non étiquetées sans coupler la politique à des architectures de génération vidéo lourdes.

• Apprentissage par Renforcement Visuel (WMRL) à l'intérieur du RLA-WM :

  • La première démonstration d'un RL visuel entraîné entièrement à l'intérieur d'un modèle du monde appris uniquement à partir de vidéos hors ligne.
  • Aucune interaction en ligne, aucune récompense conçue à la main, ni aucune perte BC auxiliaire n'est utilisée pendant la phase de RL.
  • Récompense Alignée Vidéo (VAR) : La récompense est définie comme la distance L1 négative entre les tokens DINO du déploiement du modèle et les tokens DINO de vérité terrain d'une vidéo de référence hors ligne (synchronisée dans le temps).
  • La politique (initialisée à partir de BC) est optimisée via PPO en utilisant des déploiements générés entièrement au sein du RLA-WM.

Contributions Clés

1. Action Latente Résiduelle (RLA) : Une représentation d'action latente simple et compacte apprise à partir de résidus DINO, qui est prédictive, généralisable et encode une topologie temporelle.
2. RLA-WM : Un modèle du monde basé sur des caractéristiques visuelles de l'état de l'art qui prédit la RLA via l'appariement de flux. Il surpasse les modèles basés sur les caractéristiques existants et les modèles de diffusion vidéo en précision tout en étant plusieurs ordres de grandeur plus rapide.
3. Nouvelles Applications :
   • Un WAM minimaliste qui améliore l'apprentissage par imitation à partir de vidéos sans action.
   • Un cadre de RL visuel (WMRL) qui optimise les politiques via des déploiements dans un modèle du monde appris sans interaction en ligne ni ingénierie manuelle de récompenses.

Résultats Expérimentaux

Qualité de Prédiction

Évalué sur les ensembles de données ManiSkill (simulation) et IWS (manipulation bimanuelle réelle) :

• Métriques : RLA-WM a obtenu des performances supérieures en LPIPS, SSIM et distance L1 DINO par rapport aux bases de référence, notamment DINO-WM (régression), RAE (diffusion dans l'espace des caractéristiques), FM-WM (appariement de flux dans l'espace des caractéristiques) et Vid2World (diffusion vidéo).
• Efficacité : RLA-WM a nécessité 3,5 T FLOPs par inférence, nettement inférieur à Vid2World (1,1 P FLOPs) et RAE/FM-WM (14,3 T FLOPs).
• Qualitatif : RLA-WM a produit des prédictions ancrées dans la réalité physique avec un minimum d'hallucinations, tandis que les modèles de diffusion vidéo (Vid2World) ont souvent généré des trajectoires nettes mais physiquement incohérentes.

Performance de l'Apprentissage Robotique

• Apprentissage à partir de Vidéos sans Action : Dans les tâches d'apprentissage par imitation n'utilisant que 5 % de données étiquetées par action, le WAM basé sur RLA a atteint un taux de réussite moyen de 35,6 %, surpassant la base de référence BC (27,2 %) et d'autres méthodes d'action latente (par exemple, AdaWorld à 33,7 %). L'amélioration a été la plus prononcée sur la tâche difficile Push-T (+15,2 % par rapport à la base de référence).
• RL Visuel (WMRL) : Sur 1500 graines d'évaluation, WMRL a considérablement amélioré les performances de la politique sur les robots XArm et UR10e par rapport à la politique BC initiale. En moyenne, il a réalisé un gain de +1,1 % par rapport à BC sur cinq tâches sans utiliser aucune donnée d'interaction réelle supplémentaire.

Signification et Revendications

L'article affirme que RLA-WM représente une étape importante en avant pour les modèles du monde basés sur des caractéristiques visuelles. En passant de la régression directe ou de la modélisation générative lourde à la modélisation d'action latente résiduelle, les auteurs atteignent un équilibre entre efficacité, précision et évolutivité qui était auparavant inatteignable dans les espaces de caractéristiques de haute dimension.

La signification réside dans :

1. Surmonter le Fléau de la Dimensionnalité : Démontrer que la modélisation générative dans l'espace des caractéristiques est réalisable si la cible est une représentation résiduelle compacte plutôt que le vecteur de caractéristiques complet.
2. Permettre un RL Hors Ligne Uniquement : Prouver qu'un RL visuel de haute qualité est possible en utilisant uniquement des données vidéo hors ligne et un modèle du monde appris, éliminant le besoin d'exploration en ligne coûteuse ou de fonctions de récompense conçues à la main.
3. Évolutivité vers des Données Non Étiquetées : Fournir une voie pour exploiter la vaste quantité de données vidéo robotiques non étiquetées pour l'amélioration des politiques grâce à l'approche WAM minimaliste.

Les auteurs reconnaissent des limitations, telles que la difficulté à gérer les mouvements de fond non pertinents pour la tâche (par exemple, les secousses de la caméra) et l'observabilité partielle (occlusions), suggérant un travail futur sur une RLA indépendante de la vue en 3D et une conditionnement multi-images. Cependant, ils soulignent que leurs résultats sur des ensembles de données à petite échelle isolent les gains pilotés par la méthode de la simple mise à l'échelle des données, établissant une base solide pour les futures applications en monde ouvert.