PWM / PWM_code /reward_model /scripts /launch_3pointr_stackcan_canplate.sh

Add files using upload-large-folder tool

dd253d4 verified 2 months ago

1.98 kB

#!/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

	#!/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