| #!/bin/bash |
| |
| |
| set -e |
|
|
| DATA_ROOT="${DATA_ROOT:?set DATA_ROOT}" |
| TASK="${TASK:?set TASK}" |
| MODEL_NAME="${MODEL_NAME:?set MODEL_NAME}" |
| SUBSAMPLE="${SUBSAMPLE:-150}" |
| USE_LABELS="${USE_LABELS:-0}" |
| 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] ${TASK}: ${TOTAL} eps -> train=${N_TRAIN} val=${N_VAL} subsample=${SUBSAMPLE} labels=${USE_LABELS} -> ${MODEL_NAME}" |
|
|
| 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" |
|
|
| LABEL_FLAG=""; [ "$USE_LABELS" = "1" ] && LABEL_FLAG="--use_point_labels" |
| 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 ${LABEL_FLAG} --no_img_cond \ |
| --dataset "${DATA_ROOT}/" --model_name "${MODEL_NAME}" --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 |
|
|
