#!/usr/bin/env bash
# Launch ONE v4 (DINOv3) training on AR-view (smith300 UMI) data.
# Mirrors yams's launch_train_v4.sh with smithbot-specific tweaks:
#   - Data root points at an AR-view dataset (state/*.npz + episodes.json)
#   - grip_lo/grip_hi cover the smith300 UMI gripper_rad range (~-2.0 to +0.5)
#   - frame_stride=1 (data is already 4 fps native)
#   - ckpts land in ~/smithbot/checkpoints/ (parallel to yams's ~/yam_para/checkpoints/)
#
# Usage (on puget):
#     CUDA_VISIBLE_DEVICES=0 bash launch_train_smithbot_v4.sh <name> <data_dir> <iters> [extra_args]
set -euo pipefail
if [[ $# -lt 3 ]]; then
  echo "usage: $0 <name> <data_dir> <iters> [extra args...]" >&2
  exit 2
fi
NAME="$1"; DATA_DIR="$2"; ITERS="$3"; shift 3

REPO_DIR="$HOME/smithbot_v4"
PYTHON="$HOME/yam_para/.venv/bin/python"    # reuse yams's venv on puget
# DINOv3 hub load — v4 backbone; hardcoded to the ViT-S/16+ weights + puget-local repo.
export DINO_REPO_DIR="$HOME/yam_para/dinov3"
export DINO_WEIGHTS_PATH="$HOME/yam_para/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth"
RUN_DIR="$HOME/smithbot/runs/${NAME}"
CKPT_DIR="$HOME/smithbot/checkpoints/${NAME}"
LOG="$HOME/smithbot/logs/exp_${NAME}.log"
mkdir -p "$RUN_DIR" "$CKPT_DIR" "$(dirname "$LOG")"
export WANDB_DIR="$RUN_DIR"

echo "── smithbot_v4 train ${NAME} on $(hostname -s) (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}) ──"
echo "  data:    $DATA_DIR"
echo "  iters:   $ITERS"
echo "  run dir: $RUN_DIR  (cwd; wandb writes here)"
echo "  ckpt:    $CKPT_DIR"
echo "  log:     $LOG"
echo

cd "$RUN_DIR"
exec "$PYTHON" "$REPO_DIR/train.py" \
  --name "$NAME" \
  --data_root "$DATA_DIR" \
  --max_iters "$ITERS" \
  --v4 \
  --multiview scene,wrist \
  --batch_size 6 \
  --lr 3e-4 --lr_schedule cosine --lr_warmup 500 \
  --backbone_lr_mult 0.1 \
  --n_window 32 --n_height_bins 64 --n_gripper_bins 32 --n_rot_clusters 256 \
  --pred_size 64 \
  --feat_dim 32 --d_model 256 --mlp_layers 1 --mlp_hidden_mult 2.0 \
  --height_enc_dim 32 --voxel_mlp_hidden 64 \
  --cls_fusion concat --per_voxel_grip_rot \
  --temporal_head flat_mlp \
  --mv_past_n 8 --mv_past_enc_dim 64 \
  --cam_margin_m 0.01 \
  --img_size 504 --frame_stride 1 \
  --grip_lo -2.0 --grip_hi 0.5 \
  --ckpt_every_iters 200 --vis_every_iters 200 \
  "$@"