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

NVIDIA introduced Cosmos 3 Edge, a 4B omni-model with a 2B Nemotron-based reasoner, designed to run on-device on Jetson Thor. The model was pretrained on the same physical-world data as Cosmos 3 Nano and Cosmos 3 Super, giving it a grounding in object motion and interaction. Post-training it on the Cosmos3-DROID dataset produces a manipulation policy that generates action chunks in about 1.53 seconds on a Jetson AGX Thor T5000, enabling continuous real-time control without a data-center GPU in the loop.

Post-training pipeline

The released policy trains on the nvidia/Cosmos3-DROID dataset, which contains 76k successful teleoperated trajectories, roughly 350 hours across 86 tasks and 564 scenes, collected with a Franka Panda arm and Robotiq gripper. The dataset comes in LeRobotDataset v3.0 format at 640x360 resolution. Preparation involves filtering idle and non-task frames, selecting successful demonstrations, and applying random cropping, rescaling, and color jitter during training.

The post-training configuration uses an 8-D absolute joint position action space (7 joints plus gripper), proprioception enabled, and a 3-camera observation canvas: a wrist view at 360x640 over two exterior views at 180x320 each, combined to 540x640. The model predicts 32 future actions at 15 Hz per chunk. Freshly initialized action heads use a 5x learning rate multiplier, with vision flow-matching weighted to match action loss. The run uses a global batch of 8192 with a long-cosine decay schedule across 10K iterations, checkpointing every 1K.

This is foundation model post-training, not a single-GPU fine-tune. The validated run uses 64 nodes of 4x GB200 for 60K iterations, roughly 68 hours, about 17.4K GB200-hours. Training hardware requires an NVIDIA DGX Station with GB200 Grace Blackwell Superchip or GB300 Grace Blackwell Ultra Desktop Superchip, CUDA 13.0, and NGC 26.06-py3.

Deployment on Jetson Thor

The policy is served by a WebSocket server speaking the OpenPI protocol, the same protocol used across the DROID policy ecosystem. At roughly 9 GB in BF16, the weights fit in Thor's on-board memory, so both the policy server and control client run on the robot. The server runs natively on Jetson Thor with eager execution, since stock Triton wheels lack sm_110a kernels.

A smoke test without a robot sends zeroed joint and gripper positions to verify the server returns a well-formed action chunk. In actual operation, the loop reads fresh cameras and measured joint and gripper positions, executes a prefix of the chunk, then replans. 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.

Evaluation results

In closed-loop RoboLab tasks, the post-trained policy reaches 22.9% success across 120 language-conditioned manipulation tasks. The model generates each action chunk in about 1.53 seconds at 640x540 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. The policy supports continuous streaming on-device by generating action chunks and replanning after each inference cycle, not after every observation.

Cosmos 3 supports several embodiments, including dual-arm Franka, UR, WidowX 250, and LeRobot SO101. Users can bring their own robot data by converting it to LeRobot Dataset v3 format, with the primary change being the dataset path for a DROID-like Franka setup. Different embodiments need their own experiment configuration defining action space, dimensionality, camera layout, and normalization settings. The cosmos-framework repository provides reproducible post-training and deployment scripts, with the released checkpoint available on HuggingFace.