NVIDIA Cosmos Reason2

NVIDIA Cosmos-Reason2 is no longer under active development and will receive only limited maintenance. For the latest state-of-the-art Physical AI capabilities, including stronger reasoning, world prediction, simulation, and action generation, migrate to Cosmos 3.

Future models and documentation are hosted at the main NVIDIA Cosmos repository: https://github.com/NVIDIA/Cosmos.

The repository provides several task-specific prompt templates organized into two main categories: Captioning and Embodied Reasoning. These templates can be found as .yaml files within the prompts/ directory.

### Captioning Templates
* Caption (`caption.yaml`)
* Temporal Localization (`temporal_localization.yaml`)
* Describe Anything (`describe_anything.yaml`)
* 2D Grounding (`2d_grounding.yaml`)

### Embodied Reasoning Templates
* Embodied Reasoning (`embodied_reasoning.yaml`)
* Robot COT (`robot_cot.yaml`)
* AV COT (`av_cot.yaml`)
* MVP Bench (`mvp_bench.yaml`)
* Causal VQA (`causal_vqa.yaml`)

For deeper technical guidance and specific use cases, refer to the following resources:

• Troubleshooting: Detailed solutions for common issues in docs/troubleshooting.md.
• Example Prompts: A collection of prompts to help you get started with the model in prompts/README.md.

Cosmos-Reason2 is built upon the Qwen3-VL architecture. You may find relevant architectural details and implementation guides in the Qwen3-VL Repository, Qwen3-VL vLLM documentation, or the official Qwen3 Documentation.

The notebooks in this repository are written in the jupytext py:percent format. This allows them to exist as both .ipynb and .py files. To synchronize changes between the .py and .ipynb versions, use the just command runner:

just notebooks-sync

Cosmos-RL uses a TOML configuration for training. For LLaVA-style datasets, you must configure the annotation_path and media_path in your config file (e.g., configs/llava_sft.toml). If your annotations use absolute paths, media_path can be left as an empty string ("").

To download the Llava-Instruct-150K dataset and COCO images for testing, use the following commands:

DATASET_DIR="/tmp/cosmos_reason2/cosmos_rl/data/llava_sft"
mkdir -p $DATASET_DIR
wget https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K/resolve/main/detail_23k.json -O $DATASET_DIR/annotations.json
wget http://images.cocodataset.org/zips/train2017.zip -O $DATASET_DIR/media.zip && unzip -q $DATASET_DIR/media.zip -d $DATASET_DIR

Supervised Fine-Tuning (SFT) improves model capability on specific task distributions.

Minimum Requirements:

• 4 GPUs with 80GB of memory each.

Configuration: Adjust the dp_shard_size in the [policy.parallelism] section of your TOML config based on your GPU count:

• For 4 GPUs: dp_shard_size = 4
• For 8 GPUs: dp_shard_size = 8

Execution Command: Run the SFT script using uv run and the cosmos-rl CLI:

uv run cosmos-rl --config configs/llava_sft.toml --log-dir outputs/llava_sft scripts/llava_sft.py

Use Docker to run Cosmos-Reason2. Ensure you have the NVIDIA Container Toolkit installed.

1. Build the container: Specify the CUDA version using --build-arg. For DGX Spark and Jetson AGX, use 13.0.0.

image_tag=$(docker build -f Dockerfile --build-arg=CUDA_VERSION=12.8.1 -q .)

2. Run the container:

docker run -it --gpus all --ipc=host --rm -v .:/workspace -v /workspace/.venv -v /workspace/examples/cosmos_rl/.venv -v /root/.cache:/root/.cache -e HF_TOKEN="$HF_TOKEN" $image_tag

Optional Arguments:

• --ipc=host: Uses host shared memory (recommended for parallel torchrun).
• -v /root/.cache:/root/.cache: Mounts host cache to avoid re-downloading models.
• -e HF_TOKEN="$HF_TOKEN": Passes your Hugging Face token for authentication.

Cosmos-RL is an async post-training framework for Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL). To set up the example environment, follow these steps:

1. Install Redis: Install the redis-server system dependency using conda:

   conda install -c conda-forge redis-server

2. Install the Example: Navigate to the example directory and use uv to sync the environment:

   cd examples/cosmos_rl
   uv sync

3. Monitor with Weights & Biases (Optional): To monitor training, install and login to wandb:

   uv tool install -U wandb
   wandb login

   During training, the controller log will provide a direct link to your wandb run.

To run the provided Jupyter notebooks in this directory on your local machine, follow these steps:

1. Prerequisites: Ensure you have completed the initial project setup as described in the main README.
2. Environment Setup: Use uv to synchronize the environment and activate the virtual environment:

   cd examples/notebooks
   uv sync
   source .venv/bin/activate

3. Execution:
   • You can run the notebooks as standard Python scripts using python <notebook>.py.
   • Alternatively, you can use VS Code to run them as interactive Jupyter notebooks.

Since Cosmos-Reason2 is based on the Qwen3-VL architecture, you can leverage vLLM for high-performance serving and inference. Key vLLM capabilities applicable to this model include:

• Online Serving: Deploying an OpenAI-compatible server.
• Offline Inference: Running batch inference locally.
• Multimodal Inputs: Handling the vision-language capabilities of the model.
• LoRA: Applying Low-Rank Adaptation for efficient fine-tuning.