Cosmos 3: Omnimodal World Models for Physical AI

Paper Link: https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf

Model Architecture §

Cosmos 3 is capable of processing multimodal inputs and generating multimodal outputs

Cosmos 3 treats action as a core modality, the action tokens bridge the physical world with language-based reasoning and video-based world modeling

Cosmos 3 uses modality-specific encoders to project different modalities into a unified representation space, which is then processed by a Mixture-of-Transformers (MoT) backbone

Language tokens are generated autoregressively, while other modalities are generated through iterative denoising

Encoders §

Cosmos 3 adopts 2 separate encoders for visual input

For visual understanding a ViT encoder (16 x 16 patches) pre-trained with vision-language alignment is used

For visual generated a VAE encoder from Wan2.2-TI2V-5B is used

ViT encoder uses 16 x 16 patch size followed by a two-layer MLP that merges 2 x 2 tokens and projects them into the latent space of the transformer

For audio generation they adopt audio VAE architecture

Action modeling is supported across diverse embodiments, since each domain exposes its own native control space, the actions are mapped into a unified action interface that enables consistent multimodal reasoning. generation, and policy learning across domains

Action Representation §

Cosmos 3 treats action as the change from one frame to the next

Instead of seeing before and after the model gets an action token that says what movement caused that change

Cosmos 3 splits action into reusable pieces:

1. Ego pose - eg. car / camera motion & head-camera motion
2. Effector pose - eg. for robot this’d be the gripper
3. Grasp state - eg. manipulation state, robot gripper: open, half-closed, closed

Cosmos 3 avoids learning embodiment-specific controller details like PID parameters

SE(3) - math notation for a 3D pose (3D position + 3D orientation)

6D representation is six numbers arranged as two 3D direction vectors

Action tokenization §

Purpose is to map embodiment data into a shared action space, while preserving specific structure and semantics

Cosmos uses domain-aware input and output projection layers with separate weight matrices for each embodiment domain

input projection:

Output projection:

Generation §

Action

• Forward dynamics - predicts future visual states conditioned on observed context and clean action tokens
• Inverse dynamics - infers the action tokens that explain visual transition
• Join video-action prediction - jointly predicts action and video tokens

“clean” refers to tokens that are known/provided and model doesn’t have to guess

AI-Judge quality filtering §

Cosmos 3 utilizes Gemma-4 as the vision-language judge (specifically Gemma-4-31B-it). The judge is prompted as a training-data auditor and assigns rubric-based integer scores from 1 to 5 across three primary quality dimensions:

• Faithfulness: whether all response claims are grounded in the provided image, video, or textual context
• Completeness: whether the response fully addresses the instruction without important omissions
• Correctness: whether the response is factually, logically, and task-level accurate

Supervised Fine-Tuning §

General spatial understanding §

In the supervised fine-tuning stage the general spatial and temporal understanding capabilities are enhanced

General spatial understanding: Enhance general spatial understanding through 2D and 3D grounding augmented with both real and simulation data

To read:

• Locate Anything
• Omni 3d: A large benchmark and model for 3D object detection in the wild
• Grounded 3D-Aware Spatial Vision-Language Modeling
• MimicGen
• RoboSpatial
• Can multimodal large language models understand spatial relations

Together these cover object references, free space, cross-view correspondence, camera motion, size, distance, direction, routes, counting, and room-scale reasoning

General temporal understanding §

Temporal capabilities are augmented along three axes: temporal event understanding, physical plausibility judgement, and structured spatiotemporal scene upsampling

• First human annotators create dense temporal actions: egocentric videos of every tasks are labeled with human action descriptions & start/end timestamps
• Second to increase question diversity training data is curated with FoundationMotion pipeline
• Annotate camera motion patterns such as panning and zooming so the model can learn ego-camera movement

Generator Data §

Training §

Cosmos 3 is trained in 2 main phases

The reasoner is pre-trained on large-scale image-text and video-text corpora and subsequently fine-tuned on a curated Physical AI mixture (producing a strong multimodal backbone)

Reasoner and Generator share the same transformer block architecture hence the Reasoner weights are used to initialize the Generator

Cosmos 3 generation is optimized under a rectified flow matching objective across all modalities. For a target latent from any modality a noisy latent is constructed via straight-line interpolation

High level idea is to:

• Take a clean target
• Add noise
• Train model to move from clean target toward noise
  • Reverse velocity during inference
• noise clean sample

Robot Policy Post-Training §

Cosmos3-Nano-Policy-DROID is trained by resuming mid-trained Cosmos3-Nano, with a freshly initialized action encoder, action-decoding MLP, and action embedding tokens

5x learning-rate multiplier is added to action-related parameters to facilitate faster adaptation