#!/usr/bin/env bash
# ARoPE 臂正式训练：8 卡 ZeRO-2，40k 步，每 1000 步存权重（含 --save_state 的优化器状态，state_A/B 两槽轮流、
# <output>/state 软链指向最新完整槽，约 140 GB），每 500 步验证。
# 用法：bash scripts/train_arope.sh [额外参数...]   崩溃续训：bash scripts/train_arope.sh --resume outputs/arope_40k/state
set -euo pipefail
ROOT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE
export DIFFSYNTH_SKIP_DOWNLOAD=True
export PYTHONPATH="$ROOT"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
GA="${GA:-2}"   # 梯度累积（全局 batch = 8 卡 × GA）；yaml 随之选 accelerate_zero2_ga${GA}.yaml（GA=1 用 accelerate_zero2.yaml）
CFG=$([ "$GA" = 1 ] && echo accelerate_zero2.yaml || echo "accelerate_zero2_ga${GA}.yaml")
RUN="${RUN:-arope_40k}"
mkdir -p "$ROOT/outputs/$RUN"
"$ROOT/.venv/bin/accelerate" launch --config_file "$ROOT/configs/$CFG" \
  "$ROOT/actionrope/train.py" \
  --arm arope --mask_channel --new_weight 2.0 --scene_dropout 0.1 \
  --lr 1e-5 --weight_decay 0.01 --warmup_steps 500 \
  --max_steps 40000 --save_every 1000 --save_state --val_every 500 --val_n 16 --val_timesteps 200,500,800 \
  --min_rt_ratio 0.8 --val_min_rt_ratio 0.9 \
  --grad_accum "$GA" --num_workers 2 --seed 0 \
  --output "$ROOT/outputs/$RUN" "$@" 2>&1 | tee -a "$ROOT/outputs/$RUN/train.log"