Fine-tuned Cosmos3-Nano world model for Unitree G1 humanoid robot manipulation. Given an initial observation image and a task description, the model jointly generates: (1) a video of the robot executing the task, and (2) a 26-dimensional joint-angle trajectory at 15 Hz.

Training: Policy SFT on the G1 BrainCo apple-picking dataset — 10,000 iterations on 7× A100 80 GB GPUs using FSDP.
Base model: nvidia/Cosmos3-Nano (8B params, Qwen3-VL-8B backbone)

What This Model Does

Input:  image (initial robot observation) + text prompt (task description)
Output: video frames (480p, 15 fps) + 26D joint-angle trajectory at 15 Hz

This is a policy model: give it a photo of what the robot currently sees and a description of what it should do — it predicts both how the robot moves (video) and what joint angles it should command (actions).

[your_image.jpg]  +  "Pick up the apple from the table"
        │
        ▼
  Cosmos3-Nano Policy SFT
        │
        ├── rollout.mp4          ← robot video (what it will see)
        └── actions_raw.npy      ← joint angles [T × 26] in radians @ 15 Hz

Actions are at 15 Hz and can be sent directly to the G1 robot controller. The model runs autoregressively in 16-frame chunks — the last frame of each chunk feeds into the next, so you can generate arbitrarily long rollouts.

This is a policy model: it predicts both how the robot moves (video) and what joint angles it should command (actions). Actions are at 15 Hz, matching the video frame rate.

Supported tasks (pre-trained):

pickapple — Pick up an apple from the table
grasporeo — Grasp an Oreo cookie from the table
grasprubikscube — Grasp a Rubik's cube from the table
pickcharger — Pick up a phone charger from the table
pickdoll — Pick up a doll from the table
pickdrink — Pick up a drink bottle from the table
picktissues — Pick up a tissue box from the table
picktoothpaste — Pick up a toothpaste tube from the table

Checkpoint Format

This model uses PyTorch DCP (Distributed Checkpoint) format — the same format used during FSDP training. The model/ directory contains 7 .distcp shards.

The fine-tuned checkpoint only stores the 4 trained adapter modules:

moe_gen — MoE generation router
time_embedder — Timestep embedder
vae2llm — VAE latent → LLM token bridge
llm2vae — LLM output → VAE latent bridge

The visual encoder and VLM backbone (Qwen3-VL-8B) are frozen during training and must be loaded from the base nvidia/Cosmos3-Nano checkpoint first. See Two-Stage Loading below.

Installation

Requires the Cosmos3 framework (NVIDIA internal):

bash

# Clone and install
cd cosmos/packages/cosmos3
uv sync
source .venv/bin/activate

Policy Inference (Image + Prompt → Video + Actions)

The core capability of this model: give it any image of what the robot sees + a text description of the task, and it outputs a robot video and the joint-angle commands to execute it.

Quick Start — Built-in Tasks

bash

# Run chunked autoregressive rollout for a built-in task
torchrun --nproc_per_node=1 examples/policy_rollout.py \
    --checkpoint-dir /path/to/iter_000010000 \
    --base-checkpoint-dir examples/checkpoints/Cosmos3-Nano \
    --output-dir /tmp/rollout_output \
    --n-chunks 5 \
    --tasks pickapple

This generates:

rollout.mp4 — 80-frame video (5 chunks × 16 frames @ 15 fps ≈ 5.3s)
actions_raw.npy — shape [80, 26] joint angles in radians @ 15 Hz
actions_normalized.npy — shape [80, 26] normalized to [-1, 1]
actions_raw.json — same as above in JSON
rollout_meta.json — task metadata

Your Own Image + Custom Prompt

You can use any image as the starting observation — a live camera frame from the robot, a photo of your workspace, etc.:

python

import torch, json, torchvision
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint import FileSystemReader
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
from torch.distributed.checkpoint.state_dict import get_model_state_dict
from pathlib import Path
from cosmos_framework.inference.model import Cosmos3OmniModel, Cosmos3OmniConfig
from cosmos_framework.inference.action import build_action_batch
from cosmos_framework.inference.args import ModelMode
from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig

# ── 1. Load model (2-stage: base → overlay fine-tuned) ──────────────────────
BASE_CKPT = Path("examples/checkpoints/Cosmos3-Nano")
SFT_CKPT  = Path("/path/to/iter_000010000")

import attrs
config = Cosmos3OmniConfig(**json.load(open(BASE_CKPT / "model/config.json")))
config.parallelism = attrs.asdict(ParallelismConfig(enable_inference_mode=True,
    data_parallel_shard_degree=1, data_parallel_replicate_degree=1))
config.compile = attrs.asdict(CompileConfig(enabled=False))

Cosmos3OmniModel.before_load_model()
wrapper = Cosmos3OmniModel(config)
state_dict = get_model_state_dict(wrapper.model)
dcp.load(
    state_dict=state_dict,
    storage_reader=FileSystemReader(str(SFT_CKPT / "model")),
    planner=DefaultLoadPlanner(allow_partial_load=True),
)
Cosmos3OmniModel.after_load_model(wrapper.model)
model = wrapper.model.eval().cuda()

# ── 2. Load YOUR image ───────────────────────────────────────────────────────
# Any image works: robot camera frame, photo, etc.
# Must be provided as a 1-frame video tensor [C, 1, H, W]
import torchvision.transforms.functional as TF
from PIL import Image

img = Image.open("your_robot_observation.jpg").resize((640, 480))
frame = TF.to_tensor(img)                    # [3, H, W] float [0,1]
frame_uint8 = (frame * 255).to(torch.uint8)
init_frame = frame_uint8.unsqueeze(1)        # [3, 1, H, W]

# ── 3. Write your task prompt ────────────────────────────────────────────────
task = "Pick up the red apple from the table"

prompt = json.dumps({
    "subjects": [{"description": "A Unitree G1 humanoid robot", "action": task}],
    "background_setting": "An indoor workspace",
    "cinematography": {"camera_motion": "static", "framing": "ego-perspective", "camera_angle": "ego"},
    "temporal_caption": f"A Unitree G1 humanoid robot performs: {task}.",
    "resolution": {"H": 480, "W": 640},
    "fps": 15.0,
})

# ── 4. Run one 16-frame chunk ────────────────────────────────────────────────
batch = build_action_batch(
    video=init_frame,
    action=torch.zeros(16, 64, dtype=torch.float32),
    raw_action_dim=26,
    prompt=prompt,
    view_point="ego_view",
    domain_name="g1_brainco",
    model_mode=ModelMode.POLICY,
    action_chunk_size=16,
    fps=15,
    resolution="480",
    input_video_key=model.input_video_key,
    batch_size=1,
    device="cuda",
)

with torch.no_grad():
    outputs = model.generate_samples_from_batch(batch, seed=[42], num_steps=30)

# ── 5. Get results ───────────────────────────────────────────────────────────
pixels   = model.decode(outputs["vision"][0])          # [1, C, T, H, W] float [-1,1]
actions  = outputs["action"][0][:16, :26].cpu()        # [16, 26] normalized [-1,1]

# Denormalize actions to joint angles (radians)
import numpy as np
stats = json.load(open("action_stats.json"))
q01 = np.array(stats["q01"])
q99 = np.array(stats["q99"])
actions_rad = q01 + (actions.numpy() + 1.0) / 2.0 * (q99 - q01)  # [16, 26] radians
# → send actions_rad[t] to G1 robot controller at 15 Hz

# ── 6. Run multiple chunks for a longer rollout ──────────────────────────────
# Take the last predicted frame as the next chunk's input:
last_frame = ((pixels[0, :, -1:] + 1) * 127.5).clamp(0, 255).to(torch.uint8)  # [C,1,H,W]
# ... repeat build_action_batch + generate with last_frame

Init Frames for Built-in Tasks

The rollout script reads an initial observation frame from:

/tmp/cosmos_outputs/g1_inference/<task>_init.mp4

Provide a short clip (or single frame as a video) of the robot in its starting position.

Chunked Autoregressive Rollout

The model runs in chunks of 16 frames. The last frame of each chunk becomes the conditioning frame for the next:

chunk 0: [init_frame] → model → [frame_1 ... frame_16] + [action_1 ... action_16]
chunk 1: [frame_16]   → model → [frame_17 ... frame_32] + [action_17 ... action_32]
...

Increase --n-chunks for longer rollouts (memory scales linearly).

Programmatic Usage

python

import torch
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint import FileSystemReader
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
from torch.distributed.checkpoint.state_dict import get_model_state_dict
from pathlib import Path
import json

from cosmos_framework.inference.args import ModelMode
from cosmos_framework.inference.model import Cosmos3OmniModel, Cosmos3OmniConfig
from cosmos_framework.inference.action import build_action_batch
from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig

BASE_CKPT = Path("examples/checkpoints/Cosmos3-Nano")
SFT_CKPT  = Path("/path/to/iter_000010000")

# ---- Load model ----
config = Cosmos3OmniConfig(**json.load(open(BASE_CKPT / "model/config.json")))
config.parallelism = {"enable_inference_mode": True,
                      "data_parallel_shard_degree": 1,
                      "data_parallel_replicate_degree": 1}
config.compile = {"enabled": False}

Cosmos3OmniModel.before_load_model()
wrapper = Cosmos3OmniModel(config)

# Load SFT weights (allow_partial_load skips visual encoder keys not in DCP)
state_dict = get_model_state_dict(wrapper.model)
dcp.load(
    state_dict=state_dict,
    storage_reader=FileSystemReader(str(SFT_CKPT / "model")),
    planner=DefaultLoadPlanner(allow_partial_load=True),
)
Cosmos3OmniModel.after_load_model(wrapper.model)

model = wrapper.model.eval().cuda()

# ---- Run one chunk ----
import torchvision
frames, _, _ = torchvision.io.read_video("init_frame.mp4", pts_unit="sec")
init_frame = frames[0].permute(2, 0, 1).unsqueeze(1)  # [C,1,H,W]

batch = build_action_batch(
    video=init_frame,
    action=torch.zeros(16, 64, dtype=torch.float32),
    raw_action_dim=26,
    prompt=your_task_prompt_json,
    view_point="ego_view",
    domain_name="g1_brainco",
    model_mode=ModelMode.POLICY,
    action_chunk_size=16,
    fps=15,
    resolution="480",
    input_video_key=model.input_video_key,
    batch_size=1,
    device="cuda",
)

with torch.no_grad():
    outputs = model.generate_samples_from_batch(batch, seed=[42], num_steps=30)

video_latent = outputs["vision"][0]
actions_norm = outputs["action"][0][:16, :26]  # [16, 26] normalized [-1,1]

pixels = model.decode(video_latent)            # [1, C, T, H, W] float [-1,1]

Action Denormalization

Actions are output as normalized values in [-1, 1]. Convert to raw joint angles (radians) using action_stats.json:

python

import json, numpy as np

with open("examples/data/g1_brainco/action_stats.json") as f:
    stats = json.load(f)

q01 = np.array(stats["q01"])  # [26]
q99 = np.array(stats["q99"])  # [26]

# actions_norm: numpy array [T, 26] in [-1, 1]
actions_raw = q01 + (actions_norm + 1.0) / 2.0 * (q99 - q01)  # radians

Video-to-Video Transfer (V2V)

The base Cosmos3-Nano supports V2V transfer: style-transfer from a reference video (e.g. human arm picking apple → G1 robot).

Two-Stage Loading for V2V with Fine-Tuned Weights

python

from cosmos_framework.inference.args import (
    OmniSampleOverrides, OmniSetupOverrides,
    EdgeTransferOverrides, PresetEdgeThreshold,
    BlurTransferOverrides, PresetBlurStrength,
)
from cosmos_framework.inference.inference import OmniInference, get_sample_data

BASE_CKPT      = "examples/checkpoints/Cosmos3-Nano"
FINETUNED_CKPT = "/path/to/iter_000010000"
CONFIG_YAML    = "cosmos_framework/inference/configs/model/Cosmos3-Nano.yaml"

# Stage 1: load base model (gets visual encoder + all base weights)
setup = OmniSetupOverrides(
    checkpoint_path=BASE_CKPT,
    config_file=CONFIG_YAML,
    output_dir="/tmp/v2v_out",
    guardrails=False,
).build_setup()
pipe = OmniInference.create(setup)

# Stage 2: overlay fine-tuned weights
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.filesystem import FileSystemReader
from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
from torch.distributed.checkpoint.state_dict import get_model_state_dict

state_dict = get_model_state_dict(pipe.model)
dcp.load(
    state_dict=state_dict,
    storage_reader=FileSystemReader(f"{FINETUNED_CKPT}/model"),
    planner=DefaultLoadPlanner(allow_partial_load=True),
)

# Run V2V with blur+edge dual conditioning (recommended for human→robot)
sample = OmniSampleOverrides(
    name="g1_output",
    output_dir="/tmp/v2v_out/g1_output",
    prompt="A Unitree G1 humanoid robot with five articulated fingers picking up a red apple...",
    vision_path="path/to/input_video.mp4",
    blur=BlurTransferOverrides(preset_blur_strength=PresetBlurStrength.MEDIUM),
    edge=EdgeTransferOverrides(preset_edge_threshold=PresetEdgeThreshold.MEDIUM),
    num_frames=121,
    fps=15,
    resolution="256",
).build_sample(model_config=pipe.model_config)

pipe.generate_batch([sample], get_sample_data(sample, model=pipe.model))
# Output: /tmp/v2v_out/g1_output/vision.mp4

Conditioning Strategies

Conditioning	When to use	Notes
Blur only	Scene re-styling, preserve motion loosely	Soft background color signal
Canny Edge only	Lock precise arm trajectory	No color context → flat background
Blur + Edge ⭐	Human → robot transfer	Best combo: background + motion
Blur + Edge + Depth	Maximum spatial control	Depth needs pre-computed `control_path`

For depth/segmentation conditioning, pre-compute the control video first:

python

# Pre-compute depth (Intel DPT-Large)
from transformers import pipeline as hf_pipeline
depth_estimator = hf_pipeline("depth-estimation", model="Intel/dpt-large")
# ... process frame by frame, save as depth.mp4

# Then use with TransferDataOverrides
from cosmos_framework.inference.args import TransferDataOverrides
depth=TransferDataOverrides(control_path="/path/to/depth.mp4")

Two-Stage Loading

Why is this needed?

The fine-tuned DCP only saves weights for the 4 trained modules (moe_gen, time_embedder, vae2llm, llm2vae). The 356 visual encoder keys are absent from the DCP because they were frozen during training.

Direct loading fails:

RuntimeError: Missing key in checkpoint state_dict: net.language_model.visual.blocks.0.attn.proj.bias.

Solution: Load the base model first (which includes the visual encoder), then overlay only the fine-tuned weights:

python

# ✅ Correct: 2-stage load
# Stage 1: base model loads visual encoder + all base weights
model = load_base_cosmos3_nano()

# Stage 2: overlay adapter weights, skip missing visual encoder keys
dcp.load(
    state_dict=get_model_state_dict(model),
    storage_reader=FileSystemReader("/path/to/iter_000010000/model"),
    planner=DefaultLoadPlanner(allow_partial_load=True),  # ← key flag
)

python

# ❌ Wrong: direct DCP load
dcp.load(state_dict=..., storage_reader=FileSystemReader(sft_ckpt))
# → RuntimeError: Missing key in checkpoint state_dict: net.language_model.visual...

Model Configuration

See Cosmos3-Nano-inference.yaml in this repo for the full inference config. Key parameters:

Parameter	Value
Architecture	Cosmos3-Nano (8B)
VLM Backbone	Qwen3-VL-8B
Action dimension	26 (G1 BrainCo joints)
Action frequency	15 Hz
Action chunk size	16 frames
Resolution	480p (policy), 256p (V2V)
Trained modules	moe_gen, time_embedder, vae2llm, llm2vae
Frozen modules	Qwen3-VL visual encoder + text backbone
Training iterations	10,000
Final loss	~3.8 (avg across ranks)

Training Details

Parameter	Value
Dataset	G1 BrainCo (apple-picking manipulation)
Training mode	Policy SFT (image + text → video + actions)
Optimizer	AdamW, lr=5e-6, wd=0, β=[0.9, 0.95]
Scheduler	LambdaCosine, 100-step warmup
Batch	1/GPU, grad_accum=2, effective=14
Hardware	7× A100 80GB (FSDP), ~22s/iter
Epochs	10,000 iterations (~61 hours total)
Checkpointing	Every 500 iters, DCP format

Files

model/
├── __0_0.distcp   ─┐
├── __1_0.distcp    │
├── __2_0.distcp    │  7× FSDP shards (~13 GB each = ~91 GB total)
├── __3_0.distcp    │  contains net.* and net_ema.* for 4 adapter modules
├── __4_0.distcp    │
├── __5_0.distcp    │
└── __6_0.distcp   ─┘

action_stats.json              — G1 BrainCo action normalization stats (q01, q99 per joint)
Cosmos3-Nano-inference.yaml    — Full model config YAML for inference

Citation

If you use this model, please cite the original Cosmos work:

bibtex

@misc{cosmos3nano,
  title={Cosmos3-Nano: An 8B Omni World Model},
  author={NVIDIA},
  year={2026},
  url={https://huggingface.co/nvidia/Cosmos3-Nano}
}

License

This model is released under the OpenMDW-1.1 license, inherited from the base Cosmos3-Nano model.

JeffrinSam/Cosmos3-Nano-G1-BrainCo-PolicySFT

Cosmos3-Nano — G1 BrainCo Policy SFT

What This Model Does

Checkpoint Format

Installation

Policy Inference (Image + Prompt → Video + Actions)

Quick Start — Built-in Tasks

Your Own Image + Custom Prompt

Init Frames for Built-in Tasks

Chunked Autoregressive Rollout

Programmatic Usage

Action Denormalization

Video-to-Video Transfer (V2V)

Two-Stage Loading for V2V with Fine-Tuned Weights

Conditioning Strategies

Two-Stage Loading

Model Configuration

Training Details

Files

Citation

License

JeffrinSam/
Cosmos3-Nano-G1-BrainCo-PolicySFT