le-wm

LeWorldModel (LeWM) est le code officiel de l'article « LeWorldModel : Stable End-to-End Joint-Embedding Predictive Architecture from Pixels » de Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun et Randall Balestriero. Il implémente une architecture prédictive à embeddings joints (JEPA) qui, selon le résumé, s'entraîne de manière stable de bout en bout à partir de pixels bruts en utilisant seulement deux termes de perte : une perte de prédiction du prochain embedding et un régularisateur imposant des embeddings latents distribués selon une loi gaussienne. Les auteurs affirment que cela réduit les hyperparamètres de perte réglables de six à un par rapport à la seule alternative de bout en bout existante, et que le modèle compte environ 15 millions de paramètres, entraînables sur un seul GPU en quelques heures. Le dépôt s'appuie sur deux projets externes : stable-worldmodel pour la gestion de l'environnement, la planification et l'évaluation, et stable-pretraining pour l'entraînement. L'objectif déclaré est de réduire ce dépôt à sa contribution essentielle, l'architecture du modèle et l'objectif d'entraînement, implémentés dans jepa.py en tant que module PyTorch. L'installation utilise uv avec Python 3.10 et installe stable-worldmodel avec les extras train et env. Les ensembles de données sont distribués au format HDF5 via Hugging Face et décompressés avec tar --zstd ; les fichiers .h5 extraits sont placés sous $STABLEWM_HOME (par défaut ~/.stable-wm/), qui peut être remplacé par une variable d'environnement. Les noms des ensembles de données sont référencés sans l'extension .h5 dans les fichiers de configuration Hydra sous config/train/. L'entraînement est lancé avec python train.py data=pusht après avoir défini l'entité et le projet WandB dans config/train/lewm.yaml. Les points de contrôle sont enregistrés dans $STABLEWM_HOME à la fin. Les configurations d'évaluation se trouvent sous config/eval/, et le champ policy doit pointer vers un chemin de point de contrôle relatif à $STABLEWM_HOME sans le suffixe _object.ckpt. Les points de contrôle LeWM pré-entraînés sont mis en miroir sur le Hub Hugging Face pour les environnements pusht, cube, tworooms et reacher, ainsi que les ensembles de données dans la même collection. Une suite plus large de points de contrôle de référence (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop) est disponible via Google Drive, avec un tableau indiquant quelles méthodes couvrent quels environnements. Les points de contrôle se présentent sous deux formes : un objet Python sérialisé _object.ckpt utilisé par eval.py et l'API stable_worldmodel, et un dictionnaire d'état uniquement des poids _weight.ckpt. Le chargement se fait via swm.policy.AutoCostModel avec un nom d'exécution relatif à $STABLEWM_HOME ; le module renvoyé est en mode évaluation et expose .state_dict(). Pour les miroirs Hugging Face, le README fournit un script de conversion qui reconstruit le modèle JEPA à partir de config.json et weights.pt et enregistre un point de contrôle objet. Le README fournit également un lien vers l'article, un site web du projet et la collection Hugging Face, et invite à soumettre des problèmes et à collaborer par e-mail. Il rapporte que LeWM planifie jusqu'à 48 fois plus vite que les modèles de monde basés sur des modèles de fondation tout en restant compétitif sur diverses tâches de contrôle 2D et 3D, que son espace latent encode la structure physique comme le montre le sondage de quantités physiques, et que l'évaluation de la surprise détecte des événements physiquement invraisemblables. Ce sont des affirmations tirées du résumé des auteurs, non vérifiées indépendamment ici.