SuperAPI AI 工具集

cosmos3-nano-sim-bimanual-post-training

智能体与模型

查看项目

更新于 10-23 00:10

技能介绍

该技能用于后训练 Cosmos3-Nano-Sim-Bimanual:以 Cosmos3-Nano 为初始化模型,训练一个动作条件世界模型,使其根据任务描述、RGB 观测和对齐的动作轨迹,预测双臂动作带来的视觉后果。模型输出视频而非机器人动作。整个流程从 Cosmos3-Nano 开始,按 Stage 0 至 Stage 3 四个原生阶段依次训练。

适用场景与起点:需要全新四阶段训练时,从 Cosmos3-Nano 和已完成的本地训练 bundle 开始,走 Stage 0 Quickstart;需要延续已有训练链时,使用已完成的阶段检查点和保存的配置继续训练;只需生成视频时,使用已发布的 Bimanual 模型或自行导出的 Stage 3 学生模型进行推理。修改配方前应先阅读数据契约(动作含义、归一化、时序与字幕)和参数指南。

环境要求:使用独立的 Framework checkout、Python 3.13、文档所述的 CUDA/FFmpeg 系统库及 GPU 分配,并设置 FRAMEWORK_DIR。在 checkout 根目录激活准备好的训练环境后运行训练脚本;不要清除 LD_LIBRARY_PATH。训练 extra 包含 mujoco==3.10.0 用于校准的动作/FK 预处理。注册布局使用 FSDP shard degree 8 和 context-parallel degree 4,八进程命令需在支持该布局和内存需求的分配内运行,多节点部署遵循 Framework 训练方式。

数据准备:将 BIMANUAL_BUNDLE 指向已完成的用户本地 bundle,其中包含 manifest.json、samples.json、本地媒体/动作资源和 resources/tokenizer/;仅有推理 JSONL 或动作 sidecar 不构成训练 bundle,Framework 不提供从原始机器人数据到 bundle 的转换命令。可从 Sim-Bimanual 检查点下载带版本的单源机器人示例以理解 bundle 布局和原生 dataloader 交接,该示例是小规模集成样例而非收敛数据集。准备数据时应在创建窗口前划分源 episode,并保持同一 episode 的所有窗口在同一 split;按任务和内存选择机器人窗口长度,保持 RGB 与动作对齐(T+1 观测对应 T 个转移);保留坐标帧、rot6d 约定、校准夹爪宽度、domain ID、有效性掩码和固定的全局 asinh 统计量。训练通过原生 transform 物化原始采集,推理 sidecar 中已归一化的值不是原始训练动作。应以一个受支持的机器人源开始,高级用户可提供多个受支持源,Framework 会归一化其声明的相对权重;Framework 内部会应用阶段特定的渲染,因此不要创建单独的源分辨率分层。检查代表性解码样本和实际打包批次,包括补齐到 64 的 59 个语义动作值及其掩码、单位和语义。绑定准备好的 bundle 后,先校验文件,再在加载模型权重前读取一个打包的 Stage 0 样本,两条命令均须报告 passed: true,第二条还须报告 materialized: true、batch 样本/ token 数和张量形状/类型。原生 make_bimanual_local_dataloader 负责阶段特定的流选择与打包,并安装 loader 状态续训支持。

模型与 VAE 准备:将 NANO_REVISION 设为选定的不可变 nvidia/Cosmos3-Nano 版本,使用 hf download 下载 Cosmos3-Nano 和 Wan2.2_VAE.pth,再用 convert_model_to_dcp 转换模型并记录模型/VAE 身份和转换版本。此步骤产出模型权重而非优化器或训练器状态。原生 Stage 0 配方在全新加载时跳过 net_ema. 和三个新动作模块,常规权重加载保持严格,并通过内置回调将已加载/全新常规权重复制到 EMA,不要添加第二个 EMA 回调;同一运行续训则恢复已保存的 EMA。

Stage 0 配置与启动:设置 IMAGINAIRE_OUTPUT_ROOT 为所有 worker 上相同绝对路径的可写存储,选择唯一 RUN_NAME,使用 examples/toml/sft_config/sim_bimanual_stage0_bidir.toml。先以 --dryrun 解析配置,确认本地 bundle/tokenizer、Cosmos3-Nano 加载路径、Stage 0 EMA 回调、并行度和 job.wandb_mode=disabled;训练不需要发布版推理检查点或 W&B key。随后用 torchrun --standalone --nproc-per-node=8 启动。可选 12 次迭代冒烟测试:使用独立 run 名并加 trainer.max_iter=12 和 checkpoint.save_iter=12,确认损失与梯度有限且第 12 次迭代完成检查点保存;正式训练时省略这两个覆盖项。监控损失、梯度和批次语义,保持作业运行至检查点保存完成;中断后以相同 run 名、输出根、bundle、拓扑和父检查点重跑同一命令即可恢复迭代、优化器、调度器、EMA 和 loader 状态。

四阶段继续训练:原生参考预算为 Stage 0–3 分别 35,000、8,000、4,000 和 1,500 次训练器迭代。四个 TOML 均位于 examples/toml/sft_config/:Stage 0(表示适配)为 sim_bimanual_stage0_bidir.toml,需 COSMOS_SIM_BIMANUAL_GA_CHECKPOINT;Stage 1(teacher forcing)为 sim_bimanual_stage1_tf.toml,需 COSMOS_SIM_BIMANUAL_STAGE0_CHECKPOINT;Stage 2(TF-dCM)为 sim_bimanual_stage2_tfdcm.toml,学生和冻结教师均需 COSMOS_SIM_BIMANUAL_STAGE1_CHECKPOINT;Stage 3(SF-DMD)为 sim_bimanual_stage3_sfdmd.toml,学生需 COSMOS_SIM_BIMANUAL_STAGE2_CHECKPOINT,参考/critic 需 COSMOS_SIM_BIMANUAL_STAGE0_CHECKPOINT。接受 Stage 0 后绑定其精确的已完成迭代根目录而非 model/,重复配置检查和启动,各阶段使用不同 run 名。每个新阶段从全新优化状态开始,同阶段续训恢复现有运行。Stage 3 保持四步调度:每六个训练器迭代更新一次学生,其余五次更新 critic,因此 12 次迭代包含两次学生更新和十次 critic 更新,应检查两个优化器而非仅迭代计数;冻结的教师/参考和 critic 权重不是推理产物。

导出与推理检查:仅在有限更新、完成保存和全新进程回读后接受各阶段,保留解析配置和不可变检查点身份,用留出生成判断模型质量。导出可移植的 Stage 3 学生模型时,将 STAGE3_CHECKPOINT_ROOT 设为其完成的迭代根、STAGE3_CONFIG 设为该运行的 config.yaml、STUDENT_EXPORT 设为新导出目录,运行 export_model 并显式选择学生 EMA 权重;保留所选变体和保存的四步调度,检查导出的动作元数据和本地路径清理,随后使用该导出进行推理。已发布的 schema-5 示例同样使用四步,应保留每个导出实际保存的调度阶梯,不要假定各调度完全相同。