Upload scripts/train.sh

scripts/train.sh

@@ -0,0 +1,112 @@
+#!/bin/bash
+# GUI-Shift GRPO Training Script
+
+set -e
+
+# Default values
+MODEL_NAME="Qwen/Qwen2.5-VL-7B-Instruct"
+DATA_DIR="./data/gui_transition/filtered"
+OUTPUT_DIR="./checkpoints/gui-shift"
+K=1
+NUM_GPUS=8
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model_name_or_path)
+            MODEL_NAME="$2"
+            shift 2
+            ;;
+        --data_dir)
+            DATA_DIR="$2"
+            shift 2
+            ;;
+        --output_dir)
+            OUTPUT_DIR="$2"
+            shift 2
+            ;;
+        --k)
+            K="$2"
+            shift 2
+            ;;
+        --num_gpus)
+            NUM_GPUS="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+echo "=== GUI-Shift GRPO Training ==="
+echo "Model: $MODEL_NAME"
+echo "Data: $DATA_DIR"
+echo "Output: $OUTPUT_DIR"
+echo "K value: $K"
+echo "GPUs: $NUM_GPUS"
+echo ""
+
+# Find data file for the specified K
+DATA_FILE="${DATA_DIR}/k${K}_transition_filtered.jsonl"
+if [ ! -f "$DATA_FILE" ]; then
+    DATA_FILE="${DATA_DIR}/k${K}_transition.jsonl"
+    if [ ! -f "$DATA_FILE" ]; then
+        echo "Error: Could not find data file for k=$K in $DATA_DIR"
+        exit 1
+    fi
+fi
+
+# Image folder (relative to trajectory data)
+IMAGE_FOLDER="${DATA_DIR}/images"
+
+echo "Using data file: $DATA_FILE"
+echo "Using image folder: $IMAGE_FOLDER"
+echo ""
+
+# GRPO training hyperparameters (from paper Appendix A)
+# learning_rate: 1e-6, temperature: 0.9, num_generations: 8
+# num_train_epochs: 4, max_prompt_length: 1024, max_completion_length: 256
+# per_device_train_batch_size: 2, gradient_accumulation_steps: 8
+# epsilon: 0.2, beta: 0.04
+
+export EXP_NAME="gui-shift-k${K}"
+mkdir -p "runs/${EXP_NAME}/log"
+
+# Note: This requires the VLM-R1 framework to be installed
+# See: https://github.com/om-ai-lab/VLM-R1
+
+torchrun --nproc_per_node="$NUM_GPUS" \
+    --nnodes="1" \
+    --node_rank="0" \
+    --master_addr="127.0.0.1" \
+    --master_port="12349" \
+    src/training/gui_grpo_trainer.py \
+    --model_name_or_path "$MODEL_NAME" \
+    --data_file_paths "$DATA_FILE" \
+    --image_folders "$IMAGE_FOLDER" \
+    --output_dir "$OUTPUT_DIR" \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 8 \
+    --gradient_checkpointing true \
+    --logging_steps 1 \
+    --num_train_epochs 4 \
+    --max_steps -1 \
+    --bf16 \
+    --attn_implementation flash_attention_2 \
+    --run_name "$EXP_NAME" \
+    --save_steps 400 \
+    --num_generations 8 \
+    --max_completion_length 256 \
+    --reward_funcs format accuracy \
+    --beta 0.04 \
+    --learning_rate 1.0e-6 \
+    --lr_scheduler_type cosine \
+    --warmup_ratio 0.1 \
+    --report_to wandb \
+    --dataset-name not_used \
+    --freeze_vision_modules true \
+
+echo ""
+echo "Training complete! Model saved to: $OUTPUT_DIR"