Simulation / Modeling / Design

Post-Train NVIDIA Cosmos 3 in One Day Using Agent Skills

Boost accuracy to over 90% for video question answering with NVIDIA TAO for LoRA and AutoML

What if autonomous coding AI agents could push your vision reasoning models above 90% accuracy with almost no manual effort? When adapting vision reasoning models to production video tasks, developers often lose days to data formatting, container setup, training scripts, baseline evaluation, and hyperparameter sweeps before they even know whether post-training improves accuracy.

Combining the open physical AI world foundation model NVIDIA Cosmos 3 with NVIDIA TAO agent skills offers a solution to this challenge. Cosmos 3 connects understanding, generation, simulation, and action through a shared omnimodal world model. It unifies text, image, video, ambient sound, and action tracking under a mixture-of-transformers (MoT) architecture. While the model offers exceptional out-of-the-box vision reasoning, every real-world deployment requires domain specialization to handle unique camera angles, edge cases, and environments. This is where post-training becomes important.

This post demonstrates how you can seamlessly post-train the NVIDIA Cosmos 3 Nano model for video question answering using a coding agent and the NVIDIA TAO library of vision model agentic fine-tuning skills. 

The NVIDIA experiments show that by using Low-Rank Adaptation (LoRA), the workflow efficiently adapted the model, instantly boosting the zero-shot baseline from 54.41% exact-match accuracy (4-way multiple-choice) to an impressive 87.14% in a single run. And by leveraging TAO AutoML to systematically sweep configurations and eliminate guesswork, it pushed peak accuracy to 93.35%, condensing what was a multiday engineering effort into a single day of automated execution, driven by a few natural language prompts.

Video 1. Learn step by step how to post-train Cosmos 3 Nano with NVIDIA TAO agent skills, from setup through the AutoML sweep and final report

What are the benefits of the Cosmos 3 MoT architecture?

NVIDIA Cosmos 3 MoT architecture uses an autoregressive transformer for robust reasoning and planning, paired with a diffusion transformer to accurately predict future world states and actions. Available in multiple sizes including Super 64B and Nano 16B, Cosmos 3 consistently ranks #1 among open models on benchmarks, including: 

Figure 1 shows the unified dual-tower design split into an autoregressive reasoning pathway and an iterative diffusion generation pathway. Modality-specific inputs are tokenized and processed through separate LayerNorm and MLP blocks, interacting through a cross-stream connection where key-value states (KAR, VAR) are passed on to the generative full-attention block for the textual and visual context.

What post-training method is optimal for Cosmos 3 Nano? 

Foundation models such as Cosmos 3 learn generalized patterns from massive, diverse datasets. Post-training is essential to help the models understand specialized domain tasks, vocabulary, and specific camera perspectives better.

When it comes to post-training vision language reasoning models such as Cosmos 3 Nano, developers typically choose between two methods:

  • Full-Parameter Supervised Fine-Tuning (SFT): Best when the domain shift is massive or the task requires structural model changes. However, since SFT updates all model weights, it requires extensive computing resources and may show regression on general knowledge.
  • Low-Rank Adaptation (LoRA): Ideal for fast iteration. LoRA freezes the base model weights and injects trainable rank-decomposition matrices.

For the example detailed in this post, LoRA post-training required ~7x fewer GPU hours compared to full-parameter SFT for Cosmos 3 Nano, making a one-day post-training turnaround a reality for engineering teams.

What are the options for how to post-train Cosmos 3 Nano? 

There are two options for how to post-train Cosmos 3 Nano:

  • Open framework: The fully open Cosmos 3 post-training framework exposes training recipes, dataset formats, and configuration files directly, so you can build and own every step of the post-training pipeline. This is a good fit if you need custom training logic or want to integrate with existing infrastructure.
  • TAO agent skills: The skills sit on top of that same capability and automate it. A coding agent reasons through the workflow, patches data issues, runs the training container, and sweeps hyperparameters for you. This automation is what compresses a multiday setup into the one-day, two-prompt experience outlined in the following sections.

NVIDIA TAO introduces agent skills that help coding agents execute vision model post-training workflows for supported model families. These skills package the task-specific knowledge required to train, evaluate, optimize, and serve models, including framework details, launcher behavior, config structure, data loading, and evaluation workflows.

Instead of manually assembling every command and config file, TAO skills allow a coding agent to reason through the workflow and generate the correct steps.

The Cosmos 3 architecture separates its capabilities into a Reasoner tower (a vision language model for multimodal understanding and logic) and a Generator tower (for video and action generation). When using TAO agent skills, the workflow automatically handles the separation of these weights, ensuring that your post-training specifically targets and optimizes only the Reasoner weights for your downstream tasks. 

How to install TAO skills for a coding agent

TAO skills set can be installed for use by any coding agent. This example uses Codex. Similar setup instructions are also available for Claude in the NVIDIA-TAO/tao-skill-bank

GitHub repo. The fastest setup path is to use the automated installation script from the TAO skill bank repo. 

curl -fsSL https://raw.githubusercontent.com/NVIDIA-TAO/tao-skills-bank/main/scripts/install-codex-agents.sh | bash

If you’d rather drive each step yourself, the repo also breaks down the script into steps.

Before starting the experiment, populate your terminal environment with the required keys:

export HUGGINGFACE_TOKEN="your_hf_token" #To pull missing models
export NGC_API_KEY="your_ngc_key" #For TAO Docker containers
export AUTOML_LLM_API_KEY="your_llm_key" #Required for LLM-guided hyperparameter sweeps in AutoML

Dataset

This example shows how to post-train Cosmos 3 Nano on the Woven Traffic Safety (WTS) dataset from Toyota. The experiment focuses on a four-way multiple-choice video question-answering task, which includes more than 8,000 training and validation samples.

This dataset is useful for demonstrating Cosmos 3 post-training because traffic safety requires detailed real-world scene understanding. The model must reason about complex road layouts, signals, vehicles, pedestrians, and event context from video (Figure 3). But the same workflow applies to any downstream visual-reasoning task, such as warehouse monitoring, autonomous vehicle perception, robot workcells, and more. 

Dataset example

Question: What is the formation of the road?

  • A: Single road (right curve)
  • B: Single road (straight line)
  • C: Intersection (without signal)
  • D: Intersection (with signal)

Expected answer: D

How to run LoRA post-training with a single prompt

In the same environment where you exported the credentials, you can kick off an entire end-to-end pipeline with a single Codex prompt:

Perform LoRA post-training of the Cosmos 3 model on the Woven Traffic Safety dataset.
Training data: /home/…/WTS_dataset/wts_data_train
Validation data: /home/…/WTS_dataset/wts_data_val
Base model on Hugging Face: nvidia/Cosmos3-Nano
Also perform a baseline evaluation first, to compare with the post-trained model.

Behind the scenes, Codex queries the TAO library and selects the specialized Cosmos-reason skill. The agent autonomously handles the underlying framework and data loaders and executes the follow sequentially:

  1. Automated error patching: The agent scans the dataset annotations, flags a missing video FPS parameter, and applies a configuration patch on the fly.
  2. Model caching: It uses your Hugging Face token to securely pull the Cosmos 3 weights.
  3. Baseline evaluation: Before altering any weights, it runs a zero-shot baseline evaluation. The base Cosmos 3 model scores an initial 54.41% accuracy.
  4. LoRA pipeline execution: The agent generates the optimized LoRA training configurations and triggers the TAO training container.

When post-training is complete, the agent evaluates the LoRA-adapted model and compares the results against the baseline.

In just one run, and around 30 minutes of training time on eight NVIDIA A100 Tensor Core GPUs, the model jumps to 87.14% accuracy, a massive improvement of 32 percentage points, achieved completely hands-free.

How to optimize a vision model with TAO AutoML

LoRA provides a strong first result, but post-training performance often depends on selecting the right configuration. Learning rate, LoRA rank, dropout, batch size, scheduler settings, and other hyperparameters can all affect final accuracy.

Instead of manually trying one configuration at a time, TAO AutoML lets the agent run a structured sweep across important post-training parameters.

TAO AutoML supports a range of parameter search strategies, including but not limited to:

  • Bayesian optimization
  • Hyperband
  • Bayesian optimization and Hyperband
  • Batch-first Bayesian optimization
  • LLM-guided search proposes hyperparameters using an LLM brain (NVIDIA NIM, OpenAI, or any OpenAI-compatible endpoint). Requires NVIDIA_API_KEY or AUTOML_LLM_API_KEY. For this experiment, Gemini API endpoint was used.

For a brief explanation of each algorithm, see the TAO AutoML algorithm documentation.

With another prompt in the same coding agent chat, you can ask agent to run AutoML sweep:

Run an AutoML sweep to improve the LoRA result. Let TAO choose suitable search strategies 
and tune the important training hyperparameters. Optimize validation accuracy and summarize 
the best models.

The agent generates candidate trials using multiple strategies, launches each experiment, tracks the hyperparameters, evaluates the results, and records the best configurations.

What accuracy improvements did LoRA and AutoML deliver?

By combining the automation of TAO agent skills with the optimization power of TAO AutoML, the model achieved exceptional accuracy gains over the zero-shot baseline. The entire experiment progressed from an unaligned base model to a highly specialized traffic safety expert through just two prompts.

Model variantStrategyValidation accuracyImprovement versus baseline
Cosmos 3 Nano (Base)Zero-shot baseline54.41%Reference
Cosmos 3 Nano + LoRASingle-prompt baseline LoRA87.14%+32.73 percentage points
Cosmos 3 Nano + AutoMLBayesian optimization93.35%+38.94 percentage points
Table 1. Cosmos 3 Nano validation accuracy improvements across different post-training strategies 

What hardware and runtime should be expected?

Cloud-based NVIDIA hardware was used to test the multiconfiguration experiment sweeps of the TAO skills. Thanks to the efficiency of the Cosmos 3 MoT architecture, a single LoRA post-training epoch takes roughly 30 minutes on a node of NVIDIA A100 (80 GB) GPUs. This means that in under an hour, you can quickly spin up a highly accurate model that delivers 87.14% accuracy.

However, as a crucial clarification, if you want to push your accuracy past this threshold, finding the optimal configuration requires a TAO AutoML hyperparameter sweep. In this experiment environment, this took 19.5 hours to process the 43 parallel trials, running completely in parallel across multiple nodes of A100 GPUs hosted on the Oracle Cloud Infrastructure (OCI).

For comparison, a full-parameter SFT run took 3 hours and 34 minutes on NVIDIA H100 GPUs.

What are the qualitative post-training results?

Numbers only tell part of the story. The true value of domain specialization shows up when evaluating complex, ambiguous real-world edge cases. Figures 5 and 6 show qualitative comparisons before and after model post-training.

How to deploy a post-trained Cosmos 3 LoRA adapter  

The Cosmos 3 Reasoner NIM offers the fastest path to a production-grade, OpenAI-compatible Reasoner endpoint. It features an isolated and optimized autoregressive reasoner built for efficient on-device reasoning. NIM microservices ship as prebuilt, optimized containers, enabling you to entirely skip manual vLLM dependencies or intricate CUDA-pairing configurations.

It natively serves text outputs from text, image, and video inputs. To deploy your post-trained model, simply pass your post-trained LoRA adapter directory directly to the NIM environment during deployment.

Watch the video tutorial How to Run NVIDIA Cosmos 3 Reasoner NIM for Video Reasoning for step-by-step instructions on serving the Cosmos 3 model using NIM. To incorporate the LoRA adapters, refer to the NVIDIA NIM documentation.

Get started post-training Cosmos 3

The transition from general-purpose AI to highly specialized physical AI no longer requires wrestling with infrastructure. By combining the structural reasoning power of the NVIDIA Cosmos 3 Reasoner tower with the automation of NVIDIA TAO agent skills, you can transition from raw video data to a deployment-ready, specialized vision-language model in a single day. 

With two natural language prompts, the coding agent handles framework complexity, patches data bugs, executes baseline benchmarks, and runs advanced hyperparameter tuning via TAO AutoML entirely hands-free, letting engineering teams focus on solving core physical AI and robotics challenges.

Ready to post-train your own vision reasoning models? Check out these resources to implement the workflow in your local development environment:

Discuss (0)

Tags