使用 Cosmos3-Nano-Policy-DROID 加速镜像训练机器人策略模型

更新时间：2026-07-01

概述

Cosmos3-Nano-Policy-DROID 加速镜像面向机器人策略模型训练场景，内置 Cosmos Framework 训练环境，适用于基于 DROID 数据集进行 Cosmos3-Nano-Policy 模型训练与调优。

通过百舸平台，您可以直接选择已适配的训练镜像、GPU 资源规格、数据集和模型权重，快速启动分布式训练任务，无需手动拉起 Docker 容器或自行维护底层运行环境。

适用场景

使用 DROID 数据集进行机器人策略模型训练。
基于 Cosmos3-Nano 模型进行行动策略微调。
在百舸平台上验证单机或多机 GPU 训练性能。
需要使用预置训练镜像快速复现实验流程。

环境与资源要求

项目	推荐配置
GPU 规格	推荐使用 Pro6000 规格
单机训练	单机 8 卡
多机训练	双机 16 卡，可按需扩展
存储空间	如果您需要下载资源测试，建议不少于 600G
CUDA / 驱动	建议使用满足平台镜像要求的默认驱动环境
工作目录	`/root/workspace/cosmos-framework`
数据集目录	`/root/workspace/datasets/droid_lerobot_v30`
模型目录	`/root/workspace/models`
输出目录	`/root/workspace/outputs`

说明：本文中的路径均以百舸训练任务内的 /root/workspace 为例。实际使用时，可根据任务挂载目录调整。

准备工作

如果您所在网络环境访问 HuggingFace 较慢，建议提前配置 HuggingFace 镜像源，或配置代理加速。

huggingface 镜像配置：

                Bash
                
                export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_ENABLE_HF_TRANSFER=1

百舸平台已提供数据集和模型权重的 BOS 地址，您也可以直接从 BOS 下载到训练任务的工作目录，避免依赖外部网络。

步骤二：准备数据集

方式一：通过 HuggingFace 下载

如果您选择从 HuggingFace 下载数据集，请先配置镜像源，然后执行下载命令。

                Bash
                
                export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_ENABLE_HF_TRANSFER=1

mkdir -p /root/workspace/datasets/droid_lerobot_v30
huggingface-cli download lerobot/droid_1.0.1 \
  --repo-type dataset \
  --local-dir /root/workspace/datasets/droid_lerobot_v30

方式二：通过 BOS 下载

如果已提供 DROID 数据集的 BOS 地址，建议优先通过 BOS 下载。

Bash

1mkdir -p /root/workspace/datasets/droid_lerobot_v30
2bcecmd bos cp -r bos:/aihc-datasets/huggingface.co/datasets/lerobot/droid_1.0.1/v3.0/ /root/workspace/datasets/droid_lerobot_v30/

可选：规范化 tasks.parquet

部分 LeRobot v3 格式的 DROID 数据集可能会将任务文本写入 __index_level_0__ 列。当前训练代码已兼容 task 与 __index_level_0__ 两种 schema，因此该步骤不是必需项，仅用于将数据集元信息统一为更清晰的结构。

                Bash
                
                cd /root/workspace/cosmos-framework
python -m cosmos_framework.scripts.normalize_droid_tasks_parquet \
  /root/workspace/datasets/droid_lerobot_v30 \
  --json

如果数据集已符合规范，脚本会自动跳过处理。

验证数据集

下载完成后，可以检查数据集关键文件是否存在。

Bash

1python3 - <<'PY'
2from pathlib import Path
3root = Path('/root/workspace/datasets/droid_lerobot_v30')
4assert (root / 'meta/info.json').exists(), 'meta/info.json 不存在，请检查数据集路径'
5data_files = list((root / 'data').glob('chunk-*/file-*.parquet'))
6video_files = list((root / 'videos').glob('**/*.mp4'))
7print(f'data parquets: {len(data_files)}')
8print(f'video mp4s: {len(video_files)}')
9PY

步骤三：准备模型权重

训练需要准备 Cosmos3-Nano 权重和 Wan2.2-TI2V-5B 相关权重。

方式一：通过 HuggingFace 下载并转换

                Bash
                
            

                export HF_ENDPOINT=https://hf-mirror.com 
export HF_HUB_ENABLE_HF_TRANSFER=1
export HF_HOME=/root/workspace/models/huggingface
export HF_HUB_CACHE=$HF_HOME/hub
export TRANSFORMERS_CACHE=$HF_HOME/transformers
export HF_DATASETS_CACHE=$HF_HOME/datasets

cd /root/workspace/cosmos-framework
python -m cosmos_framework.scripts.convert_model_to_dcp \
  --checkpoint-path Cosmos3-Nano \
  -o /root/workspace/models/Cosmos3-Nano-DCP
            

方式二：通过 BOS 下载

如果已提供转换后的模型权重 BOS 地址，建议直接下载到 /workspace/models 下。

                Bash
                
                mkdir -p /root/workspace/models/Cosmos3-Nano/
bcecmd bos cp -r bos:/aihc-models-bj/nvidia/Cosmos3-Nano/ /root/workspace/models/Cosmos3-Nano/

mkdir -p /root/workspace/models/Wan2.2-TI2V-5B
bcecmd bos cp -r bos:/aihc-models-bj/Wan-AI/Wan2.2-TI2V-5B/ /root/workspace/models/Wan2.2-TI2V-5B/

mkdir -p /root/workspace/models/Qwen3-VL-8B-Instruct
bcecmd bos cp -r bos:/aihc-models-bj/Qwen/Qwen3-VL-8B-Instruct/ /root/workspace/models/Qwen3-VL-8B-Instruct/

因为原权重转换脚本会从 huggingface 在线下载权重，此处可以通过以下代码生成本地读取脚本进行替换：

                Bash
                
            

                cat > /root/workspace/cosmos-framework/convert_local.py << CONVERT_SCRIPT
import os
os.environ["COSMOS_DEVICE"] = "cpu"
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["HF_HOME"] = "/root/workspace/models/huggingface"

from cosmos_framework.inference.common.init import init_script
init_script(env={"COSMOS_DEVICE": "cpu"})

import cosmos_framework.utils.checkpoint_db as ckpt_db
_original_hf_download = ckpt_db._hf_download

LOCAL_FILE_MAP = {
    ("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth"): "/root/workspace/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
}
LOCAL_DIR_MAP = {
    "Qwen/Qwen3-VL-8B-Instruct": "/root/workspace/models/Qwen3-VL-8B-Instruct",
    "nvidia/Cosmos3-Nano": "/root/workspace/models/Cosmos3-Nano",
}

def _patched_hf_download(cmd_args):
    repo_id = cmd_args[0] if cmd_args else None
    filename = None
    for arg in cmd_args[1:]:
        if not arg.startswith("--") and "." in arg:
            filename = arg
            break
    if (repo_id, filename) in LOCAL_FILE_MAP:
        local = LOCAL_FILE_MAP[(repo_id, filename)]
        print(f"[PATCH] Redirecting file {repo_id}/{filename} -> {local}")
        return local
    if repo_id in LOCAL_DIR_MAP:
        local = LOCAL_DIR_MAP[repo_id]
        for i, arg in enumerate(cmd_args):
            if arg == "--include" and i + 1 < len(cmd_args):
                subdir = cmd_args[i + 1].split("/")[0] if "/" in cmd_args[i + 1] else ""
                if subdir and os.path.isdir(os.path.join(local, subdir)):
                    local = os.path.join(local, subdir)
                    break
        print(f"[PATCH] Redirecting dir {repo_id} -> {local}")
        return local
    print(f"[WARN] No local mapping for {repo_id}, attempting original download...")
    return _original_hf_download(cmd_args)

ckpt_db._hf_download = _patched_hf_download

import cosmos_framework.configs.base.defaults.vlm as vlm_mod
_original_download_tokenizer = vlm_mod.download_tokenizer_files

def _patched_download_tokenizer(model_name, config_variant):
    if model_name in LOCAL_DIR_MAP:
        print(f"[PATCH] Using local tokenizer: {LOCAL_DIR_MAP[model_name]}")
        return LOCAL_DIR_MAP[model_name]
    return _original_download_tokenizer(model_name, config_variant)

vlm_mod.download_tokenizer_files = _patched_download_tokenizer

from cosmos_framework.scripts.convert_model_to_dcp import main
main()
CONVERT_SCRIPT
            

执行以下脚本即可实现本地权重转换：

                Bash
                
                cd /root/workspace/cosmos-framework && source .venv/bin/activate
PYTHONPATH=. python convert_local.py \
  --checkpoint-path /root/workspace/models/Cosmos3-Nano \
  -o /root/workspace/models/Cosmos3-Nano-DCP

步骤四：配置训练环境变量

在百舸训练任务启动命令中，配置以下环境变量。路径可根据实际挂载位置调整。

                Bash
                
            

                export WORKSPACE_DIR=/root/workspace/cosmos-framework
export DROID_ROOT=/root/workspace/datasets/droid_lerobot_v30
export BASE_CHECKPOINT_PATH=/root/workspace/models/Cosmos3-Nano-DCP
export WAN_VAE_PATH=/root/workspace/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
export IMAGINAIRE_OUTPUT_ROOT=/root/workspace/outputs

export HF_HOME=/root/workspace/models/huggingface
export HF_HUB_OFFLINE=1

# 避免部分 Inductor/Triton 编译场景下出现资源不足问题
export TORCHINDUCTOR_MIX_ORDER_REDUCTION=0
export COSMOS_GPU_COLOR_JITTER=1
export COSMOS_AC_LAYER_POLICY=full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,full,none,none,none,none,none,none,none,none,none,none,none,none

export GPUS_PER_NODE=8
            

多机训练时，百舸平台通常会注入或提供分布式训练所需的节点信息。请根据平台任务环境设置以下变量。

                Bash
                
                export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
export MASTER_PORT=${MASTER_PORT:-"6100"}
export NNODES=${WORLD_SIZE:-"1"}
export NODE_RANK=${RANK:-"0"}

步骤五：启动训练

在百舸训练任务的启动命令中使用以下命令。

                Bash
                
            

                cd /root/workspace/cosmos-framework
source .venv/bin/activate

PYTHONPATH=. torchrun \
  --nnodes=$NNODES \
  --node_rank=$NODE_RANK \
  --master_addr=$MASTER_ADDR \
  --master_port=$MASTER_PORT \
  --nproc_per_node=$GPUS_PER_NODE \
  -m cosmos_framework.scripts.train \
  --sft-toml examples/toml/sft_config/action_policy_droid_repro.toml \
  -- trainer.max_iter=10000 \
     model.config.parallelism.data_parallel_replicate_degree=$NNODES \
     trainer.logging_iter=50 \
     checkpoint.save_iter=1000 \
     dataloader_train.max_samples_per_batch=32 \
     ++dataloader_train.dataloader.num_workers=8 \
     ++dataloader_train.dataloader.prefetch_factor=16 \
     ++dataloader_train.dataloader.pin_memory=true \
     ++dataloader_train.dataloader.persistent_workers=true \
     model.config.compile.enabled=true \
     model.config.compile.compiled_region=language \
     model.config.compile.compile_dynamic=true \
     model.config.compile.use_cuda_graphs=false \
     model.config.activation_checkpointing.mode=full \
     trainer.callbacks.compile_tokenizer.enabled=true \
     trainer.callbacks.compile_tokenizer.compile_after_iterations=3 \
     'trainer.callbacks.compile_tokenizer.warmup_resolutions=["480"]' \
     'trainer.callbacks.compile_tokenizer.aspect_ratio=3\,4' \
     trainer.callbacks.compile_tokenizer.mode=null \
     job.wandb_mode=disabled
            

性能参考

在 Pro6000 单机 8 卡测试环境下，参考性能如下。

指标	参考值
Iter Time	11.89 s
SPS / GPU	2.69
Global SPS	21.5
MFU	约 42%
显存占用	约 84.7 GiB
数据集初始化	约 25 s

以上数据为特定测试环境下的参考值。实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。

常见问题

训练过程中出现 Inductor 或 Triton 资源不足

如果训练过程中出现与 Inductor 或 Triton 编译相关的资源不足问题，可以优先确认是否已设置以下环境变量。

                Bash
                
                export TORCHINDUCTOR_MIX_ORDER_REDUCTION=0

如果仍然出现显存不足，可以尝试适当降低单批次样本数。

Bash

1dataloader_train.max_samples_per_batch=16

后续操作

训练任务启动后，可以在百舸平台查看任务日志、资源利用率和输出目录。训练产物默认输出到：

Plain Text

1/root/workspace/outputs

如需保留训练结果，可在任务结束后将输出目录同步到 BOS 或其他持久化存储。

评价此篇文章

有帮助没帮助

使用 OpenPI-0.5 加速镜像进行训练

常见问题

百度智能云

百度百舸 · AI计算平台

百度百舸 · AI计算平台

使用 Cosmos3-Nano-Policy-DROID 加速镜像训练机器人策略模型

概述

适用场景

环境与资源要求

准备工作

步骤二：准备数据集

方式一：通过 HuggingFace 下载

方式二：通过 BOS 下载

可选：规范化 tasks.parquet

验证数据集

步骤三：准备模型权重

方式一：通过 HuggingFace 下载并转换

方式二：通过 BOS 下载

步骤四：配置训练环境变量

步骤五：启动训练

性能参考

常见问题

训练过程中出现 Inductor 或 Triton 资源不足

后续操作