DEBUG=true
if [ "$DEBUG" = true ]; then
  GPUS=1
  GPUS_PER_NODE=${GPUS}
  wandb_enable=false
  ACCELERATE_ARGS="--num_machines 1 --num_processes ${GPUS} --mixed_precision=bf16 --dynamo_backend=no"
  num_workers=0
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:-8}
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"}

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

lr=1e-4
steps=$((GPUS * 100000))
chunk_size=10
normalize_fast_actions=true
max_action_dim=7

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

accelerate launch $ACCELERATE_ARGS pi0fast/train_pi0fast.py \
  --policy.path=../pretrained/pi0fast_base \
  --fast_tokenizer_path=../pretrained/fast_libero \
  --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} \
  --state_mode Q01_Q99