Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning

Abstract: Recent video generation models demonstrate remarkable ability to capture complex physical interactions and scene evolution over time. To leverage their spatiotemporal priors, robotics works have adapted video models for policy learning but introduce complexity by requiring multiple stages of post-training and new architectural components for action generation. In this work, we introduce Cosmos Policy, a simple approach for adapting a large pretrained video model (Cosmos-Predict2) into an effective robot policy through a single stage of post-training on the robot demonstration data collected on the target platform, with no architectural modifications. Cosmos Policy learns to directly generate robot actions encoded as latent frames within the video model's latent diffusion process, harnessing the model's pretrained priors and core learning algorithm to capture complex action distributions. Additionally, Cosmos Policy generates future state images and values (expected cumulative rewards), which are similarly encoded as latent frames, enabling test-time planning of action trajectories with higher likelihood of success. In our evaluations, Cosmos Policy achieves state-of-the-art performance on the LIBERO and RoboCasa simulation benchmarks (98.5% and 67.1% average success rates, respectively) and the highest average score in challenging real-world bimanual manipulation tasks, outperforming strong diffusion policies trained from scratch, video model-based policies, and state-of-the-art vision-language-action models fine-tuned on the same robot demonstrations. Furthermore, given policy rollout data, Cosmos Policy can learn from experience to refine its world model and value function and leverage model-based planning to achieve even higher success rates in challenging tasks. We release code, models, and training data at https://research.nvidia.com/labs/dir/cosmos-policy/

• The paper introduces a unified framework that adapts large-scale video diffusion models for direct prediction of robot actions, future states, and value estimates through latent frame injection.
• The methodology repurposes a transformer-based denoiser to integrate multimodal inputs without structural modifications, achieving high success rates across simulation and real-world tasks.
• Experimental results demonstrate state-of-the-art performance with up to 98.5% success in simulations and 93.6% on real robots, with additional gains from on-policy fine-tuning for enhanced planning.

Cosmos Policy: Direct Adaptation of Video Foundation Models for Unified Visuomotor Policy, World Model, and Planning

Introduction and Motivation

The paper "Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning" (2601.16163) presents an architecture and methodology for adapting large-scale pretrained video diffusion models to the domain of robot policy learning and control. The core proposition is that video foundation models encode rich spatiotemporal priors fundamental for modeling physical interactions, temporal dynamics, and causal regularities—properties often lacking in models pretrained purely on static data or semantically focused datasets. Cosmos Policy leverages the NVIDIA Cosmos-Predict2-2B model, converting it into a multi-modal policy head that enables joint prediction of robot actions, future states (latent observations), and value estimates, all through a single-stage fine-tuning with no architectural modifications.

Figure 1: Cosmos Policy is a robot control policy derived from the Cosmos-Predict2-2B foundation video model, supporting multimodal inputs and multiple camera views, and predicting action chunks, future states, and value estimates through a unified video diffusion objective.

Methodology and Model Architecture

Cosmos Policy builds on the EDM-trained spatiotemporal latent diffusion transformer Cosmos-Predict2-2B, originally designed for video frame prediction conditioned on text and images. The adaptation to the robotics domain is enabled by a latent frame injection mechanism, where robot proprioceptive input, action sequences, and value estimates are mapped directly into additional latent frames within the diffusion sequence. This approach supports seamless integration of arbitrary robot-specific modalities without requiring custom modules or decoder heads.

The overall tokenization and injection process is detailed as follows: input images across various camera views and blank placeholders are VAE-tokenized into a latent sequence, after which the blank slots are overwritten with normalized robot proprioception, action chunk, and value data. This permutation allows for flexible composition based on the robot setup and task requirements.

Figure 2: Cosmos Policy implements latent frame injection, directly incorporating non-image modalities and supporting joint prediction of robot actions, future state, and value within the video diffusion pipeline.

Notably, the same transformer-based denoiser network is repurposed for autoregressive and parallel decoding of all relevant outputs. Auxiliary supervision is applied: policy training samples require the model to jointly predict the next action, future state, and value; world model samples require generation of future state and value from $(s, a)$; value function samples demand value prediction conditioned on $(s, a, s')$. This joint objective is managed by balancing batch samples during training.

The conventional approach for adapting internet-scale video models to robotics requires multiple fine-tuning stages, explicit policy heads, or substantial model surgery. Cosmos Policy, in contrast, achieves high-level integration by operating purely in the shared latent space and leveraging the alignment induced by video prediction pretraining.

Experiments and Evaluation

Multidomain Benchmarks and Real-World Results

Experiments are executed across three domains: LIBERO single-arm simulation, RoboCasa kitchen manipulation, and real-world ALOHA dual-arm bimanual manipulation. Cosmos Policy is trained using single-stage fine-tuning on robot demonstration data, with additional policy rollout fine-tuning for challenging planning tasks. The key components evaluated are:

• Direct policy execution: Feedforward action generation using parallel decoding.
• Model-based planning: Generating candidate action sequences, simulating rollouts, and selecting those with maximum value via best-of-N ensemble sampling.

  Figure 3: Cosmos Policy is validated on the ALOHA bimanual robotic platform, handling dexterous, multimodal, and long-horizon manipulation tasks.

Numerical Outcomes

• LIBERO Benchmark: Cosmos Policy attains a 98.5% average success rate, surpassing state-of-the-art imitation policies (including pretrained and fine-tuned VLA and video-action models) in spatial, object, goal, and long-horizon tasks.
• RoboCasa Benchmark: Achieves 67.1% average success using only 50 demonstrations per task (outperforming models trained on 300–3000+ demos).
• ALOHA Real Robot: Delivers a 93.6% average completion score across four complex tasks, outperforming OpenVLA-OFT+, $\pi_{0.5}$, and other strong diffusion policies, especially in tasks requiring high action multimodality or precision.

  Figure 4: On real-world ALOHA robot tasks, Cosmos Policy achieves the highest average scores, particularly excelling in complex, high-multimodality scenarios.

Failure modalities of state-of-the-art VLA policies are highlighted:

Figure 5: Failure cases in $\pi_{0.5}$ and OpenVLA-OFT+ include unsuccessful high-precision grasps and unreliable handling of multimodal action distributions.

Planning with Model-Enriched Policy

Training solely on demonstration data limits generalization to out-of-distribution states. Cosmos Policy overcomes this via policy rollout experience; fine-tuning on on-policy data substantially refines the transition and value predictions, enabling model-based search.

• Model-based planning (with $V(s')$ estimation) yields a 12.5-point performance gain in the most challenging ALOHA tasks, outperforming direct policy execution and model-free ($Q(s, a)$) planning.
• Fine-tuned world models more accurately capture error modes and their consequences, directly supporting effective search-based planning.

  Figure 6: Post-rollout fine-tuning, the world model captures previously missed errors, improving future state prediction and supporting successful multistep planning.

  

  Figure 7: On challenging ALOHA tasks, model-based planning utilizing $V(s')$ delivers superior performance over alternative planning formulations.

Ablations and Design Analysis

Ablation studies confirm several core architectural hypotheses:

• Auxiliary joint supervision (policy predicting $(a, s', V(s'))$) yields a 1.5% success rate boost over action-only training.
• Pretrained initialization with video diffusion models is critical; from-scratch training degrades success rates by 3.9–18.7 points, and produces unstable or jerky behaviors.
• Multistep denoising during inference can be traded for inference latency with only minor performance degradation.

Implications, Limitations, and Future Directions

Cosmos Policy demonstrates that video diffusion foundation models encode strong transferable priors for robotic visuomotor control. The latent injection/conditioning mechanism provides a path for modular integration of arbitrary robot or task-specific modalities without necessitating structural changes or custom heads. The joint learning of policy, world model, and value within a unified latent diffusion space benefits both policy robustness and sample efficiency for planning-centric RL.

Practical limitations include elevated inference latency in model-based planning modes (several seconds per action chunk using multi-GPU parallelization), constraining applicability to slow or deliberative tasks. Extension to more dynamic control scenarios will require architectural or algorithmic acceleration (e.g., reducing denoising steps, compressed conditioning, or more efficient search). Additionally, accurate planning beyond the demonstration distribution hinges on sufficient quantities of on-policy rollout data; improving generalization with fewer rollouts remains an open problem.

Future work should explore deeper planning horizons, alternative sampling strategies, and cross-domain transfer to more diverse robots and tasks, as well as integration with open-vocabulary instruction following.

Conclusion

Cosmos Policy establishes that video diffusion foundation models can be efficiently and directly adapted for high-fidelity visuomotor control, world modeling, and value inference in robotics, within a single-stage, architecture-preserving paradigm. Its latent frame injection and joint training regime avoids the typical complexities of multi-stage adaptation and enables strong zero-shot and generalization performance across varied domains. This approach delineates a scalable and generalizable path for leveraging internet-scale video pretraining in embodied AI and robotics.