Post-Train NVIDIA Cosmos 3 Edge for On-Device Robot Control

NVIDIA released Cosmos 3 Edge, a 4B omni-model with a 2B Nemotron-based reasoner, designed for on-device robot control on Jetson Thor. Post-trained on the nvidia/Cosmos3-DROID dataset (76k trajectories, ~350 hours across 86 tasks and 564 scenes), the policy achieves 22.9% success in closed-loop RoboLab tasks and generates action chunks in about 1.53 seconds on Jetson AGX Thor T5000, enabling real-time control without a data-center GPU.

Robots need policies that can adapt to their sensors, environments, and tasks while running on onboard computing hardware. World models https://www.nvidia.com/en-us/glossary/world-models/ offer a foundation for learning physical interactions, but their size can make on-device deployment difficult. This changes with the new NVIDIA Cosmos 3 Edge. Cosmos 3 Edge https://huggingface.co/nvidia/Cosmos3-Edge is a 4B omni-model https://www.nvidia.com/en-us/glossary/omni-model/ with a 2B NVIDIA Nemotron-based reasoner in the Cosmos 3 family. It was pretrained on the same physical-world data as NVIDIA Cosmos 3 Nano and NVIDIA Cosmos 3 Super and starts with the same grounding in how objects move and interact. Plus, the model is small enough to run on-device on NVIDIA Jetson Thor https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-thor/ . What you’ll build By the end of this tutorial, you’ll have a post-trained Cosmos 3 Edge manipulation policy that can run on Jetson Thor and be evaluated in closed-loop simulation. You’ll learn how to: - Post-train Cosmos 3 Edge to predict robot actions. - Serve the resulting policy on Jetson Thor. - Run inference inside a receding-horizon control loop. - Evaluate policy behavior in closed-loop simulation. Each step is reproducible from the open cosmos-framework https://github.com/NVIDIA/cosmos-framework repo, and the released checkpoint is available on HuggingFace https://huggingface.co/nvidia/Cosmos3-Edge . Why post-train Cosmos 3 Edge for on-device robot manipulation? World foundation models are pretrained on large multimodal datasets that capture patterns in object motion and physical interaction, such as how objects fall, slide, and respond to contact. Cosmos natively generates actions based on physical understanding and prediction capabilities. This prior knowledge provides a useful starting point for robot-policy training rather than requiring the policy to learn all physical relationships from task-specific demonstrations. Deploying these models on a physical robot, however, introduces two practical constraints: Device memory: The model and its runtime state must fit within the memory available on the robot. Control latency: The complete inference pipeline must be fast enough to support the robot’s required control frequency. Post-training Cosmos 3 Edge addresses each of these constraints. The resulting policy model fits within Jetson Thor’s memory, so inference runs directly on the robot rather than being offloaded to a data-center GPU. The remaining factor is control latency. The DROID action policy runs in real time, directly on the robot. On an NVIDIA Jetson AGX Thor T5000, it generates each action chunk in about 1.53 seconds running at 640×540 resolution and 15 Hz , while a single chunk covers roughly 2.13 seconds of robot motion. Because the next chunk is ready before the current one finishes, the arm moves continuously, with no data-center GPU in the loop. The policy supports continuous streaming on-device by generating action chunks and replanning after each inference cycle. It doesn’t replan after every observation. In closed-loop RoboLab https://github.com/NVLabs/RoboLab tasks, the post-trained policy reaches 22.9% success. Together, these results show that a 4B world foundation model can serve as a practical, real-time, on-device policy backbone. This model fits on a Jetson AGX Thor and runs entirely on the robot. What data does the policy train on? The release policy in this tutorial trains on the nvidia/Cosmos3-DROID https://huggingface.co/datasets/nvidia/Cosmos3-DROID dataset. It contains 76k successful teleoperated trajectories, approximately 350 hours across 86 tasks and 564 scenes, collected with a Franka Panda arm and Robotiq gripper. The dataset is packaged in LeRobotDataset v3.0 format at 640 × 360 resolution. Prepare it in three stages: - Filter idle and non-task frames. - Select successful demonstrations for training. - Apply random cropping, rescaling, and color jitter during training. hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir Cosmos3-DROID Bring your own robot data Convert your data to the LeRobot Dataset v3 per-frame camera video, joint states, gripper state, actions, task instruction . For a DROID-like Franka setup, the primary change is the dataset path. A different embodiment needs its own experiment configuration that defines the action space, dimensionality, camera layout, and normalization settings. Cosmos 3 supports several embodiments, including dual-arm Franka, UR, WidowX 250, LeRobot SO101. See the Cosmos 3 Edge model card https://huggingface.co/nvidia/Cosmos3-Edge for the embodiment list. Prerequisites Before starting, confirm that you have: - The latest Cosmos framework http://github.com/NVIDIA/cosmos-framework release. - The validated training hardware: NVIDIA DGX Station with the NVIDIA GB200 Grace Blackwell Superchip or the NVIDIA GB300 Grace Blackwell Ultra Desktop Superchip. - The supported NVIDIA CUDA and container versions: CUDA 13.0 cu130 , NGC 26.06-py3. - Access to the Cosmos 3 Edge base checkpoint https://huggingface.co/nvidia/Cosmos3-Edge . - Access to the Cosmos3-DROID dataset https://huggingface.co/datasets/nvidia/Cosmos3-DROID . - A Hugging Face access token. This is foundation model post-training, not a single-GPU fine-tune. The validated run uses 64 nodes of 4× GB200 for 60K iterations, roughly 68 hours ~17.4K GB200-hours . Plan compute accordingly. Setting | Value | |---|---| | Initialization | Cosmos3 Edge checkpoint | use state=true Table 1. Post-training configuration and hyperparameters How to run post-training At a high level, post-training is broken into four steps: - Download the dataset. - Convert the base checkpoint to distributed-checkpoint DCP format. - Apply the curation filter. - Launch 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 Edit and launch post-training script Find this file in the repo: cosmos framework/configs/base/experiment/action/posttrain config/action policy droid nano.py This launcher registers the recipe for Cosmos 3 Nano. To train Cosmos 3 Edge, make these three edits before launching: - Replace the NANO MODEL CONFIG import with EDGE MODEL CONFIG from configs/base/experiment/sft/models/edge model config.py . - Set BASE CHECKPOINT PATH to the Cosmos 3 Edge DCP checkpoint from the conversion step. - Rename examples/launch sft action policy droid nano.sh to examples/launch sft action policy droid edge.sh Everything else, including dataset, action space, curation filter, and training schedule, stays the same. Launch post-training bash examples/launch sft action policy droid edge.sh For the latest end-to-end instructions, including checkpoint conversion, environment configuration, and curation-filter setup, see the DROID post-training reproduction guide and model card https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID . These resources are updated alongside the repository. How to deploy and run the policy on Jetson Thor The policy is served by a WebSocket policy server speaking the OpenPI protocol, the same protocol used across the DROID policy ecosystem. The client sends an observation dictionary; the server returns an action chunk. For Edge, the server runs natively on Jetson Thor . At roughly 9 GB in BF16, the weights fit in Thor’s on-board memory, so both the policy server and the control client run on the robot with no data-center GPU in the loop. Start the policy server on the robot 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 Because the server runs on Thor, host="localhost" , the request never leaves the robot. On device, the policy runs in real time. This is the on-device inference that Cosmos Edge enables. Smoke test no robot needed The DROID policy is state-conditioned. Meaning that joint position and gripper position are real model inputs, not boilerplate. With no arm in the loop, zeros are the honest placeholder; this call just proves the server returns a well-formed action chunk. python 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 Running on the robot The same call sits in a replanning loop. Each cycle reads fresh cameras and the arm’s measured joint and gripper positions replacing the zeros above , executes a prefix of the chunk, then re-asks: python 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 Executing only a prefix and replanning is standard practice: the policy self-corrects every few seconds, and because it is state-conditioned, each replan starts from where the arm actually is rather than where the previous chunk assumed it would be. Starting the server with video decoding enabled returns the policy’s imagined rollout alongside the actions, so you can inspect what the world model predicts for the chunk it emitted. Evaluate the policy in a closed-loop simulation You don’t need a physical robot to validate a policy. In fact, evaluating a policy in simulation is recommended before evaluating directly on a physical robot to avoid unexpected behaviours before they affect developers or robots. RoboLab https://github.com/NVLabs/RoboLab —the Isaac Lab-Arena https://developer.nvidia.com/isaac/lab-arena benchmark behind the RoboLab leaderboard—is open, and its client connects to the same policy server. It executes each action chunk in physics and streams rendered observations back for a true closed loop across 120 language-conditioned manipulation tasks. 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 Swap --task for any of the 120 tasks to build a broader picture of policy behavior. Each run produces viewport and robot-camera videos plus a success log, so you can inspect both the outcome and the trajectory that produced it. In closed-loop RoboLab evaluation, the post-trained Edge policy reaches 22.9% success across the task suite. That result is earned at a fraction of the inference compute of larger Cosmos3 variants Nano reaches 36.8% , which is precisely the trade Edge is designed to make: real-time, fully on-robot autonomy at competitive success rates. Note: NVIDIA Isaac Sim 5.x requires the NVIDIA RTX Server Driver Release 580 or later. Confirm the supported driver version in the RoboLab and Isaac Sim documentation before building. Post-training beyond robot control Post-training techniques that improve robot control also extend to other capabilities, such as synthetic data generation for robot training. Developers can create specialized world models that generate high-quality synthetic data tailored to their environments and tasks for both indoor and outdoor robot training. Aigen, for example, post-trained Cosmos to generate diverse synthetic crop and weed variations, enabling an autonomous weeding system to achieve strong performance using just 1% real-world data for training. More broadly, Cosmos’s open weights and framework, with the OpenMDW1.1 license, make post-training a powerful, flexible, and easy way to create specialized, highly performant, and accurate custom models for physical AI.