Skip to content

README.md

Latest commit

 

History

History
169 lines (122 loc) · 10.5 KB

README.md

File metadata and controls

169 lines (122 loc) · 10.5 KB
Raw

Cosmos3-Nano Action-Policy Fine-Tuning (SFT)

This example demonstrates supervised fine-tuning (SFT) of Cosmos3-Nano into a robot action policy, using the action-policy recipe from cosmos-framework. Two embodiments are covered, each reproducing a Cosmos3 paper result:

  • DROID — reproduces Cosmos3-Nano-Policy-DROID: trained on real-robot DROID data, evaluated on the RoboLab simulation benchmark.
  • LIBERO-10 — reproduces the Cosmos3 paper's LIBERO-10 results: trained and evaluated on the LIBERO-10 simulation benchmark.
Recipe Launch shell Base model Dataset
Policy-DROID SFT (Nano) launch_sft_action_policy_droid_nano.sh Cosmos3-Nano Cosmos3-DROID success split
Policy-LIBERO-10 SFT (Nano) launch_sft_action_policy_libero_10_nano.sh Cosmos3-Nano LIBERO_LeRobot_v3 libero_10
Policy-LIBERO-all SFT (Nano) launch_sft_action_policy_libero_all_nano.sh Cosmos3-Nano LIBERO_LeRobot_v3 all 4 suites

The DROID recipe uses the registered action_policy_droid_nano experiment: joint_pos 8-D actions, proprioceptive state, concat_view 480p video, chunk length 32, episode-shuffle streaming, JSON-formatted action prompts (format_prompt_as_json=True), and the optional keep_ranges_1_0_1.json window filter. The reference reproduction runs lr 2e-4 (cosine, cycle 100000), generator loss_scale 10, global batch 8192 (HSDP 32x8 = 256 ranks; GB200 reference, 64 nodes x 4), for 10000 iters. The action prompt is serialized as JSON at both train and eval time, so evaluation must use the matching JSON prompt format.

The LIBERO recipe uses frame_wise_relative rot6d 10-D actions, quantile_rot normalization, concat_view (third-person + wrist) at 20 fps, lr 5e-5 / warmup 500 / cycle 16000, global batch 2048 (HSDP 2x8). To match the LIBERO-10 results reported in Cosmos3, we provide two presets:

  • libero-10 — action_policy_libero_nano + launch_sft_action_policy_libero_10_nano.sh; trains on libero_10 alone (max_iter 2000).
  • libero-all — action_policy_libero_all_nano + launch_sft_action_policy_libero_all_nano.sh; equal mix of all 4 LIBERO suites, which needs longer training (max_iter 5000).

For a runnable egocentric hand-pose data conversion example, see README_egocentric_hand_action.md. It converts a sample video and 3D hand-pose annotation pair into the raw 57D hand Action format used by the dataset path.

The recipe uses [job].task = "vfm" with the registered action_policy_droid_nano experiment. It trains a DROID policy model with joint_pos 8-D actions, proprioceptive state, concat_view 480p video, chunk length 32, episode-shuffle streaming, JSON-formatted action prompts, and the optional keep_ranges_1_0_1.json window filter.

Prerequisites

  1. Install cosmos-framework. This recipe drives cosmos_framework.scripts.train, so install a cosmos-framework checkout first — follow the shared cosmos-framework setup (clone into packages/cosmos3, then uv sync --all-extras --group=cu130-train; use cu128-train on a CUDA 12.x driver).
  2. Recommended container. For a curated CUDA + PyTorch base, NVIDIA recommends starting from the NGC PyTorch container nvcr.io/nvidia/pytorch:25.09-py3 (CUDA 13; use :25.06-py3 for a CUDA 12.8 driver). See the cosmos-framework setup guide.
  3. Activate the cosmos-framework venv so cosmos_framework is importable: source <path-to>/packages/cosmos3/.venv/bin/activate.
  4. Hugging Face access. The base model, Wan2.2 VAE, Cosmos3-DROID dataset, and keep_ranges_1_0_1.json filter are hosted on Hugging Face — authenticate once with uvx hf@latest auth login (or export HF_TOKEN) and accept any model/dataset terms first.
  5. Weights & Biases. The TOML follows the canonical cosmos-framework example and defaults to wandb_mode="online". Export WANDB_API_KEY for online logging, or add job.wandb_mode=disabled to EXTRA_TAIL_OVERRIDES for a local/no-W&B run.
  6. Prepare the Cosmos3-DROID dataset. The full Cosmos3-DROID dataset is large, so this cookbook expects you to stage it explicitly. By default the launcher looks for data/Cosmos3-DROID/success/meta/info.json under this folder. Set DATASET_PATH=/path/to/Cosmos3-DROID/success to use another filesystem.
  7. Run from this directory (cookbooks/cosmos3/generator/action/finetune/). Converted checkpoints, the Wan2.2 VAE, the keep_ranges_1_0_1.json filter, and run outputs default to checkpoints/, data/, and outputs/ here (all git-ignored).

Quick start

Stage the dataset, then run the launcher:

# One possible layout. The dataset is large; put it on a filesystem with enough space.
uvx hf@latest download \
    nvidia/Cosmos3-DROID \
    --repo-type dataset \
    --local-dir data/Cosmos3-DROID

bash launch_sft_action_policy_droid_nano.sh

The launcher is a complete local wrapper for the public cookbook:

  • checks that the Cosmos3-DROID success split exists
  • downloads Wan2.2_VAE.pth if needed
  • converts Cosmos3-Nano to a local DCP checkpoint if needed
  • downloads keep_ranges_1_0_1.json if needed
  • launches training with action_policy_droid_nano.toml

The script intentionally stays close to the cosmos-framework example launcher: DATASET_PATH is bridged to DROID_ROOT, BASE_CHECKPOINT_PATH and WAN_VAE_PATH are exported for the TOML, EXTRA_TAIL_OVERRIDES is appended after --, and NPROC_PER_NODE / NNODES / NODE_RANK / MASTER_ADDR / MASTER_PORT map directly to torchrun. The extra cookbook behavior is the reasoner-style local prep above and enabling the released keep_ranges_1_0_1.json filter by default.

Paths are fixed at the top of the script (under this git-ignored folder) — edit them there or export env vars to relocate data and checkpoints:

export DATASET_PATH=/scratch/Cosmos3-DROID/success
export BASE_CHECKPOINT_PATH=/scratch/checkpoints/Cosmos3-Nano
export WAN_VAE_PATH=/scratch/checkpoints/wan22_vae/Wan2.2_VAE.pth
export FILTER_PATH=/scratch/droid/keep_ranges_1_0_1.json
bash launch_sft_action_policy_droid_nano.sh

The committed TOML pins the GB200 reference shape (HSDP 32x8 = 256 ranks, global batch 8192). To run on a single 8-GPU node — e.g. a short smoke test — drop the replicate degree to 1 alongside the iteration/batch knobs:

export EXTRA_TAIL_OVERRIDES=" \
  job.wandb_mode=disabled \
  trainer.max_iter=10 \
  checkpoint.save_iter=10 \
  model.config.parallelism.data_parallel_replicate_degree=1 \
  dataloader_train.max_samples_per_batch=32 \
"
bash launch_sft_action_policy_droid_nano.sh

To reproduce the reference at full global batch 8192 on fewer GPUs, keep data_parallel_replicate_degree=1 and raise trainer.grad_accum_iter (32 on one 8-GPU node) instead of shrinking the batch.

LIBERO quick start

Each launcher stages its dataset (auto-downloaded if missing), downloads the Wan VAE, converts the base checkpoint, and trains.

libero-10:

bash launch_sft_action_policy_libero_10_nano.sh
  • downloads nvidia/LIBERO_LeRobot_v3 libero_10 to data/LIBERO_LeRobot_v3/libero_10 if missing
  • launches training with action_policy_libero_10_nano.toml (max_iter 2000)

libero-all:

bash launch_sft_action_policy_libero_all_nano.sh
  • downloads all 4 suites of nvidia/LIBERO_LeRobot_v3 to data/LIBERO_LeRobot_v3 if missing
  • launches training with action_policy_libero_all_nano.toml (max_iter 5000; it needs longer to converge)

Both download Wan2.2_VAE.pth and convert Cosmos3-Nano to a local DCP checkpoint if needed. Relocate inputs via env vars, or run a short smoke test:

export LIBERO_ROOT=/scratch/LIBERO_LeRobot_v3/libero_10   # libero-all preset: the parent dir, /scratch/LIBERO_LeRobot_v3
export EXTRA_TAIL_OVERRIDES="job.wandb_mode=disabled trainer.max_iter=10 checkpoint.save_iter=10 dataloader_train.max_samples_per_batch=32"
bash launch_sft_action_policy_libero_10_nano.sh

Checkpoints are saved every 500 iters.

Outputs

Training writes to outputs/train/<project>/<group>/<name>/:

  • checkpoints/iter_<N>/ — DCP checkpoint (model / optim / scheduler / trainer state); checkpoints/latest_checkpoint.txt names the newest.
  • config.yaml, launch metadata, logs, and one directory per registered callback.

Export to Hugging Face safetensors

RUN_DIR=outputs/train/<project>/<group>/<name>
CKPT=$RUN_DIR/checkpoints/$(cat "$RUN_DIR/checkpoints/latest_checkpoint.txt")
python -m cosmos_framework.scripts.export_model \
    --checkpoint-path "$CKPT" \
    --config-file "$RUN_DIR/config.yaml" \
    -o "$RUN_DIR/model"

Use the exported $RUN_DIR/model with the Cosmos3-Nano-Policy-DROID inference cookbook.

Convert to Diffusers

Convert the export into a Diffusers pipeline:

python -m cosmos_framework.scripts.convert_model_to_diffusers \
    --checkpoint-path "$RUN_DIR/model" -o "$RUN_DIR/diffusers"

The input is the Hugging Face directory produced by export_model above (not a raw DCP checkpoint), so run the export first. See the Export and Convert Checkpoints overview for details.

Advanced configuration

These recipes are intentionally minimal. For the full post-training reference — raw torchrun, resuming, every TOML field, the keep_ranges_1_0_1.json filter, and advanced HSDP/multi-node parallelism — see the canonical cosmos-framework docs:

SFT here is a multi-GPU torchrun job, so this cookbook ships as a launch script + this README rather than a one-click notebook.