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 é um modelo de ação de mundo para controle de robôs, construído sobre o Cosmos 3, a família de modelos de mundo omnimodais da NVIDIA. A partir das observações da câmera e de uma instrução em linguagem natural, a política prevê conjuntamente os quadros de vídeo futuros e o próximo conjunto de ações do robô. A implementação no LeRobot é adaptada da política de ação nativa do cosmos-framework e executa os modelos Cosmos3 Nano e Edge liberados, convertidos para o formato do LeRobot.

Visão geral do modelo

Cosmos 3 é um Mixture-of-Transformers (MoT). Cada camada do decodificador possui dois caminhos com suas próprias normalizações de camada, projeções de atenção e MLPs, e os dois caminhos se encontram apenas em uma autoatenção compartilhada:

  • O caminho de compreensão (o raciocinador) lê os tokens de texto do prompt e, opcionalmente, os tokens de imagem do ViT, com atenção causal.
  • O caminho de geração (o gerador) lê os tokens contínuos da subsequência de difusão: os latentes do VAE do Wan2.2 do vídeo e os tokens de ação. Ele realiza atenção bidirecional sobre ambos os caminhos e é treinado com flow matching.
Mixture-of-Transformers architecture of Cosmos 3

Arquitetura Mixture-of-Transformers (MoT) do Cosmos 3 (Figura 5 do relatório técnico). Esquerda: a subsequência autoregressiva (AR) passa pelo reasoner e a subsequência de difusão (DM) pelo generator, e as duas se encontram em uma atenção multimodal compartilhada. Direita: a máscara de atenção, causal para AR e completa para difusão.

Como política, o Cosmos3 é executado no modo de modelo de ação do mundo das receitas de ação nativas:

  • Entradas: as vistas da câmera são compostas em um único canvas de vídeo (uma vista como está, duas lado a lado, três com a primeira no topo e as outras duas abaixo dela). O prompt carrega a tarefa, a taxa de quadros, a resolução do canvas, uma descrição do layout do canvas e o número de quadros ociosos do chunk.
  • Training: a política codifica o clipe do frame atual e os chunk_size frames futuros com o VAE, adiciona ruído aos latentes futuros e ao chunk de ação e aprende a desruidar ambos. A perda é uma soma ponderada de uma perda de vídeo e uma perda de ação.
  • Inferência: o quadro de condição é limpo, o vídeo futuro e as ações são amostrados conjuntamente com o solver UniPC, e a política retorna o chunk de ações. O vídeo previsto pode ser salvo durante a avaliação (veja Avaliação).
  • Por padrão, apenas o caminho de geração, as camadas ao seu redor e os módulos de ação são treinados (train_gen_stream_only=true, freeze_vision_encoder=true), conforme as receitas nativas de pós-treinamento.

O LeRobot fornece os dois tamanhos lançados como dois tipos de política:

Cosmos3 NanoCosmos3 Edge
Tipo de políticacosmos3_nanocosmos3_edge
Parâmetros (treinados por padrão)16.5B (7.0B)4.6B (1.4B)
BackboneMoT inicializado a partir do Qwen3-VL, 36 camadas, tamanho oculto 4096MoT com um transformer denso de 2B, 28 camadas, tamanho oculto 2048
Tokenizernvidia/Cosmos3-Nanonvidia/Cosmos3-Edge
Checkpoint baselerobot/Cosmos3-Nanolerobot/Cosmos3-Edge

Requisitos de instalação

  1. Instale o LeRobot seguindo nosso Guia de Instalação.

  2. Instale as dependências do Cosmos3:

    pip install -e ".[cosmos3]"

    Se você instalou o LeRobot a partir do PyPI:

    pip install 'lerobot[cosmos3]'
  3. Adicione os extras do seu fluxo de trabalho: training para lerobot-train, libero para treinar ou avaliar no LIBERO e peft para LoRA. Por exemplo:

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

O Cosmos3 usa o backend de atenção de produto escalar do PyTorch (sdpa) por padrão. Se o seu ambiente tiver uma instalação compatível do flash_attn, você pode solicitar --policy.attn_implementation=flash_attention_2. Não o combine com --policy.compile_model=true: as camadas do decodificador compiladas são recompiladas para cada camada com flash attention e param com um erro de recompilação.

Checkpoints

CheckpointTamanhoDados de treinamentoUsar
lerobot/Cosmos3-NanoNanoLançamento de nvidia/Cosmos3-NanoBase para ajuste fino
lerobot/Cosmos3-EdgeBordaLançamento de nvidia/Cosmos3-EdgeBase para ajuste fino
lerobot/Cosmos3-Nano-Policy-LIBERONanolerobot/LIBERO-Cosmos3Avaliação LIBERO
lerobot/Cosmos3-Edge-Policy-LIBEROBordalerobot/LIBERO-Cosmos3Avaliação LIBERO
lerobot/Cosmos3-Nano-Policy-LIBERO_10Nanolerobot/LIBERO_10-Cosmos3Avaliação LIBERO-10
lerobot/Cosmos3-Edge-Policy-LIBERO_10Bordalerobot/LIBERO_10-Cosmos3Avaliação LIBERO-10
  • Os checkpoints base contêm os pesos liberados no formato LeRobot. Seus módulos de ação (action2llm, llm2action, action_modality_embed) ainda não foram treinados para um robô, portanto, inicie o fine-tuning com --policy.load_action_modules=false (veja Fine-tuning no LIBERO).
  • O config.json dos checkpoints base contém a receita nativa de pós-treino do DROID, e os padrões de classe dos tipos de política são a receita nativa do LIBERO.
  • Os checkpoints do LIBERO contêm os pesos EMA de sua execução de treinamento.
  • Os pesos são armazenados em bfloat16. config.json define dtype: float32 propositalmente: o treinamento com precisão mista a partir desses checkpoints mantém pesos mestres em float32, assim como o treinamento nativo. Para inferência e avaliação, passe --policy.dtype=bfloat16.

Existem duas maneiras de começar a partir de um checkpoint:

  • --policy.path=<checkpoint> carrega a configuração armazenada, os pesos e os processadores, e em seguida aplica a linha de comando por cima.
  • --policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint> (ou cosmos3_edge) carrega apenas os pesos e os processadores, e a configuração inicia a partir dos padrões da classe, a receita LIBERO.

Fine-tuning no LIBERO

Os padrões da classe correspondem à receita LIBERO do relatório técnico: blocos de 16 ações, embodiment domain_id=5, sem ação de estado, normalização por quantil das ações, UniPC com 30 passos e sem classifier-free guidance na inferência, e uma taxa de aprendizado de 5e-5 com 500 passos de aquecimento e um decaimento linear até zero no passo 16000. Partindo de um checkpoint base com --policy.type, portanto, é necessário apenas --policy.load_action_modules=false.

A receita treina nos dados do LIBERO convertidos para o layout do Cosmos3:

Ambos provêm de nvidia/LIBERO_LeRobot_v3, os dados da receita nativa. As ações são 10D [delta position (3), rot6d (6), gripper] em vez das ações de 7D axis angle do LIBERO, as duas câmeras observation.images.image (terceira pessoa) e observation.images.image2 (punho) são 256x256 a 20 FPS, e as estatísticas das ações são as da receita nativa. Durante a avaliação, o ambiente LIBERO converte as ações de volta (veja Evaluation).

Comando de treinamento

A receita nativa do LIBERO treina com um lote global de 2048 amostras por 2000 iterações no LIBERO-10 (o melhor checkpoint está próximo da iteração 1500) e por 5000 iterações nos quatro conjuntos (o melhor próximo de 4500). O comando abaixo reproduz isso em uma máquina com 8 GPUs: FSDP2 sobre as 8 GPUs, 32 amostras por GPU e 8 passos de acúmulo de gradiente, resultando em 32 x 8 x 8 = 2048 amostras por atualização do otimizador. Cada contagem de iterações nativa é multiplicada pelos 8 passos de acúmulo: 2000 iterações se tornam --steps=16000, um checkpoint a cada 500 iterações --save_freq=4000, e o aquecimento de 500 e o fim do decaimento de 16000 se tornam 4000 e 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

Para as quatro suítes, use --dataset.repo_id=lerobot/LIBERO-Cosmos3 e --steps=40000 (5000 iterações). Para o Cosmos3 Edge, use --policy.type=cosmos3_edge --policy.pretrained_path=lerobot/Cosmos3-Edge.

Principais parâmetros de treinamento

ParâmetroDescrição
--policy.load_action_modules=falseInicia os módulos de ação a partir da sua inicialização, em vez de carregá-los, como fazem as receitas nativas de pós-treino. Defina-o apenas ao fazer fine-tuning a partir de um checkpoint base (veja as observações abaixo).
--parallelism.dp_shard=8FSDP2 sobre as 8 GPUs. A política declara suas unidades de wrap do FSDP2, portanto --accelerator.fsdp.wrap_modules não é necessário. Consulte o guia de treinamento multi GPU.
--accelerator.fsdp.mixed_precision_policy.*cálculo e redução de gradientes em bfloat16 com pesos mestres em float32, como no treinamento nativo. Mantenha o padrão --policy.dtype (float32) com ele: uma política em bfloat16 treinaria pesos mestres em bfloat16.
--accelerator.gradient_accumulation.stepsO lote global é --batch_size x GPUs x etapas de acúmulo. sync_each_batch=true reduz os gradientes após cada micro lote, de modo que eles permanecem fragmentados e não necessitam de uma cópia extra não fragmentada.
--steps, --save_freq, --policy.scheduler_*_stepsLeRobot conta passos de loop (micro lotes), não atualizações do otimizador. Multiplique as contagens de iterações da receita nativa (2000 ou 5000 iterações, warmup 500, decay 16000, checkpoints a cada 500) pelos passos de acúmulo. Mantenha --steps no máximo --policy.scheduler_decay_steps: a taxa de aprendizado é zero depois disso.
--ema.type=power --ema.sigma_rel=0.1O EMA da receita nativa. Cada checkpoint também contém os pesos do EMA em pretrained_model_ema/, que são os pesos a serem avaliados e publicados.
--tf32=falseMantém as multiplicações de matrizes em float32 em precisão total, assim como no treinamento nativo.
--policy.gradient_checkpointing=trueRecalcula as ativações das camadas do decodificador e dos blocos de visão na passagem reversa para economizar memória.
--policy.compile_model=trueOpcional: torch.compile das camadas do decodificador e do codificador do VAE.

Em uma única GPU ou com DDP, substitua as flags do FSDP2 (--parallelism.dp_shard e --accelerator.fsdp.mixed_precision_policy.*) por --accelerator.mixed_precision=bf16. Por exemplo, Cosmos3 Edge em uma GPU, com 32 amostras por passo e 64 passos de acúmulo para o mesmo lote 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 é redefinido para true assim que o checkpoint é carregado, de modo que a execução salva carrega novamente seus módulos de ação treinados. Dois casos mantêm os módulos de ação em sua inicialização por engano:

  • Ao construir uma política sem --policy.pretrained_path ou --policy.path: nada é carregado, portanto a flag permanece false e uma avaliação ou retomada posterior dessa execução ignora seus módulos de ação treinados.
  • PEFT com os módulos de ação fora de --peft.full_training_modules: os módulos de ação então nunca são treinados.

Defina --policy.load_action_modules=false apenas ao fazer fine-tuning a partir de um checkpoint base. Com PEFT, mantenha os action modules em --peft.full_training_modules (o padrão).

Requisitos de hardware

Medido em GPUs A800 de 80 GB com os comandos desta página e gradient checkpointing:

ConfiguraçãoGPUsLote por GPUMemória de pico por GPU
Nano, ajuste fino completo (FSDP2)83238 GB
Nano, LoRA1464 GB
Edge, ajuste fino completo13247 GB
Edge, LoRA1821 GB
Avaliação em bfloat16, Nano e Edge110 ambientes39 GB e 15 GB
  • Cosmos3 Nano requer GPUs de 80 GB, e pelo menos 4 delas para o fine-tuning completo: antes de o FSDP2 dividir o modelo, cada GPU mantém o modelo inteiro em float32 (cerca de 66 GB).
  • O comando LIBERO acima leva cerca de 45 segundos por atualização do otimizador em 8 GPUs A800, ou seja, cerca de 25 horas para o LIBERO-10.

Ajuste fino em seus próprios dados

Cosmos3 recebe a receita de treinamento da linha de comando e o conjunto de dados: a política não contém lógica específica de robô. Qualquer conjunto de dados LeRobot funciona, desde que siga os pontos abaixo.

Requisitos do conjunto de dados

  • Câmeras: de uma a três funcionalidades observation.images.*, na ordem do info.json do conjunto de dados. Cada visão é redimensionada para --policy.view_size (por padrão, o tamanho da primeira câmera) e composta em um único canvas de vídeo: uma visão como está, duas lado a lado, três com a primeira no topo e as outras duas lado a lado abaixo dela. Coloque as câmeras na ordem desse layout, por exemplo [third person, wrist] para LIBERO e [wrist, exterior 1, exterior 2] para DROID.
  • Ver descrição: --policy.view_description é a frase do prompt que descreve o layout do canvas. Os padrões descrevem o layout do LIBERO (padrões da classe) e o layout do DROID (base config.json). Escreva a frase do seu próprio layout: uma descrição incorreta não é detectada e apenas enfraquece o prompt.
  • Representação da ação: a política aprende as ações do dataset como estão, sem etapa de conversão. Converta as ações para a representação que deseja aprender antes do treinamento, como o dados do LIBERO faz com rotações rot6d. Se o ambiente esperar outra representação, converta as ações de volta nos processadores do ambiente (o ambiente LIBERO faz isso para o Cosmos3).
  • Largura da ação: no máximo --policy.max_action_dim (64) canais. A largura da saída vem do conjunto de dados.
  • Taxa de quadros: fps vem dos metadados do conjunto de dados e deve ser a taxa de controle real: ele é escrito no prompt e define o eixo temporal das posições do vídeo e os limiares dos quadros ociosos.
  • Estatísticas: a normalização padrão de ação é QUANTILES, que requer as estatísticas q01 e q99 da ação em meta/stats.json (conjuntos de dados gravados pelo LeRobot as possuem). A receita nativa DROID usa IDENTITY em vez disso.
  • Estado: com --policy.use_state=true, o estado atual é colocado à frente das ações como uma ação de estado limpo. observation.state então precisa ter a mesma largura e significado de canal que a ação. O estado e a ação são normalizados com suas próprias estatísticas.
  • Idioma: a task de cada quadro é a instrução do prompt.
  • Embodiment: --policy.domain_id seleciona a linha das projeções de ação, uma de 32. As receitas nativas usam 5 para LIBERO, 8 para DROID e 30 para RoboCasa. Quando os módulos de ação começam a partir da inicialização (--policy.load_action_modules=false), cada linha começa não treinada e qualquer valor funciona.
  • Chunk: --policy.chunk_size ações por chunk (16 para LIBERO, 32 para DROID) e --policy.n_action_steps ações executadas por chunk.

Os campos deixados como None (fps, view_size) e os nomes dos canais de ação vêm do dataset. Todos os outros campos vêm da configuração, portanto, defina-os na linha de comando ou em um arquivo --config_path.

Nomes dos canais de ação e estado. O prompt indica quantos frames do chunk estão ociosos (o robô não se move), e Cosmos3IdleFrameProcessorStep os conta a partir dos nomes dos canais de ação em meta/info.json. Nomes no formato [<prefix>_]<type>_<index> indicam os tipos de canal: pos, uma rotação (rot6d, rot9d, axisangle, euler_xyz, quat_xyzw, quat_wxyz), gripper ou joint, com um prefixo para distinguir os braços (por exemplo, left_rot6d_0). Os dados do LIBERO usam pos_0..2, rot6d_0..5, gripper_0, e o layout do DROID usa joint_0..6, gripper_0.

Esse nome não é obrigatório. Sem ele, a política apenas registra um aviso e a contagem de prompts não inclui nenhum frame ocioso, o que é exatamente o que a receita LIBERO observa de qualquer forma (seus episódios não possuem frames ociosos sob a regra nativa). Siga o nome apenas se você quiser a contagem de frames ociosos da receita DROID nativa.

Receita DROID

O config.json dos checkpoints base é a receita nativa de pós-treino do DROID: embodiment domain_id=8, chunks de 32 ações, o estado atual como ação de estado e frames ociosos contados sobre ele (use_state=true, idle_state=true), sem normalização (IDENTITY), UniPC com 4 passos e uma escala de classifier-free guidance de 3 na inferência, e uma taxa de aprendizado de 8e-5 (a taxa nativa de 2e-4 multiplicada por 0,4) sem warmup e com decaimento linear até zero no passo 100000. Começar a partir dele com --policy.path treina com essa receita.

A receita nativa treina nos episódios bem-sucedidos de nvidia/Cosmos3-DROID. Um conjunto de dados DROID no layout Cosmos3 tem:

  • As três câmeras observation.images.wrist_left, observation.images.exterior_1_left e observation.images.exterior_2_left, nesta ordem (a de pulso no topo e as duas vistas externas abaixo dela), com 360x640 nos dados nativos.
  • Ações de posição conjunta 8D nomeadas de joint_0 a joint_6 e gripper_0, e observation.state com o mesmo layout.
  • Uma taxa de quadros de 15 FPS.

A receita nativa treina com um lote global de 8192 amostras por 10000 iterações, com um checkpoint a cada 1000 iterações. Em uma máquina com 8 GPUs, 8 amostras por GPU e 128 passos de acúmulo resultam no mesmo lote global (8 x 8 x 128 = 8192), e a contagem de cada iteração é multiplicada por 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

A receita nativa DROID foi executada em 256 GPUs. Em uma máquina com 8 GPUs A800, uma atualização do otimizador deste comando leva cerca de 33 minutos, portanto, a receita completa requer muitas máquinas. Execute o mesmo comando em cada máquina com as flags de multi nó do accelerate launch (ou torchrun), faça o sharding dentro de cada máquina e a replicação entre máquinas com --parallelism.dp_replicate=<number of machines> --parallelism.dp_shard=8, e divida --accelerator.gradient_accumulation.steps e as contagens de passos pelo número de máquinas. Consulte HSDP no guia de treinamento multi GPU.

Um novo robô

Para um novo robô, comece pelos padrões da classe com --policy.type e defina os campos dos seus dados na linha de comando, por exemplo:

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 \
  ...

Configuração do backbone

A configuração do backbone --policy.mot_backbone é um único valor de linha de comando. Passe um dicionário YAML com apenas os campos a serem alterados, por exemplo, um modelo com menos camadas para um teste rápido sem pesos pré-treinados:

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

Os campos omitidos mantêm os padrões de tamanho ou os valores da configuração carregada com --policy.path. Escreva números decimais como 1.0e-6 (o YAML lê 1e-6 como uma string) e forneça por completo os campos cujo valor é um dicionário, como rope_parameters. Os campos de arquitetura devem corresponder aos pesos que você carrega.

Fine-tuning com PEFT (LoRA)

Cosmos3 suporta LoRA por meio da integração PEFT do LeRobot (pip install -e ".[cosmos3,training,peft]"). Por padrão, o LoRA adapta as projeções de atenção do caminho de geração (q_proj_gen, k_proj_gen, v_proj_gen, o_proj_gen), e os módulos de ação (action2llm, llm2action, action_modality_embed) são treinados por completo e salvos junto com o adaptador.

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, a taxa de aprendizado e o decaimento de peso são os valores da receita nativa de LoRA do Cosmos3. Sem --peft.lora_alpha, o PEFT usa 8.
  • Com --policy.load_action_modules=false, os módulos de ação podem ter formas diferentes (--policy.max_action_dim, --policy.num_embodiment_domains). Isso funciona apenas para fine-tuning completo: ao recarregar uma execução PEFT, o modelo base é carregado, o que requer as mesmas formas.
  • --policy.train_gen_stream_only e --policy.freeze_vision_encoder não têm efeito com PEFT: apenas os adaptadores e os módulos de ação são treinados.
  • Não adicione as projeções de ação (action2llm, llm2action) a --peft.target_modules: elas são tabelas por embodiment lidas diretamente por uma multiplicação de matriz agrupada, portanto o LoRA não as alteraria. Elas já são treinadas por completo.
  • Mantenha o padrão float32 --policy.dtype com PEFT, para que os módulos de ação que treinam por completo mantenham pesos mestres em float32.
  • As receitas nativas não definem tamanho de lote ou agendamento para LoRA. O comando mantém o agendamento LIBERO (decaimento linear até zero no passo 16000) com 4 amostras por passo.
  • Cosmos3 Nano com LoRA treina em uma GPU de 80 GB com gradient checkpointing, cerca de 64 GB com 4 amostras por GPU (veja Requisitos de hardware).

Avaliação

Avalie um checkpoint do LIBERO nos quatro conjuntos:

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
  • Para um checkpoint de sua própria execução de treinamento, passe os pesos EMA: --policy.path=outputs/train/cosmos3_nano_libero_10/checkpoints/last/pretrained_model_ema.
  • --env.observation_height=256 --env.observation_width=256 renderiza as câmeras na resolução de 256x256 dos dados de treinamento.
  • O ambiente LIBERO adiciona automaticamente a conversão de ações do Cosmos3 aos seus processadores: as ações rot6d de 10D são convertidas de volta para as ações axis angle de 7D do LIBERO. Ao executar a política fora do lerobot-eval e de um robô simulado, construa também os processadores do ambiente com make_env_pre_post_processors(env_cfg, policy_cfg) de lerobot.envs.factory.
  • Consulte o guia do LIBERO para o protocolo de avaliação e as outras opções.

Vídeo previsto

Cosmos3 prevê os quadros futuros junto com as ações. Passe --policy.save_predicted_video=true para decodificar o vídeo previsto de cada chunk com o VAE e gravar pred_episode_*.mp4 ao lado dos vídeos renderizados eval_episode_*.mp4.

Cosmos3 Edge na tarefa LIBERO-10 “coloque a tigela preta na gaveta inferior do armário e feche-a”. Esquerda: as observações em terceira pessoa e do pulso do rollout. Direita: o vídeo previsto pela política, amostrado novamente a cada 16 ações a partir da observação atual.

Resultados

LIBERO

ModelosLibero-SpatialLibero-ObjectLibero-GoalLibero-10Média
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%

Taxas de sucesso em loop fechado no LIBERO. Cosmos3-Nano-Policy-LIBERO e Cosmos3-Edge-Policy-LIBERO são treinados nas quatro suítes do LIBERO (lerobot/Cosmos3-Nano-Policy-LIBERO, lerobot/Cosmos3-Edge-Policy-LIBERO). Cosmos3-Nano-Policy-LIBERO_10 e Cosmos3-Edge-Policy-LIBERO_10 são treinados apenas no LIBERO-10 (lerobot/Cosmos3-Nano-Policy-LIBERO_10, lerobot/Cosmos3-Edge-Policy-LIBERO_10), e a configuração de treinamento do Cosmos3-Nano-Policy-LIBERO_10 segue a receita do LIBERO-10 relatada no relatório técnico. Cosmos-Nano-TechReport é a taxa de sucesso no LIBERO-10 do Cosmos3 Nano treinado apenas no LIBERO-10, conforme relatado na Tabela 20 do relatório técnico (MT-init, iteração 2000, 500 rollouts).

Citação

@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}
}

Licença

A implementação do LeRobot segue a Apache 2.0 License do LeRobot. Os pesos do Cosmos3 e o código nativo são liberados sob a OpenMDW-1.1 License.

Update on GitHub