LeRobot documentation

Cosmos3

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.7.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Cosmos3

Cosmos3은 Cosmos 3 기반의 로봇 제어용 월드 액션 모델입니다. Cosmos 3는 NVIDIA의 옴니모달 월드 모델 제품군입니다. 카메라 관측값과 언어 지시를 입력으로 받아, 정책은 미래 비디오 프레임과 다음 로봇 액션 청크를 동시에 예측합니다. LeRobot 구현은 네이티브 cosmos-framework 액션 정책에 기반하여 적응되었으며, LeRobot 형식으로 변환된 공개된 Cosmos3 Nano 및 Edge 모델을 실행합니다.

모델 개요

Cosmos 3은 Mixture-of-Transformers(MoT)입니다. 각 decoder 레이어에는 자체 layer norm, attention projection 및 MLP를 가진 두 개의 경로가 있으며, 두 경로는 공유 self-attention에서만 만납니다:

  • 이해 경로(추론기)는 프롬프트의 텍스트 토큰과 선택적으로 ViT 이미지 토큰을 인과적 어텐션(causal attention)으로 읽습니다.
  • 생성 경로(생성기)는 확산 서브시퀀스의 연속 토큰, 즉 비디오의 Wan2.2 VAE 잠재 변수와 액션 토큰을 읽습니다. 두 경로 모두에 대해 양방향으로 주의를 기울이며(flow matching)로 학습됩니다.
Mixture-of-Transformers architecture of Cosmos 3

Cosmos 3의 Mixture-of-Transformers (MoT) 아키텍처(기술 보고서의 Figure 5). 왼쪽: 오토레그레시브(AR) 서브시퀀스는 reasoner를, 디퓨전(DM) 서브시퀀스는 generator를 통과한 후 공유 멀티모달 어텐션에서 합쳐집니다. 오른쪽: AR은 causal, 디퓨전은 full인 어텐션 마스크입니다.

정책에 따라 Cosmos3는 네이티브 액션 레시피의 월드 액션 모델 모드에서 실행됩니다:

  • 입력: 카메라 뷰는 하나의 비디오 캔버스로 구성됩니다(하나의 뷰는 그대로, 두 개의 뷰는 나란히, 세 개의 뷰는 첫 번째가 위에 있고 나머지 두 개가 아래에 배치). 프롬프트에는 작업, 프레임 레이트, 캔버스 해상도, 캔버스 레이아웃 설명, 해당 청크의 유휴 프레임 수가 포함됩니다.
  • 학습: 정책은 현재 프레임의 클립과 chunk_size개의 미래 프레임을 VAE로 인코딩하고, 미래 잠재 변수와 액션 청크에 노이즈를 추가한 뒤, 이 둘 모두를 디노이즈하도록 학습합니다. 손실 함수는 비디오 손실과 액션 손실의 가중 합입니다.
  • Inference: 조건 프레임은 클린 상태이며, 미래 비디오와 액션은 UniPC 솔버와 함께 공동으로 샘플링되고, 정책은 액션 청크를 반환합니다. 평가 중 예측된 비디오를 저장할 수 있습니다(자세한 내용은 Evaluation 참조).
  • 기본적으로 네이티브 포스트 트레이닝 레시피와 마찬가지로 생성 경로, 그 주변의 레이어 및 액션 모듈만 학습됩니다(train_gen_stream_only=true, freeze_vision_encoder=true).

LeRobot는 두 가지 공개된 크기를 두 가지 정책 유형으로 제공합니다:

Cosmos3 NanoCosmos3 Edge
정책 유형cosmos3_nanocosmos3_edge
매개변수 (기본적으로 학습됨)16.5B (7.0B)4.6B (1.4B)
BackboneQwen3-VL에서 초기화된 MoT, 36개 레이어, hidden size 4096밀집 2B transformer, 28개 레이어, hidden size 2048을 갖춘 MoT
토크나이저nvidia/Cosmos3-Nanonvidia/Cosmos3-Edge
기본 체크포인트lerobot/Cosmos3-Nanolerobot/Cosmos3-Edge

설치 요구 사항

  1. LeRobot은 설치 가이드를 참고하여 설치하세요.

  2. Cosmos3의 종속성을 설치합니다:

    pip install -e ".[cosmos3]"

    LeRobot을 PyPI에서 설치한 경우:

    pip install 'lerobot[cosmos3]'
  3. 워크플로우에 필요한 추가 항목을 추가하세요: lerobot-train을 위한 training, LIBERO에서 학습 또는 평가를 위한 libero, LoRA를 위한 peft입니다. 예를 들어:

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

Cosmos3은 기본적으로 PyTorch의 scaled dot product attention 백엔드(sdpa)를 사용합니다. 환경에 호환되는 flash_attn 설치 파일이 있는 경우 --policy.attn_implementation=flash_attention_2를 요청할 수 있습니다. --policy.compile_model=true와 함께 사용하지 마세요. 컴파일된 디코더 레이어는 flash attention을 사용하는 각 레이어마다 다시 컴파일되어 재컴파일 오류로 중단됩니다.

체크포인트

체크포인트크기학습 데이터사용
lerobot/Cosmos3-NanoNanonvidia/Cosmos3-Nano 릴리스파인튜닝의 기반
lerobot/Cosmos3-EdgeEdgenvidia/Cosmos3-Edge 릴리스파인튜닝의 기반
lerobot/Cosmos3-Nano-Policy-LIBERONanolerobot/LIBERO-Cosmos3LIBERO 평가
lerobot/Cosmos3-Edge-Policy-LIBEROEdgelerobot/LIBERO-Cosmos3LIBERO 평가
lerobot/Cosmos3-Nano-Policy-LIBERO_10Nanolerobot/LIBERO_10-Cosmos3LIBERO-10 평가
lerobot/Cosmos3-Edge-Policy-LIBERO_10Edgelerobot/LIBERO_10-Cosmos3LIBERO-10 평가
  • 기본 체크포인트에는 LeRobot 형식으로 공개된 가중치가 포함되어 있습니다. 해당 액션 모듈(action2llm, llm2action, action_modality_embed)은 아직 로봇을 위해 학습되지 않았으므로, --policy.load_action_modules=false를 사용하여 파인튜닝을 시작하세요(LIBERO에서 파인튜닝 참고).
  • 기본 체크포인트의 config.json에는 네이티브 DROID 후학습 레시피가 포함되어 있으며, 정책 유형의 클래스 기본값은 네이티브 LIBERO 레시피입니다.
  • LIBERO 체크포인트에는 해당 학습 실행의 EMA 가중치가 포함되어 있습니다.
  • 가중치는 bfloat16로 저장됩니다. config.json은 의도적으로 dtype: float32를 설정합니다. 이러한 체크포인트에서 혼합 정밀도 학습을 수행하면 네이티브 학습과 마찬가지로 float32 마스터 가중치를 유지합니다. 추론 및 평가 시에는 --policy.dtype=bfloat16를 전달하세요.

checkpoint에서 시작하는 방법은 두 가지가 있습니다:

  • --policy.path=<checkpoint>는 저장된 구성, 가중치 및 프로세서를 로드한 다음, 그 위에 커맨드 라인을 적용합니다.
  • --policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint> (또는 cosmos3_edge)는 가중치와 프로세서만 로드하며, 설정은 클래스 기본값인 LIBERO 레시피에서 시작됩니다.

LIBERO에서 Fine-tuning

클래스 기본값은 기술 보고서의 LIBERO 레시피입니다: 16개 동작의 청크, domain_id=5 embodiment, 상태 동작 없음, 동작의 quantile normalization, 추론 시 30단계 UniPC 및 classifier-free guidance 없음, 그리고 500 warmup 단계와 16000단계에서 0으로 선형 감쇠하는 5e-5의 학습률. 따라서 --policy.type이 있는 기본 체크포인트에서 시작하는 경우 --policy.load_action_modules=false만 필요합니다.

이 레시피는 Cosmos3 레이아웃으로 변환된 LIBERO 데이터로 학습합니다:

두 데이터셋 모두 네이티브 레시피의 데이터인 nvidia/LIBERO_LeRobot_v3에서 가져온 것입니다. 액션은 LIBERO의 7D axis angle 액션 대신 10D [delta position (3), rot6d (6), gripper]이며, 두 카메라 observation.images.image(3인칭)와 observation.images.image2(wrist)는 20 FPS에서 256x256 해상도입니다. 또한 액션 통계는 네이티브 레시피의 것입니다. 평가 시 LIBERO 환경이 액션을 원래 형태로 변환합니다(자세한 내용은 Evaluation 참조).

학습 명령

네이티브 LIBERO 레시피는 LIBERO-10에서 2000회 반복(최적 체크포인트는 1500회 반복 부근) 및 네 가지 스위트에서 5000회 반복(최적은 4500회 부근) 동안 전역 배치 2048 샘플로 학습합니다. 아래 명령은 8개 GPU가 있는 단일 머신에서 이를 재현합니다: 8개 GPU에 대한 FSDP2, GPU당 32 샘플, 8단계 그라디언트 축적으로, 32 x 8 x 8 = 2048 샘플이 옵티마이저 업데이트마다 사용됩니다. 모든 네이티브 반복 횟수는 8단계 축적 단계로 곱해집니다: 2000회 반복은 --steps=16000이 되고, 500회 반복마다의 체크포인트는 --save_freq=4000이 되며, 500의 워밍업과 16000의 감쇠 종료는 각각 4000과 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

네 가지 스위트에는 --dataset.repo_id=lerobot/LIBERO-Cosmos3와 --steps=40000(5000회 반복)을 사용하세요. Cosmos3 Edge에는 --policy.type=cosmos3_edge --policy.pretrained_path=lerobot/Cosmos3-Edge를 사용하세요.

주요 학습 매개변수

매개변수설명
--policy.load_action_modules=false…
--parallelism.dp_shard=88개 GPU에 대한 FSDP2. 정책이 FSDP2 래핑 단위를 선언하므로 --accelerator.fsdp.wrap_modules는 필요하지 않습니다. 멀티 GPU 학습 가이드를 참조하세요.
--accelerator.fsdp.mixed_precision_policy.*네이티브 학습과 마찬가지로 bfloat16로 연산 및 그래디언트 리덕션을 수행하고 float32 마스터 가중치를 사용합니다. 기본 --policy.dtype(float32)를 그대로 유지하세요. bfloat16 정책을 사용하면 bfloat16 마스터 가중치로 학습됩니다.
--accelerator.gradient_accumulation.steps전체 배치(global batch)는 --batch_size x GPU 수 x 누적 단계(accumulation steps)입니다. sync_each_batch=true를 사용하면 각 마이크로 배치(micro batch)마다 그라디언트가 감소(reduce)되므로, 샤드(sharded)된 상태로 유지되며 추가적인 언샤드(unsharded)된 사본이 필요하지 않습니다.
--steps, --save_freq, --policy.scheduler_*_stepsLeRobot는 옵티마이저 업데이트가 아닌 루프 스텝(마이크로 배치)을 카운트합니다. 네이티브 레시피의 반복 횟수(2000 또는 5000 iterations, warmup 500, decay 16000, 500마다 체크포인트)를 accumulation steps로 곱하세요. --steps는 --policy.scheduler_decay_steps 이하로 유지하세요. 해당 값 이후에는 학습률이 0이 됩니다.
--ema.type=power --ema.sigma_rel=0.1네이티브 레시피의 EMA입니다. 모든 체크포인트에는 pretrained_model_ema/에 EMA 가중치가 함께 저장되며, 이는 평가 및 공개에 사용되는 가중치입니다.
--tf32=falsefloat32 행렬 곱셈을 네이티브 학습과 동일하게 전체 정밀도로 유지합니다.
--policy.gradient_checkpointing=true메모리를 절약하기 위해 역전파 과정에서 디코더 레이어와 비전 블록의 활성화 값을 다시 계산합니다.
--policy.compile_model=true선택 사항: 디코더 레이어와 VAE 인코더의 torch.compile.

단일 GPU 또는 DDP를 사용하는 경우, FSDP2 플래그(--parallelism.dp_shard 및 --accelerator.fsdp.mixed_precision_policy.*)를 --accelerator.mixed_precision=bf16로 대체합니다. 예를 들어, 단일 GPU에서 Cosmos3 Edge를 1스텝당 32개의 샘플과 동일한 2048의 전역 배치 크기를 유지하기 위한 64개의 누적 스텝으로 실행하는 경우:

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는 체크포인트가 로드된 후 true로 리셋되므로, 저장된 실행은 학습된 액션 모듈을 다시 로드합니다. 다음 두 가지 경우에서 액션 모듈이 실수로 초기화 상태로 유지됩니다:

  • --policy.pretrained_path 또는 --policy.path 없이 정책을 빌드하는 경우: 아무것도 로드되지 않으므로 플래그는 false로 유지되며, 해당 실행의 후속 평가나 재개 시 학습된 액션 모듈이 건너뜁니다.
  • --peft.full_training_modules 외부의 action 모듈과 함께 PEFT를 사용할 경우: action 모듈은 절대 학습되지 않습니다.

--policy.load_action_modules=false는 베이스 체크포인트에서 파인튜닝할 때만 설정하세요. PEFT를 사용하는 경우, 액션 모듈을 --peft.full_training_modules(기본값)에 유지하세요.

하드웨어 요구 사항

gradient checkpointing를 활성화한 상태에서 이 페이지의 명령을 사용하여 A800 80 GB GPU에서 측정했습니다:

설정GPUGPU당 배치GPU별 피크 메모리
Nano, 전체 미세 조정(FSDP2)83238 GB
Nano, LoRA1464 GB
Edge, 전체 미세 조정13247 GB
Edge, LoRA1821 GB
bfloat16, Nano 및 Edge에서의 평가110개의 환경39 GB 및 15 GB
  • Cosmos3 Nano는 80 GB GPU가 필요하며, 전체 fine-tuning을 위해서는 최소 4개가 필요합니다. FSDP2가 모델을 sharding하기 전에 각 GPU는 전체 float32 모델(약 66 GB)을 보유해야 합니다.
  • 위의 LIBERO 명령은 8개의 A800 GPU에서 옵티마이저 업데이트당 약 45초가 소요되므로, LIBERO-10의 경우 약 25시간이 걸립니다.

자체 데이터로 Fine-tuning

Cosmos3는 명령줄과 데이터셋에서 학습 레시피를 가져옵니다. 정책에는 로봇에 특화된 로직이 포함되어 있지 않습니다. 아래 항목을 준수하는 모든 LeRobot 데이터셋이 작동합니다.

데이터셋 요구 사항

  • 카메라: 데이터셋 info.json의 순서대로 하나에서 세 개까지의 observation.images.* 특징. 모든 뷰는 --policy.view_size(기본값은 첫 번째 카메라의 크기)로 리사이즈되어 하나의 비디오 캔버스로 구성됩니다: 하나의 뷰는 그대로, 두 개는 나란히, 세 개는 첫 번째가 위에 있고 나머지 두 개가 아래에 나란히 배치됩니다. 카메라를 해당 레이아웃의 순서대로 배치하세요. 예를 들어 LIBERO의 경우 [third person, wrist], DROID의 경우 [wrist, exterior 1, exterior 2]입니다.
  • 설명 보기: --policy.view_description는 캔버스 레이아웃을 설명하는 프롬프트의 문장입니다. 기본값은 LIBERO 레이아웃(클래스 기본값)과 DROID 레이아웃(기본 config.json)을 설명합니다. 자신의 레이아웃에 맞는 문장을 작성하세요. 잘못된 설명은 감지되지 않으며 프롬프트의 성능만 약화시킵니다.
  • 액션 표현: 정책은 데이터셋의 액션을 그대로 학습하며 변환 단계가 없습니다. 학습하려는 표현으로 액션을 변환한 후 학습을 시작해야 하며, LIBERO 데이터는 rot6d 회전으로 이를 수행합니다. 환경이 다른 표현을 기대하는 경우, 환경 프로세서에서 액션을 원래 표현으로 되돌립니다(Cosmos3를 위한 LIBERO 환경이 이를 수행합니다).
  • 액션 폭: 최대 --policy.max_action_dim(64) 채널. 출력 폭은 데이터셋에서 가져옵니다.
  • 프레임 속도: fps는 데이터셋 메타데이터에서 가져오며 실제 제어 속도가 되어야 합니다. 프롬프트에 기록되고, 비디오 위치의 시간 축과 유휴 프레임의 임계값을 설정합니다.
  • 통계: 기본 동작 정규화는 QUANTILES이며, meta/stats.json에 있는 동작의 q01 및 q99 통계가 필요합니다(LeRobot으로 작성된 데이터셋에는 해당 통계가 포함되어 있습니다). 네이티브 DROID 레시피는 대신 IDENTITY를 사용합니다.
  • State: --policy.use_state=true를 사용하면 현재 상태가 액션 앞에 clean state action으로 배치됩니다. 이 경우 observation.state는 액션과 동일한 폭과 채널 의미를 가져야 합니다. 상태와 액션은 각각 고유한 통계로 정규화됩니다.
  • Language: 모든 프레임의 task는 프롬프트의 지시입니다.
  • Embodiment: --policy.domain_id은 32개 중 하나인 action projections의 행을 선택합니다. 네이티브 레시피는 LIBERO에 5, DROID에 8, RoboCasa에 30을 사용합니다. action 모듈이 초기화 상태에서 시작하는 경우(--policy.load_action_modules=false), 모든 행이 미학습 상태로 시작하므로 어떤 값이든 사용할 수 있습니다.
  • Chunk: --policy.chunk_size는 chunk당 실행할 action 수(LIBERO는 16, DROID는 32)이고, --policy.n_action_steps는 chunk당 실행할 action 수입니다.

None으로 남겨진 필드(fps, view_size)와 액션 채널 이름은 데이터셋에서 가져옵니다. 나머지 필드는 모두 구성에서 가져오므로, 커맨드 라인 또는 --config_path 파일에서 설정해야 합니다.

액션 및 상태 채널 이름. 프롬프트는 청크의 프레임 중 몇 프레임이 유휴 상태(로봇이 움직이지 않는 상태)인지 명시하며, Cosmos3IdleFrameProcessorStep은 meta/info.json의 액션 채널 이름에서 이를 계산합니다. [<prefix>_]<type>_<index> 형태의 이름은 채널 유형을 나타내며, pos, 회전(rot6d, rot9d, axisangle, euler_xyz, quat_xyzw, quat_wxyz), gripper 또는 joint이며, 팔을 구분하기 위한 접두사가 포함될 수 있습니다(예: left_rot6d_0). LIBERO 데이터는 pos_0..2, rot6d_0..5, gripper_0를, DROID 레이아웃은 joint_0..6, gripper_0를 사용합니다.

이러한 명명은 필수 사항이 아닙니다. 명명을 사용하지 않으면 정책은 경고만 기록하고 프롬프트는 유휴 프레임을 카운트하지 않으며, 이는 LIBERO 레시피가 기본적으로 관찰하는 내용과 정확히 일치합니다(네이티브 규칙 하에서 해당 에피소드에는 유휴 프레임이 없으므로). 네이티브 DROID 레시피의 유휴 프레임 카운트를 원할 경우에만 해당 명명을 따르세요.

DROID 레시피

베이스 체크포인트의 config.json는 네이티브 DROID 후학습 레시피입니다. 구체적으로, embodiment domain_id=8, 32개 액션의 청크, 현재 상태를 상태 액션으로 사용하며 해당 상태에 대해 유휴 프레임을 카운트하는 설정(use_state=true, idle_state=true), 정규화 없음(IDENTITY), 추론 시 4단계 UniPC와 classifier-free guidance scale 3, 그리고 warmup 없이 100000단계에서 0으로 선형 감쇠하는 학습률 8e-5(네이티브 2e-4에 배율 0.4를 적용한 값)를 사용합니다. 이 설정을 --policy.path로 시작하면 해당 레시피로 학습이 진행됩니다.

네이티브 레시피는 nvidia/Cosmos3-DROID의 성공 에피소드로 학습합니다. Cosmos3 레이아웃의 DROID 데이터셋에는 다음이 있습니다:

  • 세 개의 카메라 observation.images.wrist_left, observation.images.exterior_1_left, observation.images.exterior_2_left는 이 순서(위쪽은 wrist, 그 아래에 두 개의 exterior 뷰)로 배치되며, 원본 데이터의 해상도는 360x640입니다.
  • joint_0에서 joint_6 및 gripper_0로 명명된 8D 관절 위치 동작과 동일한 레이아웃을 가진 observation.state
  • 15 FPS의 프레임 레이트.

네이티브 레시피는 10000회 반복으로 전역 배치 8192개의 샘플을 학습하며, 1000회 반복마다 체크포인트를 저장합니다. GPU 8개가 장착된 단일 머신에서 GPU당 8개의 샘플과 128개의 누적 단계를 사용하면 동일한 전역 배치(8 x 8 x 128 = 8192)를 얻으며, 모든 반복 횟수는 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

네이티브 DROID 레시피는 256개의 GPU에서 실행되었습니다. 8개의 A800 GPU가 장착된 한 대의 머신에서 이 명령의 옵티마이저 업데이트 1회가 약 33분 소요되므로, 전체 레시피를 실행하려면 여러 대의 머신이 필요합니다. 모든 머신에서 accelerate launch(또는 torchrun)의 멀티 노드 플래그를 사용하여 동일한 명령을 실행하고, 각 머신 내에서 샤딩을 수행하며 머신 간에는 복제를 수행하도록 --parallelism.dp_replicate=<number of machines> --parallelism.dp_shard=8를 설정하고, --accelerator.gradient_accumulation.steps와 스텝 수를 머신 수로 나누세요. 멀티 GPU 트레이닝 가이드의 HSDP 섹션을 참조하세요.

새로운 로봇

새 로봇의 경우 --policy.type를 사용하여 클래스 기본값에서 시작하고, 데이터 필드를 명령줄에서 설정합니다. 예를 들어:

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

Backbone 구성

백본 구성 --policy.mot_backbone는 단일 커맨드 라인 값입니다. 변경할 필드만 포함된 YAML 딕셔너리를 전달하세요. 예를 들어, 사전 학습된 가중치 없이 빠른 테스트를 위해 레이어 수가 적은 모델을 사용할 수 있습니다:

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

누락된 필드는 크기의 기본값 또는 --policy.path로 로드된 구성의 값을 유지합니다. 부동소수점 수는 1.0e-6 형식으로 작성하세요(YAML은 1e-6을 문자열로 읽습니다). 또한 rope_parameters처럼 값이 dict인 필드는 전체를 명시해야 합니다. 아키텍처 필드는 로드하는 가중치와 일치해야 합니다.

PEFT(LoRA)를 활용한 Fine-tuning

Cosmos3는 LeRobot PEFT 통합(pip install -e ".[cosmos3,training,peft]")을 통해 LoRA를 지원합니다. 기본적으로 LoRA는 생성 경로의 어텐션 투영(q_proj_gen, k_proj_gen, v_proj_gen, o_proj_gen)을 적응시키고, 액션 모듈(action2llm, llm2action, action_modality_embed)은 전체 학습되며 어댑터와 함께 저장됩니다.

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, 학습률과 가중치 감쇠는 네이티브 Cosmos3 LoRA 레시피의 값입니다. --peft.lora_alpha가 없으면 PEFT는 8을 사용합니다.
  • --policy.load_action_modules=false를 사용하면 액션 모듈의 모양이 변경될 수 있습니다(--policy.max_action_dim, --policy.num_embodiment_domains). 이는 전체 파인튜닝에서만 작동합니다. PEFT 실행을 다시 로드하면 베이스가 로드되는데, 이는 동일한 모양이 필요합니다.
  • --policy.train_gen_stream_only와 --policy.freeze_vision_encoder는 PEFT와 함께 사용할 때 효과가 없습니다. 어댑터와 액션 모듈만 학습됩니다.
  • 액션 투영(action2llm, llm2action)을 --peft.target_modules에 추가하지 마세요. 이들은 각 embodiment별로 그룹화된 행렬 곱셈을 통해 직접 읽히는 테이블이므로 LoRA를 적용해도 변경되지 않습니다. 이미 전체적으로 학습됩니다.
  • PEFT를 사용할 때는 기본 float32 --policy.dtype를 유지하여, 전체 학습을 수행하는 action 모듈이 float32 마스터 가중치를 유지하도록 합니다.
  • native recipes는 LoRA 배치 크기나 스케줄을 정의하지 않습니다. 해당 명령은 LIBERO 스케줄(스텝 16000에서 0으로 선형 감쇠)을 유지하며, 스텝당 4개의 샘플을 사용합니다.
  • Cosmos3 Nano는 LoRA를 사용하여 80 GB GPU 1대에서 gradient checkpointing을 통해 학습할 수 있으며, GPU당 샘플 4개 기준 약 64 GB가 필요합니다(하드웨어 요구 사항 참조).

평가

LIBERO 체크포인트를 네 가지 스위트에서 평가합니다:

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
  • 자체 학습 실행의 체크포인트의 경우, EMA 가중치를 전달하세요: --policy.path=outputs/train/cosmos3_nano_libero_10/checkpoints/last/pretrained_model_ema.
  • --env.observation_height=256 --env.observation_width=256는 카메라를 학습 데이터의 256x256 해상도로 렌더링합니다.
  • LIBERO 환경은 Cosmos3 액션 변환을 프로세서에 자동으로 추가합니다: 10D rot6d 액션은 LIBERO의 7D axis angle 액션으로 다시 변환됩니다. lerobot-eval과 시뮬레이션 로봇 외부에서 정책을 실행하는 경우에도 lerobot.envs.factory의 make_env_pre_post_processors(env_cfg, policy_cfg)를 사용하여 환경 프로세서를 구성해야 합니다.
  • 평가 프로토콜 및 기타 옵션에 대해서는 LIBERO 가이드를 참조하세요.

예측된 비디오

Cosmos3은 미래 프레임과 동작을 함께 예측합니다. --policy.save_predicted_video=true를 전달하면 VAE로 각 chunk의 예측된 비디오를 디코딩하여 렌더링된 eval_episode_*.mp4 비디오 옆에 pred_episode_*.mp4를 작성합니다.

Cosmos3 Edge의 LIBERO-10 작업 “검은 그릇을 캐비닛의 아래 서랍에 넣고 닫기” 결과. 왼쪽: 롤아웃의 3인칭 및 손목 관측값. 오른쪽: 정책이 예측한 비디오로, 현재 관측값에서 16개 동작마다 다시 샘플링되었습니다.

결과

LIBERO

모델Libero-SpatialLibero-ObjectLibero-GoalLibero-10평균
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%

LIBERO의 폐쇄 루프 성공률. Cosmos3-Nano-Policy-LIBERO와 Cosmos3-Edge-Policy-LIBERO는 네 가지 LIBERO 스위트(lerobot/Cosmos3-Nano-Policy-LIBERO, lerobot/Cosmos3-Edge-Policy-LIBERO)로 학습되었습니다. Cosmos3-Nano-Policy-LIBERO_10과 Cosmos3-Edge-Policy-LIBERO_10은 LIBERO-10만으로 학습되었으며(lerobot/Cosmos3-Nano-Policy-LIBERO_10, lerobot/Cosmos3-Edge-Policy-LIBERO_10), Cosmos3-Nano-Policy-LIBERO_10의 학습 설정은 기술 보고서에 보고된 LIBERO-10 레시피를 따릅니다. Cosmos-Nano-TechReport는 기술 보고서의 Table 20에 보고된 바와 같이, LIBERO-10만으로 학습된 Cosmos3 Nano의 LIBERO-10 성공률입니다(MT-init, iteration 2000, 500 rollouts).

인용

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

라이선스

LeRobot 구현은 LeRobot의 Apache 2.0 License를 따릅니다. Cosmos3 가중치와 네이티브 코드는 OpenMDW-1.1 License 하에 공개됩니다.

Update on GitHub