Skip to content

Cosmos3 DROID Action-Policy Post-Training

This example post-trains Cosmos3-Nano into a DROID action policy on the Cosmos3-DROID dataset, reproducing Cosmos3-Nano-Policy-DROID. The model predicts absolute joint-position actions conditioned on proprioceptive state and video observations at a resolution of 480p (640×360). Experiment action_policy_droid_nano, launcher launch_sft_action_policy_droid_nano.sh.

Two external inputs are required: (1) a pre-downloaded Cosmos3-DROID dataset in LeRobotDataset v3.0 format, and (2) a DCP base checkpoint converted from the chosen base model (Cosmos3-Nano or Cosmos3-Edge).

The recipe runs multi-node via HSDP (single node / 8 GPUs and beyond).

Prerequisites

  • Setup — clone the repo, install the training extras (uv sync --all-extras --group=cu130-train), and activate the environment.
  • Environment Variables — set environment variables.
  • FAQ — troubleshooting (OOM during SFT, defaults) and common pitfalls.

The runnable artifacts (TOML recipe, paired launch shell) live in examples/; all commands below run from the repo root with the environment activated.

Inputs You Provide

This package ships the training stack — the registered action_policy_droid_nano experiment, the DROID action dataset class with the recipe knobs (action_space=joint_pos, use_state, concat_view), and the EMA warm-start in checkpoint/dcp.py. Two inputs are external and must be provided per environment:

  1. Cosmos3-DROID dataset (in LeRobotDataset v3.0 format) — pre-download the dataset into a directory named droid_plus_lerobot_640x360_20260412 (the loader resolves the dataset schema from the directory name) and point DROID_ROOT at that directory — the parent containing success/, not .../success itself.
  2. DCP base checkpoint — convert Cosmos3-Nano to DCP and point BASE_CHECKPOINT_PATH at it (see Full Reproduction). Action heads are not loaded from it (they init fresh).

Recipe

knob value
init Cosmos3-Nano (public Hugging Face repo)
action space joint_pos (absolute joint position, 8-D incl. gripper)
state use_state=true (proprioception; valid only with joint_pos)
task mode policy (single-task; the joint multi-task default is avoided)
resolution 480
viewpoint / video concat_view / video_mode=null
chunk length 32 (tokenizer encode_exact_durations=[33])
sequence packing max_num_tokens_after_packing=-1 (full vision sequence per step)
shuffle episode-shuffle stream (decorrelates the per-step global batch)
window filter keep_ranges_1_0_1.json (KarlP/droid) — trains the curated ≈74% window set
lr 2e-4
global batch 8192 = max_samples_per_batch × world size × grad_accum_iter (reduce the first / raise the last to fit GPU memory)
eval disabled for the reproduction run

Full Reproduction

The OSS flow mirrors the other recipes (see docs/training.md):

# Step 1: download Cosmos3-DROID into a directory named droid_plus_lerobot_640x360_20260412
# (the loader resolves the dataset schema from the directory name), then point
# DATASET_PATH at that parent directory — the one containing success/, not .../success.
hf download nvidia/Cosmos3-DROID --repo-type dataset \
  --local-dir /path/to/droid_plus_lerobot_640x360_20260412
export DATASET_PATH=/path/to/droid_plus_lerobot_640x360_20260412

# Step 2: convert the base checkpoint -> $BASE_CHECKPOINT_PATH
python -m cosmos_framework.scripts.convert_model_to_dcp \
  --checkpoint-path Cosmos3-Nano \
  -o $BASE_CHECKPOINT_PATH

# Step 3: download the keep_ranges_1_0_1.json window filter (drops idle/non-task frames -> trains
# the curated ~74% window set, matching the released model).
export FILTER_DIR=/path/to/droid_filters
hf download KarlP/droid keep_ranges_1_0_1.json --local-dir $FILTER_DIR

# Step 4: launch. The TOML selects the experiment + scalars; the dataset/action
# knobs come from the registered experiment.
export BASE_CHECKPOINT_PATH=/path/to/base_checkpoint
export WAN_VAE_PATH=/path/to/Wan2.2_VAE.pth
export NPROC_PER_NODE=8
# Enable the keep_ranges_1_0_1.json filter via EXTRA_TAIL_OVERRIDES (space-separated Hydra
# overrides; an exported string survives `bash <wrapper>`). $FILTER_DIR is expanded HERE,
# at export time — it must already be set in this shell, or filter_dict_path silently
# becomes /keep_ranges_1_0_1.json and the dataset fails with FileNotFoundError.
: "${FILTER_DIR:?set FILTER_DIR (Step 3) before composing EXTRA_TAIL_OVERRIDES}"
export EXTRA_TAIL_OVERRIDES=" \
  dataloader_train.dataloader.datasets.droid.dataset.use_filter_dict=True \
  dataloader_train.dataloader.datasets.droid.dataset.filter_dict_path=$FILTER_DIR/keep_ranges_1_0_1.json \
"
bash examples/launch_sft_action_policy_droid_nano.sh

The recipe TOML (action_policy_droid_nano.toml) sets the scalar knobs (max_iter, save_iter, grad_clip, parallelism, wandb); the dataset/action knobs (joint_pos, use_state, concat_view, 480p, chunk 32, count-based batch) live in the registered action_policy_droid_nano experiment per the schema's design. For multi-node HSDP, set model.config.parallelism.data_parallel_replicate_degree = <num_nodes> (intra-node shard stays 8).

The keep_ranges_1_0_1.json filter maps each DROID trajectory key to a list of [start, end] frame ranges; only windows whose start falls inside a kept range are trained on (episodes absent from the dict are dropped). To train on the full window set instead, leave EXTRA_TAIL_OVERRIDES unset.

Checkpoints

  • Saved every save_iter iters (1000 in the validated run) to the object store, at <bucket>/<project>/<group>/<job.name>/checkpoints/iter_<N>/.
  • The run is resumable from the latest checkpoint (re-launch with the same job.name).
  • Export to HF safetensors via cosmos_framework.scripts.export_model (see docs/training.md).