#!/bin/bash
# SO-101 PickOrange StarVLA training launcher — westc single RTX 4080 SUPER 32G.
# No `conda activate` (hangs in non-interactive SSH) -> full env-binary paths.
# Self-logs to /root/starvla_train.log. Set RESUME=1 to resume from latest ckpt.
exec > /root/starvla_train.log 2>&1
set -o pipefail

ENV=/root/autodl-tmp/envs/starvla
REPO=/root/autodl-tmp/starVLA
cd $REPO

export CUDA_VISIBLE_DEVICES=0
export TORCH_CUDA_ARCH_LIST=8.9
export TOKENIZERS_PARALLELISM=false
export WANDB_MODE=disabled
export PYTHONUNBUFFERED=1
export HF_HOME=/root/autodl-tmp/hf_cache
export PATH=$ENV/bin:/usr/local/cuda-12.4/bin:$PATH
# training is offline (local base_vlm + local dataset) -> no proxy needed

RESUME_FLAG=""
if [ "${RESUME:-0}" = "1" ]; then
  RESUME_FLAG="--trainer.is_resume True"
  echo "### RESUME mode: will pick up latest checkpoint"
fi

echo "=== StarVLA SO-101 training start ==="; date
$ENV/bin/python -c "import torch;print('torch',torch.__version__,'cuda',torch.cuda.is_available())"

$ENV/bin/accelerate launch \
  --config_file starVLA/config/deepseeds/deepspeed_zero2.yaml \
  --num_processes 1 \
  --main_process_port 29521 \
  starVLA/training/train_starvla.py \
  --config_yaml examples/SO101_PickOrange/train_files/config_so101_qwengr00t.yaml \
  --framework.name QwenGR00T \
  --framework.qwenvl.base_vlm /root/autodl-tmp/models/Qwen3-VL-4B-Instruct \
  --datasets.vla_data.data_root_dir /root/autodl-tmp/datasets \
  --datasets.vla_data.data_mix so101_pickorange \
  --datasets.vla_data.per_device_batch_size 8 \
  --trainer.freeze_modules qwen_vl_interface \
  --trainer.max_train_steps 30000 \
  --trainer.save_interval 3000 \
  --trainer.keep_last_checkpoints 2 \
  --trainer.logging_frequency 10 \
  --run_root_dir /root/autodl-tmp/starvla-outputs \
  --run_id so101_pickorange_qwengr00t \
  $RESUME_FLAG

echo "=== training exited ==="; date