#!/bin/bash
# Train the can+plate (label-conditioned, 2-object) 3-PoinTr on Stack-Can.
# Mirrors the push_box_boxmark recipe (subsample=150, use_point_labels, point_dim=2,
# no_img_cond). 98 both-object demos (ep042 dropped). Runs on GPU0 (pointr env).
set -e

DATA_ROOT="/data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate"
TASK="stack_can_canplate"
SUBSAMPLE=150
SPLIT_NAME="${TASK}_flow"
EPOCHS="${EPOCHS:-5000}"
BATCH_SIZE="${BATCH_SIZE:-4}"
GPU="${GPU:-0}"

cd /home/hongyi/PWM/reward_model/3PoinTr_repo
source /home/hongyi/miniconda3/etc/profile.d/conda.sh
conda activate pointr

TOTAL=$(ls -d "${DATA_ROOT}/no_actions/${TASK}"/[0-9][0-9][0-9][0-9] 2>/dev/null | wc -l)
[ "$TOTAL" -lt 5 ] && { echo "ERROR: only ${TOTAL} episodes built"; exit 1; }
N_VAL=$((TOTAL / 10)); [ "$N_VAL" -lt 1 ] && N_VAL=1
N_TRAIN=$((TOTAL - N_VAL))
echo "[launch] $TOTAL episodes -> train=$N_TRAIN, val=$N_VAL  (subsample=$SUBSAMPLE)"

TASK_SPLITS="{\"${TASK}\": [0, 0, 0, ${N_TRAIN}, ${N_VAL}, 0]}"
python scripts_training/splitter.py \
  --dataset "${DATA_ROOT}/" \
  --task-splits "${TASK_SPLITS}" \
  --out-dir "${DATA_ROOT}/" \
  --name "${SPLIT_NAME}"
SPLIT_FILE="${DATA_ROOT}/${SPLIT_NAME}.json"
echo "[launch] split -> ${SPLIT_FILE}"

mkdir -p /home/hongyi/PWM/reward_model/3pointr_models
CUDA_VISIBLE_DEVICES=$GPU python scripts_training/pretrain_flow.py \
  --device cuda:0 \
  --epochs "${EPOCHS}" \
  --batch_size "${BATCH_SIZE}" \
  --lr 1e-4 \
  --arch_dim 256 --arch_depth 2 --arch_heads 4 \
  --arch_query_depth 1 --arch_query_heads 1 \
  --subsample "${SUBSAMPLE}" \
  --point_dim 2 \
  --use_point_labels \
  --no_img_cond \
  --dataset "${DATA_ROOT}/" \
  --model_name "${TASK}_${SUBSAMPLE}pts_384" \
  --tasks "${TASK}" \
  --save_dir "/home/hongyi/PWM/reward_model/3pointr_models" \
  --split_file "${SPLIT_FILE}" \
  --flow_data_only \
  --eval_every 50 \
  --save_vis_every 250 \
  --weight_decay 1e-4 \
  --seed 42 \
  --num_workers 4 \
  --flow_loss_w 10.0