#!/bin/bash
LLM=qwen2_vl
ACTION_HEAD=unet_diffusion_policy
TASK=aloha_robotwin_place

ROOT=/data/private/liuza/robotiwin/policy/TinyVLA/TinyVLA-v2
mnop=/data/private/liuza/robotiwin/policy/TinyVLA/TinyVLA-v2/model_param/InternVL3-1B/
BS=128
LR=2e-5
noise_samples=8
OUTPUT=${ROOT}/${ACTION_HEAD}_results/${TASK}-${BS}BS-${LR}LR-${noise_samples}noise_samples
if [ -d "$OUTPUT" ]; then
   echo 'output exists'
else
   echo '!!output not exists!!'
   mkdir -p $OUTPUT
fi

mkdir -p $OUTPUT/src
cp -r ./aloha_scripts $OUTPUT/src/
cp -r ./scripts $OUTPUT/
cp -r ./data_utils $OUTPUT/src/
cp -r ./vla $OUTPUT/src/
cp -r ./policy_heads $OUTPUT/src/

deepspeed --master_port 29604 --num_gpus=8 --num_nodes=1 ./train_vla.py \
  --deepspeed scripts/zero2.json \
  --action_dim 14 \
  --state_dim 14 \
  --flash_attn True \
  --chunk_size 16 \
  --noise_samples ${noise_samples} \
  --policy_head_type $ACTION_HEAD \
  --episode_first False \
  --task_name $TASK \
  --model_name_or_path $mnop \
  --freeze_vision_tower False \
  --freeze_backbone False \
  --bf16 True \
  --output_dir $OUTPUT \
  --max_steps 60000 \
  --per_device_train_batch_size ${BS} \
  --gradient_accumulation_steps 1 \
  --save_strategy "steps" \
  --save_steps 10000 \
  --save_total_limit 50 \
  --learning_rate ${LR} \
  --weight_decay 0. \
  --warmup_ratio 0. \
  --lr_scheduler_type "cosine" \
  --logging_steps 5 \
  --tf32 True \
  --model_max_length 2048 \
  --gradient_checkpointing True \
  --dataloader_num_workers 8 \
  --report_to tensorboard \
  --logging_dir $OUTPUT/log | tee $OUTPUT/log.log

echo $OUTPUT
