#!/bin/bash
LLM=qwen2_vl   #qwen2_vl  paligemma
LLM_MODEL_SIZE=2B #3B
# LLM_MODEL_SIZE=2_8B
# lora only vit and tune adapter
ACTION_HEAD=dit_diffusion_policy  #act #unet_diffusion_policy dit_diffusion_policy

echo '7.5h'
#sleep 7.5h
ROOT=/home/jovyan/tzb # /home/jovyan/tzb || /gpfs/private/tzb
DIT_ROOT=/home/share # /home/share || /gpfs/share/share

#PRETRAIN=${ROOT}/wjj/model_param/multi_head2/${ACTION_HEAD}_results/checkpoint_all/${LLM}_${LLM_MODEL_SIZE}_pure/vanilla_aloha_${LLM}_vla_pt_f_vit/qwen2_vl_all_data_1200_align_frozen_dit_lora_chunk_50/checkpoint-40000 # non substeps DIT
#PRETRAIN=${ROOT}/wjj/model_param/multi_head2/${ACTION_HEAD}_results/checkpoint_all/${LLM}_${LLM_MODEL_SIZE}/vanilla_aloha_${LLM}_vla_pt_f_vit/qwen2_vl_all_data_1200_combine_constant_pretrain_DIT_H_full_param/checkpoint-60000 # with substeps DIT
#PRETRAIN=${ROOT}/wjj/model_param/multi_head2/${ACTION_HEAD}_results/checkpoint_all/${LLM}_${LLM_MODEL_SIZE}/vanilla_aloha_${LLM}_vla_pt_f_vit/qwen2_vl_4_cameras_1_12_all_data_pretrain_DiT_XH_full_param_stage_1_50/checkpoint-60000 # with substeps DIT
#PRETRAIN=${ROOT}/wjj/model_param/multi_head2/${ACTION_HEAD}_results/checkpoint_all/${LLM}_${LLM_MODEL_SIZE}/vanilla_aloha_${LLM}_vla_pt_f_vit/qwen2_3_cameras_1_17_all_data_pretrain_DiT_H_full_param_stage_1_50/checkpoint-60000 # with substeps DIT
PRETRAIN=${ROOT}/wjj/model_param/multi_head2/${ACTION_HEAD}_results/checkpoint_all/${LLM}_${LLM_MODEL_SIZE}/vanilla_aloha_${LLM}_vla_pt_f_vit/qwen2_vl_3_cameras_1_17_all_data_pretrain_6w_DiT_H_Non_EMA_full_param_stage_1_50/checkpoint-60000 # with substeps DIT

#DIT_PRETRAIN=${DIT_ROOT}/ljm/model_param/scaledp/resnet50_with_film_nosubreason/fold_t_shirt_easy_version_all_add_clean_table_1_0_4_DiT-H_320_240_32_1e-4_numsteps_40000_sub_0_2025_01_04_17_38_19/policy_step_40000_2025-01-05_13-30-34.ckpt # non substeps DIT
DIT_PRETRAIN=${DIT_ROOT}/ljm/model_param/scaledp/resnet50_with_film_subreason/fold_t_shirt_easy_version_all_add_clean_table_1_0_4_DiT-H_320_240_32_1e-4_numsteps_40000_sub_1_2025_01_04_17_26_23/policy_step_40000_2025-01-05_12-40-45.ckpt # with substeps DIT


if [ "${LLM}" == "paligemma" ]; then
  echo "Using PaliGemma"
  mnop=${ROOT}/wjj/model_param/PaliGemma/paligemma/pixel_224/vla-paligemma-3b-pt-224
else
  mnop=${ROOT}/wjj/model_param/Qwen2-VL-${LLM_MODEL_SIZE}-Instruct
fi

mnop=$PRETRAIN # pretrain ckpt as base
TASK_NAME="folding_two_shirts_by_drag"

OUTPUT=${ROOT}/wjj/train_results/dexvla_lerobot_results/${LLM}_${LLM_MODEL_SIZE}/${task_name}_Stage3
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 ./qwen2_vla $OUTPUT/src/
cp -r ./policy_heads $OUTPUT/src/

# tinyvla set "use_reasoning with_llm_head load_pretrain using_film" false
# paligemma flash_attn False

deepspeed --master_port 29604 --num_gpus=8 --num_nodes=1 ./train_vla.py \
  --deepspeed scripts/zero2.json \
  --use_reasoning True \
  --lora_enable False \
  --action_dim 14 \
  --state_dim 14 \
  --flash_attn True \
  --chunk_size 50 \
  --lora_module "vit llm" \
  --load_pretrain False \
  --history_images_length 1 \
  --model_pretrain $PRETRAIN \
  --load_pretrain_dit False \
  --pretrain_dit_path $DIT_PRETRAIN \
  --ground_truth_reasoning False \
  --using_all_reasoning_hidden False \
  --using_film True \
  --using_ema False \
  --policy_head_type $ACTION_HEAD \
  --policy_head_size "DiT_H" \
  --with_llm_head True \
  --image_size_stable "(320,240)" \
  --image_size_wrist "(320,240)" \
  --lora_r 64 \
  --lora_alpha 256 \
  --episode_first False \
  --task_name $TASK_NAME \
  --model_name_or_path $mnop \
  --version v0 \
  --tune_mm_mlp_adapter True \
  --freeze_vision_tower False \
  --freeze_backbone False \
  --mm_use_im_start_end False \
  --mm_use_im_patch_token False \
  --image_aspect_ratio pad \
  --group_by_modality_length False \
  --bf16 True \
  --output_dir $OUTPUT \
  --max_steps 20000 \
  --per_device_train_batch_size 12 \
  --gradient_accumulation_steps 1 \
  --save_strategy "steps" \
  --save_steps 10000 \
  --save_total_limit 50 \
  --learning_rate 2e-5 \
  --weight_decay 0. \
  --warmup_ratio 0.01 \
  --lr_scheduler_type "cosine" \
  --logging_steps 50 \
  --tf32 True \
  --model_max_length 2048 \
  --gradient_checkpointing True \
  --dataloader_num_workers 8 \
  --lazy_preprocess True \
  --policy_class $ACTION_HEAD \
  --concat "token_cat" \
  --report_to tensorboard \
  --logging_dir $OUTPUT/log | tee $OUTPUT/log.log

for dir in "$OUTPUT"/*/ ; do
    # 检查文件夹名称是否包含'checkpoint'
    if [[ "$(basename "$dir")" == *"checkpoint"* ]]; then
        cp ${mnop}/preprocessor_config.json $dir
        cp ${mnop}/chat_template.json $dir
        # cp $OUTPUT/non_lora_trainables.bin $dir
    fi
done

mv ./60030.log $OUTPUT
echo $OUTPUT
