LeRobot documentation

Cosmos3

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Cosmos3

Cosmos3 est un modèle d’action du monde pour le contrôle de robots, construit sur Cosmos 3, la famille de modèles du monde omnimodaux de NVIDIA. À partir des observations de la caméra et d’une instruction en langage, la politique prédit conjointement les futures images vidéo et le prochain bloc d’actions du robot. L’implémentation LeRobot est adaptée de la politique d’action native cosmos-framework et exécute les modèles Cosmos3 Nano et Edge publiés, convertis au format LeRobot.

Aperçu du modèle

Cosmos 3 est un Mixture-of-Transformers (MoT). Chaque couche de décodeur possède deux chemins disposant chacun de leurs propres normalisations de couche, projections d'attention et MLP, et les deux chemins ne se rejoignent que dans une auto-attention partagée :

  • Le chemin de compréhension (le raisonneur) lit les jetons de texte du prompt, et éventuellement les jetons d’image ViT, avec une attention causale.
  • La voie de génération (le générateur) lit les jetons continus de la sous-séquence de diffusion : les latents VAE de Wan2.2 pour la vidéo et les jetons d'action. Elle effectue une attention bidirectionnelle sur les deux voies et est entraînée par correspondance de flux.
Mixture-of-Transformers architecture of Cosmos 3

Architecture Mixture-of-Transformers (MoT) de Cosmos 3 (figure 5 du rapport technique). À gauche : la sous-séquence autorégressive (AR) passe par le raisonneur et la sous-séquence de diffusion (DM) par le générateur, et les deux convergent dans une attention multimodale partagée. À droite : le masque d'attention, causal pour l'AR et complet pour la diffusion.

En tant que politique, Cosmos3 s’exécute en mode de modèle d’action du monde des recettes d’action natives :

  • Entrées : les vues de la caméra sont composées en un seul canevas vidéo (une vue telle quelle, deux côte à côte, trois avec la première en haut et les deux autres en dessous). Le prompt contient la tâche, le taux d’images par seconde, la résolution du canevas, une description de la disposition du canevas et le nombre d’images inactives du bloc.
  • Entraînement : la politique encode le clip de l'image courante et les chunk_size images futures avec le VAE, ajoute du bruit aux latents futurs et au chunk d'action, et apprend à débruitiser les deux. La perte est une somme pondérée d'une perte vidéo et d'une perte d'action.
  • Inférence : le cadre de condition est propre, la vidéo future et les actions sont échantillonnées conjointement avec le solveur UniPC, et la politique renvoie le bloc d’actions. La vidéo prédite peut être enregistrée lors de l’évaluation (voir Évaluation).
  • Par défaut, seuls le chemin de génération, les couches qui l’entourent et les modules d’action sont entraînés (train_gen_stream_only=true, freeze_vision_encoder=true), comme dans les recettes de post-entraînement natives.

LeRobot fournit les deux tailles publiées en tant que deux types de politique :

Cosmos3 NanoCosmos3 Edge
Type de politiquecosmos3_nanocosmos3_edge
Paramètres (entraînés par défaut)16.5B (7.0B)4.6B (1.4B)
BackboneMoT initialisé à partir de Qwen3-VL, 36 couches, taille cachée 4096MoT avec un transformer dense de 2B, 28 couches, taille cachée 2048
Tokenizernvidia/Cosmos3-Nanonvidia/Cosmos3-Edge
Point de contrôle de baselerobot/Cosmos3-Nanolerobot/Cosmos3-Edge

Prérequis d'installation

  1. Installez LeRobot en suivant notre Guide d'installation.

  2. Installer les dépendances de Cosmos3 :

    pip install -e ".[cosmos3]"

    Si vous avez installé LeRobot depuis PyPI :

    pip install 'lerobot[cosmos3]'
  3. Ajoutez les extras de votre flux de travail : training pour lerobot-train, libero pour entraîner ou évaluer sur LIBERO, et peft pour LoRA. Par exemple :

    pip install -e ".[cosmos3,training,libero]"

Cosmos3 utilise par défaut le backend d’attention à point produit mis à l’échelle de PyTorch (sdpa). Si votre environnement dispose d’une installation compatible de flash_attn, vous pouvez demander --policy.attn_implementation=flash_attention_2. Ne le combinez pas avec --policy.compile_model=true : les couches décodeur compilées se recompilent pour chaque couche avec flash attention et s’arrêtent avec une erreur de recompilation.

Points de contrôle

Point de contrôleTailleDonnées d'entraînementUtiliser
lerobot/Cosmos3-NanoNanoSortie de nvidia/Cosmos3-NanoBase pour l'ajustement fin
lerobot/Cosmos3-EdgeEdgeSortie de nvidia/Cosmos3-EdgeBase pour l'ajustement fin
lerobot/Cosmos3-Nano-Policy-LIBERONanolerobot/LIBERO-Cosmos3Évaluation LIBERO
lerobot/Cosmos3-Edge-Policy-LIBEROEdgelerobot/LIBERO-Cosmos3Évaluation LIBERO
lerobot/Cosmos3-Nano-Policy-LIBERO_10Nanolerobot/LIBERO_10-Cosmos3Évaluation LIBERO-10
lerobot/Cosmos3-Edge-Policy-LIBERO_10Edgelerobot/LIBERO_10-Cosmos3Évaluation LIBERO-10
  • Les points de contrôle de base contiennent les poids publiés au format LeRobot. Leurs modules d’action (action2llm, llm2action, action_modality_embed) n’ont pas encore été entraînés pour un robot, il convient donc de commencer l’ajustement fin avec --policy.load_action_modules=false (voir Ajustement fin sur LIBERO).
  • Le config.json des points de contrôle de base contient la recette de post-entraînement native de DROID, et les valeurs par défaut de la classe des types de politique correspondent à la recette native de LIBERO.
  • Les points de contrôle LIBERO contiennent les poids EMA de leur exécution d'entraînement.
  • Les poids sont stockés en bfloat16. config.json définit dtype: float32 volontairement : l’entraînement en précision mixte à partir de ces points de contrôle conserve les poids maîtres en float32, comme lors de l’entraînement natif. Pour l’inférence et l’évaluation, passez --policy.dtype=bfloat16.

Il existe deux façons de démarrer à partir d'un point de contrôle :

  • --policy.path=<checkpoint> charge la configuration stockée, les poids et les processeurs, puis applique la ligne de commande par-dessus.
  • --policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint> (ou cosmos3_edge) charge uniquement les poids et les processeurs, et la configuration démarre des valeurs par défaut de la classe, la recette LIBERO.

Fine-tuning sur LIBERO

Les valeurs par défaut de la classe correspondent à la recette LIBERO du rapport technique : des blocs de 16 actions, l’incarnation domain_id=5, aucune action d’état, une normalisation par quantiles des actions, UniPC avec 30 étapes et sans guidage sans classifieur à l’inférence, ainsi qu’un taux d’apprentissage de 5e-5 avec 500 étapes de montée en charge et un décroissement linéaire jusqu’à zéro à l’étape 16000. En partant d’un point de contrôle de base avec --policy.type, il suffit donc d’ajouter --policy.load_action_modules=false.

La recette entraîne le modèle sur les données LIBERO converties au format Cosmos3 :

Les deux proviennent de nvidia/LIBERO_LeRobot_v3, les données de la recette native. Les actions sont des [delta position (3), rot6d (6), gripper] 10D au lieu des actions axis angle 7D de LIBERO, les deux caméras observation.images.image (troisième personne) et observation.images.image2 (poignet) sont en 256x256 à 20 FPS, et les statistiques des actions sont celles de la recette native. Lors de l'évaluation, l'environnement LIBERO convertit les actions en retour (voir Évaluation).

Commande d'entraînement

La recette native LIBERO entraîne avec un lot global de 2048 échantillons pendant 2000 itérations sur LIBERO-10 (le meilleur point de contrôle se situe vers l'itération 1500) et pendant 5000 itérations sur les quatre suites (meilleur vers 4500). La commande ci-dessous la reproduit sur une machine avec 8 GPU : FSDP2 sur les 8 GPU, 32 échantillons par GPU et 8 étapes d'accumulation de gradient, soit 32 x 8 x 8 = 2048 échantillons par mise à jour de l'optimiseur. Chaque nombre d'itérations natif est multiplié par les 8 étapes d'accumulation : 2000 itérations deviennent --steps=16000, un point de contrôle toutes les 500 itérations --save_freq=4000, et le warmup de 500 et la fin du déclin de 16000 deviennent 4000 et 128000.

accelerate launch --num_processes=8 $(which lerobot-train) \
  --policy.type=cosmos3_nano \
  --policy.pretrained_path=lerobot/Cosmos3-Nano \
  --policy.load_action_modules=false \
  --policy.gradient_checkpointing=true \
  --policy.repo_id=${HF_USER}/Cosmos3-Nano-Policy-LIBERO_10 \
  --dataset.repo_id=lerobot/LIBERO_10-Cosmos3 \
  --parallelism.dp_shard=8 \
  --accelerator.fsdp.mixed_precision_policy.param_dtype=bfloat16 \
  --accelerator.fsdp.mixed_precision_policy.reduce_dtype=bfloat16 \
  --accelerator.fsdp.mixed_precision_policy.cast_forward_inputs=false \
  --accelerator.gradient_accumulation.steps=8 \
  --accelerator.gradient_accumulation.sync_each_batch=true \
  --tf32=false \
  --ema.type=power \
  --ema.sigma_rel=0.1 \
  --batch_size=32 \
  --steps=16000 \
  --save_freq=4000 \
  --policy.scheduler_warmup_steps=4000 \
  --policy.scheduler_decay_steps=128000 \
  --output_dir=outputs/train/cosmos3_nano_libero_10 \
  --job_name=cosmos3_nano_libero_10 \
  --trackio.enable=true

Pour les quatre suites, utilisez --dataset.repo_id=lerobot/LIBERO-Cosmos3 et --steps=40000 (5000 itérations). Pour Cosmos3 Edge, utilisez --policy.type=cosmos3_edge --policy.pretrained_path=lerobot/Cosmos3-Edge.

Paramètres d'entraînement clés

ParamètreDescription
--policy.load_action_modules=falseDémarre les modules d'action depuis leur initialisation au lieu de les charger, comme le font les recettes natives de post-entraînement. Ne le définissez que lors de l'ajustement fin à partir d'un point de contrôle de base (voir les notes ci-dessous).
--parallelism.dp_shard=8FSDP2 sur les 8 GPU. La politique déclare ses unités d’emballage FSDP2, donc --accelerator.fsdp.wrap_modules n’est pas nécessaire. Consultez le guide d’entraînement multi-GPU.
--accelerator.fsdp.mixed_precision_policy.*Calcul et réduction des gradients en bfloat16 avec des poids maîtres en float32, comme dans l’entraînement natif. Conservez la valeur par défaut de --policy.dtype (float32) : une politique en bfloat16 entraînerait des poids maîtres en bfloat16.
--accelerator.gradient_accumulation.stepsLe lot global est --batch_size x GPU x étapes d'accumulation. sync_each_batch=true réduit les gradients après chaque micro-lot, de sorte qu'ils restent fragmentés et n'ont pas besoin d'une copie supplémentaire non fragmentée.
--steps, --save_freq, --policy.scheduler_*_stepsLeRobot compte les étapes de boucle (micro-lots), et non les mises à jour de l'optimiseur. Multipliez les nombres d'itérations de la recette native (2000 ou 5000 itérations, warmup 500, decay 16000, points de contrôle tous les 500) par les étapes d'accumulation. Gardez --steps au plus égal à --policy.scheduler_decay_steps : le taux d'apprentissage est nul après cela.
--ema.type=power --ema.sigma_rel=0.1L’EMA de la recette native. Chaque point de contrôle contient également les poids EMA dans pretrained_model_ema/, qui sont les poids à évaluer et à publier.
--tf32=falseConserve les multiplications de matrices en float32 en pleine précision, comme le fait l'entraînement natif.
--policy.gradient_checkpointing=trueRecalcule les activations des couches du décodeur et des blocs visuels lors de la rétropropagation afin d'économiser de la mémoire.
--policy.compile_model=trueFacultatif : torch.compile des couches du décodeur et de l'encodeur VAE.

Sur un GPU unique ou avec DDP, remplacez les drapeaux FSDP2 (--parallelism.dp_shard et --accelerator.fsdp.mixed_precision_policy.*) par --accelerator.mixed_precision=bf16. Par exemple, Cosmos3 Edge sur un GPU unique, avec 32 échantillons par étape et 64 étapes d'accumulation pour un lot global de 2048 :

lerobot-train \
  --policy.type=cosmos3_edge \
  --policy.pretrained_path=lerobot/Cosmos3-Edge \
  --policy.load_action_modules=false \
  --policy.gradient_checkpointing=true \
  --policy.repo_id=${HF_USER}/Cosmos3-Edge-Policy-LIBERO_10 \
  --dataset.repo_id=lerobot/LIBERO_10-Cosmos3 \
  --accelerator.mixed_precision=bf16 \
  --accelerator.gradient_accumulation.steps=64 \
  --tf32=false \
  --ema.type=power \
  --ema.sigma_rel=0.1 \
  --batch_size=32 \
  --steps=128000 \
  --save_freq=32000 \
  --policy.scheduler_warmup_steps=32000 \
  --policy.scheduler_decay_steps=1024000 \
  --output_dir=outputs/train/cosmos3_edge_libero_10

--policy.load_action_modules=false est réinitialisé à true une fois le point de contrôle chargé, de sorte que l’exécution sauvegardée recharge ses modules d’action entraînés. Deux cas conservent par erreur les modules d’action à leur initialisation :

  • Construire une politique sans --policy.pretrained_path ni --policy.path : rien n’est chargé, le drapeau reste donc false et une évaluation ou une reprise ultérieure de cette exécution ignore ses modules d’action entraînés.
  • PEFT avec les modules d'action en dehors de --peft.full_training_modules : les modules d'action ne s'entraînent alors jamais.

Définissez --policy.load_action_modules=false uniquement lors de l’ajustement fin à partir d’un point de contrôle de base. Avec PEFT, conservez les modules d’action dans --peft.full_training_modules (la valeur par défaut).

Exigences matérielles

Mesuré sur des GPU A800 80 Go avec les commandes de cette page et le gradient checkpointing :

ConfigurationGPULot par GPUMémoire maximale par GPU
Nano, fine-tuning complet (FSDP2)83238 Go
Nano, LoRA1464 Go
Edge, fine-tuning complet13247 Go
Edge, LoRA1821 Go
Évaluation en bfloat16, Nano et Edge110 environnements39 Go et 15 Go
  • Cosmos3 Nano nécessite des GPU de 80 Go, et au moins 4 d'entre eux pour un fine-tuning complet : avant que FSDP2 ne découpe le modèle, chaque GPU contient le modèle entier en float32 (environ 66 Go).
  • La commande LIBERO ci-dessus prend environ 45 secondes par mise à jour de l'optimiseur sur 8 GPU A800, soit environ 25 heures pour LIBERO-10.

Fine-tuning sur vos propres données

Cosmos3 prend la recette d'entraînement depuis la ligne de commande et le jeu de données : la politique ne contient aucune logique spécifique au robot. Tout jeu de données LeRobot fonctionne dès lors qu'il respecte les points ci-dessous.

Exigences relatives aux jeux de données

  • Caméras : une à trois fonctionnalités observation.images.*, dans l’ordre du fichier info.json du jeu de données. Chaque vue est redimensionnée à --policy.view_size (par défaut, la taille de la première caméra) et composée en une seule toile vidéo : une vue telle quelle, deux côte à côte, trois avec la première en haut et les deux autres côte à côte en dessous. Placez les caméras dans l’ordre de cette disposition, par exemple [third person, wrist] pour LIBERO et [wrist, exterior 1, exterior 2] pour DROID.
  • Afficher la description : --policy.view_description est la phrase du prompt qui décrit la disposition du canevas. Les valeurs par défaut décrivent la disposition LIBERO (valeurs par défaut de la classe) et la disposition DROID (fichier config.json de base). Rédigez la phrase correspondant à votre propre disposition : une description erronée n'est pas détectée et ne fait que affaiblir le prompt.
  • Représentation des actions : la politique apprend les actions du jeu de données telles qu’elles sont, sans étape de conversion. Convertissez les actions dans la représentation que vous souhaitez apprendre avant l’entraînement, comme le fait les données LIBERO avec les rotations rot6d. Si l’environnement s’attend à une autre représentation, convertissez les actions en retour dans les processeurs d’environnement (l’environnement LIBERO le fait pour Cosmos3).
  • Largeur d'action : au maximum --policy.max_action_dim (64) canaux. La largeur de sortie provient du jeu de données.
  • Taux d’images par seconde : fps provient des métadonnées du jeu de données et doit correspondre au taux de contrôle réel : il est inscrit dans le prompt et définit l’axe temporel des positions de la vidéo ainsi que les seuils des images inactives.
  • Statistiques : la normalisation d’action par défaut est QUANTILES, qui nécessite les statistiques q01 et q99 de l’action dans meta/stats.json (les jeux de données écrits par LeRobot les contiennent). La recette DROID native utilise IDENTITY à la place.
  • État : avec --policy.use_state=true, l'état actuel est placé devant les actions en tant qu'action d'état propre. observation.state doit alors avoir la même largeur et la même signification de canal que l'action. L'état et l'action sont normalisés avec leurs propres statistiques.
  • Langue : la task de chaque image est l'instruction du prompt.
  • Embodiment : --policy.domain_id sélectionne la ligne des projections d’action, parmi 32. Les recettes natives utilisent 5 pour LIBERO, 8 pour DROID et 30 pour RoboCasa. Lorsque les modules d’action démarrent depuis leur initialisation (--policy.load_action_modules=false), chaque ligne démarre non entraînée et n’importe quelle valeur fonctionne.
  • Chunk : --policy.chunk_size actions par chunk (16 pour LIBERO, 32 pour DROID) et --policy.n_action_steps actions exécutées par chunk.

Les champs laissés à None (fps, view_size) ainsi que les noms des canaux d’action proviennent du jeu de données. Tous les autres champs proviennent de la configuration, il faut donc les définir sur la ligne de commande ou dans un fichier --config_path.

Noms des canaux d’action et d’état. Le prompt indique combien de frames du chunk sont inactives (le robot ne bouge pas), et Cosmos3IdleFrameProcessorStep les compte à partir des noms de canaux d’action dans meta/info.json. Les noms de la forme [<prefix>_]<type>_<index> indiquent les types de canaux : pos, une rotation (rot6d, rot9d, axisangle, euler_xyz, quat_xyzw, quat_wxyz), gripper ou joint, avec un préfixe pour distinguer les bras (par exemple left_rot6d_0). Les données LIBERO utilisent pos_0..2, rot6d_0..5, gripper_0, la disposition DROID joint_0..6, gripper_0.

Ce nommage n’est pas obligatoire. Sans lui, la politique n’enregistre qu’un avertissement et le prompt ne compte aucune image inerte, ce qui correspond exactement à ce que la recette LIBERO observe de toute façon (ses épisodes ne contiennent pas d’images inertes selon la règle native). Suivez ce nommage uniquement si vous souhaitez obtenir le nombre d’images inertes de la recette DROID native.

Recette DROID

Le config.json des points de contrôle de base correspond à la recette de post-entraînement native de DROID : incarnation domain_id=8, blocs de 32 actions, l’état actuel en tant qu’action d’état et les images inactives comptées sur celui-ci (use_state=true, idle_state=true), aucune normalisation (IDENTITY), UniPC avec 4 étapes et une échelle de guidage sans classifieur de 3 à l’inférence, ainsi qu’un taux d’apprentissage de 8e-5 (le taux natif de 2e-4 multiplié par son facteur de 0,4) sans échauffement et avec un décroissement linéaire jusqu’à zéro à l’étape 100000. En démarrant à partir de celui-ci avec --policy.path, l’entraînement est effectué avec cette recette.

La recette native entraîne le modèle sur les épisodes réussis de nvidia/Cosmos3-DROID. Un jeu de données DROID au format Cosmos3 comporte :

  • Les trois caméras observation.images.wrist_left, observation.images.exterior_1_left et observation.images.exterior_2_left, dans cet ordre (la caméra de poignet en haut, les deux vues extérieures en dessous), 360x640 dans les données natives.
  • Actions de position articulaire 8D nommées joint_0 à joint_6 et gripper_0, et observation.state avec la même structure.
  • Un taux de rafraîchissement de 15 FPS.

La recette native entraîne avec un lot global de 8192 échantillons pendant 10000 itérations, avec un point de contrôle toutes les 1000 itérations. Sur une machine équipée de 8 GPU, 8 échantillons par GPU et 128 étapes d'accumulation donnent le même lot global (8 x 8 x 128 = 8192), et chaque nombre d'itérations est multiplié par 128 :

accelerate launch --num_processes=8 $(which lerobot-train) \
  --policy.path=lerobot/Cosmos3-Nano \
  --policy.load_action_modules=false \
  --policy.gradient_checkpointing=true \
  --policy.repo_id=${HF_USER}/Cosmos3-Nano-Policy-DROID \
  --dataset.repo_id=${HF_USER}/DROID-Cosmos3 \
  --parallelism.dp_shard=8 \
  --accelerator.fsdp.mixed_precision_policy.param_dtype=bfloat16 \
  --accelerator.fsdp.mixed_precision_policy.reduce_dtype=bfloat16 \
  --accelerator.fsdp.mixed_precision_policy.cast_forward_inputs=false \
  --accelerator.gradient_accumulation.steps=128 \
  --accelerator.gradient_accumulation.sync_each_batch=true \
  --tf32=false \
  --ema.type=power \
  --ema.sigma_rel=0.1 \
  --batch_size=8 \
  --steps=1280000 \
  --save_freq=128000 \
  --policy.scheduler_decay_steps=12800000 \
  --output_dir=outputs/train/cosmos3_nano_droid \
  --job_name=cosmos3_nano_droid \
  --trackio.enable=true

La recette DROID native a été exécutée sur 256 GPU. Sur une machine équipée de 8 GPU A800, une mise à jour de l’optimiseur de cette commande prend environ 33 minutes, si bien que la recette complète nécessite de nombreuses machines. Exécutez la même commande sur chaque machine avec les drapeaux multi-nœuds de accelerate launch (ou torchrun), effectuez le sharding au sein de chaque machine et la réplication entre les machines avec --parallelism.dp_replicate=<number of machines> --parallelism.dp_shard=8, et divisez --accelerator.gradient_accumulation.steps ainsi que les nombres d’étapes par le nombre de machines. Consultez HSDP dans le guide d’entraînement multi-GPU.

Un nouveau robot

Pour un nouveau robot, partez des valeurs par défaut de la classe avec --policy.type et définissez les champs de vos données sur la ligne de commande, par exemple :

lerobot-train \
  --policy.type=cosmos3_edge \
  --policy.pretrained_path=lerobot/Cosmos3-Edge \
  --policy.load_action_modules=false \
  --policy.domain_id=0 \
  --policy.chunk_size=32 \
  --policy.n_action_steps=32 \
  --policy.view_description='The left half shows the front camera; the right half shows the wrist-mounted camera.' \
  --dataset.repo_id=${HF_USER}/my_dataset \
  ...

Configuration du backbone

La configuration du backbone --policy.mot_backbone est une valeur unique en ligne de commande. Transmettez un dictionnaire YAML ne contenant que les champs à modifier, par exemple un modèle avec moins de couches pour un test rapide sans poids préentraînés :

--policy.type=cosmos3_edge --policy.mot_backbone='{text_config: {num_hidden_layers: 2}, vision_config: {num_hidden_layers: 1}}'

Les champs omis conservent les valeurs par défaut de la taille, ou les valeurs de la configuration chargée avec --policy.path. Écrivez les nombres flottants sous la forme 1.0e-6 (YAML interprète 1e-6 comme une chaîne de caractères) et fournissez en entier les champs dont la valeur est un dictionnaire, tels que rope_parameters. Les champs d’architecture doivent correspondre aux poids que vous chargez.

Fine-tuning avec PEFT (LoRA)

Cosmos3 prend en charge LoRA via l'intégration PEFT de LeRobot (pip install -e ".[cosmos3,training,peft]"). Par défaut, LoRA adapte les projections d'attention du chemin de génération (q_proj_gen, k_proj_gen, v_proj_gen, o_proj_gen), tandis que les modules d'action (action2llm, llm2action, action_modality_embed) sont entraînés en entier et sauvegardés avec l'adaptateur.

lerobot-train \
  --policy.type=cosmos3_nano \
  --policy.pretrained_path=lerobot/Cosmos3-Nano \
  --policy.load_action_modules=false \
  --policy.gradient_checkpointing=true \
  --policy.repo_id=${HF_USER}/Cosmos3-Nano-Policy-LIBERO_10-LoRA \
  --peft.method_type=lora \
  --peft.r=16 \
  --peft.lora_alpha=32 \
  --policy.optimizer_lr=5e-4 \
  --policy.optimizer_weight_decay=0 \
  --dataset.repo_id=lerobot/LIBERO_10-Cosmos3 \
  --accelerator.mixed_precision=bf16 \
  --tf32=false \
  --ema.type=power \
  --ema.sigma_rel=0.1 \
  --batch_size=4 \
  --steps=16000 \
  --output_dir=outputs/train/cosmos3_nano_libero_10_lora
  • --peft.lora_alpha=32, le taux d’apprentissage et le weight decay correspondent aux valeurs de la recette LoRA native de Cosmos3. Sans --peft.lora_alpha, PEFT utilise 8.
  • Avec --policy.load_action_modules=false, les modules d’action peuvent changer de forme (--policy.max_action_dim, --policy.num_embodiment_domains). Cela fonctionne uniquement pour l’affinage complet : le rechargement d’une exécution PEFT charge le modèle de base, qui nécessite les mêmes formes.
  • --policy.train_gen_stream_only et --policy.freeze_vision_encoder n’ont aucun effet avec PEFT : seuls les adaptateurs et les modules d’action sont entraînés.
  • N’ajoutez pas les projections d’action (action2llm, llm2action) à --peft.target_modules : ce sont des tables par incarnation lues directement par une multiplication matricielle groupée, de sorte que LoRA ne les modifierait pas. Elles s’entraînent déjà en entier.
  • Conservez le --policy.dtype par défaut float32 avec PEFT, afin que les modules d'action entraînés en entier conservent des poids maîtres en float32.
  • Les recettes natives ne définissent ni taille de lot LoRA ni calendrier. La commande conserve le calendrier LIBERO (décroissance linéaire jusqu'à zéro à l'étape 16000) avec 4 échantillons par étape.
  • Cosmos3 Nano avec LoRA s'entraîne sur un seul GPU de 80 Go avec le gradient checkpointing, environ 64 Go pour 4 échantillons par GPU (voir Exigences matérielles).

Évaluation

Évaluer un point de contrôle LIBERO sur les quatre suites :

lerobot-eval \
  --policy.path=lerobot/Cosmos3-Nano-Policy-LIBERO \
  --policy.dtype=bfloat16 \
  --tf32=false \
  --env.type=libero \
  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
  --env.observation_height=256 \
  --env.observation_width=256 \
  --eval.n_episodes=10 \
  --eval.batch_size=10 \
  --output_dir=outputs/eval/cosmos3_nano_libero
  • Pour un point de contrôle de votre propre exécution d'entraînement, passez les poids EMA : --policy.path=outputs/train/cosmos3_nano_libero_10/checkpoints/last/pretrained_model_ema.
  • --env.observation_height=256 --env.observation_width=256 rend les caméras à la résolution 256x256 des données d'entraînement.
  • L’environnement LIBERO ajoute automatiquement la conversion d’actions Cosmos3 à ses processeurs : les actions rot6d 10D sont converties en actions axis angle 7D de LIBERO. Lorsque vous exécutez la politique en dehors de lerobot-eval et d’un robot simulé, construisez également les processeurs d’environnement avec make_env_pre_post_processors(env_cfg, policy_cfg) depuis lerobot.envs.factory.
  • Consultez le guide LIBERO pour le protocole d'évaluation et les autres options.

Vidéo prédite

Cosmos3 prédit les futures images ainsi que les actions. Passez --policy.save_predicted_video=true pour décoder la vidéo prédite de chaque chunk avec le VAE et écrire pred_episode_*.mp4 à côté des vidéos eval_episode_*.mp4 rendues.

Cosmos3 Edge sur la tâche LIBERO-10 « mettre le bol noir dans le tiroir inférieur du meuble et le fermer ». À gauche : les observations en tierce personne et au poignet du déroulement. À droite : la vidéo prédite par la politique, rééchantillonnée toutes les 16 actions à partir de l’observation actuelle.

Résultats

LIBERO

ModèlesLibero-SpatialLibero-ObjectLibero-GoalLibero-10Moyenne
Cosmos3-Nano-Policy-LIBERO98.0% (490/500)99.4% (497/500)97.4% (487/500)96.8% (484/500)97.90%
Cosmos3-Edge-Policy-LIBERO98.6% (493/500)99.4% (497/500)97.2% (486/500)96.6% (483/500)97.95%
Cosmos3-Nano-Policy-LIBERO_10N/AN/AN/A97.6% (488/500)97.60%
Cosmos3-Edge-Policy-LIBERO_10N/AN/AN/A95.2% (476/500)95.20%
Cosmos-Nano-TechReportN/AN/AN/A97.4%97.40%

Taux de succès en boucle fermée sur LIBERO. Cosmos3-Nano-Policy-LIBERO et Cosmos3-Edge-Policy-LIBERO sont entraînés sur les quatre suites LIBERO (lerobot/Cosmos3-Nano-Policy-LIBERO, lerobot/Cosmos3-Edge-Policy-LIBERO). Cosmos3-Nano-Policy-LIBERO_10 et Cosmos3-Edge-Policy-LIBERO_10 sont entraînés uniquement sur LIBERO-10 (lerobot/Cosmos3-Nano-Policy-LIBERO_10, lerobot/Cosmos3-Edge-Policy-LIBERO_10), et la configuration d’entraînement de Cosmos3-Nano-Policy-LIBERO_10 suit la recette LIBERO-10 décrite dans le rapport technique. Cosmos-Nano-TechReport correspond au taux de succès sur LIBERO-10 de Cosmos3 Nano entraîné uniquement sur LIBERO-10, tel que rapporté dans le tableau 20 du rapport technique (initialisation MT, itération 2000, 500 rollouts).

Citation

@techreport{nvidia2026cosmos3,
  title  = {Cosmos 3: Omnimodal World Models for physical AI},
  author = {{NVIDIA Cosmos Team}},
  year   = {2026},
  url    = {https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf}
}

Licence

L'implémentation LeRobot suit la licence Apache 2.0 de LeRobot. Les poids Cosmos3 et le code natif sont publiés sous la licence OpenMDW-1.1.

Update on GitHub