#!/bin/bash
#SBATCH --account=nvr_lpr_rvp
#SBATCH --partition=polar4,polar3,polar,grizzly
#SBATCH --nodes=1
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=56
#SBATCH --time=3:50:00
#SBATCH --job-name=xvla_spatial_object
#SBATCH --output=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/logs/xvla_spatial_object_%j.out
#SBATCH --error=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/logs/xvla_spatial_object_%j.err
#SBATCH --comment=fact_off

trap 'scontrol requeue ${SLURM_JOB_ID}; exit 15' SIGTERM

ACCELERATE=/lustre/fsw/portfolios/nvr/users/jtremblay/conda_envs/XVLA/bin/accelerate
XVLA_DIR=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/X-VLA
DATA_ROOT=/lustre/fsw/portfolios/nvr/users/jtremblay/yu/conflict_maniskill/demo_conflict
TRAIN_META="${DATA_ROOT}/spatial_object/300/huggingface_data/spatial_object/conflict/meta/info.json"
OUTPUT_DIR="${XVLA_DIR}/output/spatial_object"

export HF_HOME=/lustre/fsw/portfolios/nvr/users/jtremblay/hugging_face
export TRANSFORMERS_CACHE=/lustre/fsw/portfolios/nvr/users/jtremblay/hugging_face/transformers_cache

cd ${XVLA_DIR}

# Find latest checkpoint to resume from
LATEST_CKPT=$(/lustre/fsw/portfolios/nvr/users/jtremblay/conda_envs/XVLA/bin/python3 -c "
import os, re, sys
output_dir = sys.argv[1]
if not os.path.isdir(output_dir):
    print('')
    sys.exit(0)
ckpts = []
for d in os.listdir(output_dir):
    m = re.match(r'ckpt-(\\d+)', d)
    if m and os.path.isdir(os.path.join(output_dir, d)):
        ckpts.append(int(m.group(1)))
if ckpts:
    print(f'ckpt-{max(ckpts)}')
else:
    print('')
" "${OUTPUT_DIR}" 2>/dev/null)

if [ -n "$LATEST_CKPT" ]; then
    LATEST_STEP=$(echo "$LATEST_CKPT" | grep -oP '\d+')
    echo "Resuming from $LATEST_CKPT (step $LATEST_STEP)"
    START_MODEL="${OUTPUT_DIR}/${LATEST_CKPT}"
    REMAINING_ITERS=$((50000 - LATEST_STEP))
    FREEZE_STEPS=0
    WARMUP_STEPS=0
else
    echo "Starting from pretrained X-VLA-Pt"
    START_MODEL="${XVLA_DIR}/deploy/X-VLA-Pt"
    REMAINING_ITERS=50000
    FREEZE_STEPS=1000
    WARMUP_STEPS=2000
fi

if [ "$REMAINING_ITERS" -le 0 ]; then
    echo "Training already complete at 50000 steps. Exiting."
    exit 0
fi

$ACCELERATE launch \
    --mixed_precision bf16 \
    --num_processes 4 \
    --num_machines 1 \
    train.py \
    --models "${START_MODEL}" \
    --train_metas_path "${TRAIN_META}" \
    --output_dir "${OUTPUT_DIR}" \
    --learning_rate 1e-4 \
    --learning_coef 0.1 \
    --batch_size 16 \
    --iters ${REMAINING_ITERS} \
    --freeze_steps ${FREEZE_STEPS} \
    --warmup_steps ${WARMUP_STEPS} \
    --save_interval 5000