LeRobot documentation
Cosmos3
Cosmos3
Cosmos3 là một mô hình hành động thế giới cho điều khiển robot, được xây dựng dựa trên Cosmos 3, hệ thống mô hình thế giới đa phương thức của NVIDIA. Từ quan sát của camera và một chỉ dẫn ngôn ngữ, chính sách đồng thời dự đoán các khung hình video trong tương lai và khối hành động robot tiếp theo. Bản triển khai LeRobot được điều chỉnh từ chính sách hành động gốc của cosmos-framework và chạy các mô hình Cosmos3 Nano và Edge đã phát hành, được chuyển đổi sang định dạng LeRobot.
- 📄 Báo cáo kỹ thuật Cosmos 3 (arXiv:2606.02800)
- 💻 cosmos-framework và Cosmos cookbook
- 🤗 Bộ sưu tập mô hình Cosmos 3
Tổng quan về Mô hình
Cosmos 3 là một Mixture-of-Transformers (MoT). Mỗi lớp giải mã có hai nhánh riêng biệt với các layer norm, phép chiếu attention và MLP của chúng, và hai nhánh này chỉ gặp nhau trong một self-attention dùng chung:
- Đường dẫn hiểu (bộ suy luận) đọc các token văn bản của prompt và, tùy chọn, các token ảnh ViT, với sự chú ý nhân quả.
- Đường dẫn sinh (bộ sinh) đọc các token liên tục của chuỗi con khuếch tán: các latent của VAE Wan2.2 cho video và các token hành động. Nó chú ý hai chiều trên cả hai đường dẫn và được huấn luyện bằng flow matching.
Kiến trúc Mixture-of-Transformers (MoT) của Cosmos 3 (Hình 5 trong báo cáo kỹ thuật). Bên trái: chuỗi con tự hồi quy (AR) đi qua bộ suy luận và chuỗi con khuếch tán (DM) đi qua bộ tạo, hai chuỗi gặp nhau trong một cơ chế chú ý đa phương thức dùng chung. Bên phải: mặt nạ chú ý, mang tính nhân quả cho AR và đầy đủ cho khuếch tán.
Theo chính sách, Cosmos3 chạy ở chế độ mô hình hành động thế giới của các công thức hành động gốc:
- Đầu vào: các góc nhìn từ camera được ghép thành một khung video duy nhất (một góc nhìn giữ nguyên, hai góc nhìn đặt cạnh nhau, ba góc nhìn với góc nhìn đầu tiên ở trên và hai góc nhìn còn lại ở bên dưới). Prompt chứa nhiệm vụ, tần số khung hình, độ phân giải của khung hình, mô tả bố cục khung hình và số lượng khung hình trống của khối.
- Đào tạo: chính sách mã hóa clip của khung hình hiện tại và
chunk_sizekhung hình tương lai bằng VAE, thêm nhiễu vào các latent tương lai và vào action chunk, và học cách khử nhiễu cho cả hai. Tổn thất là tổng có trọng số của tổn thất video và tổn thất hành động. - Suy luận: khung điều kiện sạch, video tương lai và các hành động được lấy mẫu đồng thời bằng bộ giải UniPC, và chính sách trả về cụm hành động. Video dự đoán có thể được lưu lại trong quá trình đánh giá (xem Đánh giá).
- Mặc định, chỉ có đường dẫn sinh, các lớp xung quanh nó và các module hành động được huấn luyện (
train_gen_stream_only=true,freeze_vision_encoder=true), giống như trong các công thức hậu huấn luyện gốc.
LeRobot cung cấp hai kích thước đã phát hành dưới dạng hai loại chính sách:
| Cosmos3 Nano | Cosmos3 Edge | |
|---|---|---|
| Loại chính sách | cosmos3_nano | cosmos3_edge |
| Tham số (được huấn luyện theo mặc định) | 16.5B (7.0B) | 4.6B (1.4B) |
| Backbone | MoT được khởi tạo từ Qwen3-VL, 36 lớp, kích thước ẩn 4096 | MoT với transformer 2B dạng dense, 28 lớp, kích thước ẩn 2048 |
| Tokenizer | nvidia/Cosmos3-Nano | nvidia/Cosmos3-Edge |
| Checkpoint gốc | lerobot/Cosmos3-Nano | lerobot/Cosmos3-Edge |
Yêu cầu cài đặt
Cài đặt LeRobot bằng cách làm theo Hướng dẫn cài đặt của chúng tôi.
Cài đặt các phụ thuộc của Cosmos3:
pip install -e ".[cosmos3]"Nếu bạn đã cài đặt LeRobot từ PyPI:
pip install 'lerobot[cosmos3]'Thêm các phần mở rộng cho quy trình làm việc của bạn:
trainingcholerobot-train,liberođể huấn luyện hoặc đánh giá trên LIBERO, vàpeftcho LoRA. Ví dụ:pip install -e ".[cosmos3,training,libero]"
Cosmos3 sử dụng backend chú ý tích vô hướng có tỷ lệ của PyTorch (sdpa) theo mặc định. Nếu môi trường của bạn có cài đặt flash_attn tương thích, bạn có thể yêu cầu --policy.attn_implementation=flash_attention_2. Không kết hợp nó với --policy.compile_model=true: các lớp decoder đã biên dịch sẽ biên dịch lại cho từng lớp có flash attention và dừng lại với lỗi biên dịch lại.
Điểm kiểm tra
| Checkpoint | Kích thước | Dữ liệu huấn luyện | Sử dụng |
|---|---|---|---|
| lerobot/Cosmos3-Nano | Nano | Phát hành nvidia/Cosmos3-Nano | Nền tảng để tinh chỉnh |
| lerobot/Cosmos3-Edge | Edge | Phát hành nvidia/Cosmos3-Edge | Nền tảng để tinh chỉnh |
| lerobot/Cosmos3-Nano-Policy-LIBERO | Nano | lerobot/LIBERO-Cosmos3 | Đánh giá LIBERO |
| lerobot/Cosmos3-Edge-Policy-LIBERO | Edge | lerobot/LIBERO-Cosmos3 | Đánh giá LIBERO |
| lerobot/Cosmos3-Nano-Policy-LIBERO_10 | Nano | lerobot/LIBERO_10-Cosmos3 | Đánh giá LIBERO-10 |
| lerobot/Cosmos3-Edge-Policy-LIBERO_10 | Edge | lerobot/LIBERO_10-Cosmos3 | Đánh giá LIBERO-10 |
- Các checkpoint cơ bản chứa các trọng số đã phát hành ở định dạng LeRobot. Các module hành động của chúng (
action2llm,llm2action,action_modality_embed) chưa được huấn luyện cho robot nào, vì vậy hãy bắt đầu tinh chỉnh với--policy.load_action_modules=false(xem Tinh chỉnh trên LIBERO). config.jsoncủa các checkpoint gốc chứa công thức huấn luyện sau (post-training) DROID gốc, và các giá trị mặc định của lớp cho các kiểu policy là công thức LIBERO gốc.- Các checkpoint của LIBERO chứa trọng số EMA của quá trình huấn luyện tương ứng.
- Trọng số được lưu trữ ở bfloat16.
config.jsoncố ý đặtdtype: float32: huấn luyện độ chính xác hỗn hợp từ các checkpoint này giữ lại trọng số master float32, giống như huấn luyện gốc. Đối với suy luận và đánh giá, hãy truyền--policy.dtype=bfloat16.
Có hai cách để bắt đầu từ một checkpoint:
--policy.path=<checkpoint>tải cấu hình đã lưu, trọng số và các bộ xử lý, sau đó áp dụng lệnh dòng lệnh lên trên.--policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint>(hoặccosmos3_edge) chỉ tải trọng số và các bộ xử lý, và cấu hình bắt đầu từ các giá trị mặc định của lớp, công thức LIBERO.
Fine-tuning trên LIBERO
Các giá trị mặc định của lớp là công thức LIBERO của báo cáo kỹ thuật: các khối 16 hành động, domain_id=5, không có hành động trạng thái, chuẩn hóa phân vị của các hành động, UniPC với 30 bước và không có hướng dẫn không phân loại viên (classifier-free guidance) khi suy luận, và tốc độ học 5e-5 với 500 bước khởi động và suy giảm tuyến tính về không ở bước 16000. Do đó, bắt đầu từ một checkpoint cơ sở với --policy.type chỉ cần --policy.load_action_modules=false.
Công thức huấn luyện trên dữ liệu LIBERO đã được chuyển đổi sang bố cục Cosmos3:
- lerobot/LIBERO-Cosmos3: bốn bộ (Spatial, Object, Goal, 10), 1693 tập.
- lerobot/LIBERO_10-Cosmos3: chỉ LIBERO-10, 379 tập.
Cả hai đều có nguồn gốc từ nvidia/LIBERO_LeRobot_v3, dữ liệu của công thức gốc. Các hành động là 10D [delta position (3), rot6d (6), gripper] thay vì các hành động 7D axis angle của LIBERO, hai camera observation.images.image (thứ ba) và observation.images.image2 (cổ tay) có độ phân giải 256x256 ở 20 FPS, và thống kê hành động là của công thức gốc. Trong quá trình đánh giá, môi trường LIBERO sẽ chuyển đổi các hành động trở lại (xem Đánh giá).
Lệnh huấn luyện
Công thức LIBERO gốc được huấn luyện với batch toàn cục gồm 2048 mẫu trong 2000 lượt lặp trên LIBERO-10 (checkpoint tốt nhất nằm gần lượt lặp 1500) và trong 5000 lượt lặp trên bốn bộ bài kiểm tra (tốt nhất gần 4500). Lệnh dưới đây tái hiện quy trình này trên một máy với 8 GPU: FSDP2 trên 8 GPU, 32 mẫu mỗi GPU và 8 bước tích lũy gradient, do đó 32 x 8 x 8 = 2048 mẫu cho mỗi lần cập nhật optimizer. Mọi số lượt lặp gốc đều được nhân với 8 bước tích lũy: 2000 lượt lặp trở thành --steps=16000, một checkpoint mỗi 500 lượt lặp --save_freq=4000, và warmup 500 cùng điểm kết thúc suy giảm 16000 trở thành 4000 và 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Đối với bốn bộ dữ liệu, hãy sử dụng --dataset.repo_id=lerobot/LIBERO-Cosmos3 và --steps=40000 (5000 lần lặp). Đối với Cosmos3 Edge, hãy sử dụng --policy.type=cosmos3_edge --policy.pretrained_path=lerobot/Cosmos3-Edge.
Các tham số huấn luyện chính
| Tham số | Mô tả |
|---|---|
--policy.load_action_modules=false | Bắt đầu các module hành động từ trạng thái khởi tạo thay vì nạp chúng, giống như các công thức hậu huấn luyện bản địa. Chỉ đặt nó khi tinh chỉnh từ một checkpoint cơ sở (xem các ghi chú bên dưới). |
--parallelism.dp_shard=8 | FSDP2 trên 8 GPU. Chính sách khai báo các đơn vị bọc FSDP2 của nó, nên --accelerator.fsdp.wrap_modules không cần thiết. Xem hướng dẫn huấn luyện đa GPU. |
--accelerator.fsdp.mixed_precision_policy.* | tính toán bfloat16 và giảm gradient với trọng số master float32, giống như trong huấn luyện native. Giữ nguyên giá trị mặc định --policy.dtype (float32) cùng với nó: một policy bfloat16 sẽ huấn luyện các trọng số master bfloat16. |
--accelerator.gradient_accumulation.steps | Batch toàn cục là --batch_size x số GPU x số bước tích lũy. sync_each_batch=true giảm các gradient sau mỗi micro batch, do đó chúng vẫn được phân mảnh và không cần một bản sao không phân mảnh bổ sung. |
--steps, --save_freq, --policy.scheduler_*_steps | LeRobot đếm các bước lặp (micro batches), không phải các lần cập nhật optimizer. Nhân số lần lặp của công thức gốc (2000 hoặc 5000 iterations, warmup 500, decay 16000, checkpoints mỗi 500) với số bước tích lũy. Giữ --steps tối đa bằng --policy.scheduler_decay_steps: tốc độ học sẽ bằng 0 sau đó. |
--ema.type=power --ema.sigma_rel=0.1 | EMA của công thức gốc. Mỗi checkpoint cũng chứa trọng số EMA trong pretrained_model_ema/, đây là các trọng số để đánh giá và công bố. |
--tf32=false | Giữ các phép nhân ma trận float32 ở độ chính xác đầy đủ, giống như quá trình huấn luyện gốc. |
--policy.gradient_checkpointing=true | Tính lại các kích hoạt của các lớp giải mã và các khối thị giác trong bước truyền ngược để tiết kiệm bộ nhớ. |
--policy.compile_model=true | Tùy chọn: torch.compile cho các lớp giải mã và bộ mã hóa VAE. |
Trên một GPU duy nhất hoặc khi sử dụng DDP, hãy thay các cờ FSDP2 (--parallelism.dp_shard và --accelerator.fsdp.mixed_precision_policy.*) bằng --accelerator.mixed_precision=bf16. Ví dụ Cosmos3 Edge trên một GPU, với 32 mẫu mỗi bước và 64 bước tích lũy để đạt cùng kích thước lô toàn cục là 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được đặt lại thànhtruesau khi checkpoint được tải, do đó run đã lưu sẽ tải lại các action module đã được huấn luyện. Hai trường hợp sau đây giữ các action module ở trạng thái khởi tạo một cách nhầm lẫn:
- Xây dựng một policy mà không có
--policy.pretrained_pathhoặc--policy.path: không có gì được nạp, do đó cờ vẫn ở trạng tháifalsevà một lần đánh giá hoặc tiếp tục (resume) sau đó của run đó sẽ bỏ qua các module hành động đã được huấn luyện của nó.- PEFT với các module hành động nằm ngoài
--peft.full_training_modules: các module hành động khi đó sẽ không bao giờ được huấn luyện.Chỉ đặt
--policy.load_action_modules=falsekhi fine-tuning từ một checkpoint cơ sở. Với PEFT, hãy giữ các action modules trong--peft.full_training_modules(mặc định).
Yêu cầu phần cứng
Đo trên GPU A800 80 GB với các lệnh của trang này và gradient checkpointing:
| Cài đặt | GPU | Batch trên mỗi GPU | Bộ nhớ tối đa trên mỗi GPU |
|---|---|---|---|
| Nano, tinh chỉnh đầy đủ (FSDP2) | 8 | 32 | 38 GB |
| Nano, LoRA | 1 | 4 | 64 GB |
| Edge, tinh chỉnh đầy đủ | 1 | 32 | 47 GB |
| Edge, LoRA | 1 | 8 | 21 GB |
| Đánh giá trong bfloat16, Nano và Edge | 1 | 10 môi trường | 39 GB và 15 GB |
- Cosmos3 Nano cần các GPU 80 GB, và ít nhất 4 GPU để fine-tuning đầy đủ: trước khi FSDP2 chia nhỏ (shard) mô hình, mỗi GPU phải chứa toàn bộ mô hình float32 (khoảng 66 GB).
- Lệnh LIBERO ở trên mất khoảng 45 giây cho mỗi lần cập nhật optimizer trên 8 GPU A800, do đó mất khoảng 25 giờ cho LIBERO-10.
Tinh chỉnh trên dữ liệu của bạn
Cosmos3 lấy công thức huấn luyện từ dòng lệnh và bộ dữ liệu: chính sách không chứa bất kỳ logic cụ thể nào cho robot. Bất kỳ bộ dữ liệu LeRobot nào cũng hoạt động được một khi nó tuân thủ các điểm dưới đây.
Yêu cầu về bộ dữ liệu
- Camera: một đến ba đặc trưng
observation.images.*, theo thứ tự tronginfo.jsoncủa bộ dữ liệu. Mỗi góc nhìn được điều chỉnh kích thước về--policy.view_size(mặc định là kích thước của camera đầu tiên) và được ghép thành một khung video duy nhất: một góc nhìn giữ nguyên, hai góc nhìn đặt cạnh nhau, ba góc nhìn với góc nhìn đầu tiên ở trên và hai góc nhìn còn lại đặt cạnh nhau bên dưới. Hãy đặt các camera theo thứ tự của bố cục đó, ví dụ[third person, wrist]cho LIBERO và[wrist, exterior 1, exterior 2]cho DROID. - Mô tả:
--policy.view_descriptionlà câu trong prompt mô tả bố cục canvas. Giá trị mặc định mô tả bố cục LIBERO (giá trị mặc định của class) và bố cục DROID (fileconfig.jsoncơ sở). Hãy viết câu mô tả cho bố cục của bạn: một mô tả sai sẽ không được phát hiện mà chỉ làm suy yếu prompt. - Biểu diễn hành động: chính sách học các hành động của bộ dữ liệu theo đúng dạng ban đầu và không có bước chuyển đổi. Hãy chuyển đổi các hành động sang biểu diễn mà bạn muốn học trước khi huấn luyện, giống như cách dữ liệu LIBERO làm với phép quay rot6d. Nếu môi trường kỳ vọng một biểu diễn khác, hãy chuyển đổi các hành động trở lại trong các bộ xử lý môi trường (môi trường LIBERO làm điều này cho Cosmos3).
- Độ rộng hành động: tối đa
--policy.max_action_dim(64) kênh. Độ rộng đầu ra được lấy từ bộ dữ liệu. - Tốc độ khung hình:
fpsđược lấy từ metadata của dataset và phải là tốc độ điều khiển thực tế: nó được ghi vào prompt và thiết lập trục thời gian của các vị trí video cũng như ngưỡng của các khung hình nhàn rỗi. - Thống kê: hành động chuẩn hóa mặc định là
QUANTILES, cần các thống kêq01vàq99của hành động trongmeta/stats.json(các bộ dữ liệu do LeRobot ghi có sẵn). Công thức DROID bản địa sử dụngIDENTITYthay thế. - Trạng thái: với
--policy.use_state=true, trạng thái hiện tại được đặt trước các hành động dưới dạng một hành động trạng thái sạch.observation.statesau đó cần có cùng chiều rộng và ý nghĩa kênh với hành động. Trạng thái và hành động được chuẩn hóa bằng thống kê riêng của chúng. - Ngôn ngữ:
taskcủa mỗi khung hình là chỉ dẫn của prompt. - Embodiment:
--policy.domain_idchọn hàng của các phép chiếu hành động, một trong 32 hàng. Các công thức gốc sử dụng 5 cho LIBERO, 8 cho DROID và 30 cho RoboCasa. Khi các module hành động bắt đầu từ trạng thái khởi tạo (--policy.load_action_modules=false), mọi hàng đều chưa được huấn luyện và bất kỳ giá trị nào cũng hoạt động. - Chunk:
--policy.chunk_sizehành động mỗi chunk (16 cho LIBERO, 32 cho DROID) và--policy.n_action_stepshành động được thực thi mỗi chunk.
Các trường để nguyên None (fps, view_size) và tên các kênh hành động được lấy từ dataset. Tất cả các trường khác đều đến từ cấu hình, vì vậy hãy đặt chúng trên dòng lệnh hoặc trong tệp --config_path.
Tên kênh hành động và trạng thái. Prompt nêu rõ có bao nhiêu khung hình trong chunk là ở trạng thái nhàn rỗi (robot không di chuyển), và
Cosmos3IdleFrameProcessorStepđếm chúng từ tên kênh hành động trongmeta/info.json. Các tên có dạng[<prefix>_]<type>_<index>cho biết loại kênh:pos, một phép quay (rot6d,rot9d,axisangle,euler_xyz,quat_xyzw,quat_wxyz),gripperhoặcjoint, kèm tiền tố để phân biệt các cánh tay (ví dụleft_rot6d_0). Dữ liệu LIBERO sử dụngpos_0..2, rot6d_0..5, gripper_0, còn bố cục DROID dùngjoint_0..6, gripper_0.Việc đặt tên này không bắt buộc. Nếu không đặt tên, chính sách chỉ ghi nhận một cảnh báo và prompt không đếm khung hình nhàn rỗi, điều này hoàn toàn khớp với những gì công thức LIBERO quan sát được (các episode của nó không có khung hình nhàn rỗi theo quy tắc gốc). Chỉ tuân thủ quy tắc đặt tên nếu bạn muốn có số lượng khung hình nhàn rỗi của công thức DROID gốc.
Công thức DROID
config.json của các checkpoint gốc là công thức hậu huấn luyện DROID gốc: domain_id=8, các khối gồm 32 hành động, trạng thái hiện tại được dùng làm hành động trạng thái và các khung hình nhàn rỗi được tính trên đó (use_state=true, idle_state=true), không chuẩn hóa (IDENTITY), UniPC với 4 bước và thang độ dẫn hướng không phân loại (classifier-free guidance) là 3 khi suy luận, cùng tốc độ học 8e-5 (tốc độ học gốc 2e-4 nhân với hệ số 0.4) không có warmup và suy giảm tuyến tính về 0 tại bước 100000. Bắt đầu từ nó với --policy.path sẽ huấn luyện theo công thức này.
Công thức gốc huấn luyện trên các tập hợp thành công của nvidia/Cosmos3-DROID. Một tập dữ liệu DROID theo bố cục Cosmos3 có:
- Ba camera
observation.images.wrist_left,observation.images.exterior_1_leftvàobservation.images.exterior_2_left, theo thứ tự này (camera cổ tay ở trên, hai góc nhìn bên ngoài ở dưới), kích thước 360x640 trong dữ liệu gốc. - 8 hành động vị trí khớp được đặt tên từ
joint_0đếnjoint_6vàgripper_0, cùng vớiobservation.statecó cùng bố cục. - Tốc độ khung hình là 15 FPS.
Công thức gốc huấn luyện với lô toàn cục gồm 8192 mẫu trong 10000 lần lặp, với một điểm kiểm tra mỗi 1000 lần lặp. Trên một máy với 8 GPU, 8 mẫu mỗi GPU và 128 bước tích lũy cho ra cùng một lô toàn cục (8 x 8 x 128 = 8192), và mọi số lần lặp đều được nhân với 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=trueCông thức DROID gốc được chạy trên 256 GPU. Trên một máy với 8 GPU A800, một lần cập nhật optimizer của lệnh này mất khoảng 33 phút, do đó công thức đầy đủ cần nhiều máy. Chạy cùng một lệnh trên mọi máy với các cờ đa nút của accelerate launch (hoặc torchrun), chia nhỏ trong mỗi máy và nhân bản giữa các máy bằng --parallelism.dp_replicate=<number of machines> --parallelism.dp_shard=8, và chia --accelerator.gradient_accumulation.steps cùng số bước cho số lượng máy. Xem HSDP trong hướng dẫn huấn luyện đa GPU.
Một robot mới
Đối với một robot mới, hãy bắt đầu từ các giá trị mặc định của lớp với --policy.type và đặt các trường dữ liệu của bạn trên dòng lệnh, ví dụ:
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 \
...Cấu hình Backbone
Cấu hình backbone --policy.mot_backbone là một giá trị dòng lệnh duy nhất. Truyền vào một dict YAML chỉ chứa các trường cần thay đổi, ví dụ một model có ít lớp hơn để kiểm tra nhanh mà không cần trọng số được huấn luyện trước:
--policy.type=cosmos3_edge --policy.mot_backbone='{text_config: {num_hidden_layers: 2}, vision_config: {num_hidden_layers: 1}}'Các trường bị bỏ qua sẽ giữ nguyên giá trị mặc định của kích thước, hoặc giá trị của cấu hình đã tải với --policy.path. Hãy viết số thực dưới dạng 1.0e-6 (YAML đọc 1e-6 như một chuỗi), và cung cấp đầy đủ các trường có giá trị là dict, chẳng hạn như rope_parameters. Các trường kiến trúc phải khớp với trọng số bạn tải.
Tinh chỉnh với PEFT (LoRA)
Cosmos3 hỗ trợ LoRA thông qua tích hợp PEFT của LeRobot (pip install -e ".[cosmos3,training,peft]"). Mặc định, LoRA điều chỉnh các phép chiếu attention của đường dẫn sinh (q_proj_gen, k_proj_gen, v_proj_gen, o_proj_gen), trong khi các module hành động (action2llm, llm2action, action_modality_embed) được huấn luyện đầy đủ và được lưu cùng với adapter.
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, tốc độ học và hệ số suy giảm trọng số là các giá trị của công thức LoRA gốc của Cosmos3. Nếu không có--peft.lora_alpha, PEFT sẽ sử dụng 8.- Với
--policy.load_action_modules=false, các action module có thể thay đổi hình dạng (--policy.max_action_dim,--policy.num_embodiment_domains). Cách này chỉ hoạt động với full fine-tuning: việc nạp lại một run PEFT sẽ nạp base, đòi hỏi các hình dạng phải giống nhau. --policy.train_gen_stream_onlyvà--policy.freeze_vision_encoderkhông có tác dụng với PEFT: chỉ có các adapter và các module hành động được huấn luyện.- Không thêm các phép chiếu hành động (
action2llm,llm2action) vào--peft.target_modules: chúng là các bảng theo từng dạng thân (embodiment) được đọc trực tiếp bởi một phép nhân ma trận nhóm, do đó LoRA sẽ không thay đổi chúng. Chúng đã được huấn luyện đầy đủ. - Giữ nguyên
--policy.dtypemặc định là float32 khi dùng PEFT, để các module hành động được huấn luyện đầy đủ vẫn giữ trọng số master ở float32. - Các công thức gốc không xác định kích thước lô hoặc lịch trình cho LoRA. Lệnh giữ nguyên lịch trình LIBERO (giảm tuyến tính về 0 ở bước 16000) với 4 mẫu mỗi bước.
- Cosmos3 Nano với LoRA được huấn luyện trên một GPU 80 GB với gradient checkpointing, khoảng 64 GB với 4 mẫu trên mỗi GPU (xem Yêu cầu phần cứng).
Đánh giá
Đánh giá một checkpoint LIBERO trên bốn bộ bài kiểm tra:
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- Đối với checkpoint của phiên chạy huấn luyện của bạn, hãy truyền trọng số EMA:
--policy.path=outputs/train/cosmos3_nano_libero_10/checkpoints/last/pretrained_model_ema. --env.observation_height=256 --env.observation_width=256hiển thị các camera ở độ phân giải 256x256 của dữ liệu huấn luyện.- Môi trường LIBERO tự động thêm chuyển đổi hành động Cosmos3 vào các bộ xử lý của nó: các hành động rot6d 10D được chuyển đổi ngược lại thành các hành động axis angle 7D của LIBERO. Khi bạn chạy chính sách bên ngoài
lerobot-evalvà một robot mô phỏng, hãy xây dựng các bộ xử lý môi trường bằngmake_env_pre_post_processors(env_cfg, policy_cfg)từlerobot.envs.factorytương tự. - Xem hướng dẫn LIBERO để biết về giao thức đánh giá và các tùy chọn khác.
Video dự đoán
Cosmos3 dự đoán các khung hình tương lai cùng với các hành động. Truyền --policy.save_predicted_video=true để giải mã video dự đoán của mỗi chunk bằng VAE và ghi pred_episode_*.mp4 cạnh các video eval_episode_*.mp4 đã được render.
Cosmos3 Edge trên tác vụ LIBERO-10 “đặt chiếc bát đen vào ngăn kéo dưới cùng của tủ và đóng lại”. Bên trái: quan sát từ góc nhìn thứ ba và cổ tay của lần rollout. Bên phải: video được dự đoán bởi chính sách, được lấy mẫu lại mỗi 16 hành động từ quan sát hiện tại.
Kết quả
LIBERO
| Mô hình | Libero-Spatial | Libero-Object | Libero-Goal | Libero-10 | Trung bình |
|---|---|---|---|---|---|
| Cosmos3-Nano-Policy-LIBERO | 98.0% (490/500) | 99.4% (497/500) | 97.4% (487/500) | 96.8% (484/500) | 97.90% |
| Cosmos3-Edge-Policy-LIBERO | 98.6% (493/500) | 99.4% (497/500) | 97.2% (486/500) | 96.6% (483/500) | 97.95% |
| Cosmos3-Nano-Policy-LIBERO_10 | N/A | N/A | N/A | 97.6% (488/500) | 97.60% |
| Cosmos3-Edge-Policy-LIBERO_10 | N/A | N/A | N/A | 95.2% (476/500) | 95.20% |
| Cosmos-Nano-TechReport | N/A | N/A | N/A | 97.4% | 97.40% |
Tỷ lệ thành công vòng kín trên LIBERO. Cosmos3-Nano-Policy-LIBERO và Cosmos3-Edge-Policy-LIBERO được huấn luyện trên bốn bộ LIBERO (lerobot/Cosmos3-Nano-Policy-LIBERO, lerobot/Cosmos3-Edge-Policy-LIBERO). Cosmos3-Nano-Policy-LIBERO_10 và Cosmos3-Edge-Policy-LIBERO_10 chỉ được huấn luyện trên LIBERO-10 (lerobot/Cosmos3-Nano-Policy-LIBERO_10, lerobot/Cosmos3-Edge-Policy-LIBERO_10), và cấu hình huấn luyện của Cosmos3-Nano-Policy-LIBERO_10 tuân theo công thức LIBERO-10 được báo cáo trong báo cáo kỹ thuật. Cosmos-Nano-TechReport là tỷ lệ thành công LIBERO-10 của Cosmos3 Nano được huấn luyện chỉ trên LIBERO-10, như được báo cáo trong Bảng 20 của báo cáo kỹ thuật (MT-init, lặp 2000, 500 rollouts).
Trích dẫn
@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}
}Giấy phép
Cài đặt LeRobot tuân theo Giấy phép Apache 2.0 của LeRobot. Trọng số Cosmos3 và mã nguồn gốc được phát hành theo Giấy phép OpenMDW-1.1.
Update on GitHub