Persistent Robot World Models: Stabilizing Multi-Step Rollouts via Reinforcement Learning

1. Problématique : Le Biais d'Exposition et la Dégradation des Déroulements

Les modèles de monde robotiques basés sur la diffusion vidéo (conditionnés par les actions du robot) sont prometteurs pour simuler des tâches complexes difficiles à modéliser avec des moteurs physiques traditionnels. Cependant, ils souffrent d'un problème fondamental lors de leur déploiement en boucle fermée (autoregressive) :

• Biais d'exposition (Exposure Bias) : Pendant l'entraînement, le modèle est formé sur des historiques de frames réels (ground-truth). En inférence, il doit prédire la prochaine séquence vidéo en se basant sur ses propres prédictions précédentes.
• Propagation d'erreurs : Toute erreur mineure commise à une étape $t$ est encodée dans l'historique et conditionne l'étape $t+1$. Ces erreurs s'accumulent rapidement, entraînant une dégradation visuelle catastrophique en quelques secondes (dissolution des objets, dérive des effecteurs, incohérence de la scène).
• Limitation actuelle : Les modèles actuels, optimisés pour la prédiction à court terme, ne parviennent pas à maintenir la cohérence structurelle et spatiale sur des horizons longs, rendant leur utilisation comme simulateurs virtuels peu fiables.

2. Méthodologie : Post-entraînement par Apprentissage par Renforcement (RL)

L'article propose une approche de post-entraînement par RL pour aligner le modèle de monde sur ses propres déroulements générés, comblant ainsi l'écart entre l'entraînement et l'inférence.

A. Adaptation de l'Objectif RL Contrastif pour la Diffusion

Les méthodes RL standard (comme les gradients de politique) sont difficiles à appliquer aux modèles de diffusion car ces derniers ne fournissent pas de vraisemblance directe et la rétropropagation à travers le processus de débruitage est coûteuse.

• Solution : Les auteurs adaptent une méthode de RL contrastif (inspirée de DiffusionNFT) au contexte des modèles de monde robotiques.
• Mécanisme : Au lieu de maximiser la vraisemblance, le modèle génère un groupe de candidats, les note via une récompense, et met à jour les poids pour favoriser les candidats à haute récompense et pénaliser ceux à faible récompense.
• Adaptation théorique : L'article démontre mathématiquement que cette approche, initialement conçue pour des modèles de prédiction de vitesse (flow-matching), s'applique exactement aux modèles utilisant la prédiction de $x_0$ (comme Ctrl-World), sans nécessiter de termes de correction dépendants du bruit ($\sigma$).
• Formule de perte : La perte combine deux branches (positive et négative) de prédictions de données propres ($x_0$), pondérées par le poids de récompense $r$ :
  $$\mathcal{L}(\theta) = \mathbb{E} \left[ r \|\hat{x}_0^+ - x_0\|^2 + (1-r) \|\hat{x}_0^- - x_0\|^2 \right]$$
  où $\hat{x}_0^+$ et $\hat{x}_0^-$ sont des interpolations entre la prédiction du modèle actuel et celle d'une référence figée (EMA).

B. Protocole d'Entraînement pour Vidéo Autoregressive

Contrairement à la génération d'images où l'on peut échantillonner plusieurs images à partir d'un même prompt, la vidéo autoregressive n'a pas de prompt fixe.

• Stratégie de "Branchement" (Branching) :
  1. Préfixe partagé : Le modèle génère un préfixe autoregressif de longueur variable $P$ (échantillonnée aléatoirement entre 0 et 9) à partir d'une observation initiale réelle. Cela simule l'accumulation d'erreurs.
  2. Génération de candidats : À partir de l'état de l'historique figé à la fin du préfixe, $K=16$ continuations candidates indépendantes sont générées.
  3. Évaluation et Classement : Chaque candidat est noté par rapport aux frames réelles (ground-truth) correspondantes.
  4. Mise à jour : Les poids de récompense sont normalisés (Z-score) au sein du groupe pour créer un signal d'apprentissage relatif, puis utilisés pour mettre à jour le modèle via la perte contrastive.

C. Récompenses Visuelles Multi-Vues

Pour évaluer la qualité dans un contexte robotique multi-caméras, les auteurs conçoivent une récompense agrégée :

• Métriques : Combinaison de LPIPS (similarité perceptuelle), SSIM (fidélité structurelle) et PSNR (déviation des pixels).
• Agrégation : Les métriques sont moyennées temporellement sur le clip et spatialement sur les trois vues (deux caméras externes + caméra poignet).
• Robustesse : Cette combinaison évite les angles morts d'une métrique unique (ex: LPIPS pour les distorsions structurelles, PSNR pour les dérives massives).

3. Contributions Clés

1. Post-entraînement RL pour modèles de monde : Première application réussie d'un schéma de post-entraînement RL sur des modèles de monde robotiques, optimisant directement la cohérence à long terme plutôt que la précision d'un seul pas.
2. Théorie de convergence étendue : Démonstration que les garanties de convergence des méthodes RL contrastives s'appliquent exactement aux architectures de diffusion basées sur la prédiction de $x_0$, sans correction théorique supplémentaire.
3. Protocole d'entraînement adaptatif : Introduction d'une stratégie de branchement avec préfixes de longueur variable, exposant le modèle à divers régimes d'erreur (début vs fin de déroulement) pour éviter le surapprentissage à un seul type d'erreur.
4. Récompenses visuelles multi-vues : Développement de signaux de récompense denses et à faible variance combinant plusieurs métriques perceptuelles sur plusieurs vues caméra.

4. Résultats Expérimentaux

Les expériences ont été menées sur le dataset DROID (robot manipulation à grande échelle) en comparant le modèle post-entraîné (PersistWorld) avec l'état de l'art (Ctrl-World, WPE, IRASim).

• Qualité de Déroulement (Rollout Fidelity) :
  • Réduction de LPIPS : -14% sur les caméras externes et -46% par rapport à WPE.
  • Amélioration SSIM : +9.1% sur la caméra poignet (vue critique pour la manipulation fine).
  • Amélioration PSNR : +1.59 dB sur la caméra poignet.
  • Le modèle maintient l'intégrité structurelle des objets (ex: un bol ne se transforme pas en blob) et la cohérence de l'effecteur sur des horizons de ~11 secondes.
• Comparaisons Appariées : Dans 98% des cas, le modèle post-entraîné est préféré au modèle de base lors de comparaisons visuelles.
• Étude de Préférence Humaine : Une étude en aveugle avec des experts a montré un taux de préférence de 80% en faveur de PersistWorld.
• Évaluation Centrée sur la Tâche : Les gains sont particulièrement marqués sur les pixels masqués des objets et du robot, prouvant que le modèle apprend mieux les dynamiques d'interaction plutôt que de simplement préserver le fond statique.

5. Signification et Impact

Ce travail résout le problème critique de la dégradation à long terme des modèles de monde robotiques, un obstacle majeur à leur utilisation comme simulateurs fiables pour l'entraînement de politiques (RL) ou l'évaluation de stratégies (VLA).

• Persistance : Le modèle devient "persistant", capable de générer des déroulements cohérents sur de longues périodes sans s'effondrer.
• Simulateur Virtuel : Cela ouvre la voie à l'utilisation de modèles de monde vidéo comme environnements virtuels scalables et haute fidélité pour le développement de robots, réduisant le besoin de données réelles coûteuses.
• Généralité : La méthode est modulaire et indépendante de la récompense, permettant d'intégrer à l'avenir des contraintes physiques ou des signaux de tâche spécifiques.

En résumé, PersistWorld transforme les générateurs vidéo en simulateurs robotiques fiables en utilisant l'apprentissage par renforcement pour apprendre à corriger ses propres erreurs de prédiction, établissant un nouvel état de l'art pour la génération vidéo conditionnée par l'action en robotique.