[Feat] support action forward dynamics finetuning

cookbooks/cosmos3/generator/action/finetune/action_fd_droid_posttrain.md

@@ -0,0 +1,173 @@
+# 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                      | `toml/sft_config/action_fd_droid_posttrain.toml`                                                                       |
+| Launch shell              | `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](https://huggingface.co/datasets/nvidia/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 cosmos-framework as described in the shared
+  [setup instructions](../../../README.md), then activate its
+  virtual environment.
+- Authenticate with Hugging Face (`uvx hf@latest auth login` or `HF_TOKEN`) and
+  accept the terms for the required model and dataset repositories.
+- The TOML defaults to online Weights & Biases logging. Export `WANDB_API_KEY`,
+  or disable it with `EXTRA_TAIL_OVERRIDES="job.wandb_mode=disabled"`.
+- Run commands from `cookbooks/cosmos3/generator/action/finetune/`.
+- 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:
+
+1. **[Cosmos3-DROID](https://huggingface.co/datasets/nvidia/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.
+2. **DCP base checkpoint** — prepare a base DCP checkpoint and point
+   `BASE_CHECKPOINT_PATH` at it. For local smoke runs this can be
+   `checkpoints/Cosmos3-Nano`.
+3. **Wan2.2 VAE** — download `Wan2.2_VAE.pth` from
+   [`Wan-AI/Wan2.2-TI2V-5B`](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B)
+   and point `WAN_VAE_PATH` at the file.
+
+## Data Layout
+
+Set `DATASET_PATH` to the Cosmos3-DROID parent directory:
+
+```shell
+export DATASET_PATH=/path/to/Cosmos3-DROID
+```
+
+The loader uses both subsets:
+
+```text
+$DATASET_PATH/
+├── success/
+│   └── meta/info.json
+└── failure/
+    └── meta/info.json
+```
+
+or institution-sharded split roots:
+
+```text
+$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:
+
+```shell
+# 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 data/Cosmos3-DROID --quiet
+export DATASET_PATH=data/Cosmos3-DROID
+
+# Step 2: prepare the base DCP checkpoint and Wan2.2 VAE.
+python -m cosmos_framework.scripts.convert_model_to_dcp \
+  -o checkpoints/Cosmos3-Nano \
+  --checkpoint-path Cosmos3-Nano
+uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \
+  --local-dir checkpoints/wan22_vae
+export BASE_CHECKPOINT_PATH=checkpoints/Cosmos3-Nano
+export WAN_VAE_PATH=checkpoints/wan22_vae/Wan2.2_VAE.pth
+
+# Step 3: optionally choose the output root and launch.
+export OUTPUT_ROOT=/path/to/output_root
+export LD_LIBRARY_PATH=''
+
+bash 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:
+
+```shell
+PYTHONPATH=. python -m cosmos_framework.scripts.train \
+  --sft-toml toml/sft_config/action_fd_droid_posttrain.toml \
+  --dryrun
+```
+
+## Run Training
+
+Recommended paired launch shell:
+
+```shell
+bash launch_sft_action_fd_droid_posttrain.sh
+```
+
+Pass short smoke-run overrides through `EXTRA_TAIL_OVERRIDES`:
+
+```shell
+export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10"
+bash launch_sft_action_fd_droid_posttrain.sh
+```
+
+Single-node, 8 GPU:
+
+```shell
+PYTHONPATH=. torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \
+  --sft-toml toml/sft_config/action_fd_droid_posttrain.toml
+```
+
+Multi-node HSDP:
+
+```shell
+PYTHONPATH=. torchrun --nnodes=$NNODES --node_rank=$NODE_RANK --nproc_per_node=8 \
+  -m cosmos_framework.scripts.train \
+  --sft-toml 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:
+
+```text
+$IMAGINAIRE_OUTPUT_ROOT/cosmos3_action_fd/action_sft/<job.name>/
+```
+
+DCP checkpoints are saved under:
+
+```text
+$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.

cookbooks/cosmos3/generator/action/finetune/launch_sft_action_fd_droid_posttrain.sh

@@ -0,0 +1,125 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# ============================================================================
+# Structured-TOML launch for action_fd_droid_posttrain.
+#
+# This script intentionally does not source _sft_launcher_common.sh. It contains
+# the recipe-specific defaults, input checks, path anchoring, tail overrides, and
+# torchrun invocation in one file.
+#
+# Run with the cosmos-framework venv active (see action_fd_droid_posttrain.md).
+# Training starts from the installed cosmos-framework checkout root so relative
+# model config paths (e.g. Qwen3-VL JSON) resolve correctly.
+#
+# Env vars (override for your filesystem):
+#   DATASET_PATH            Cosmos3-DROID parent dir (success/ + failure/)
+#   BASE_CHECKPOINT_PATH    Base DCP checkpoint
+#   WAN_VAE_PATH            Wan2.2 VAE .pth
+#   COSMOS_FRAMEWORK_ROOT   Optional override for the framework checkout root
+#   NPROC_PER_NODE          torchrun --nproc_per_node (default 8)
+#   EXTRA_TAIL_OVERRIDES    space-separated Hydra overrides
+#   MASTER_PORT             torchrun --master_port (default 50012)
+#   NNODES                  torchrun --nnodes for multi-node
+#   NODE_RANK               torchrun --node_rank for multi-node
+#   MASTER_ADDR             torchrun --master_addr for multi-node
+#   OUTPUT_ROOT             default output root (default outputs/train)
+#   LOG_FILENAME            override log filename
+# ============================================================================
+
+set -euo pipefail
+
+# Cookbook recipe directory (TOML, logs, default relative paths).
+COOKBOOK_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+TOML_FILE="$COOKBOOK_DIR/toml/sft_config/action_fd_droid_posttrain.toml"
+
+: "${DATASET_PATH:=data/Cosmos3-DROID}"
+: "${BASE_CHECKPOINT_PATH:=checkpoints/Cosmos3-Nano}"
+: "${WAN_VAE_PATH:=checkpoints/wan22_vae/Wan2.2_VAE.pth}"
+
+[[ "$DATASET_PATH" = /* ]] || DATASET_PATH="$COOKBOOK_DIR/$DATASET_PATH"
+[[ "$BASE_CHECKPOINT_PATH" = /* ]] || BASE_CHECKPOINT_PATH="$COOKBOOK_DIR/$BASE_CHECKPOINT_PATH"
+[[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$COOKBOOK_DIR/$WAN_VAE_PATH"
+export DATASET_PATH BASE_CHECKPOINT_PATH WAN_VAE_PATH
+
+OUTPUT_ROOT="${OUTPUT_ROOT:-$COOKBOOK_DIR/outputs/train}"
+LOG_DIR="$OUTPUT_ROOT/logs"
+TOML_STEM="$(basename "$TOML_FILE" .toml)"
+LOG_FILE="$LOG_DIR/${LOG_FILENAME:-${TOML_STEM}_sft.log}"
+IMAGINAIRE_OUTPUT_ROOT="${IMAGINAIRE_OUTPUT_ROOT:-$OUTPUT_ROOT}"
+mkdir -p "$LOG_DIR"
+
+echo ">>> $(date '+%H:%M:%S') Checking inputs..."
+[[ -f "$TOML_FILE" ]] || { echo "ERROR: TOML not found: $TOML_FILE" >&2; exit 1; }
+[[ -d "$DATASET_PATH" ]] || { echo "ERROR: DATASET_PATH not found: $DATASET_PATH" >&2; exit 1; }
+[[ -d "$BASE_CHECKPOINT_PATH" ]] || { echo "ERROR: BASE_CHECKPOINT_PATH not found: $BASE_CHECKPOINT_PATH" >&2; exit 1; }
+[[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }
+
+[[ -f "$DATASET_PATH/success/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/success/*/meta/info.json")" ]] || {
+    echo "ERROR: missing Cosmos3-DROID success split under $DATASET_PATH (expected success/meta/info.json or success/*/meta/info.json)" >&2
+    exit 1
+}
+[[ -f "$DATASET_PATH/failure/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/failure/*/meta/info.json")" ]] || {
+    echo "ERROR: missing Cosmos3-DROID failure split under $DATASET_PATH (expected failure/meta/info.json or failure/*/meta/info.json)" >&2
+    exit 1
+}
+
+if [[ -n "${COSMOS_FRAMEWORK_ROOT:-}" ]]; then
+    FRAMEWORK_ROOT="$(cd "$COSMOS_FRAMEWORK_ROOT" && pwd)"
+else
+    FRAMEWORK_ROOT="$(
+        python -c 'import pathlib, cosmos_framework; print(pathlib.Path(cosmos_framework.__file__).resolve().parents[1])'
+    )"
+fi
+
+QWEN3_VL_CONFIG="$FRAMEWORK_ROOT/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"
+[[ -f "$QWEN3_VL_CONFIG" ]] || {
+    cat >&2 <<EOF
+ERROR: cannot resolve cosmos-framework model configs under:
+  $FRAMEWORK_ROOT
+
+Expected Qwen3-VL config at:
+  $QWEN3_VL_CONFIG
+
+Activate the cosmos-framework venv, or export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework.
+EOF
+    exit 1
+}
+
+TAIL_OVERRIDES=(
+    "trainer.straggler_detection.enabled=False"
+)
+if [[ -n "${EXTRA_TAIL_OVERRIDES:-}" ]]; then
+    # EXTRA_TAIL_OVERRIDES is intentionally word-split to match the framework launcher UX.
+    # shellcheck disable=SC2206
+    EXTRA_OVERRIDES_ARRAY=(${EXTRA_TAIL_OVERRIDES})
+    TAIL_OVERRIDES+=("${EXTRA_OVERRIDES_ARRAY[@]}")
+fi
+
+TRAILING_ARGS=(-- "${TAIL_OVERRIDES[@]}")
+
+TORCHRUN_ARGS=(--nproc_per_node="${NPROC_PER_NODE:-8}" --master_port="${MASTER_PORT:-50012}")
+[[ -n "${NNODES:-}" ]] && TORCHRUN_ARGS+=(--nnodes="$NNODES")
+[[ -n "${NODE_RANK:-}" ]] && TORCHRUN_ARGS+=(--node_rank="$NODE_RANK")
+[[ -n "${MASTER_ADDR:-}" ]] && TORCHRUN_ARGS+=(--master_addr="$MASTER_ADDR")
+
+cd "$FRAMEWORK_ROOT"
+echo ">>> $(date '+%H:%M:%S') cookbook:   $COOKBOOK_DIR"
+echo ">>> $(date '+%H:%M:%S') framework:  $FRAMEWORK_ROOT"
+echo ">>> $(date '+%H:%M:%S') TOML:       $TOML_FILE"
+echo ">>> $(date '+%H:%M:%S') dataset:    $DATASET_PATH"
+echo ">>> $(date '+%H:%M:%S') checkpoint: $BASE_CHECKPOINT_PATH"
+echo ">>> $(date '+%H:%M:%S') log:        $LOG_FILE"
+
+IMAGINAIRE_OUTPUT_ROOT="$IMAGINAIRE_OUTPUT_ROOT" \
+PYTHONPATH="${FRAMEWORK_ROOT}${PYTHONPATH:+:$PYTHONPATH}" \
+    torchrun "${TORCHRUN_ARGS[@]}" -m cosmos_framework.scripts.train \
+    --sft-toml="$TOML_FILE" \
+    "${TRAILING_ARGS[@]}" \
+    2>&1 | tee "$LOG_FILE"
+
+EXIT_CODE=${PIPESTATUS[0]}
+echo ">>> $(date '+%H:%M:%S') Done (exit $EXIT_CODE)"
+exit $EXIT_CODE

cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_fd_droid_posttrain.toml

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# ============================================================================
+# DROID-Merged LeRobot forward-dynamics post-training.
+#
+# Env required:
+#   DATASET_PATH=/path/to/Cosmos3-DROID
+#   BASE_CHECKPOINT_PATH=<base DCP dir or object-store path>
+#   WAN_VAE_PATH=<Wan2.2_VAE.pth>
+#   IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root
+#
+# The registered experiment owns dataset mode/action-space wiring:
+#   mode="forward_dynamics", action_space="ee_pose", chunk_length=16, res480.
+# ============================================================================
+
+[job]
+task       = "vfm"
+experiment = "action_fd_droid_posttrain"
+project    = "cosmos3_action_fd"
+group      = "action_sft"
+name       = "action_fd_droid_posttrain"
+wandb_mode = "online"
+
+[model]
+precision = "bfloat16"
+
+[model.parallelism]
+data_parallel_shard_degree     = 8
+data_parallel_replicate_degree = 1
+
+[model.activation_checkpointing]
+mode = "selective"
+
+[model.tokenizer]
+vae_path = "${oc.env:WAN_VAE_PATH}"
+
+[scheduler]
+cycle_lengths = [20000]
+
+[trainer]
+max_iter     = 20000
+logging_iter = 50
+
+[checkpoint]
+load_path = "${oc.env:BASE_CHECKPOINT_PATH}"
+save_iter = 250