PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Moo Jin Kim, Yihuai Gao, Tsung-Yi Lin, Yen-Chen Lin, Yunhao Ge, Grace Lam, Percy Liang, Shuran Song, Ming-Yu Liu, Chelsea Finn, Jinwei Gu

arXiv:2601.16163v1cs.AIcs.RO

TL;DR

Existing video-model robot policies often require multiple training stages and architectural additions for action generation. Cosmos Policy instead fine-tunes a pretrained video model in one stage without architectural changes, directly generating actions, future states, and values as latent frames. It achieves state-of-the-art simulation performance and strong real-world results, with model-based planning providing further gains.

  • Problem

    Prior adaptations of video models for robot policies often require multiple post-training stages and new architectural components for action generation.

  • Method

    Cosmos Policy single-stage fine-tunes Cosmos-Predict2 without architectural modifications to generate actions, future states, and values as latent frames within latent diffusion.

  • Results

    Cosmos Policy achieves 98.5% LIBERO and 67.1% RoboCasa average success rates, and 93.6% average success in challenging real-world bimanual manipulation tasks.

  • Takeaways & Limitations

    Policy rollout data can refine the world model and value function, enabling model-based planning that improves success on challenging tasks.

  • Takeaways & Limitations

    Model-based planning has substantially lower inference speed, taking around 5 seconds to produce one action chunk, and effective planning requires substantial rollout data.

Abstract

from arXiv · show

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/

1 INTRODUCTION

Cosmos Policy adapts a pretrained video model into a robot policy without the multi-stage training and architectural additions used by prior approaches. It directly generates actions, future states, and values, achieving strong simulation and real-world manipulation results and further gains from planning.

  • Pretrained video models provide temporal-causal priors for physically plausible and coherent scene evolution, complementing the semantic concepts learned by vision-language models.
  • Prior video-policy methods often require multiple training stages and new components, including separate action modules, to generate robot actions.
  • Cosmos Policy uses a single post-training stage on robot demonstrations, makes no architectural modifications, and directly models actions, future states, and values as latent frames.
  • 98.5% and 67.1% average success rates establish state-of-the-art performance on LIBERO and RoboCasa, respectively, while 93.6% is the highest average success rate in the reported real-world bimanual tasks.
  • 12.5 percent higher average task completion results from adding model-based planning in two challenging real-world manipulation tasks.

2 RELATED WORK

Related work adapts video and vision-language models to robotic manipulation, but existing approaches commonly add training stages, action modules, or other components. Cosmos Policy instead directly incorporates actions and additional modalities into the pretrained video model’s native latent diffusion process.

  • Video-based robot policies: Video-based robot policies use robot and action-less Internet video data, with some methods fine-tuning video models before training separate action modules.
  • Video-based robot policies: Other approaches train unified video-action models that jointly predict future frames and actions, but they do not fully leverage pretrained video-model capabilities.
  • Cosmos Policy: Cosmos Policy uses single-stage fine-tuning to generate actions, proprioceptive state, and state values within the pretrained model’s native latent diffusion process.
  • Vision-language-action models: Vision-language-action models fine-tune vision-language backbones on large-scale robotic imitation data and show strong generalization to semantic concepts unseen during robotic interaction.
  • World models and value functions: World-model methods integrate learned dynamics and planning with learning across classical model-predictive control and modern neural reinforcement-learning paradigms.

3 PRELIMINARIES

The paper formulates manipulation as finite-horizon imitation learning and builds on latent video diffusion, where clean image sequences are corrupted and denoised while conditioning information remains available. The framework also uses world models and value functions to represent future observations and expected returns for planning.

  • Cosmos video model: Cosmos-Predict2 is a latent video diffusion model that predicts subsequent frames from an initial image and text using VAE-encoded continuous tokens and EDM denoising.
  • Cosmos video model: The denoiser learns to recover a clean VAE-encoded sequence from Gaussian-corrupted latents, conditioning on text through cross-attention and noise level through adaptive layer normalization.
  • Cosmos video model: A conditioning mask keeps the first input-image latent frame clean while subsequent frames receive noise during training.
  • MDP formulation and imitation learning: Manipulation is modeled as a finite-horizon MDP with states, actions, transitions, rewards, and horizon, and policies are trained by imitation on expert state-action demonstrations.
  • MDP formulation and imitation learning: Policies predict multi-timestep action chunks to improve motion smoothness and success rates.
  • Latent diffusion sequence: Figure 2 depicts tokenizing images into latent frames, injecting modalities into the sequence, and denoising noised frames conditioned on clean frames.
  • World models and value functions: World models predict future observations from current observations and actions, while value functions represent expected discounted returns under a policy.

4 COSMOS POLICY: ADAPTING VIDEO MODEL FOR CONTROL & PLANNING

Cosmos Policy adapts Cosmos-Predict2 into a unified robot policy by encoding actions, future states, and values as latent frames without architectural changes. It supports direct action generation and model-based planning, with rollout data refining the world model and value function.

  • Latent frame injection: Cosmos Policy adapts a pretrained video model to predict robot actions, future states, and values within a unified latent diffusion sequence.The model is trained on tuples of current state, action, future state, and future value, with auxiliary supervision for policy and world-model learning.
  • Latent frame injection: An 11-frame latent sequence interleaves proprioception, camera views, action chunks, future observations, and future state values for left-to-right decoding.New modalities are normalized to [−1, +1] and duplicated across latent volumes before being inserted into the sequence.
  • Joint training and decoding: The model can decode actions, future states, and values jointly in parallel or autoregressively, trading faster inference for potentially higher-quality predictions.Direct policy evaluation requires only actions, whereas planning uses autoregressive decoding for future-state and value predictions.
  • Scope and data: Cosmos Policy achieves successful long-horizon, high-precision real-world manipulation, while demonstration data generally covers successful outcomes and can limit planning distributions.Some datasets, including LIBERO training data, also contain suboptimal trajectories caused by teleoperation errors.
  • Planning: Rollout data is used to fine-tune the model with heavier weighting on world-model and value predictions, while demonstrations remain the primary policy-training source.The rollout fine-tuning split assigns 90 percent of each batch to world-model and value-function training and 10 percent to policy training.
  • Planning: Cosmos Policy supports direct execution and best-of-N planning, selecting among sampled action proposals using predicted future states and values.Planning uses a policy model to generate proposals and a planning model to score them through predicted future state values.

5 EXPERIMENTS

Cosmos Policy is evaluated across simulation and real-world manipulation benchmarks against diffusion, video-based, and fine-tuned VLA policies, with additional analyses of planning and failure modes. It achieves strong multi-task performance, including state-of-the-art simulation results, highest overall ALOHA performance, and improved scores from model-based planning.

  • Real-world ALOHA: Cosmos Policy obtains the highest overall ALOHA score and outperforms all other methods in three of four real-world tasks.The score measures average percent completion for each task.
  • Baselines: The compared policies include diffusion policies trained from scratch, video model-based policies, and fine-tuned vision-language-action models.The evaluation compares against methods including Diffusion Policy, Dita, UVA, UWM, Video Policy, π0, π0.5, OpenVLA-OFT, CogVLA, UniVLA, DP-VLA, and GR00T-N1.5.
  • Simulation benchmarks: 98.5% and 67.1% average success rates establish new state-of-the-art results for Cosmos Policy on LIBERO and RoboCasa, respectively.LIBERO averages 500 trials per suite across three seeds; RoboCasa averages 50 trials per task across three seeds and uses 50 demonstrations versus more than 300 for prior methods.
  • Overall results: Cosmos Policy achieves the highest overall performance across LIBERO, RoboCasa, and real-world ALOHA evaluations.The ALOHA suite contains four bimanual manipulation tasks evaluated under both in-distribution and out-of-distribution conditions.
  • Failure analysis: Fine-tuned VLAs struggle with high-precision grasps and multimodal action sequences, whereas Cosmos Policy handles both more reliably.π0.5 commonly loses the ziploc-bag grip, while OpenVLA-OFT+ reaches between candies rather than directly toward one.
  • Model-based planning: 12.5 points of average score improvement are observed on two challenging ALOHA tasks when using model-based planning with the V(s′) formulation.Fine-tuning on rollout data improves future-state prediction and helps avoid errors such as losing the ziploc-bag slider grasp.

6 DISCUSSION

Cosmos Policy achieves strong performance across simulated and real-world robotic environments, while rollout data supports refining its world model and value predictions for model-based planning. Planning remains constrained by inference speed, rollout-data requirements, and shallow search depth.

  • Cosmos Policy demonstrates strong performance in LIBERO, RoboCasa, and ALOHA robot environments.
  • Policy rollout data can refine world model and value function predictions to enable effective model-based planning.
  • Around 5 seconds are required to produce one action chunk with model-based planning, which may limit use in dynamic tasks.
  • Effective planning requires substantial rollout data beyond the demonstration distribution, while the current approach uses one layer of best-of-N search.

7 REPRODUCIBILITY STATEMENT

The project releases model checkpoints, training data, and code, including training and evaluation scripts. Further training and evaluation details are provided in the appendices.

  • Model checkpoints, training data, and code, including training and evaluation scripts, are released on the project website.
  • Further training and evaluation details are provided in Appendices A.2 and A.3.

A.1 LATENT INJECTION IMPLEMENTATION DETAILS

Cosmos Policy injects proprioception, actions, and values into placeholder latent frames within a video diffusion sequence. The model denoises these injected modalities, which are later recovered directly from the generated latent frames.

  • Latent injection overwrites placeholder latent frames with normalized and duplicated robot proprioception, action chunks, and values.
  • Action chunks are normalized to [−1, +1], flattened, duplicated, reshaped into latent-frame volumes, and used to overwrite target frames.
  • Cosmos Policy denoises corrupted portions of the latent sequence while conditioning on the uncorrupted portions.
  • Generated action latents are decoded by averaging duplicated values, un-normalizing, and deploying all or part of the predicted action chunk.
  • Values are extracted by averaging their latent volume and un-normalizing, without requiring VAE decoding for injected non-image modalities.

A.2.1 COSMOS POLICY NOISE DISTRIBUTION

Cosmos Policy modifies the pretrained video model’s noise schedule to improve precision in robot action and multimodal prediction. It adds weight to higher noise levels during training and raises the minimum noise level during inference.

  • Robot policy training requires more precise generations than video generation because small action imprecisions can cause catastrophic failures.
  • Figure 9 contrasts the base log-normal noise distribution with Cosmos Policy’s adjusted hybrid distribution, which places more weight on higher noise levels.
  • Cosmos Policy uses a hybrid log-normal-uniform noise distribution with a 0.7/0.3 sampling split and greater weight on larger noise levels.
  • The hybrid noise distribution empirically improves action prediction accuracy and overall success rate.
  • At inference, σmin is raised to 4 from 0.002 while σmax remains 80, improving action, future-state, and value prediction accuracy by lower L1 loss.

A.2.2 LIBERO TRAINING DETAILS

LIBERO training fully fine-tunes Cosmos-Predict2 to predict 16-timestep action chunks alongside future-state targets.

  • Cosmos Policy is trained for 40K gradient steps using 64 H100 GPUs with global batch size 1920 over 48 hours.
  • The base Cosmos-Predict2 model is fully fine-tuned, updating all model weights.
  • The policy predicts 16-timestep action chunks and executes each full chunk before requerying.
  • After training, action, future proprioception, and future wrist-image latent L1 losses are 0.012, 0.007, and 0.068, respectively.

A.2.3 ROBOCASA TRAINING DETAILS

RoboCasa and real-world ALOHA evaluations use platform-specific training budgets, action-chunk execution, and both in-distribution and out-of-distribution trials. Cosmos Policy achieves the highest aggregate ALOHA success rates, although π0.5 performs slightly better specifically on OOD scenarios.

  • RoboCasa training: Cosmos Policy is trained for 45K gradient steps on 32 H100 GPUs with global batch size 800 over 48 hours.
  • ALOHA training: ALOHA evaluations train one policy across four tasks using 185 demonstrations, 50K gradient steps, and 8 H100 GPUs over 48 hours.
  • ALOHA training: ALOHA uses 50-timestep action chunks spanning 2 seconds at a 25 Hz controller and executes each full chunk before requerying.
  • Baseline training: For fair comparison, π0.5 and π0 receive 400K gradient steps while OpenVLA-OFT+ receives 32K under the same 48-hour, 8-GPU budget.
  • Evaluation procedure: Direct-policy evaluation predicts actions, future state, and value in parallel, then discards future-state and value predictions when planning is unused.
  • Evaluation procedure: Planning uses ensembles of 3 future-state predictions and 5 value predictions to robustify estimates.
  • Evaluation procedure: The ALOHA suite contains 101 trials across four tasks, testing both in-distribution and out-of-distribution initial conditions.
  • Scoring: ALOHA scoring ranges from 0 to 100 points and measures partial task completion rather than binary success.
Loading 2601.16163v1…