diff --git a/cookbooks/cosmos3/generator/action/finetune/README.md b/cookbooks/cosmos3/generator/action/finetune/README.md index fc5c1dd1..26ef5e6c 100644 --- a/cookbooks/cosmos3/generator/action/finetune/README.md +++ b/cookbooks/cosmos3/generator/action/finetune/README.md @@ -1,15 +1,17 @@ # Cosmos3-Nano Action-Policy Fine-Tuning (SFT) -This example demonstrates supervised fine-tuning (SFT) of [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano) into a robot action policy, using the action-policy recipe from [cosmos-framework](https://github.com/NVIDIA/cosmos-framework). Two embodiments are covered, each reproducing a Cosmos3 paper result: +This example demonstrates supervised fine-tuning (SFT) of [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano) into a robot action policy, using the action-policy recipe from [cosmos-framework](https://github.com/NVIDIA/cosmos-framework). Three embodiments are covered: - **DROID** — reproduces [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/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. +- **RoboCasa** — mobile-base kitchen manipulation: trained on any RoboCasa task split (target or pretrain, atomic or composite) and evaluated closed-loop in held-out kitchens. | Recipe | Launch shell | Base model | Dataset | | --- | --- | --- | --- | | Policy-DROID SFT (Nano) | `launch_sft_action_policy_droid_nano.sh` | Cosmos3-Nano | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) success split | | Policy-LIBERO-10 SFT (Nano) | `launch_sft_action_policy_libero_10_nano.sh` | Cosmos3-Nano | [LIBERO_LeRobot_v3](https://huggingface.co/datasets/nvidia/LIBERO_LeRobot_v3) `libero_10` | | Policy-LIBERO-all SFT (Nano) | `launch_sft_action_policy_libero_all_nano.sh` | Cosmos3-Nano | [LIBERO_LeRobot_v3](https://huggingface.co/datasets/nvidia/LIBERO_LeRobot_v3) all 4 suites | +| Policy-RoboCasa SFT (Nano) | `launch_sft_action_policy_robocasa_nano.sh` | Cosmos3-Nano | [RoboCasa](https://robocasa.ai) target or pretrain splits (LeRobot v3.0) | 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. @@ -18,12 +20,31 @@ The LIBERO recipe uses `frame_wise_relative` rot6d 10-D actions, `quantile_rot` - **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`](./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 RoboCasa recipe uses the registered `action_policy_robocasa_nano` experiment. It is the +first of the three with a **mobile base**, so the action contract is wider than the arm-only +10-D used by LIBERO. RoboCasa's native base command is passed through unchanged +(`base_encoding="raw"`, 15-D): -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. +``` +[base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] +``` + +Observation is `camera_set="left_wrist"` — `agentview_left` and `eye_in_hand` concatenated +horizontally at their native 256x256 each — plus end-effector proprioception as a clean +condition token (`use_state=True`). `action_normalization=None`, because the state token is +absolute and delta statistics do not apply. Chunk length 32 at 20 fps, lr 5e-5 / warmup 500 / +cycle 16000, global batch 2048 (HSDP 8x2 = 16 ranks; 4 GB200 nodes x 4 GPUs). + +The recipe is independent of which RoboCasa tasks it trains on. `TASK_SET` selects the split: + +| `TASK_SET` | Tasks | RoboCasa split | +| --- | --- | --- | +| `target_atomic` (default) | 18 | `target/atomic` | +| `pretrain_atomic` | 65 | `pretrain/atomic` | +| `pretrain_all` | 300 | `pretrain` atomic + composite | + +The names resolve to the task tuples in cosmos-framework's `robocasa_lerobot_dataset.py`, +which is the single source of truth. Larger splits need a longer schedule — see the examples below. ## Prerequisites @@ -126,6 +147,131 @@ bash launch_sft_action_policy_libero_10_nano.sh Checkpoints are saved every 500 iters. +## RoboCasa quick start + +Get the RoboCasa dataset from [robocasa.ai](https://robocasa.ai), convert the split you want, +then train. The example below uses `target/atomic`; `SRC_ROOT` and `TASK_SET` select any other. + +The converter ships with cosmos-framework: + +```shell +CONVERT=$(python -c 'import cosmos_framework, pathlib +print(pathlib.Path(cosmos_framework.__file__).parent / "scripts/convert_robocasa_to_lerobot_v30.sh")') + +# RoboCasa publishes LeRobot v2.1; the pinned lerobot is v3.0-only and rejects it. +# Conversion runs on copies -- the released dataset is never modified -- and is idempotent. +SRC_ROOT=/path/to/robocasa/datasets/v1.0/target/atomic \ + bash "$CONVERT" # -> data/robocasa_v30///lerobot/ + +bash launch_sft_action_policy_robocasa_nano.sh +``` + +Other splits: convert the matching source directory, point `ROBOCASA_ROOT` at it, and give the +larger corpus a longer schedule. + +```shell +# 65 pretrain/atomic tasks. The converter takes every task it finds under SRC_ROOT. +SRC_ROOT=/path/to/robocasa/datasets/v1.0/pretrain/atomic V30_ROOT=$PWD/data/robocasa_pretrain \ + bash "$CONVERT" + +TASK_SET=pretrain_atomic ROBOCASA_ROOT=$PWD/data/robocasa_pretrain \ + bash launch_sft_action_policy_robocasa_nano.sh + +# 300 pretrain tasks (atomic + composite) -- convert both splits into one root, 40k steps +SRC_ROOT=/path/to/robocasa/datasets/v1.0/pretrain/composite V30_ROOT=$PWD/data/robocasa_pretrain \ + bash "$CONVERT" + +TASK_SET=pretrain_all ROBOCASA_ROOT=$PWD/data/robocasa_pretrain \ + EXTRA_TAIL_OVERRIDES="trainer.max_iter=40000 scheduler.cycle_lengths=[40000] scheduler.warm_up_steps=[1250]" \ + bash launch_sft_action_policy_robocasa_nano.sh +``` + +The launcher resolves `TASK_SET` to its task tuple, verifies every task is staged before +starting, and names any that are missing; the loader itself skips absent tasks with a warning +rather than failing. + +Relocate inputs via env vars, or run a short smoke test: + +```shell +export ROBOCASA_ROOT=/scratch/robocasa_v30 +# The TOML pins HSDP 8x2 = 16 ranks (4 nodes x 4 GPUs). On a single 8-GPU node, drop the +# replicate degree to 1 as well, or torchrun starts 8 processes for a 16-rank layout. +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=8 \ +" +bash launch_sft_action_policy_robocasa_nano.sh +``` + +Checkpoints are saved every 1000 iters. + +### Closed-loop evaluation + +RoboCasa is scored by rolling the policy out in simulation, not by held-out loss. The evaluator +lives in cosmos-framework at `cosmos_framework/simulation/robocasa/closed_loop_eval.py`, alongside +the LIBERO one. + +It drives the simulator against a policy server +(`cosmos_framework.scripts.action_policy_server_robocasa`) over HTTP, so it needs **two Python +environments**: robosuite/robocasa and cosmos-framework have conflicting numpy/mujoco pins and +cannot share a venv. The server runs in the cosmos-framework venv; the evaluator runs in a +separate venv built per the [RoboCasa install guide](https://robocasa.ai). Both processes run +together — start the server first and leave it up: + +```shell +# --- process 1, cosmos-framework venv ------------------------------------------------- +# Serve the DCP checkpoint directly; no safetensors export is needed for evaluation. +RUN_DIR=outputs/train/// +CKPT=$RUN_DIR/checkpoints/$(cat "$RUN_DIR/checkpoints/latest_checkpoint.txt") + +python -m cosmos_framework.scripts.action_policy_server_robocasa \ + --checkpoint-path "$CKPT/model" --config-file "$RUN_DIR/config.yaml" \ + --port 8912 --num-steps 10 --guidance 3.0 --fps 20 + +# --- process 2, robosuite/robocasa venv ----------------------------------------------- +# cosmos-framework is NOT installed in this venv (that is the point of splitting them), but +# the evaluator only imports pure-python modules from it, so a checkout on PYTHONPATH is +# enough. --dataset-dir is the ORIGINAL RoboCasa export, not the converted copy: the +# simulator reads scene XML from it. +MUJOCO_GL=egl PYTHONPATH=/path/to/cosmos-framework \ +python -m cosmos_framework.simulation.robocasa.closed_loop_eval \ + --server-url http://127.0.0.1:8912 \ + --dataset-dir /path/to/robocasa/datasets/v1.0/target/atomic///lerobot \ + --output-dir eval_out --num-test-episodes 20 \ + --action-horizon 32 --camera-set left_wrist \ + --use-state --use-base-action --base-encoding raw \ + --seed 0 --image-size 256 --cam-size 256 +``` + +The evaluation contract must match the training recipe — a mismatch does not raise, it just +degrades the policy. For this recipe: `--action-horizon 32` (equals `chunk_length`), +`--camera-set left_wrist`, `--use-state`, and `--use-base-action --base-encoding raw` for the +15-D contract. + +The server derives that same 15-D width from the checkpoint's experiment config, so the two +sides agree without being configured twice; `GET /info` reports the width in effect. Pass +`--raw-action-dim` only to override it, e.g. when serving a checkpoint whose config is not +available. + +`smoke_test_robocasa_eval.sh` runs this exact pair end to end on one task for a couple of +rollouts. Use it to check the handshake before launching a full evaluation: + +```shell +RUN_DIR=outputs/train/// \ +FRAMEWORK_ROOT=/path/to/cosmos-framework \ +SIM_PYTHON=/path/to/robocasa-venv/bin/python \ +DATASET_DIR=/path/to/robocasa/datasets/v1.0/target/atomic///lerobot \ + bash smoke_test_robocasa_eval.sh +``` + +The policy server instantiates the text-safety guardrail, which downloads the gated +[`nvidia/Cosmos-Guardrail1`](https://huggingface.co/nvidia/Cosmos-Guardrail1) repo (~7 GB). Accept +its terms and authenticate (`uvx hf@latest auth login`) before the first evaluation, and do not +set `HF_HUB_OFFLINE=1` on that run. + ## Outputs Training writes to `outputs/train////`: diff --git a/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_robocasa_nano.sh b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_robocasa_nano.sh new file mode 100755 index 00000000..7354b799 --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/launch_sft_action_policy_robocasa_nano.sh @@ -0,0 +1,136 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Complete recipe: RoboCasa mobile-manipulation action-policy SFT on Cosmos3-Nano (HSDP 8x2). +# Run from this folder with the cosmos-framework venv active (see README): +# bash launch_sft_action_policy_robocasa_nano.sh +# +# Trains a MOBILE-BASE manipulation policy on RoboCasa with the native base-motion contract +# (`base_encoding="raw"`, 15-D) and the full-resolution agentview_left | eye_in_hand composite. +# See the TOML header for the full contract. +# +# The recipe is not tied to one task set. TASK_SET selects which RoboCasa split to train on: +# +# target_atomic 18 tasks target/atomic (default; the worked example) +# pretrain_atomic 65 tasks pretrain/atomic +# pretrain_all 300 tasks pretrain atomic + composite +# +# Anything beyond the default also needs ROBOCASA_ROOT pointed at the matching converted export +# and a longer schedule -- pass e.g. +# EXTRA_TAIL_OVERRIDES="trainer.max_iter=40000 scheduler.cycle_lengths=[40000]" +# +# It prepares the small dependencies, checks for the staged dataset, and trains. +# Paths are fixed under this (git-ignored) folder, matching the DROID/LIBERO wrappers, while +# the TOML and tail-overrides match the cosmos-framework example. + +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" + +TOML_FILE="${TOML_FILE:-toml/sft_config/action_policy_robocasa_nano.toml}" +# RoboCasa's released export is LeRobot v2.1; the installed lerobot is v3.0-only and rejects it. +# ROBOCASA_ROOT must point at a CONVERTED copy: ///lerobot/. +: "${ROBOCASA_ROOT:=$PWD/data/robocasa_v30}" +: "${BASE_CHECKPOINT_PATH:=$PWD/checkpoints/Cosmos3-Nano}" +: "${WAN_VAE_PATH:=$PWD/checkpoints/wan22_vae/Wan2.2_VAE.pth}" + +# Resolve TASK_SET to a task list. The tuples live in the framework's +# robocasa_lerobot_dataset module, which is the single source of truth; reading them here keeps +# the names from drifting between the recipe and the dataset. +TASK_SET="${TASK_SET:-target_atomic}" +case "$TASK_SET" in + target_atomic) TASK_CONST=DEFAULT_ALL_ATOMIC_TASKS ;; + pretrain_atomic) TASK_CONST=DEFAULT_PRETRAIN_ATOMIC_TASKS ;; + pretrain_all) TASK_CONST=DEFAULT_PRETRAIN_ALL_TASKS ;; + *) echo "TASK_SET must be one of: target_atomic pretrain_atomic pretrain_all" >&2; exit 1 ;; +esac +_mod=cosmos_framework.data.generator.action.datasets.robocasa_lerobot_dataset +_tasks="$(python -c "from ${_mod} import ${TASK_CONST} as T; print(' '.join(T))")" || { + echo "ERROR: could not read ${TASK_CONST} from ${_mod}; is the cosmos-framework venv active?" >&2 + exit 1 +} +read -ra TASKS <<< "$_tasks" +# The TOML pins the default set; anything else is selected with a hydra list override. +if [[ "$TASK_SET" != "target_atomic" ]]; then + _list="$(IFS=,; echo "${TASKS[*]}")" + EXTRA_TAIL_OVERRIDES="${EXTRA_TAIL_OVERRIDES:-} dataloader_train.dataloader.datasets.robocasa.dataset.task_names=[${_list}]" +fi +echo "task set : ${TASK_SET} (${#TASKS[@]} tasks)" + +# 1. Check the v3.0-converted dataset is staged. tasks.parquet is a v3.0-only marker, so its +# presence distinguishes a converted copy from the raw v2.1 export. +missing=() +for t in "${TASKS[@]}"; do + ls "$ROBOCASA_ROOT/$t"/*/lerobot/meta/tasks.parquet >/dev/null 2>&1 || missing+=("$t") +done +if (( ${#missing[@]} )); then + cat >&2 <///lerobot/meta/tasks.parquet +The installed lerobot is v3.0-only, so the released v2.1 export must be converted first: + SRC_ROOT=/path/to/robocasa/datasets/v1.0/target/atomic bash convert_robocasa_to_lerobot_v30.sh +or point at an existing converted root: + export ROBOCASA_ROOT=/path/to/converted/v30/root + +The loader skips absent tasks with a warning rather than failing, so this check runs up front +to avoid training on a partial task set by accident. +EOF + exit 1 +fi + +# 2. Download the Wan2.2 VAE (skipped if present). +if [[ ! -f "$WAN_VAE_PATH" ]]; then + uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth --local-dir "$(dirname "$WAN_VAE_PATH")" +fi + +# 3. Convert the base checkpoint to DCP (skipped if present). +if [[ ! -d "$BASE_CHECKPOINT_PATH" ]]; then + python -m cosmos_framework.scripts.convert_model_to_dcp -o "$BASE_CHECKPOINT_PATH" --checkpoint-path Cosmos3-Nano +fi + +# 4. Train (HSDP 8x2 = 16 ranks per the TOML; set NNODES/NODE_RANK/MASTER_ADDR per node). +# The TOML reads these paths from the environment. +export ROBOCASA_ROOT +export BASE_CHECKPOINT_PATH +export WAN_VAE_PATH + +TAIL_OVERRIDES=() +if [[ -n "${EXTRA_TAIL_OVERRIDES:-}" ]]; then + # EXTRA_TAIL_OVERRIDES is intentionally word-split to match the framework launcher UX. + # shellcheck disable=SC2206 + TAIL_OVERRIDES=(${EXTRA_TAIL_OVERRIDES}) +fi + +TORCHRUN_ARGS=(--nproc_per_node="${NPROC_PER_NODE:-8}") +TORCHRUN_ARGS+=(--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") + +OUTPUT_ROOT="${OUTPUT_ROOT:-$PWD/outputs/train}" + +# The training entrypoint does NOT chdir, but the model config loads release-root-relative +# resource files, so torchrun must run from the package root. Derive it from the installed +# package and make the TOML path absolute so it still resolves from that cwd. Every other +# path the run needs (ROBOCASA_ROOT / BASE_CHECKPOINT_PATH / WAN_VAE_PATH / OUTPUT_ROOT) is +# already absolute. +COSMOS_PKG_ROOT="$(python -c 'import cosmos_framework, pathlib; print(pathlib.Path(cosmos_framework.__file__).resolve().parent.parent)')" +TOML_ABS="$PWD/$TOML_FILE" +echo "torchrun cwd=$COSMOS_PKG_ROOT toml=$TOML_ABS" + +# torchrun workers inherit the launcher's cwd, so run it from the package root (subshell +# keeps the cd local). All paths passed in are absolute, so nothing else is affected. +if (( ${#TAIL_OVERRIDES[@]} )); then + ( cd "$COSMOS_PKG_ROOT" && IMAGINAIRE_OUTPUT_ROOT="${IMAGINAIRE_OUTPUT_ROOT:-$OUTPUT_ROOT}" \ + torchrun "${TORCHRUN_ARGS[@]}" \ + -m cosmos_framework.scripts.train --sft-toml="$TOML_ABS" \ + -- "${TAIL_OVERRIDES[@]}" ) +else + ( cd "$COSMOS_PKG_ROOT" && IMAGINAIRE_OUTPUT_ROOT="${IMAGINAIRE_OUTPUT_ROOT:-$OUTPUT_ROOT}" \ + torchrun "${TORCHRUN_ARGS[@]}" \ + -m cosmos_framework.scripts.train --sft-toml="$TOML_ABS" ) +fi diff --git a/cookbooks/cosmos3/generator/action/finetune/smoke_test_robocasa_eval.sh b/cookbooks/cosmos3/generator/action/finetune/smoke_test_robocasa_eval.sh new file mode 100755 index 00000000..73ee772e --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/smoke_test_robocasa_eval.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Smoke test for the RoboCasa closed-loop evaluation handshake. +# +# Runs the exact two-process pair the README documents -- policy server in the +# cosmos-framework venv, evaluator in the robosuite/robocasa venv -- on a single task for a +# couple of short rollouts, and checks that the width and state contracts the two sides agree +# on are the ones this recipe trains with. The evaluator sends a 15-D state token; a server +# that resolved a different width rejects every request, so the mismatch has to be caught +# here rather than an hour into a real evaluation. +# +# RUN_DIR=outputs/train/// \ +# FRAMEWORK_ROOT=/path/to/cosmos-framework \ +# SIM_PYTHON=/path/to/robocasa-venv/bin/python \ +# DATASET_DIR=/path/to/robocasa/datasets/v1.0/target/atomic///lerobot \ +# bash smoke_test_robocasa_eval.sh +# +# Run it from this folder with the cosmos-framework venv active. + +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" + +: "${RUN_DIR:?set RUN_DIR to a training run directory (it must contain checkpoints/ and config.yaml)}" +: "${DATASET_DIR:?set DATASET_DIR to an ORIGINAL RoboCasa export: ...///lerobot}" +: "${SIM_PYTHON:?set SIM_PYTHON to the python of the robosuite/robocasa venv}" +# cosmos-framework is not installed in the simulator venv; the evaluator imports only +# pure-python modules from it, so a checkout on PYTHONPATH is enough. +: "${FRAMEWORK_ROOT:?set FRAMEWORK_ROOT to a cosmos-framework checkout}" + +# The DCP checkpoint is served directly; evaluation needs no safetensors export. +if [[ -z "${CHECKPOINT_PATH:-}" ]]; then + LATEST="$RUN_DIR/checkpoints/$(cat "$RUN_DIR/checkpoints/latest_checkpoint.txt")" + CHECKPOINT_PATH="$LATEST/model" +fi +CONFIG_FILE="${CONFIG_FILE:-$RUN_DIR/config.yaml}" +PORT="${PORT:-8912}" +NUM_EPISODES="${NUM_EPISODES:-2}" +OUTPUT_DIR="${OUTPUT_DIR:-$PWD/outputs/smoke_eval}" +# The width this recipe trains with (use_base_action=True, base_encoding="raw"). +EXPECTED_ACTION_DIM="${EXPECTED_ACTION_DIM:-15}" + +for f in "$CHECKPOINT_PATH" "$CONFIG_FILE" "$DATASET_DIR" "$SIM_PYTHON" "$FRAMEWORK_ROOT"; do + [[ -e "$f" ]] || { echo "ERROR: does not exist: $f" >&2; exit 1; } +done + +mkdir -p "$OUTPUT_DIR" +SERVER_LOG="$OUTPUT_DIR/server.log" + +echo "=== 1/3 starting policy server on port $PORT (log: $SERVER_LOG)" +# No --raw-action-dim: the point of the test is that the server resolves the width itself. +python -m cosmos_framework.scripts.action_policy_server_robocasa \ + --checkpoint-path "$CHECKPOINT_PATH" --config-file "$CONFIG_FILE" \ + --port "$PORT" --num-steps 10 --guidance 3.0 --fps 20 \ + > "$SERVER_LOG" 2>&1 & +SERVER_PID=$! +trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT + +echo "=== 2/3 waiting for /info and checking the contract" +INFO="" +for _ in $(seq 1 120); do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: the policy server exited during startup; tail of $SERVER_LOG:" >&2 + tail -30 "$SERVER_LOG" >&2 + exit 1 + fi + if INFO="$(curl -sf -m 5 "http://127.0.0.1:$PORT/info")"; then + break + fi + INFO="" + sleep 5 +done +[[ -n "$INFO" ]] || { echo "ERROR: /info never came up; tail of $SERVER_LOG:" >&2; tail -30 "$SERVER_LOG" >&2; exit 1; } + +INFO="$INFO" EXPECTED_ACTION_DIM="$EXPECTED_ACTION_DIM" python - <<'PY' +import json, os, sys + +info = json.loads(os.environ["INFO"]) +expected = int(os.environ["EXPECTED_ACTION_DIM"]) +print("[smoke] /info:", json.dumps(info, indent=2)) +dim = info.get("raw_action_dim") +source = info.get("raw_action_dim_source") +if dim != expected: + sys.exit( + f"FAIL: server resolved raw_action_dim={dim} (from {source}), but this recipe is " + f"{expected}-D. The evaluator would send a {expected}-D state token and every /predict " + f"would be rejected." + ) +if not info.get("requires_state"): + sys.exit( + "FAIL: the server does not require a state token, so this checkpoint was not trained " + "with use_state=True. Serving it under --use-state silently drops a conditioning signal." + ) +print(f"[smoke] OK: raw_action_dim={dim} (from {source}), requires_state=True") +PY + +echo "=== 3/3 running $NUM_EPISODES rollouts through the evaluator" +# Same flags as the README; only the episode count is cut down. The rollout horizon comes +# from robocasa's task registry, so it is not shortened here. +MUJOCO_GL="${MUJOCO_GL:-egl}" PYTHONPATH="$FRAMEWORK_ROOT" \ +"$SIM_PYTHON" -m cosmos_framework.simulation.robocasa.closed_loop_eval \ + --server-url "http://127.0.0.1:$PORT" \ + --dataset-dir "$DATASET_DIR" \ + --output-dir "$OUTPUT_DIR" --num-test-episodes "$NUM_EPISODES" \ + --action-horizon 32 --camera-set left_wrist \ + --use-state --use-base-action --base-encoding raw \ + --seed 0 --image-size 256 --cam-size 256 + +# The rollouts are far too short to succeed; completing them without a contract error is the +# result being checked. +echo "=== PASS: the server/evaluator handshake works end to end; results in $OUTPUT_DIR" diff --git a/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_robocasa_nano.toml b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_robocasa_nano.toml new file mode 100644 index 00000000..7a11b18d --- /dev/null +++ b/cookbooks/cosmos3/generator/action/finetune/toml/sft_config/action_policy_robocasa_nano.toml @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + + +# RoboCasa mobile-base action-policy SFT on Cosmos3-Nano. +# +# All 18 `target/atomic` RoboCasa tasks with a mobile base. The action contract is 15-D and +# carries RoboCasa's native base command unchanged (`base_encoding="raw"`): +# +# [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] +# +# Observation is `camera_set="left_wrist"` — agentview_left | eye_in_hand concatenated +# horizontally at their native 256x256 each — plus EEF proprioception as a clean condition +# token (`use_state=True`). `action_normalization=None`: the state token is absolute, so delta +# statistics do not apply. 32-step action chunk at 20 fps. +# +# Global batch 2048 = 128 samples x 16 ranks (HSDP 8x2, i.e. 4 nodes x 4 GPUs). +# +# Closed-loop eval must use the same contract: +# ACTION_HORIZON=32 CAMERA_SET=left_wrist USE_STATE=1 +# USE_BASE_ACTION=1 BASE_ENCODING=raw RAW_ACTION_DIM=15 + +[job] +task = "vfm" +experiment = "action_policy_robocasa_nano" +project = "cosmos3_action_robocasa" +group = "action_sft" +name = "action_policy_robocasa" +wandb_mode = "online" + +[model] +precision = "bfloat16" +max_num_tokens_after_packing = 74000 + +[model.parallelism] +data_parallel_shard_degree = 8 +data_parallel_replicate_degree = 2 # HSDP 8x2 = 16 ranks (4 nodes); gbs = 128 x 16 = 2048 + +[model.activation_checkpointing] +mode = "selective" +save_ops_regex = ["fmha"] + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[optimizer] +lr = 5.0e-05 + +[scheduler] +cycle_lengths = [16000] # linear decay over the full run (== max_iter) +warm_up_steps = [500] + +[trainer] +max_iter = 16000 +logging_iter = 50 +grad_accum_iter = 1 + +[checkpoint] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 1000