Cosmos3 Generator Action¶

Cosmos3-Nano action-generation supports three distinct tasks:

Forward dynamics (fd) — predict future observations from a start image plus an action trajectory.
Inverse dynamics (id) — predict ego-motion trajectories from input AV or camera videos using the Cosmos3-Nano.
Policy (policy) — predict future observations and action trajectories from a start image, task instruction, and state.

Each of the modes can be explored interactively using the Action Viewer Hugging Face space. The rest of this doc shows how to run these modes on selected embodiments directly.

Overview¶

Examples are shown across native PyTorch (Cosmos Framework), Diffusers, TensorRT-LLM, and vLLM-Omni. Every backend uses the sample assets under assets/ and covers three tasks:

Environment setup for all backends is centralized in the shared Cosmos3 cookbooks environment setup guide; each backend below links to the section you need.

Generator requires the Guardrail. Request access to the gated nvidia/Cosmos-1.0-Guardrail HF repository before running these examples. To disable the guardrail, set enable_safety_checker=False (Diffusers), use_guardrails: false (TensorRT-LLM extra_params), guardrails: false (vLLM-Omni extra_params/extra_args), or --no-guardrails (Cosmos Framework).

Action Definition¶

Cosmos3 treats action as a modality whose tokens represent transitions between consecutive visual states. This cookbook follows the unified action interface defined in the paper: ego and end-effector poses are represented as 9D pose deltas, consisting of 3D translation and a 6D continuous rotation representation. Grasp state encodes the current manipulation state, using a 1D open-close value for grippers or a 15D human-hand representation with 3D state for each of 5 fingers.

Embodiment	Representation	Dimensionality	Unit	Post-processing	Generation duration
Autonomous vehicle	Ego pose (9D)	9D	Meter	Normalization	60 frames @ 10FPS
DROID	End-effector pose (9D) + gripper grasp state (1D)	10D	Meter	Multiview concatenation, `to-OpenCV`, normalization	16 frames @ 15FPS
UMI	End-effector pose (9D) + gripper grasp state (1D)	10D	Meter	Normalization	16 frames @ 20FPS
Human hand pose	Ego pose (9D) + right wrist (9D) + right fingertips (15D) + left wrist (9D) + left fingertips (15D)	57D	Meter	Wrist-frame alignment, normalization	16 frames @ 15FPS

Action data samples across different embodiments can be inspected interactively in the Cosmos3 Action Viewer Hugging Face Space.

Run with Cosmos Framework¶

Quickstart¶

Set up the environment: Cosmos Framework setup. Activate the framework venv, then run the native inference entrypoint. Forward dynamics on Nano looks like:

torchrun --nproc-per-node=1 \
  -m cosmos_framework.scripts.inference \
  --parallelism-preset=latency \
  -i <forward-dynamics input spec>.json \
  -o /tmp/cosmos3_action_fd \
  --checkpoint-path Cosmos3-Nano \
  --seed 0

The input spec pairs a start image with an action trajectory. The notebooks assemble ready-to-run specs for AV, DROID, UMI, and human hand-pose examples from the checked-in assets under assets/. Outputs are written under the framework checkout.

Cosmos Framework Walkthrough¶

The Cosmos Framework build their input spec, run inference, and visualize the generated videos:

run_fd_with_cosmos_framework.ipynb — forward dynamics for AV, DROID, UMI, and human hand-pose examples using Cosmos3-Nano.
run_id_with_cosmos_framework.ipynb — inverse dynamics, predicting ego-motion trajectories from input AV videos using Cosmos3-Nano.
run_policy_with_cosmos_framework.md - policy, predicting future observations and action trajectories for DROID robot using Cosmos3-Nano-Policy-DROID and Cosmos3-Edge-Policy-DROID.

Run with Diffusers¶

Quickstart¶

Set up the environment: Diffusers setup. Action inputs are grouped into a CosmosActionCondition. The pipeline derives the frame count from chunk_size + 1 and the conditioning canvas from resolution_tier, so height, width, and num_frames stay unset:

import json
from pathlib import Path

import torch
from diffusers import Cosmos3OmniPipeline, CosmosActionCondition
from diffusers.utils import export_to_video, load_image

action_root = Path(".agents/skills/cookbooks/cosmos3/generator/action")
raw_actions = torch.as_tensor(
    json.load(open(action_root / "assets/actions/av_traj_forward.json")), dtype=torch.float32
)

pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16)
pipe.to("cuda")

result = pipe(
    prompt="You are an autonomous vehicle planning system.",
    action=CosmosActionCondition(
        mode="forward_dynamics",
        chunk_size=60,
        domain_name="av",
        resolution_tier=480,
        raw_actions=raw_actions,
        image=load_image(str(action_root / "assets/images/av_0.jpg")),
        view_point="ego_view",
    ),
    fps=10,
    num_inference_steps=30,
    guidance_scale=1.0,
    use_system_prompt=False,
    generator=torch.Generator(device="cuda").manual_seed(0),
)
export_to_video(result.video, "/tmp/cosmos3_action_fd.mp4", fps=10, macro_block_size=1)

For inverse dynamics, pass mode="inverse_dynamics" with video= instead of image= and no raw_actions; the predicted trajectory comes back as result.action. mode="policy" conditions on a first frame like forward dynamics but takes no raw_actions, predicting the action chunk alongside the video. Cosmos3-Super also ships action_gen=True, so the forward and inverse calls work against that checkpoint too.

Diffusers Notebook Walkthrough¶

The Diffusers notebooks provision a dedicated venv with the LeRobot readers, pose helpers, and plotting extras, register it as the Cosmos3 Diffusers Action kernel, then run each task with Cosmos3OmniPipeline and write outputs under outputs/notebooks/diffusers/:

run_fd_with_diffusers.ipynb — forward dynamics for three AV ego trajectories and three camera-pose trajectories, an autoregressive multiview DROID rollout, an autoregressive UMI rollout, and a bimanual hand-pose chunk.
run_id_with_diffusers.ipynb — inverse dynamics for the checked-in AV clips plus the Bridge, AgiBotWorld-Beta, RoboMIND (Franka, Franka dual-arm, UR), UMI, and Fractal episodes, plotting each prediction as a camera or end-effector trajectory.
run_policy_with_diffusers.ipynb — policy, jointly predicting future observations and an action chunk for DROID with Cosmos3-Nano-Policy-DROID, with the predicted chunk plotted as an end-effector trajectory.

The examples that read the checked-in LeRobot episodes use the readers from the Cosmos framework, and the trajectory plots use its pose helpers, so set COSMOS3_REPO to your framework checkout before running them. The DROID reader resolves its feature layout from the name of the directory it is given, so also set COSMOS3_DROID_ROOT to a directory named after the DROID release you have available.

Run with TensorRT-LLM¶

Quickstart¶

Set up and launch the VisualGen server with the one-GPU Nano config: TensorRT-LLM setup, including the server-side NLTK data setup. Keep NLTK_DATA exported in the server's shell. Action requests upload a conditioning image as multipart image_reference or a video as video_reference, and put the mode-specific fields under extra_params. Use the same trained AV prompt as the Cosmos Framework reference; the legacy view_point field is ignored by current TensorRT-LLM.

Both notebooks' current multipart fields were exercised against TensorRT-LLM source revision bca6761ab84fbcd58fc7f914eade7de48b32e35e. Follow the source pin in the shared setup.

import json
from pathlib import Path

import requests
from safetensors.torch import load as load_safetensors

action_root = Path(".agents/skills/cookbooks/cosmos3/generator/action")
image_path = action_root / "assets/images/av_0.jpg"
actions = json.loads((action_root / "assets/actions/av_traj_forward.json").read_text())
prompt = "You are an autonomous vehicle planning system."

with image_path.open("rb") as image_file:
    response = requests.post(
        "http://localhost:8000/v1/videos/sync",
        data={
            "prompt": prompt,
            "format": "safetensors",
            "response_format": "file",
            "seed": "0",
            "extra_params": json.dumps(
                {
                    "action_mode": "forward_dynamics",
                    "domain_name": "av",
                    "action": actions,
                    "use_guardrails": True,
                }
            ),
        },
        files={"image_reference": (image_path.name, image_file, "image/jpeg")},
        headers={"Accept": "application/octet-stream"},
    )
response.raise_for_status()
payload = load_safetensors(response.content)
print(payload["video"].shape, payload["action"].shape, payload["frame_rate"].item())

The av domain preset supplies its 60-step, 9D, 480p/10 fps recipe; other recognized domains similarly fill omitted action width, chunk, resolution, and frame rate. Inverse dynamics predicts its action tensor, while forward dynamics carries the conditioned action alongside the rollout. Both supported modes use a tensor output: format=auto selects safetensors, and explicit mp4/avi is rejected rather than dropping the trajectory.

The quickstart uses the blocking POST /v1/videos/sync route with response_format=file, which returns the tensor payload bytes directly. The older /v1/videos/generations spelling is a deprecated alias. The asynchronous POST /v1/videos route returns HTTP 202; poll GET /v1/videos/{id} and download the same tensor payload from GET /v1/videos/{id}/content.

TensorRT-LLM notebook walkthrough¶

run_fd_with_trt_llm.ipynb — forward dynamics from the checked-in AV image and 60×9 action trajectory.
run_id_with_trt_llm.ipynb — inverse dynamics from the checked-in AV observation clip.

Both notebooks decode TensorRT-LLM's safetensors response, write the rollout video, and inspect the named action tensor. This response contract is not the top-level action JSON used by vLLM-Omni.

Run with vLLM-Omni¶

Quickstart¶

Set up the environment and start the server: vLLM-Omni setup (Docker recommended). From the cookbook root (.agents/skills), set export COSMOS3_WORKDIR=$PWD and export COSMOS3_HOST_PORT=8001, then run the Docker command from the env setup guide. Wait until this succeeds:

curl http://localhost:8001/v1/models

Forward-dynamics requests are multipart POSTs to /v1/videos — a start image under files={"input_reference": ...} plus an extra_params payload carrying the action trajectory. The vLLM notebooks use these diffusion defaults for action generation (see run_fd_with_vllm_omni.ipynb and run_id_with_vllm_omni.ipynb):

Field	Value
`num_inference_steps`	`30`
`guidance_scale`	`1.0`
`flow_shift`	`10.0`

The notebooks build the full request body for AV, DROID, UMI and human hand-pose examples, including autoregressive chunked generation for the robotics examples. Policy inference uses async POST /v1/videos to retrieve a rollout video plus top-level action metadata.

VLLM-Omni Notebook Walkthrough¶

The vLLM-Omni notebooks send requests through the OpenAI-compatible video API and write outputs under outputs/cosmos3_action_vllm/:

run_fd_with_vllm_omni.ipynb — forward dynamics for AV, DROID, UMI, and human hand-pose examples.
run_id_with_vllm_omni.ipynb — inverse dynamics, predicting ego-motion trajectories from input AV videos.
run_policy_with_vllm_omni.ipynb — policy inference for DROID through the async video API.

Post-Train for Cosmos3-Nano-Policy-DROID¶

To reproduce our post-training recipe for Cosmos3-Nano-Policy-DROID, use the Cosmos3-Nano-Policy-DROID SFT cookbook. It follows the same launch-script pattern as the other Cosmos3 finetune cookbooks while delegating the canonical training implementation to Cosmos Framework.

The same action-policy SFT cookbook also covers LIBERO-10 (launch_sft_action_policy_libero_10_nano.sh) — fine-tuning Cosmos3-Nano on the libero_10 simulation benchmark with the same launch-script pattern.

TODO¶

[ ] Add additional embodiment examples.