Technical Summary: Cosmos Policy

Problem Statement

Recent advancements in video generation models have demonstrated an impressive ability to capture complex physical interactions, temporal causality, and scene evolution. However, adapting these models for robotic control has traditionally required complex, multi-stage training pipelines and the introduction of new architectural components (e.g., separate action diffusers or inverse dynamics models). Existing unified video-action models often fail to leverage the powerful spatiotemporal priors of large-scale pretrained video models due to custom designs trained from scratch. Furthermore, while Vision-Language-Action (VLA) models have become popular backbones, they are primarily trained on static image-text pairs, potentially limiting their grasp of temporal dynamics and implicit physics compared to video models.

The core challenge addressed by this work is how to effectively adapt a large, pretrained video foundation model into a robot policy through a simple, single-stage process without architectural modifications, while simultaneously enabling the policy to perform model-based planning using learned world models and value functions.

Methodology

Cosmos Policy Architecture

The authors introduce Cosmos Policy, a robot policy fine-tuned from the NVIDIA Cosmos-Predict2-2B video foundation model. The approach is distinguished by its lack of architectural changes; instead, it leverages the model's native latent diffusion process to handle multimodal inputs and outputs.

1. Latent Frame Injection:
To incorporate modalities not originally supported by the video model (robot proprioception, actions, state values, and multi-view cameras), the authors propose Latent Frame Injection.

• Mechanism: Additional modalities are encoded as new latent frames and interleaved directly into the video model's latent diffusion sequence.
• Sequence Structure: For a robotic setup with multiple cameras and proprioception, the latent sequence is structured as $(s, a, s', V(s'))$, where:
  • $s$: Current state (robot proprioception + current camera images).
  • $a$: Action chunk (encoded as a latent frame).
  • $s'$: Future state (future proprioception + future camera images).
  • $V(s')$: Future state value (expected cumulative reward, encoded as a latent frame).
• Flexibility: This method allows for arbitrary robot configurations (e.g., varying numbers of cameras) by simply adding or removing latent frames in the sequence without altering the model weights or structure.

2. Joint Training Objectives:
The model is fine-tuned on robot demonstration data using a single-stage objective that jointly learns three functions:

• Policy: Predicts actions, future states, and values conditioned on the current state ($p(a, s', V(s') | s)$).
• World Model: Predicts future states and values conditioned on state and action ($p(s', V(s') | s, a)$).
• Value Function: Predicts future values conditioned on state, action, and future state ($p(V(s') | s, a, s')$).

The training data consists of a mix of expert demonstrations (50% of the batch) and policy rollouts (50% of the batch). The rollout data is used to train the world model and value function, allowing the system to learn from both successful and failed trajectories.

3. Planning with Model-Based Search:
Cosmos Policy supports two deployment modes:

• Direct Policy: Uses parallel decoding to generate actions directly from the current state.
• Model-Based Planning: Utilizes the learned world model and value function to search for optimal actions.
  • Process: The system samples multiple action proposals, uses the "planning model" (a refined checkpoint trained heavily on rollout data) to predict the resulting future states and values, and selects the action leading to the highest predicted value.
  • Ensemble Strategy: To handle multimodal distributions and variance, the system ensembles predictions by querying the world model and value function multiple times per action, aggregating results via "majority mean" to determine success likelihood.

Key Contributions

1. Simple Adaptation of Video Models: The paper demonstrates that a large pretrained video model can be adapted into a state-of-the-art robot policy through a single stage of post-training with no architectural modifications, utilizing latent frame injection to encode actions and values.
2. Unified Policy, World Model, and Value Function: The approach unifies policy learning, world modeling, and value estimation within a single latent diffusion framework, leveraging the model's pretrained priors to capture complex action distributions and temporal dynamics.
3. Planning via Rollout Refinement: The authors show that by collecting policy rollout data and fine-tuning the value and world model components, the system can perform effective model-based planning (best-of-N sampling) to improve success rates in challenging tasks.
4. State-of-the-Art Performance: The method achieves new benchmarks in simulation and real-world manipulation, outperforming policies trained from scratch, video-based policies, and fine-tuned VLAs.

Experimental Results

The authors evaluated Cosmos Policy across three benchmarks:

• LIBERO Simulation: Achieved an average success rate of 98.5%, outperforming strong baselines including Diffusion Policy, Dita, $\pi_0$, OpenVLA-OFT, and CogVLA.
• RoboCasa Simulation: Achieved an average success rate of 67.1% using only 50 human demonstrations per task (significantly fewer than the 300+ used by many competitors), surpassing methods like GR00T-N1.5 and Video Policy.
• Real-World ALOHA Tasks: In challenging bimanual manipulation tasks (e.g., folding shirts, handling ziploc bags), Cosmos Policy achieved the highest average score (93.6%).
  • It notably outperformed fine-tuned VLAs ($\pi_0.5$, OpenVLA-OFT+) in tasks requiring high-precision grasps and handling high action multimodality.
  • Planning Impact: When enhanced with model-based planning using rollout data, the policy saw a 12.5% average increase in task completion scores on the most difficult ALOHA tasks.

Ablation Studies:

• Removing auxiliary learning targets (predicting $s'$ and $V$ alongside $a$) resulted in a 1.5% drop in success rate.
• Training the model from scratch (without pretrained video weights) resulted in a 3.9% drop, highlighting the importance of video model priors.
• Model-based planning using a $V(s')$ formulation (requiring a world model) outperformed a model-free $Q(s, a)$ variant, suggesting the value of learned environment dynamics.

Significance and Claims

The paper claims that Cosmos Policy represents a significant step forward in robot learning by demonstrating that the spatiotemporal priors of large-scale video generation models are a superior foundation for low-level control compared to static image-based VLA models.

Key claims include:

• Simplicity and Efficiency: The ability to create a powerful robot policy without architectural changes or multi-stage training pipelines.
• Data Efficiency: The method achieves high performance with fewer demonstrations (50 vs. 300+) compared to other state-of-the-art methods, likely due to the strong initialization from the video foundation model.
• Planning Capability: The unified architecture naturally supports model-based planning, allowing the robot to "imagine" future states and values to select better actions, a capability that improves performance in high-precision and multimodal tasks where direct policies struggle.
• Generalization: The approach shows strong generalization in both in-distribution and out-of-distribution scenarios across simulation and real-world environments.

The authors acknowledge limitations, specifically the inference speed of model-based planning (approx. 5 seconds per action chunk) and the current need for substantial rollout data to refine the world model, identifying these as areas for future improvement.