Cosmos3
Collection
Cosmos3-Nano and Cosmos3-Edge base checkpoints converted to LeRobot format. • 8 items • Updated • 2
Cosmos3-Edge-Policy-LIBERO_10
Cosmos3-Edge fine-tuned on the LIBERO-10 as a LeRobot policy. From the task instruction and the third person and wrist camera images, the policy predicts a chunk of 16 actions together with the future video of the scene. This repository holds the EMA weights at optimizer update 2000.
| Item | Value |
|---|---|
| Base checkpoint | lerobot/Cosmos3-Edge, the LeRobot format of nvidia/Cosmos3-Edge |
| Policy type | cosmos3_edge |
| Backbone | Nemotron 3 Dense VL MoT, 28 decoder layers, hidden width 2048 |
| Weights | EMA weights, 1188 tensors: 1184 BF16 and 4 FP32 (the timestep embedding) |
| Observations | observation.images.image (third person) and observation.images.image2 (wrist), 256x256 RGB at 20 FPS, no state |
| Actions | 10D [delta position (3), rot6d (6), gripper], converted to the 7D actions of LIBERO by the environment processors |
Results on LIBERO
| Model | Libero-10 |
|---|---|
| Cosmos3-Nano-Policy-LIBERO_10 | 97.6% (488/500) |
| Cosmos3-Edge-Policy-LIBERO_10 | 95.2% (476/500) |
| Cosmos3-Nano-Policy-LIBERO_10 (TechReport) | 97.40% |
Closed-loop success rates: 50 episodes for each of the 10 tasks of a suite, 500 episodes per suite, with the evaluation setup below.
Evaluation
Command
The results were measured with MuJoCo 2.3.7. The LeRobot lock file installs MuJoCo 3.8.1, so install the older version first:
uv pip install mujoco==2.3.7
lerobot-eval \
--policy.path=lerobot/Cosmos3-Edge-Policy-LIBERO_10 \
--policy.dtype=bfloat16 \
--tf32=false \
--seed=0 \
--env.type=libero \
--env.task=libero_10 \
--env.obs_type=pixels \
--env.observation_height=256 \
--env.observation_width=256 \
--env.episode_length=510 \
--eval.n_episodes=50 \
--eval.batch_size=4 \
--eval.use_async_envs=true \
--output_dir=outputs/eval/cosmos3-edge-policy-libero_10/libero_10
Evaluation setup
| Item | Value |
|---|---|
| LeRobot commit | d0a20dc38e058a8aa3f7bc48cf32d48bb2f99e5b (huggingface/lerobot#5134) |
| MuJoCo | 2.3.7 |
| robosuite | 1.4.0 |
| hf-libero | 0.1.4 |
| Gymnasium | 1.3.0 |
| Rendering | MuJoCo EGL on the NVIDIA GPU (MUJOCO_GL=egl) |
| Cameras | third person and wrist, 256x256 RGB, 20 Hz |
| Control | relative end effector control, hard reset, the fixed initial states of each task |
| Episodes | 50 per task (initial states and environment seeds 0 to 49), 4 environments per batch |
| Policy budget | 510 for Libero-10 steps, counted after the 10 settling steps of the environment |
| Inference | bfloat16, SDPA attention, no torch.compile, TF32 disabled, cuDNN benchmarking disabled |
| Sampling | UniPC with 30 steps, guidance scale 1 (no classifier-free guidance), sampling shift 5, each chunk of 16 actions executed in full |
| Seed | 0, set once per suite before the policy is built |
- MuJoCo 3.x changes the outcome of contact rich LIBERO tasks, so results measured with it are not comparable.
- The policy budgets are those of the native Cosmos3 LIBERO evaluation. The LeRobot defaults are longer.
- The video canvas step of the preprocessor ran on CPU, and cuDNN benchmarking was disabled.
lerobot-evalruns the preprocessor on the policy device and enables cuDNN benchmarking, which changes the inputs and the sampling by small numerical amounts.
Evaluation hardware and software
| Item | Value |
|---|---|
| GPU | 1 NVIDIA A800-SXM4-80GB per suite |
| NVIDIA driver | 595.71.05 |
| CPU | 2 x Intel Xeon Platinum 8358P (128 threads) |
| Operating system | Ubuntu 24.04 (glibc 2.39), Linux 6.8, x86_64 |
| Python | 3.12.15 |
| PyTorch | 2.11.0+cu128 (CUDA 12.8, cuDNN 9.19.0, NCCL 2.28.9) |
| torchvision | 0.26.0+cu128 |
| Transformers | 5.5.4 |
| Diffusers | 0.39.0 |
| Accelerate | 1.15.0 |
| PEFT | 0.21.0 |
| NumPy | 2.2.6 |
| PyAV / torchcodec | 15.1.0 / 0.11.1 |
| safetensors | 0.8.0 |
| huggingface_hub | 1.32.0 |
| tokenizers | 0.22.2 |
| OpenCV | 4.14.0.94 |
| SciPy | 1.18.1 |
Training
Training setup
| Item | Value |
|---|---|
| Base checkpoint | lerobot/Cosmos3-Edge, action modules reinitialized (--policy.load_action_modules=false) |
| Training data | lerobot/LIBERO_10-Cosmos3: 379 episodes, 10 tasks, every episode used (no validation split) |
| Global batch | 2048 samples per optimizer update |
| Batch per GPU x GPUs x gradient accumulation | 64 x 32 x 1 = 2048 |
| GPUs | 32 NVIDIA A100-SXM4-80GB (4 nodes of 8, InfiniBand) |
| Parallelism | FSDP2 HSDP: parameters sharded over the 8 GPUs of a node, replicated over the 4 nodes |
| Optimizer updates | 4000 in the run, a checkpoint every 200 |
| This checkpoint | optimizer update 2000 (LeRobot step 2000), chosen by the LIBERO success rate among the saved checkpoints |
| Optimizer | AdamW, learning rate 5e-5 (2.5e-4 for the action modules), betas (0.9, 0.99), weight decay 0.05, gradient clipping 1.0 |
| Learning rate schedule | linear warmup over 500 updates, then linear decay that would reach 0 at update 16000 |
| Precision | float32 master weights, optimizer state and EMA, bfloat16 compute and gradient reduction, TF32 disabled |
| Gradient checkpointing | on |
| Attention / compilation | SDPA, no torch.compile |
| EMA | power EMA with sigma_rel=0.1, updated at every optimizer update |
| Seed | 42 |
| Data loading | PyAV video decoding, 8 workers per GPU, prefetch factor 1 |
| Software | the evaluation software above (MuJoCo is not used in training) |
Keep the global batch at 2048 samples on other GPUs: batch per GPU x number of GPUs x gradient accumulation steps = 2048. LeRobot counts --steps, --save_freq, --policy.scheduler_warmup_steps and --policy.scheduler_decay_steps in micro batches, so multiply the optimizer update counts by the gradient accumulation steps (1 here). Every rank builds the float32 policy on the CPU before FSDP shards it (about 18 GB for Cosmos3-Edge), so a node needs that much host memory for each of its ranks that load at the same time.
Training command
# Run on each of the 4 nodes with NODE_RANK=0 to 3.
torchrun \
--nnodes=4 --nproc_per_node=8 --node_rank=${NODE_RANK} \
--master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
$(which lerobot-train) \
--policy.type=cosmos3_edge \
--policy.pretrained_path=lerobot/Cosmos3-Edge \
--policy.load_action_modules=false \
--policy.gradient_checkpointing=true \
--policy.attn_implementation=sdpa \
--policy.compile_model=false \
--policy.push_to_hub=false \
--dataset.repo_id=lerobot/LIBERO_10-Cosmos3 \
--dataset.video_backend=pyav \
--parallelism.dp_replicate=4 \
--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=1 \
--accelerator.gradient_accumulation.sync_each_batch=true \
--tf32=false \
--ema.type=power \
--ema.sigma_rel=0.1 \
--batch_size=64 \
--steps=4000 \
--save_freq=200 \
--policy.scheduler_warmup_steps=500 \
--policy.scheduler_decay_steps=16000 \
--num_workers=8 \
--prefetch_factor=1 \
--seed=42 \
--output_dir=outputs/train/cosmos3_edge_policy_libero_10 \
--job_name=cosmos3_edge_policy_libero_10
The weights of this repository are outputs/train/cosmos3_edge_policy_libero_10/checkpoints/002000/pretrained_model_ema of this run.
Loading and fine-tuning
--policy.path=lerobot/Cosmos3-Edge-Policy-LIBERO_10 loads the configuration, the weights and the processors. config.json sets dtype: float32 on purpose: fine-tuning from this checkpoint keeps float32 master weights. For inference and evaluation pass --policy.dtype=bfloat16.
License and provenance
The NVIDIA model materials use OpenMDW 1.1. Retain the license and applicable copyright and origin notices when redistributing them. NOTICE records the original model and the attribution of this distribution. Separately installed runtime software retains its own licenses.
Consult the original NVIDIA model card and its safety, bias, privacy, and explainability notes for upstream limitations.
- Downloads last month
- -