对 NVIDIA Cosmos 3 Edge 进行后训练，实现设备端机器人控制

机器人在板载计算硬件上运行时，需要能够适应其传感器、环境和任务的策略。世界模型为学习物理交互提供了基础，但其规模可能会导致难以在设备上部署。新的 NVIDIA Cosmos 3 Edge 带来了变化。

Cosmos 3 Edge 是 Cosmos 3 系列中的 4B 全模型 (配备基于 2B NVIDIA Nemotron 的推理器) 。它使用与 NVIDIA Cosmos 3 Nano 和 NVIDIA Cosmos 3 Super 相同的物理世界数据进行预训练，并以物体移动和交互方式相同的基础进行训练。此外，该模型足够小，可以在 NVIDIA Jetson Thor 上本地运行。

您将构建的内容

本教程结束后，您将获得经过后训练的 Cosmos 3 边缘操作策略，该策略可以在 Jetson Thor 上运行，并在闭环仿真中进行评估。

您将学习如何：

1. 对 Cosmos 3 Edge 进行后训练，以预测机器人动作。
2. 在 Jetson Thor 上提供生成的策略。
3. 在水平下降控制循环中运行推理。
4. 评估闭环仿真中的策略行为。

每个步骤都可从打开的 cosmos-framework 存储库中重现，并且已发布的检查点可在 HuggingFace 上获取。

为何要对 Cosmos 3 Edge 进行后训练，以便在设备上操作机器人？

世界基础模型基于大型多模态数据集进行预训练，这些数据集可捕获物体运动和物理交互的模式，例如物体的下落、滑动和对接触的反应。Cosmos 基于物理理解和预测能力原生生成动作。这种先验知识为机器人策略训练提供了一个有用的起点，而不是要求策略从特定任务的演示中学习所有物理关系。

但是，在实体机器人上部署这些模型会带来两个实际限制：

• 设备内存：模型及其运行时状态必须适合机器人的可用内存。
• 控制延迟：整个推理工作流必须足够快，以支持机器人所需的控制频率。

后训练 Cosmos 3 Edge 可解决上述每项限制。生成的策略模型适合 Jetson Thor 的内存，因此推理直接在机器人上运行，而不是将其卸载到数据中心 GPU。

其余因素是控制延迟。DROID 动作策略直接在机器人上实时运行。在 NVIDIA Jetson AGX Thor T5000 上，它会在大约 1.53 秒内生成每个动作数据块 (分辨率为 640 × 540，15 Hz) ，而单个数据块可完成大约 2.13 秒的机器人动作。由于下一个数据块在当前数据块完成之前就已准备就绪，因此机械臂会不断移动，且循环中没有数据中心 GPU。该策略通过在每个推理周期后生成动作块和重新规划，支持在设备上持续串流。它不会在每次观察后重新进行规划。

在闭环RoboLab任务中，后训练策略的成功率为22.9%。这些结果表明，4B世界基础模型可以作为实用的实时设备端策略主干。此模型适用于Jetson AGX Thor，完全在机器人上运行。

策略基于哪些数据进行训练？

本教程中的发布策略基于 nvidia/Cosmos3-DROID 数据集进行训练。它包含 76000 个成功的远程操作轨迹，涉及 86 项任务和 564 个场景，大约 350 小时，使用 Franka Panda 手臂和 Robotiq 抓手收集。

数据集以 640 × 360 分辨率的 LeRobotDataset v3.0 格式打包。分三个阶段进行准备：

• 过滤空闲帧和非任务帧。
• 选择成功的演示进行训练。
• 在训练期间应用随机裁剪、重新缩放和颜色抖动。

hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir Cosmos3-DROID

引入您自己的机器人数据

将您的数据转换为 LeRobot 数据集 v3 (每帧摄像头视频、关节状态、抓手状态、动作、任务指令) 。

对于类似 DROID 的 Franka 设置，主要更改是数据集路径。不同的具身需要自己的实验配置来定义动作空间、维度、摄像头布局和归一化设置。

Cosmos 3 支持多种形态，包括双臂 Franka、UR、WidowX 250、LeRobot SO101。详情请参阅 Cosmos 3 Edge 模型卡。

预备知识

开始之前，请确认您拥有：

• 最新 Cosmos 框架 版本。
• 经过验证的训练硬件：搭载 NVIDIA GB200 Grace Blackwell 超级芯片的 NVIDIA DGX 工作站或 NVIDIA GB300 Grace Blackwell Ultra 桌面超级芯片。
• 受支持的 NVIDIA CUDA 和容器版本：CUDA 13.0 (cu130) 、NGC 26.06-py3。
• 访问Cosmos 3 Edge 基本检查点。
• 访问 Cosmos3-DROID 数据集。
• Hugging Face 访问令牌。

这是基础模型的后训练，而不是单个 GPU 的微调。经过验证的运行使用 64 个节点 ( 4 个 GB200 以上) 进行 6 万次迭代，大约为 68 小时 (约 17.4 K GB200 小时) 。相应地规划计算。

如何运行后训练

后期训练一般分为四个步骤：

1. 下载数据集。
2. 将基础检查点转换为分布式检查点 (DCP) 格式。
3. 应用筛选器。
4. 启动

# Download the Cosmos3-DROID dataset (success split)
hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir /path/to/Cosmos3-DROID

# Convert the base checkpoint to DCP (one-time; runs on CPU — it repacks weights, no GPU math)
python -m cosmos_framework.scripts.convert_model_to_dcp \ -o /path/to/Cosmos3-Edge-dcp --checkpoint-path Cosmos3-Edge

# Launch post-training
bash examples/launch_sft_action_policy_droid_nano.sh

编辑并启动后训练脚本

在资源库中找到此文件：

cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py

该启动器注册了 Cosmos 3 Nano 的 recipe。要训练 Cosmos 3 Edge，请在启动前进行以下三次编辑：

• 将 NANO_MODEL_CONFIG 导入替换为 EDGE_MODEL_CONFIG (来自 configs/base/experiment/sft/models/edge_model_config.py) 。
• 将 BASE_CHECKPOINT_PATH 设置为转换步骤中的 Cosmos 3 Edge DCP Checkpoint。
• 将 examples/launch_sft_action_policy_droid_nano.sh 重命名为 examples/launch_sft_action_policy_droid_edge.sh

其他一切 (包括数据集、动作空间、管护过滤器和训练时间表) 都保持不变。

# Launch post-training
bash examples/launch_sft_action_policy_droid_edge.sh

有关最新的端到端指令 (包括检查点转换、环境配置和管理过滤器设置) ，请参阅 DROID 后训练再现指南和模型卡。这些资源会与资源库一起更新。

如何在 Jetson Thor 上部署和运行策略

该策略由使用 OpenPI 协议的 WebSocket 策略服务器提供服务，OpenPI 协议与 DROID 策略生态系统中使用的协议相同。客户端发送观察字典；服务器则返回动作块。对于 Edge，服务器在 Jetson Thor 上以原生方式运行。在 BF16 中，权重大约为 9 GB，适合 Thor 的板载内存，因此策略服务器和控制客户端都在机器人上运行，没有数据中心 GPU 在循环中。

启动机器人上的策略服务器

export HF_TOKEN=<your_hf_token>
# Thor: run eager; stock Triton wheels lack sm_110a kernels
export TORCHDYNAMO_DISABLE=1
python -m cosmos_framework.scripts.action_policy_server_robolab \
  --checkpoint_path nvidia/Cosmos3-Edge-Policy-DROID \
  --port 8000 \
  --format-prompt-as-json True
# first launch downloads weights and takes a few minutes; then:
curl localhost:8000/healthz   # -> OK

由于服务器在 Thor host="localhost" 上运行，因此请求从未离开机器人。该策略在设备上实时运行。这是 Cosmos Edge 实现的本地推理。

烟雾测试 (无需机器人)

DROID 策略受状态限制。这意味着 joint_position 和 gripper_position 是真实模型输入，而非样板文件。由于循环中没有机械臂，因此零是可靠的占位符；此调用只是证明服务器返回了格式良好的动作块。

import numpy as np
from PIL import Image
from openpi_client.websocket_client_policy import WebsocketClientPolicy

client = WebsocketClientPolicy(host="localhost", port=8000)

observation = {
    "prompt": "put the marker in the basket",
    "observation/wrist_image_left": np.asarray(Image.open("wrist.jpg")),
    "observation/exterior_image_1_left": np.asarray(Image.open("left.jpg")),
    "observation/exterior_image_2_left": np.asarray(Image.open("right.jpg")),
    "observation/joint_position": np.zeros(7, dtype=np.float32),   # smoke test only
    "observation/gripper_position": np.float32(0.0),               # smoke test only
}
result = client.infer(observation)
actions = result["action"]  # [32, 8]: 7 joint positions + gripper, 15 Hz

在机器人上运行

相同的调用位于重新规划循环中。每个循环会读取新的摄像头以及测量到的机械臂关节和抓手位置 (替换上述零) ，执行数据块的前缀，然后再次询问：

def get_observation():
    return {
        "prompt": "put the marker in the basket",
        "observation/wrist_image_left": read_camera("wrist"),
        "observation/exterior_image_1_left": read_camera("left"),
        "observation/exterior_image_2_left": read_camera("right"),
        "observation/joint_position": robot.joint_positions(),   # REAL, 7 floats
        "observation/gripper_position": robot.gripper_position(),
    }

while not task_done():
    result = client.infer(get_observation())
    execute_actions(result["action"][:16], hz=15)   # first 16 of 32, then replan

仅执行前缀并重新规划是一种标准做法：策略每几秒钟就会自我校正一次，而且由于它是受状态条件限制的，因此每次重新规划都从 ARM 实际所在的位置开始，而不是从前一个数据块假设的位置开始。启用视频解码后启动服务器会返回策略的预期部署以及操作，因此您可以检查世界模型对其发出的数据块的预测。

在闭环仿真中评估策略

您不需要物理机器人来验证策略。事实上，建议先在仿真中评估策略，然后再直接在实体机器人上进行评估，以避免意外行为影响开发者或机器人。 RoboLab 是 Isaac Lab-Arena 排行榜背后的 Isaac Lab-Arena 基准测试，现已开放，其客户端连接到同一策略服务器。它执行物理中的每个动作块，并将渲染的观察结果流式传输回，从而在 120 个受语言限制的操作任务中实现真正的闭环。

# Clone RoboLab and fetch the scene assets
git clone https://github.com/NVlabs/RoboLab.git && cd RoboLab
git lfs pull

# Build the sim image and run one task against the policy server
./docker/build_docker.sh latest
./docker/run_docker.sh latest
python policies/cosmos3/run.py --task BananaInBowlTask

# Or run headless across many envs for success-rate statistics
python policies/cosmos3/run.py --task BananaInBowlTask --num-envs 10 --headless

将 --task 替换为 120 个任务中的任何一个，以构建更广泛的策略行为概览。每次运行都会生成视口和机器人摄像头视频以及成功日志，因此您可以检查结果和生成结果的轨迹。

在闭环 RoboLab 评估中，经过后训练的边缘策略在整个任务套件中的成功率为 22.9%。获得这一结果所需的推理计算量只是更大型 Cosmos3 变体的一小部分 ( Nano 达到 36.8%) ，而这正是 trade Edge 旨在实现的：以极具竞争力的成功率实现实时、完全在机器人上的自主性。

注意：NVIDIA Isaac Sim 5.x 需要 NVIDIA RTX 服务器驱动版本 580 或更高版本。构建前，请在 RoboLab 和 Isaac Sim 文档中确认受支持的驱动版本。

超出机器人控制的后训练

改善机器人控制的后训练技术也扩展到其他功能，例如用于机器人训练的合成数据生成。开发者可以创建专门的世界模型，生成针对其环境和任务定制的高质量合成数据，用于室内外机器人训练。

例如，Aigen 利用经过后训练的 Cosmos 生成多样化的合成作物和杂草变种，使自主除草系统仅使用 1% 的真实数据进行训练，即可实现强劲性能。

更广泛地说，Cosmos 的开放权重和框架以及 OpenMDW1.1 许可证使后训练成为一种强大、灵活和轻松的方法，可用于为物理 AI 创建专业、高性能和准确的自定义模型。

开始使用

• 下载 Cosmos Edge 策略权重和数据
• 观看 8 月 20 日的 Cosmos Labs 直播
• 在 Jetson 上探索 Cosmos Edge
• 下载 Cosmos 3 Checkpoint
• 在 GitHub 上获取代码
• 试用 Cosmos 3 Nano Reasoner 和 Cosmos 3 Nano