Cosmos3

Cosmos3 是一个用于机器人控制的世界动作模型，基于 NVIDIA 的全模态世界模型系列 Cosmos 3 构建。根据相机观测和语言指令，该策略联合预测未来的视频帧和下一段机器人动作。LeRobot 的实现改编自原生的 cosmos-framework 动作策略，并运行已发布的、转换为 LeRobot 格式的 Cosmos3 Nano 和 Edge 模型。

• 📄 Cosmos 3 技术报告 (arXiv:2606.02800)
• 💻 cosmos-framework 和 Cosmos cookbook
• 🤗 Cosmos 3 模型集合

模型概览

Cosmos 3 是一种 Transformer 混合体（Mixture-of-Transformers，MoT）。每个解码器层都有两条路径，各自拥有独立的层归一化、注意力投影和 MLP，这两条路径仅在共享的自注意力中汇合：

• 理解路径（推理器）使用因果注意力读取提示的文本标记，以及可选的 ViT 图像标记。
• 生成路径（生成器）读取扩散子序列的连续 token：视频的 Wan2.2 VAE 潜变量以及动作 token。它对两条路径进行双向注意力计算，并使用 flow matching 进行训练。

Cosmos 3 的 Mixture-of-Transformers (MoT) 架构（技术报告中的图 5）。左侧：自回归（AR）子序列经过推理器，扩散（DM）子序列经过生成器，两者在共享的多模态注意力中汇合。右侧：注意力掩码，AR 部分为因果掩码，扩散部分为全掩码。

作为策略，Cosmos3 以原生动作配方中的世界动作模型模式运行：

• 输入：将各个相机视角组合成一个视频画布（一个视角保持原样，两个并排，三个则第一个在上方，另外两个在下方）。提示词包含任务、帧率、画布分辨率、画布布局描述以及该数据块中空闲帧的数量。
• 训练：策略使用 VAE 对当前帧的片段以及 chunk_size 个未来帧进行编码，为未来潜在变量和动作块添加噪声，并学习对两者进行去噪。损失是视频损失和动作损失的加权和。
• 推理：条件帧是干净的，未来视频和动作通过 UniPC 求解器联合采样，策略返回动作块。在评估期间可以保存预测的视频（参见 评估）。
• 默认情况下，仅生成路径、其周围的层以及动作模块参与训练（train_gen_stream_only=true、freeze_vision_encoder=true），这与原生后训练方案一致。

LeRobot 将发布的两种尺寸作为两种策略类型提供：

安装要求

1. 按照我们的安装指南安装 LeRobot。

2. 安装 Cosmos3 依赖项：

   已复制

   pip install -e ".[cosmos3]"

   如果您从 PyPI 安装了 LeRobot：

   已复制

   pip install 'lerobot[cosmos3]'

3. 添加工作流所需的附加组件：training 用于 lerobot-train，libero 用于在 LIBERO 上进行训练或评估，以及 peft 用于 LoRA。例如：

   已复制

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

Cosmos3 默认使用 PyTorch 的缩放点积注意力后端（sdpa）。如果您的环境中安装了兼容的 flash_attn，可以请求 --policy.attn_implementation=flash_attention_2。请勿将其与 --policy.compile_model=true 结合使用：编译后的解码器层会为每个使用 flash attention 的层重新编译，并因重新编译错误而停止。

检查点

• 基础检查点以 LeRobot 格式存储了发布的权重。其动作模块（action2llm、llm2action、action_modality_embed）尚未针对机器人进行训练，因此请从使用 --policy.load_action_modules=false 开始微调（参见 在 LIBERO 上进行微调）。
• 基础检查点的 config.json 包含原生的 DROID 后训练配方，而策略类型的类默认值则是原生的 LIBERO 配方。
• LIBERO 检查点保存了其训练运行中的 EMA 权重。
• 权重以 bfloat16 格式存储。config.json 有意将 dtype: float32 设置为 float32：从这些检查点进行混合精度训练时，会保留 float32 主权重，这与原生训练的行为一致。在进行推理和评估时，请传入 --policy.dtype=bfloat16。

从检查点开始有两种方式：

• --policy.path=<checkpoint> 会加载已存储的配置、权重和处理器，然后在其上应用命令行参数。
• --policy.type=cosmos3_nano --policy.pretrained_path=<checkpoint>（或 cosmos3_edge）仅加载权重和处理器，配置则从类默认值，即 LIBERO 配方开始。

在 LIBERO 上进行微调

该类的默认值对应技术报告中的 LIBERO 配方：16 个动作的块，具身 domain_id=5，无状态动作，对动作进行分位数归一化，推理时使用 30 步的 UniPC 且不使用无分类器引导，学习率为 5e-5，预热 500 步，并在第 16000 步时线性衰减至零。因此，从基础检查点开始，使用 --policy.type 时只需设置 --policy.load_action_modules=false。

该配方在转换为 Cosmos3 布局的 LIBERO 数据上进行训练：

• lerobot/LIBERO-Cosmos3：四个套件（Spatial、Object、Goal、10），共 1693 个回合。
• lerobot/LIBERO_10-Cosmos3：仅包含 LIBERO-10，共 379 个回合。

两者均来自 nvidia/LIBERO_LeRobot_v3，即原生配方（native recipe）的数据。动作为 10 维 [delta position (3), rot6d (6), gripper]，而非 LIBERO 的 7 维轴角（axis angle）动作；两个相机 observation.images.image（第三人称视角）和 observation.images.image2（腕部视角）的分辨率为 256x256，帧率为 20 FPS；动作统计量也采用原生配方的数值。在评估期间，LIBERO 环境会将动作转换回原始格式（参见 评估）。

训练命令

原生 LIBERO 配方在 LIBERO-10 上以 2048 的全局批次大小训练 2000 次迭代（最佳检查点位于第 1500 次迭代附近），并在四个套件上训练 5000 次迭代（最佳位于第 4500 次迭代附近）。以下命令可在配备 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。

关键训练参数

在单 GPU 或使用 DDP 时，请将 FSDP2 标志（--parallelism.dp_shard 和 --accelerator.fsdp.mixed_precision_policy.*）替换为 --accelerator.mixed_precision=bf16。例如，在单 GPU 上运行 Cosmos3 Edge，每步 32 个样本，64 个累积步，以保持相同的 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 在检查点加载完成后会被重置为 true，因此保存的运行会重新加载其训练好的 action modules。以下两种情况会错误地使 action modules 保持其初始化状态：

• 在不使用 --policy.pretrained_path 或 --policy.path 的情况下构建策略：不会加载任何内容，因此该标志保持为 false，之后对该运行进行评估或恢复时会跳过其已训练的动作模块。
• 在 --peft.full_training_modules 之外使用 PEFT 的动作模块：此时动作模块将永远不会被训练。

仅在从基础检查点进行微调时，才设置 --policy.load_action_modules=false。在使用 PEFT 时，请将 action modules 保留在 --peft.full_training_modules 中（默认值）。

硬件要求

使用本页中的命令和梯度检查点，在 A800 80 GB GPU 上测量：

• Cosmos3 Nano 需要 80 GB 的 GPU，并且全量微调至少需要 4 个：在 FSDP2 对模型进行分片之前，每个 GPU 都会持有完整的 float32 模型（约 66 GB）。
• 上述 LIBERO 命令在 8 块 A800 GPU 上每次优化器更新大约需要 45 秒，因此完成 LIBERO-10 大约需要 25 小时。

使用您自己的数据进行微调

Cosmos3 从命令行和训练数据集中获取训练配方：策略中不包含任何与特定机器人相关的逻辑。任何 LeRobot 数据集，只要遵循以下要点，即可使用。

数据集要求

• 相机：一到三个 observation.images.* 特征，顺序与数据集 info.json 中的顺序一致。每个视角都会调整为 --policy.view_size 大小（默认为第一个相机的尺寸），并组合成一个视频画布：一个视角原样显示，两个并排显示，三个则第一个在上方，另外两个在下方并排显示。请按该布局顺序放置相机，例如 LIBERO 使用 [third person, wrist]，DROID 使用 [wrist, exterior 1, exterior 2]。
• 查看描述：--policy.view_description 是提示词中描述画布布局的句子。默认值分别描述了 LIBERO 布局（类默认值）和 DROID 布局（基础 config.json）。请编写描述您自己布局的句子：错误的描述不会被检测到，只会削弱提示词的效果。
• 动作表示：策略按数据集原样学习动作，没有转换步骤。在训练前，请将动作转换为您希望学习的表示形式，例如 LIBERO 数据使用 rot6d 旋转。如果环境期望另一种表示形式，请在环境处理器中将动作转换回去（LIBERO 环境对 Cosmos3 执行此操作）。
• 动作宽度：最多为 --policy.max_action_dim（64）个通道。输出宽度由数据集决定。
• 帧率：fps 来自数据集元数据，且必须是真实的控制速率：它会被写入提示词，并用于设置视频位置的时间轴以及空闲帧的阈值。
• 统计信息：默认的动作归一化方式为 QUANTILES，这需要 meta/stats.json 中动作的 q01 和 q99 统计值（由 LeRobot 写入的数据集包含这些统计值）。原生 DROID 配方则使用 IDENTITY。
• 状态：使用 --policy.use_state=true 时，当前状态会作为干净的状态动作放置在动作之前。observation.state 需要与动作具有相同的宽度和通道含义。状态和动作会分别使用各自的统计量进行归一化。
• 语言：每一帧的 task 都是提示词的指令。
• 具身：--policy.domain_id 用于选择动作投影的行，共 32 行。原生配方中，LIBERO 使用 5，DROID 使用 8，RoboCasa 使用 30。当动作模块从初始化状态开始（--policy.load_action_modules=false）时，每一行都未经验证，因此任意值均可使用。
• Chunk：每个 --policy.chunk_size 包含的动作数（LIBERO 为 16，DROID 为 32），以及每个 --policy.n_action_steps 执行的动作数。

保持为 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 原生后训练配方：具身 domain_id=8，每块 32 个动作，将当前状态作为状态动作，并在其上计算空闲帧（use_state=true、idle_state=true），不进行归一化（IDENTITY），推理时使用 4 步的 UniPC 和 3 的无分类器引导尺度，学习率为 8e-5（原生 2e-4 乘以 0.4 的系数），无预热，并在第 100000 步线性衰减至零。使用 --policy.path 从此处开始训练时，将采用此配方。

原生配方在 nvidia/Cosmos3-DROID 的成功回合上进行训练。采用 Cosmos3 布局的 DROID 数据集包含：

• 三个摄像头 observation.images.wrist_left、observation.images.exterior_1_left 和 observation.images.exterior_2_left，按此顺序排列（腕部摄像头在上方，两个外部视角在下方），原始数据中的分辨率为 360x640。
• 8 维关节位置动作，命名为 joint_0 至 joint_6 以及 gripper_0，并且 observation.state 具有相同的布局。
• 15 FPS 的帧率。

原生配方使用 8192 个样本的全局批次训练 10000 次迭代，每 1000 次迭代保存一次检查点。在配备 8 个 GPU 的单台机器上，每个 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 的单台机器上，执行此命令的一次优化器更新大约需要 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 \
  ...

骨干网络配置

骨干网络配置 --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），请完整提供。架构字段必须与您加载的权重相匹配。

使用 PEFT (LoRA) 进行微调

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 主权重。
• 原生配方未定义 LoRA 的批处理大小或调度策略。该命令保留了 LIBERO 的调度策略（在第 16000 步时线性衰减至零），每步使用 4 个样本。
• Cosmos3 Nano 配合 LoRA 可在单张 80 GB GPU 上训练，启用梯度检查点时，每 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 动作转换添加进去：10 维 rot6d 动作会被转换回 LIBERO 的 7 维轴角动作。当您在 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 任务“将黑色碗放入橱柜的底层抽屉并关上”中的表现。左图：rollout 过程中的第三人称视角和手腕视角观测。右图：策略预测的视频，从当前观测开始，每 16 个动作重新采样一次。

结果

LIBERO

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 是 Cosmos3 Nano 仅在 LIBERO-10 上训练时的 LIBERO-10 成功率，如技术报告表 20 中所述（MT-init，迭代 2000，500 次 rollout）。

引用

@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 许可证。Cosmos3 权重和原生代码在 OpenMDW-1.1 许可证 下发布。