File size: 1,979 Bytes
dd253d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | #!/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
|