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 es un modelo de acción mundial para el control de robots, construido sobre Cosmos 3, la familia de modelos de mundo omnimodales de NVIDIA. A partir de las observaciones de la cámara y una instrucción en lenguaje, la política predice conjuntamente los futuros fotogramas de video y el siguiente bloque de acciones del robot. La implementación de LeRobot se adapta de la política de acción nativa de cosmos-framework y ejecuta los modelos Cosmos3 Nano y Edge publicados, convertidos al formato de LeRobot.

Resumen del modelo

Cosmos 3 es una Mixture-of-Transformers (MoT). Cada capa del decodificador tiene dos vías con sus propias normalizaciones de capa, proyecciones de atención y MLPs, y las dos vías solo se unen en una autoatención compartida:

  • La vía de comprensión (el razonador) lee los tokens de texto del prompt y, opcionalmente, los tokens de imagen de ViT, con atención causal.
  • La vía de generación (el generador) lee los tokens continuos del subsecuencia de difusión: los latentes del VAE de Wan2.2 del video y los tokens de acción. Presta atención de forma bidireccional sobre ambas vías y se entrena con flow matching.
Mixture-of-Transformers architecture of Cosmos 3

Arquitectura Mixture-of-Transformers (MoT) de Cosmos 3 (Figura 5 del informe técnico). Izquierda: la subsecuencia autoregresiva (AR) pasa por el razonador y la subsecuencia de difusión (DM) por el generador, y ambas se encuentran en una atención multimodal compartida. Derecha: la máscara de atención, causal para AR y completa para difusión.

Como política, Cosmos3 se ejecuta en el modo de modelo de acción del mundo de las recetas de acción nativas:

  • Entradas: las vistas de la cámara se componen en un único lienzo de video (una vista tal cual, dos lado a lado, tres con la primera arriba y las otras dos debajo). El prompt contiene la tarea, la tasa de fotogramas, la resolución del lienzo, una descripción del diseño del lienzo y el número de fotogramas inactivos del bloque.
  • Entrenamiento: la política codifica el clip del cuadro actual y los chunk_size cuadros futuros con el VAE, añade ruido a los latentes futuros y al trozo de acción, y aprende a desruidar ambos. La pérdida es una suma ponderada de una pérdida de vídeo y una pérdida de acción.
  • Inferencia: el marco de condición es limpio, el video futuro y las acciones se muestrean conjuntamente con el solver UniPC, y la política devuelve el trozo de acciones. El video predicho puede guardarse durante la evaluación (ver Evaluación).
  • Por defecto, solo se entrenan la ruta de generación, las capas que la rodean y los módulos de acción (train_gen_stream_only=true, freeze_vision_encoder=true), tal como en las recetas de entrenamiento posterior nativas.

LeRobot ofrece los dos tamaños publicados como dos tipos de política:

Cosmos3 NanoCosmos3 Edge
Tipo de políticacosmos3_nanocosmos3_edge
Parámetros (entrenados por defecto)16.5B (7.0B)4.6B (1.4B)
BackboneMoT inicializado desde Qwen3-VL, 36 capas, tamaño oculto 4096MoT con un transformer denso de 2B, 28 capas, tamaño oculto 2048
Tokenizernvidia/Cosmos3-Nanonvidia/Cosmos3-Edge
Punto de control baselerobot/Cosmos3-Nanolerobot/Cosmos3-Edge

Requisitos de instalación

  1. Instala LeRobot siguiendo nuestra Guía de instalación.

  2. Instale las dependencias de Cosmos3:

    pip install -e ".[cosmos3]"

    Si instalaste LeRobot desde PyPI:

    pip install 'lerobot[cosmos3]'
  3. Agrega los extras de tu flujo de trabajo: training para lerobot-train, libero para entrenar o evaluar en LIBERO y peft para LoRA. Por ejemplo:

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

Cosmos3 utiliza el backend de atención de producto punto escalado de PyTorch (sdpa) de forma predeterminada. Si su entorno tiene una instalación compatible de flash_attn, puede solicitar --policy.attn_implementation=flash_attention_2. No lo combine con --policy.compile_model=true: las capas del decodificador compiladas se recompilan para cada capa con flash attention y se detienen con un error de recompilación.

Puntos de control

CheckpointTamañoDatos de entrenamientoUsar
lerobot/Cosmos3-NanoNanoLanzamiento de nvidia/Cosmos3-NanoBase para el ajuste fino
lerobot/Cosmos3-EdgeEdgeLanzamiento de nvidia/Cosmos3-EdgeBase para el ajuste fino
lerobot/Cosmos3-Nano-Policy-LIBERONanolerobot/LIBERO-Cosmos3Evaluación de LIBERO
lerobot/Cosmos3-Edge-Policy-LIBEROEdgelerobot/LIBERO-Cosmos3Evaluación de LIBERO
lerobot/Cosmos3-Nano-Policy-LIBERO_10Nanolerobot/LIBERO_10-Cosmos3Evaluación de LIBERO-10
lerobot/Cosmos3-Edge-Policy-LIBERO_10Edgelerobot/LIBERO_10-Cosmos3Evaluación de LIBERO-10
  • Los puntos de control base contienen los pesos publicados en el formato LeRobot. Sus módulos de acción (action2llm, llm2action, action_modality_embed) aún no se han entrenado para un robot, por lo que comience el ajuste fino con --policy.load_action_modules=false (consulte Ajuste fino en LIBERO).
  • El config.json de los puntos de control base contiene la receta nativa de posentrenamiento de DROID, y los valores predeterminados de la clase de los tipos de política son la receta nativa de LIBERO.
  • Los puntos de control de LIBERO contienen los pesos EMA de su ejecución de entrenamiento.
  • Los pesos se almacenan en bfloat16. config.json establece dtype: float32 a propósito: el entrenamiento con precisión mixta desde estos puntos de control mantiene los pesos maestros en float32, al igual que el entrenamiento nativo. Para la inferencia y la evaluación, pase --policy.dtype=bfloat16.

Hay dos formas de comenzar desde un punto de control:

  • --policy.path=<checkpoint> carga la configuración almacenada, los pesos y los procesadores, y luego aplica la línea de comandos encima.
  • --policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint> (o cosmos3_edge) carga solo los pesos y los procesadores, y la configuración parte de los valores predeterminados de la clase, la receta LIBERO.

Ajuste fino en LIBERO

Los valores predeterminados de la clase son la receta LIBERO del informe técnico: bloques de 16 acciones, encarnación domain_id=5, sin acción de estado, normalización cuantílica de las acciones, UniPC con 30 pasos y sin guía libre de clasificador en la inferencia, y una tasa de aprendizaje de 5e-5 con 500 pasos de calentamiento y un decaimiento lineal hasta cero en el paso 16000. Por lo tanto, al partir de un punto de control base con --policy.type solo se necesita --policy.load_action_modules=false.

La receta entrena con los datos de LIBERO convertidos al formato de Cosmos3:

Ambos provienen de nvidia/LIBERO_LeRobot_v3, los datos de la receta nativa. Las acciones son de 10D [delta position (3), rot6d (6), gripper] en lugar de las acciones de 7D axis angle de LIBERO, las dos cámaras observation.images.image (tercera persona) y observation.images.image2 (muñeca) son de 256x256 a 20 FPS, y las estadísticas de las acciones son las de la receta nativa. Durante la evaluación, el entorno LIBERO convierte las acciones de vuelta (ver Evaluation).

Comando de entrenamiento

La receta nativa de LIBERO entrena con un lote global de 2048 muestras durante 2000 iteraciones en LIBERO-10 (el mejor checkpoint está cerca de la iteración 1500) y durante 5000 iteraciones en los cuatro conjuntos (el mejor cerca de 4500). El comando siguiente lo reproduce en una máquina con 8 GPUs: FSDP2 sobre las 8 GPUs, 32 muestras por GPU y 8 pasos de acumulación de gradientes, por lo que 32 x 8 x 8 = 2048 muestras por actualización del optimizador. Cada número de iteraciones nativo se multiplica por los 8 pasos de acumulación: 2000 iteraciones se convierten en --steps=16000, un checkpoint cada 500 iteraciones --save_freq=4000, y el warmup de 500 y el fin del decaimiento de 16000 se convierten en 4000 y 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 los cuatro conjuntos de datos, use --dataset.repo_id=lerobot/LIBERO-Cosmos3 y --steps=40000 (5000 iteraciones). Para Cosmos3 Edge, use --policy.type=cosmos3_edge --policy.pretrained_path=lerobot/Cosmos3-Edge.

Parámetros clave de entrenamiento

ParámetroDescripción
--policy.load_action_modules=falseInicia los módulos de acción desde su inicialización en lugar de cargarlos, tal como lo hacen las recetas nativas de posentrenamiento. Establézcalo solo al realizar un ajuste fino desde un punto de control base (consulte las notas a continuación).
--parallelism.dp_shard=8FSDP2 en las 8 GPU. La política declara sus unidades de envoltura de FSDP2, por lo que no se necesita --accelerator.fsdp.wrap_modules. Consulte la guía de entrenamiento multi GPU.
--accelerator.fsdp.mixed_precision_policy.*cálculo y reducción de gradientes en bfloat16 con pesos maestros en float32, como en el entrenamiento nativo. Mantenga el valor predeterminado de --policy.dtype (float32): una política en bfloat16 entrenaría pesos maestros en bfloat16.
--accelerator.gradient_accumulation.stepsEl lote global es --batch_size x GPUs x pasos de acumulación. sync_each_batch=true reduce los gradientes después de cada micro lote, de modo que permanecen fragmentados y no requieren una copia adicional no fragmentada.
--steps, --save_freq, --policy.scheduler_*_stepsLeRobot cuenta los pasos del bucle (micro lotes), no las actualizaciones del optimizador. Multiplique los recuentos de iteraciones de la receta nativa (2000 o 5000 iteraciones, warmup 500, decay 16000, checkpoints cada 500) por los pasos de acumulación. Mantenga --steps como máximo en --policy.scheduler_decay_steps: la tasa de aprendizaje es cero después de eso.
--ema.type=power --ema.sigma_rel=0.1La EMA de la receta nativa. Cada punto de control también contiene los pesos de la EMA en pretrained_model_ema/, que son los pesos que se deben evaluar y publicar.
--tf32=falseMantiene las multiplicaciones de matrices en float32 con precisión completa, tal como lo hace el entrenamiento nativo.
--policy.gradient_checkpointing=trueRecalcula las activaciones de las capas del decodificador y de los bloques visuales en la pasada inversa para ahorrar memoria.
--policy.compile_model=trueOpcional: torch.compile de las capas del decodificador y del codificador VAE.

En una sola GPU o con DDP, reemplace las banderas de FSDP2 (--parallelism.dp_shard y --accelerator.fsdp.mixed_precision_policy.*) por --accelerator.mixed_precision=bf16. Por ejemplo, Cosmos3 Edge en una GPU, con 32 muestras por paso y 64 pasos de acumulación para el mismo 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 se restablece a true una vez que se carga el checkpoint, por lo que la ejecución guardada vuelve a cargar sus módulos de acción entrenados. Dos casos mantienen los módulos de acción en su inicialización por error:

  • Construir una política sin --policy.pretrained_path ni --policy.path: no se carga nada, por lo que la bandera permanece en false y una evaluación o reanudación posterior de esa ejecución omite sus módulos de acción entrenados.
  • PEFT con los módulos de acción fuera de --peft.full_training_modules: los módulos de acción nunca se entrenan.

Establezca --policy.load_action_modules=false solo al realizar fine-tuning desde un checkpoint base. Con PEFT, mantenga los action modules en --peft.full_training_modules (el valor predeterminado).

Requisitos de hardware

Medido en GPUs A800 de 80 GB con los comandos de esta página y gradient checkpointing:

ConfiguraciónGPUsLote por GPUMemoria máxima por GPU
Nano, ajuste fino completo (FSDP2)83238 GB
Nano, LoRA1464 GB
Edge, ajuste fino completo13247 GB
Edge, LoRA1821 GB
Evaluación en bfloat16, Nano y Edge110 entornos39 GB y 15 GB
  • Cosmos3 Nano requiere GPUs de 80 GB y, al menos, 4 de ellas para el ajuste fino completo: antes de que FSDP2 divida el modelo, cada GPU almacena el modelo completo en float32 (aproximadamente 66 GB).
  • El comando LIBERO anterior tarda aproximadamente 45 segundos por actualización del optimizador en 8 GPUs A800, por lo que toma unas 25 horas para LIBERO-10.

Ajuste fino con sus propios datos

Cosmos3 toma la receta de entrenamiento desde la línea de comandos y el conjunto de datos: la política no contiene lógica específica de robots. Cualquier conjunto de datos de LeRobot funciona una vez que cumple con los puntos a continuación.

Requisitos del conjunto de datos

  • Cámaras: de una a tres características observation.images.*, en el orden del info.json del conjunto de datos. Cada vista se redimensiona a --policy.view_size (por defecto, el tamaño de la primera cámara) y se compone en un único lienzo de video: una vista tal cual, dos lado a lado, y tres con la primera arriba y las otras dos lado a lado debajo. Coloca las cámaras en el orden de ese diseño, por ejemplo [third person, wrist] para LIBERO y [wrist, exterior 1, exterior 2] para DROID.
  • Ver descripción: --policy.view_description es la oración del prompt que describe el diseño del lienzo. Los valores predeterminados describen el diseño de LIBERO (valores predeterminados de la clase) y el diseño de DROID (archivo base config.json). Escriba la oración de su propio diseño: una descripción incorrecta no se detecta y solo debilita el prompt.
  • Representación de la acción: la política aprende las acciones del tal cual son y no tiene un paso de conversión. Convierta las acciones a la representación que desee aprender antes del entrenamiento, como hace los datos de LIBERO con las rotaciones rot6d. Si el entorno espera otra representación, convierta las acciones de vuelta en los procesadores del entorno (el entorno de LIBERO hace esto para Cosmos3).
  • Ancho de la acción: como máximo --policy.max_action_dim (64) canales. El ancho de la salida proviene del conjunto de datos.
  • Frecuencia de fotogramas: fps proviene de los metadatos del conjunto de datos y debe ser la tasa de control real: se escribe en el prompt y establece el eje temporal de las posiciones del video y los umbrales de los fotogramas inactivos.
  • Estadísticas: la normalización de acciones predeterminada es QUANTILES, que requiere las estadísticas q01 y q99 de la acción en meta/stats.json (los conjuntos de datos escritos por LeRobot las incluyen). La receta nativa de DROID utiliza IDENTITY en su lugar.
  • Estado: con --policy.use_state=true el estado actual se coloca delante de las acciones como una acción de estado limpia. observation.state entonces necesita el mismo ancho y significado de canal que la acción. El estado y la acción se normalizan con sus propias estadísticas.
  • Idioma: la task de cada fotograma es la instrucción del prompt.
  • Embodiment: --policy.domain_id selecciona la fila de las proyecciones de acción, una de 32. Las recetas nativas usan 5 para LIBERO, 8 para DROID y 30 para RoboCasa. Cuando los módulos de acción parten de su inicialización (--policy.load_action_modules=false), cada fila comienza sin entrenar y cualquier valor es válido.
  • Chunk: --policy.chunk_size acciones por chunk (16 para LIBERO, 32 para DROID) y --policy.n_action_steps acciones ejecutadas por chunk.

Los campos que se dejan en None (fps, view_size) y los nombres de los canales de acción provienen del conjunto de datos. Todos los demás campos provienen de la configuración, por lo que debes establecerlos en la línea de comandos o en un archivo --config_path.

Nombres de los canales de acción y estado. El prompt indica cuántos fotogramas del bloque están inactivos (el robot no se mueve), y Cosmos3IdleFrameProcessorStep los cuenta a partir de los nombres de los canales de acción en meta/info.json. Los nombres de la forma [<prefix>_]<type>_<index> indican los tipos de canal: pos, una rotación (rot6d, rot9d, axisangle, euler_xyz, quat_xyzw, quat_wxyz), gripper o joint, con un prefijo para distinguir los brazos (por ejemplo, left_rot6d_0). Los datos de LIBERO usan pos_0..2, rot6d_0..5, gripper_0, y el esquema de DROID usa joint_0..6, gripper_0.

Este nombre no es obligatorio. Sin él, la política solo registra una advertencia y el prompt no cuenta ningún fotograma inactivo, que es exactamente lo que ve la receta de LIBERO de todos modos (sus episodios no tienen fotogramas inactivos bajo la regla nativa). Siga el nombre solo si desea la cantidad de fotogramas inactivos de la receta nativa de DROID.

Receta DROID

El config.json de los puntos de control base es la receta nativa de post-entrenamiento de DROID: encarnación domain_id=8, bloques de 32 acciones, el estado actual como acción de estado y los fotogramas inactivos contados sobre él (use_state=true, idle_state=true), sin normalización (IDENTITY), UniPC con 4 pasos y una escala de guía libre de clasificador de 3 en la inferencia, y una tasa de aprendizaje de 8e-5 (la nativa 2e-4 multiplicada por su factor de 0.4) sin calentamiento y con un decaimiento lineal hasta cero en el paso 100000. Empezar desde él con --policy.path entrena con esta receta.

La receta nativa entrena en los episodios exitosos de nvidia/Cosmos3-DROID. Un conjunto de datos DROID en el formato Cosmos3 tiene:

  • Las tres cámaras observation.images.wrist_left, observation.images.exterior_1_left y observation.images.exterior_2_left, en este orden (la muñeca arriba y las dos vistas exteriores debajo), con una resolución nativa de 360x640 en los datos.
  • Acciones de posición conjunta 8D denominadas joint_0 a joint_6 y gripper_0, y observation.state con el mismo diseño.
  • Una tasa de fotogramas de 15 FPS.

La receta nativa entrena con un lote global de 8192 muestras durante 10000 iteraciones, con un punto de control cada 1000 iteraciones. En una máquina con 8 GPUs, 8 muestras por GPU y 128 pasos de acumulación producen el mismo lote global (8 x 8 x 128 = 8192), y cada recuento de iteraciones se multiplica 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

La receta nativa de DROID se ejecutó en 256 GPUs. En una máquina con 8 GPUs A800, una actualización del optimizador de este comando tarda aproximadamente 33 minutos, por lo que la receta completa requiere muchas máquinas. Ejecute el mismo comando en cada máquina con las banderas de multi nodo de accelerate launch (o torchrun), fragmente dentro de cada máquina y replique entre máquinas con --parallelism.dp_replicate=<number of machines> --parallelism.dp_shard=8, y divida --accelerator.gradient_accumulation.steps y los recuentos de pasos entre el número de máquinas. Consulte HSDP en la guía de entrenamiento multi GPU.

Un robot nuevo

Para un robot nuevo, comience con los valores predeterminados de la clase con --policy.type y establezca los campos de sus datos en la línea de comandos, por ejemplo:

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

Configuración del backbone

La configuración del backbone --policy.mot_backbone es un único valor de línea de comandos. Pasa un diccionario YAML con solo los campos que deseas cambiar, por ejemplo, un modelo con menos capas para una prueba rápida sin pesos preentrenados:

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

Los campos omitidos conservan los valores predeterminados del tamaño o los valores de la configuración cargada con --policy.path. Escribe los decimales como 1.0e-6 (YAML interpreta 1e-6 como una cadena de texto) y proporciona por completo los campos cuyo valor es un diccionario, como rope_parameters. Los campos de arquitectura deben coincidir con los pesos que cargues.

Ajuste fino con PEFT (LoRA)

Cosmos3 admite LoRA a través de la integración de PEFT de LeRobot (pip install -e ".[cosmos3,training,peft]"). De forma predeterminada, LoRA adapta las proyecciones de atención de la ruta de generación (q_proj_gen, k_proj_gen, v_proj_gen, o_proj_gen), y los módulos de acción (action2llm, llm2action, action_modality_embed) se entrenan por completo y se guardan junto con el 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, la tasa de aprendizaje y el decaimiento de peso son los valores de la receta nativa de LoRA de Cosmos3. Sin --peft.lora_alpha, PEFT utiliza 8.
  • Con --policy.load_action_modules=false, los módulos de acción pueden cambiar de forma (--policy.max_action_dim, --policy.num_embodiment_domains). Esto solo funciona para el ajuste fino completo: al recargar una ejecución de PEFT se carga la base, que necesita las mismas formas.
  • --policy.train_gen_stream_only y --policy.freeze_vision_encoder no tienen efecto con PEFT: solo se entrenan los adaptadores y los módulos de acción.
  • No añada las proyecciones de acción (action2llm, llm2action) a --peft.target_modules: son tablas por encarnación que se leen directamente mediante una multiplicación de matrices agrupada, por lo que LoRA no las modificaría. Ya se entrenan por completo.
  • Mantenga el --policy.dtype predeterminado de float32 con PEFT, para que los módulos de acción que se entrenan por completo conserven los pesos maestros en float32.
  • Las recetas nativas no definen el tamaño de lote ni el calendario de LoRA. El comando mantiene el calendario de LIBERO (decaimiento lineal hasta cero en el paso 16000) con 4 muestras por paso.
  • Cosmos3 Nano con LoRA se entrena en una GPU de 80 GB con gradient checkpointing, aproximadamente 64 GB con 4 muestras por GPU (ver Requisitos de hardware).

Evaluación

Evalúe un punto de control de LIBERO en los cuatro 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 un punto de control de tu propia ejecución de entrenamiento, pasa los pesos EMA: --policy.path=outputs/train/cosmos3_nano_libero_10/checkpoints/last/pretrained_model_ema.
  • --env.observation_height=256 --env.observation_width=256 renderiza las cámaras con la resolución de 256x256 de los datos de entrenamiento.
  • El entorno LIBERO añade automáticamente la conversión de acciones de Cosmos3 a sus procesadores: las acciones rot6d de 10D se convierten de nuevo a las acciones axis angle de 7D de LIBERO. Cuando ejecutas la política fuera de lerobot-eval y un robot simulado, también construye los procesadores del entorno con make_env_pre_post_processors(env_cfg, policy_cfg) de lerobot.envs.factory.
  • Consulte la guía de LIBERO para conocer el protocolo de evaluación y las demás opciones.

Video predicho

Cosmos3 predice los fotogramas futuros junto con las acciones. Pasa --policy.save_predicted_video=true para decodificar el video predicho de cada chunk con el VAE y escribir pred_episode_*.mp4 junto a los videos eval_episode_*.mp4 renderizados.

Cosmos3 Edge en la tarea LIBERO-10 “coloca el bol negro en el cajón inferior del armario y ciérralo”. Izquierda: las observaciones de tercera persona y de la muñeca del despliegue. Derecha: el video predicho por la política, muestreado nuevamente cada 16 acciones desde la observación actual.

Resultados

LIBERO

ModelosLibero-SpatialLibero-ObjectLibero-GoalLibero-10Promedio
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%

Tasas de éxito en bucle cerrado en LIBERO. Cosmos3-Nano-Policy-LIBERO y Cosmos3-Edge-Policy-LIBERO se entrenan en las cuatro suites de LIBERO (lerobot/Cosmos3-Nano-Policy-LIBERO, lerobot/Cosmos3-Edge-Policy-LIBERO). Cosmos3-Nano-Policy-LIBERO_10 y Cosmos3-Edge-Policy-LIBERO_10 se entrenan solo en LIBERO-10 (lerobot/Cosmos3-Nano-Policy-LIBERO_10, lerobot/Cosmos3-Edge-Policy-LIBERO_10), y la configuración de entrenamiento de Cosmos3-Nano-Policy-LIBERO_10 sigue la receta de LIBERO-10 reportada en el informe técnico. Cosmos-Nano-TechReport es la tasa de éxito en LIBERO-10 de Cosmos3 Nano entrenado solo en LIBERO-10, tal como se reporta en la Tabla 20 del informe técnico (MT-init, iteración 2000, 500 rollouts).

Cita

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

Licencia

La implementación de LeRobot sigue la Apache 2.0 License de LeRobot. Los pesos de Cosmos3 y el código nativo se publican bajo la OpenMDW-1.1 License.

Update on GitHub