Scheduled Commit

scripts/train/dexmg_lifttray_finetune_wan22.sh

@@ -0,0 +1,111 @@
+#!/bin/bash
+# DreamZero fine-tuning on DexMimicGen bimanual_panda_hand.LiftTray (Wan2.2-TI2V-5B backbone),
+# warm-started from the T-Rex LoRA checkpoint (checkpoints/dreamzero_trex_wan22_lora).
+#
+# DexMG LiftTray: two Panda arms + Inspire dexterous hands, 1000 generated demos.
+#   state 40-dim / action 24-dim (per side: delta-EEF pos+rot 6 + hand 6) -> max_action_dim=64
+#   3 views: ego_view, left_wrist_view, right_wrist_view (256x256 @ 20fps, resized to 320x160)
+#
+# Usage:
+#   bash scripts/train/dexmg_lifttray_finetune_wan22.sh
+#
+# Prerequisites:
+#   1) LiftTray LeRobot dataset downloaded, e.g. to
+#      /scratch1/home/zhicao/dexmimicgen/datasets/lerobot/bimanual_panda_hand.LiftTray
+#   2) GEAR metadata generated on top of it (run once):
+#      python scripts/data/convert_lerobot_to_gear.py \
+#          --dataset-path /scratch1/home/zhicao/dexmimicgen/datasets/lerobot/bimanual_panda_hand.LiftTray \
+#          --embodiment-tag robocasa_bimanual_panda_inspire_hand \
+#          --action-horizon 24
+#   3) Wan2.2-TI2V-5B / Wan2.1 CLIP / umt5-xxl weights (already under ./checkpoints)
+
+export HYDRA_FULL_ERROR=1
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then
+    :
+elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then
+    DREAMZERO_ROOT="$SCRIPT_REPO_ROOT"
+else
+    echo "ERROR: Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/."
+    exit 1
+fi
+
+# ============ USER CONFIGURATION ============
+DEXMG_DATA_ROOT=${DEXMG_DATA_ROOT:-"/scratch1/home/zhicao/dexmimicgen/datasets/lerobot/bimanual_panda_hand.LiftTray"}
+OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_dexmg_lifttray_wan22_lora"}
+
+# Warm start: T-Rex LoRA checkpoint (7+22 joints per side -> 6+6 here; both pad to 64)
+PRETRAINED_CKPT=${PRETRAINED_CKPT:-"$DREAMZERO_ROOT/checkpoints/dreamzero_trex_wan22_lora/checkpoint-16000"}
+
+WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"}
+IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"}
+TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"}
+# =============================================
+
+# Validate dataset exists and is GEAR-converted
+if [ ! -d "$DEXMG_DATA_ROOT" ]; then
+    echo "ERROR: DexMG LiftTray dataset not found at $DEXMG_DATA_ROOT"
+    exit 1
+fi
+if [ ! -f "$DEXMG_DATA_ROOT/meta/embodiment.json" ]; then
+    echo "ERROR: $DEXMG_DATA_ROOT/meta/embodiment.json missing - run convert_lerobot_to_gear.py first (see header)"
+    exit 1
+fi
+if [ ! -d "$PRETRAINED_CKPT" ]; then
+    echo "ERROR: warm-start checkpoint not found at $PRETRAINED_CKPT"
+    exit 1
+fi
+
+EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py"
+cd "$DREAMZERO_ROOT"
+
+torchrun --nproc_per_node=gpu --standalone "$EXPERIMENT_PY" \
+    report_to=wandb \
+    data=dreamzero/dexmg_lifttray_wan22 \
+    wandb_project=dreamzero \
+    train_architecture=lora \
+    num_frames=33 \
+    action_horizon=24 \
+    max_action_dim=64 \
+    ++action_head_cfg.config.diffusion_model_cfg.action_dim=64 \
+    num_views=3 \
+    model=dreamzero/vla \
+    model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
+    model/dreamzero/transform=dreamzero_cotrain \
+    num_frame_per_block=2 \
+    num_action_per_block=24 \
+    num_state_per_block=1 \
+    seed=42 \
+    training_args.learning_rate=1e-5 \
+    training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \
+    save_steps=4000 \
+    training_args.warmup_ratio=0.05 \
+    output_dir="$OUTPUT_DIR" \
+    per_device_train_batch_size=1 \
+    max_steps=20000 \
+    weight_decay=1e-5 \
+    save_total_limit=5 \
+    upload_checkpoints=false \
+    bf16=true \
+    tf32=true \
+    eval_bf16=true \
+    dataloader_pin_memory=false \
+    dataloader_num_workers=1 \
+    save_lora_only=true \
+    max_chunk_size=4 \
+    save_strategy=steps \
+    enable_wandb_video_reconstruction=true \
+    wandb_video_reconstruction_steps=1000 \
+    wandb_video_reconstruction_episode=0 \
+    wandb_video_reconstruction_num_chunks=4 \
+    dexmg_data_root="$DEXMG_DATA_ROOT" \
+    pretrained_model_path="$PRETRAINED_CKPT" \
+    ++action_head_cfg.config.defer_lora_injection=true \
+    dit_version="$WAN22_CKPT_DIR" \
+    text_encoder_pretrained_path="$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \
+    image_encoder_pretrained_path="$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \
+    vae_pretrained_path="$WAN22_CKPT_DIR/Wan2.2_VAE.pth" \
+    tokenizer_path="$TOKENIZER_DIR" \
+    "$@"