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

TL;DR

Cosmos Policy presents a single-stage fine-tuning pipeline that repurposes a pretrained video diffusion model (Cosmos-Predict2-2B) into a robot policy by latent-frame injection of actions, future observations, and values. The approach jointly trains a policy, world model, and value function within the model’s latent diffusion framework, enabling both direct control and model-based planning via best-of-N search. With rollout-data refinement, the planning capability improves performance on challenging manipulation tasks, achieving state-of-the-art results on LIBERO and RoboCasa and strong real-world ALOHA performance, while highlighting trade-offs in inference speed. The work demonstrates that pretrained video priors can substantially boost visuomotor control and planning without architectural changes, offering practical benefits for data-efficient planning-based robotics.

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/

Figures (12)

• Figure 1: We present Cosmos Policy, a state-of-the-art robot policy fine-tuned from the NVIDIA Cosmos-Predict2-2B video foundation model. Cosmos Policy handles multimodal inputs and multi-view camera images and predicts (1) a robot action chunk, (2) future state (represented by robot proprioception and image observations), and (3) value (expected rewards-to-go at the future state). No architectural changes are made to the base video model, and all modalities are jointly modeled through the video diffusion learning objective.
• Figure 2: The latent diffusion sequence of Cosmos Policy. We illustrate latent frame injection---the primary mechanism for adapting the pretrained Cosmos-Predict2 into a policy that can predict robot actions, future states, and values without architectural changes. First, raw images are tokenized into latent frames (first row). Then, additional modalities are inserted directly into the latent frame sequence of the video diffusion model (second row). The model is then tasked to denoise the noised latent frames conditioned on the clean frames (third row). See Section \ref{['sec:latent_injection']} for more details. (Note: For simplicity, this figure does not depict certain implementation details; see Figure \ref{['fig:cosmos_policy_latent_diffusion_sequence_detailed_version']} for a more detailed visualization.)
• Figure 3: Cosmos Policy in the ALOHA robot tasks. Cosmos Policy can successfully execute real-world robotic control tasks that require long-horizon, high-precision manipulation and have high action multimodality.
• Figure 4: Real-world ALOHA robot evaluation results. We evaluate state-of-the-art policies on a suite of four tasks and measure the score, which represents average percent completion of each task. Cosmos Policy achieves highest overall score, outperforming all other methods in three of four tasks.
• Figure 5: Common failure modes of $\pi_{0.5}$ and OpenVLA-OFT+ on two challenging ALOHA robot tasks.Left:$\pi_{0.5}$ struggles to execute a high-precision grasp and loses grip of the ziploc bag. Right: OpenVLA-OFT+ reaches between two candies rather than towards one, suggesting difficulty with modeling the highly multimodal action distribution.