Auto-sync: 2026-06-28 07:19:58

scripts/slurm/export_retrieval_residual_policy_targets.sbatch

@@ -0,0 +1,66 @@
+#!/bin/bash
+#SBATCH --job-name=export_resid_targets
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=2
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=12G
+#SBATCH --time=00:30:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
+DATASET="${DATASET:-$SCRATCH_ROOT/experiments/six_task_h16_collection}"
+CHECKPOINT="${CHECKPOINT:?Set CHECKPOINT to a trained DoVLA checkpoint}"
+OUT="${OUT:?Set OUT to the target-map JSON path}"
+SPLIT="${SPLIT:-all}"
+RETRIEVAL_NEIGHBORS="${RETRIEVAL_NEIGHBORS:-1}"
+RETRIEVAL_METRIC="${RETRIEVAL_METRIC:-raw}"
+RETRIEVAL_RESIDUAL_SCALE="${RETRIEVAL_RESIDUAL_SCALE:-0.35}"
+EXCLUDE_TYPES="${EXCLUDE_TYPES:-residual_random_negative,residual_wrong_direction,residual_near_miss}"
+if [[ -n "${EXCLUDE_TYPES_COLON:-}" ]]; then
+  EXCLUDE_TYPES="${EXCLUDE_TYPES_COLON//:/,}"
+fi
+MAX_GROUPS="${MAX_GROUPS:-}"
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$(dirname "$OUT")"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+EXTRA_ARGS=()
+if [[ -n "$MAX_GROUPS" ]]; then
+  EXTRA_ARGS+=(--max-groups "$MAX_GROUPS")
+fi
+if [[ "${NO_LEAVE_ONE_OUT:-0}" == "1" ]]; then
+  EXTRA_ARGS+=(--no-leave-one-out)
+fi
+if [[ "${NO_CLIP_ACTIONS:-0}" == "1" ]]; then
+  EXTRA_ARGS+=(--no-clip-actions)
+fi
+
+apptainer exec --nv \
+  --env "OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1,DOVLA_TORCH_THREADS=1,PYTHONDONTWRITEBYTECODE=1" \
+  -B "$PROJECT_DIR:$PROJECT_DIR" \
+  -B "/scratch/$USER:/scratch/$USER" \
+  "$SIF" "$PYTHON" scripts/export_retrieval_residual_policy_targets.py \
+    --checkpoint "$CHECKPOINT" \
+    --dataset "$DATASET" \
+    --out "$OUT" \
+    --device cuda \
+    --split "$SPLIT" \
+    --retrieval-neighbors "$RETRIEVAL_NEIGHBORS" \
+    --retrieval-metric "$RETRIEVAL_METRIC" \
+    --retrieval-residual-scale "$RETRIEVAL_RESIDUAL_SCALE" \
+    --exclude-types "$EXCLUDE_TYPES" \
+    "${EXTRA_ARGS[@]}"