DEBUG=true
if [ "$DEBUG" = true ]; then
  GPUS=2
  GPUS_PER_NODE=${GPUS}
  wandb_enable=false
  ACCELERATE_ARGS="--num_machines 1 --num_processes ${GPUS} --mixed_precision=bf16 --dynamo_backend=no"
  num_workers=8
fi

# distributed settings
GPUS=${GPUS:-8}
GPUS_PER_NODE=${GPUS_PER_NODE:-8}
NODES=$((GPUS / GPUS_PER_NODE))
PER_DEVICE_BATCH_SIZE=${PER_DEVICE_BATCH_SIZE:-16}
wandb_enable=${wandb_enable:-true}
num_workers=${num_workers:-8}

# set environments
source scripts/env.sh
conda activate vla_next

ACCELERATE_ARGS=${ACCELERATE_ARGS:-"--main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT \
  --num_machines ${NODES} --num_processes=${GPUS} --multi_gpu \
  --mixed_precision=no --dynamo_backend=no"}
  # --deepspeed_config_file configs/zero1.json

# dataset=configs/data-demo.yaml
dataset=libero_spatial_no_noops_1.0.0_lerobot
# dataset=bridge_orig_1.0.0_lerobot

cfg=$(echo $dataset | sed 's/^\([a-zA-Z]\+\).*/\1/')

lr=1e-4 # bridge 5e-5 / libero 1e-4
steps=$((GPUS * 30000)) # bridge 100k / libero 30k (>1 epoch)
chunk_size=50

# * fine-tuning
job_name=pi0_${dataset}_ck${chunk_size}_gpu${GPUS}_lr${lr}_bs${PER_DEVICE_BATCH_SIZE}_s$((steps / 1000))k

accelerate launch $ACCELERATE_ARGS pi0/train_pi0.py \
  --attention_implementation fa \
  --policy.path=../pretrained/pi0 \
  --dataset.repo_id=${dataset} \
  --dataset.image_transforms.enable=true \
  --num_workers=${num_workers} \
  --policy_optimizer_lr=${lr} \
  --chunk_size=${chunk_size} \
  --steps=${steps} \
  --batch=${PER_DEVICE_BATCH_SIZE} \
  --save_freq=20000 \
  --log_freq=1000 \
  --job_name=${job_name} \
  --wandb.enable=${wandb_enable} \
  --wandb.disable_artifact=true \
  --wandb.project=${WANDB_PROJECT} \
  --wandb.entity=${WANDB_ENTITY}