Overview
Cosmos3’s policy path is a different animal from text-to-video. It is the part of the model that looks at a scene, reads a task, and works out what to do. Give it an observation and a prompt and it predicts an action chunk. Give it an action and it rolls out a plausible future. Give it a transition that already happened and it infers the action in between. This page uses Cosmos3-Nano-Policy-DROID. Do not swap in the generalCosmos3-Nano generation checkpoint when you want actions; that path lives in Cosmos3 Generation Mode.
One plugin, cosmos3_policy, serves three modes on one GPU or several:
examples/cosmos3/run_cosmos3_policy.py wires all three modes. It runs with decode_video=True, saves the action as JSON, and writes a rollout mp4 whenever pixels come back. --cfg 2 or --tp N moves the same run onto managed workers. It handles one request at a time: a demo, not a server.Architecture
The policy path shares the Cosmos3 transformer with the generation path. Its request adds an action latent, a domain id, and a mode. Video and action travel through the same denoising loop, and the mode only decides which of them is a clean condition and which is generated.phyai/src/phyai/models/cosmos3
main_cosmos3_policy.py
scheduler_cosmos3_policy.py
model_runner_policy_cosmos3.py
model_runner_vae_cosmos3.py
modeling_cosmos3.py
vae_wan.py
sampler_unipc.py
The three modes
policy has the shape of a control loop: observation plus task in, action chunk out. The first observation frame is the clean condition; everything after it, video and action alike, is generated from noise. The question it answers is “given this scene, what should the robot do?”
forward_dynamics hands the model an observation and a known action and asks for the video. The action is the clean condition and the video is the target. The question is “if the robot does this, what happens next?” This mode needs --action-file.
inverse_dynamics runs the other way. You give it a video and it recovers the action chunk that would explain the change. The whole video is clean; the action is recovered from noise. The question is “what moved the scene from A to B?”
Input contract
Cosmos3PolicyProcessor.preprocess() takes a dict. The example script builds it from the CLI arguments:
Images become
(1, 3, T, H, W) in [-1, 1]. With --video the script reads the first action_chunk_size + 1 frames and repeats the last one when the clip is short.
Domain and action dimensions
Actions have two widths.action_dim is the model’s internal width, 64 by default. raw_action_dim is the real width of the robot’s action space. The processor pads conditioning actions up to action_dim and slices the output back down to raw_action_dim.
An integer
domain_id carries no width, so pass --raw-action-dim alongside it.
Run path
1
Prepare weights
Download Cosmos3-Nano-Policy-DROID. The policy path needs at least:
2
Construct the engine
The plugin is
"cosmos3_policy". With decode_video=True the VAE is loaded and decoded rollout pixels come back with the action.use_karras_sigmas=None reads the schedule from the checkpoint; False switches to linear flow with flow_shift.3
Construct the processor
The processor resizes and pads the observation, tokenizes the prompt, pads the action, resolves the domain id, and later slices the output.
4
Preprocess input
processed.video_shape is still in pixels. Convert it with pixel_to_latent_shape when you build the request.5
Build the request
6
Step and postprocess
action always comes back, shaped (1, action_chunk, raw_action_dim). pixels appears when the engine was built with decode_video=True, in [0, 1].Script examples
Predict an action from a single image:.cache/: cosmos3_policy_out_action.json with the action chunk, and cosmos3_policy_out.mp4 when pixels were decoded.
Roll out a video from a known action:
action.json takes either of these shapes; the numbers are placeholders. DROID actions are 10 wide, and a chunk shorter than action_chunk_size is padded by repeating its last step.
--condition-frames, an image conditions on frame 0 and a video on frames 0,1.
Output postprocessing
Cosmos3PolicyProcessor.postprocess() pulls action out of the result, slices it to raw_action_dim, and, when you gave it action_stats_path, denormalizes it back to physical units:
Without stats the action stays in the model’s normalized scale. Switching embodiment means switching weights, domain, and stats together.
Multi-GPU execution
The policy plugin takes the sameParallelConfig fields as the generation plugin. tp_size shards the transformer and is the knob that matters here: the examples run with guidance_scale=1.0, so cfg_size=2 would compute an unconditional branch only to discard it (the scheduler warns when it sees this). A run needs cfg_size * tp_size GPUs, picked with CUDA_VISIBLE_DEVICES on the launching process.
parallel block to EngineConfig and build the processor with device="cpu", so the request stays on the CPU while the workers own the GPUs. Keep the engine behind an if __name__ == "__main__": guard, because workers start with spawn and re-import the main module. The deployment argument is optional and only extends the startup timeout here.
engine.step() returns CUDA-IPC views of the output rank’s tensors; the postprocessor moves the action, and the pixels when there are any, to the CPU. tp_size must divide both the attention heads and the KV heads (1, 2, 4, or 8 for this checkpoint), with dense and attention TP kept equal. Replicas and external launchers are covered in Parallel serving.