Action FD DROID Post-Training¶

This document describes how to run the action_fd_droid_posttrain experiment. It trains Cosmos3-DROID forward dynamics in cosmos_framework.

Overview¶

Piece	Value
Experiment	`action_fd_droid_posttrain`
TOML	`examples/toml/sft_config/action_fd_droid_posttrain.toml`
Launch shell	`examples/launch_sft_action_fd_droid_posttrain.sh`
Config module	`cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py`
Dataset wrapper	`cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py`
Dataset root	Cosmos3-DROID parent root containing `success/` and `failure/`
Task mode	`forward_dynamics`
Action space	`ee_pose`: 10-D `[pos_delta, rot6d_delta, gripper]`
Chunk length / resolution	`16` frames at `480`

Prerequisites¶

Install the training environment as described in docs/setup.md.
Run commands from the repository root.
In NGC / PyTorch containers, set LD_LIBRARY_PATH='' before Python commands.

Inputs You Provide¶

This package ships the training stack — the registered action_fd_droid_posttrain experiment, the dataset class, and the paired TOML/launch shell. Three inputs are external and must be provided per environment:

Cosmos3-DROID dataset (in LeRobotDataset v3.0 format) — pre-download the dataset and point DATASET_PATH at the resulting .../Cosmos3-DROID parent directory. This experiment trains on both success/ and failure/ subsets.
DCP base checkpoint — prepare a base DCP checkpoint and point BASE_CHECKPOINT_PATH at it. For local smoke runs this can be examples/checkpoints/Cosmos3-Nano.

Data Layout¶

Set DATASET_PATH to the Cosmos3-DROID parent directory:

export DATASET_PATH=/path/to/Cosmos3-DROID

The loader uses both subsets:

$DATASET_PATH/
├── success/
│   └── meta/info.json
└── failure/
    └── meta/info.json

or institution-sharded split roots:

$DATASET_PATH/
├── success/<institution>/meta/info.json
└── failure/<institution>/meta/info.json

The launch shell passes DATASET_PATH through to the training config.

Full Reproduction¶

The flow mirrors the other SFT recipes:

# Step 1: prepare Cosmos3-DROID -> $DATASET_PATH.
# If you do not already have the dataset locally, download it from Hugging Face.
uvx hf@latest download --repo-type dataset nvidia/Cosmos3-DROID \
  --local-dir examples/data/Cosmos3-DROID --quiet
export DATASET_PATH=examples/data/Cosmos3-DROID

# Step 2: point to the base DCP checkpoint and Wan2.2 VAE.
export BASE_CHECKPOINT_PATH=examples/checkpoints/Cosmos3-Nano
export WAN_VAE_PATH=/path/to/Wan2.2_VAE.pth

# Step 3: choose the output root and launch.
export IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root
export LD_LIBRARY_PATH=''

bash examples/launch_sft_action_fd_droid_posttrain.sh

BASE_CHECKPOINT_PATH should point at the base DCP checkpoint. When using the paired launch shell, it must be a local directory because the shared launcher validates it before starting torchrun.

Validate The Config¶

Use --dryrun before launching workers:

PYTHONPATH=. python -m cosmos_framework.scripts.train \
  --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml \
  --dryrun

Run Training¶

Recommended paired launch shell:

bash examples/launch_sft_action_fd_droid_posttrain.sh

Pass short smoke-run overrides through EXTRA_TAIL_OVERRIDES:

export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10"
bash examples/launch_sft_action_fd_droid_posttrain.sh

Single-node, 8 GPU:

PYTHONPATH=. torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \
  --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml

Multi-node HSDP:

PYTHONPATH=. torchrun --nnodes=$NNODES --node_rank=$NODE_RANK --nproc_per_node=8 \
  -m cosmos_framework.scripts.train \
  --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml \
  -- model.parallelism.data_parallel_replicate_degree=$NNODES

Keep data_parallel_shard_degree=8 and set model.parallelism.data_parallel_replicate_degree to the number of nodes.

Outputs¶

Training outputs land under:

$IMAGINAIRE_OUTPUT_ROOT/cosmos3_action_fd/action_sft/<job.name>/

DCP checkpoints are saved under:

$RUN_DIR/checkpoints/iter_<N>/

The run is resumable by relaunching with the same output directory and job name.

Notes¶

This recipe uses mode="forward_dynamics", so actions are conditioning and the model trains video prediction from the first frame plus action sequence.
DROIDMergedLeRobotDataset uses action_space="ee_pose" for the 10-D end-effector pose action path.