Upload training and eval scripts

openpi/scripts/run_twin_dual_push_128_packed_5k.sh

@@ -0,0 +1,110 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+VENV="$ROOT/.venv/bin/activate"
+RUN_LOG_DIR="/workspace/run_logs"
+mkdir -p "$RUN_LOG_DIR"
+
+export HF_HOME=/workspace/.hf
+export HF_HUB_CACHE=/workspace/.hf/hub
+export HF_DATASETS_CACHE=/workspace/.hf/datasets
+export HUGGINGFACE_HUB_CACHE=/workspace/.hf/hub
+export XDG_CACHE_HOME=/workspace/.cache
+export OPENPI_LEROBOT_HOME=/workspace/lerobot
+export OPENPI_TORCH_COMPILE_SAMPLE_ACTIONS=0
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+cd "$ROOT"
+source "$VENV"
+
+BASELINE_CONFIG="pi05_twin_dual_push_128_packed_baseline_pytorch_5k"
+PARALLEL_CONFIG="pi05_twin_dual_push_128_packed_parallel_pytorch_5k"
+BASELINE_EXP="dual_push_128_packed_baseline_5k"
+PARALLEL_EXP="dual_push_128_packed_parallel_5k"
+VAL_REPO="lsnu/twin_dual_push_128_val"
+INTERMEDIATE_VAL_BATCHES=50
+FINAL_VAL_BATCHES=100
+SAMPLE_VAL_BATCHES=16
+SAMPLE_NUM_STEPS="4,10"
+RUN_WARMSTART_CHECK="${RUN_WARMSTART_CHECK:-0}"
+
+BASELINE_CKPT_ROOT="$ROOT/checkpoints/$BASELINE_CONFIG/$BASELINE_EXP"
+PARALLEL_CKPT_ROOT="$ROOT/checkpoints/$PARALLEL_CONFIG/$PARALLEL_EXP"
+
+BASELINE_LOG="$RUN_LOG_DIR/${BASELINE_EXP}.log"
+PARALLEL_LOG="$RUN_LOG_DIR/${PARALLEL_EXP}.log"
+WARMSTART_LOG="$RUN_LOG_DIR/warmstart_dual_push_128_5k.log"
+
+wait_for_checkpoint_dir() {
+  local ckpt_dir="$1"
+  while [[ ! -d "$ckpt_dir" ]]; do
+    echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] waiting for checkpoint dir: $ckpt_dir"
+    sleep 30
+  done
+}
+
+run_eval() {
+  local config_name="$1"
+  local ckpt_dir="$2"
+  local num_batches="$3"
+  local log_path="$4"
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] eval start config=$config_name ckpt=$ckpt_dir batches=$num_batches"
+  python scripts/eval_twin_val_loss_pytorch.py \
+    --config_name "$config_name" \
+    --checkpoint_dir "$ckpt_dir" \
+    --repo_id "$VAL_REPO" \
+    --num_batches "$num_batches" \
+    --num_workers 0 \
+    --sample_num_batches "$SAMPLE_VAL_BATCHES" \
+    --sample_num_steps "$SAMPLE_NUM_STEPS" \
+    >"$log_path" 2>&1
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] eval done log=$log_path"
+}
+
+run_train_and_evals() {
+  local config_name="$1"
+  local exp_name="$2"
+  local ckpt_root="$3"
+  local train_log="$4"
+
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] train start config=$config_name exp=$exp_name"
+  torchrun --standalone --nproc_per_node=4 scripts/train_pytorch.py \
+    "$config_name" \
+    --exp_name "$exp_name" \
+    --overwrite \
+    >"$train_log" 2>&1
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] train done config=$config_name exp=$exp_name"
+
+  for step in 1000 2000 5000; do
+    wait_for_checkpoint_dir "$ckpt_root/$step"
+    local num_batches="$INTERMEDIATE_VAL_BATCHES"
+    if [[ "$step" == "5000" ]]; then
+      num_batches="$FINAL_VAL_BATCHES"
+    fi
+    run_eval \
+      "$config_name" \
+      "$ckpt_root/$step" \
+      "$num_batches" \
+      "$RUN_LOG_DIR/${exp_name}_val_${step}.log"
+  done
+}
+
+echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] packed dual-push 128 5k runner started"
+if [[ "$RUN_WARMSTART_CHECK" == "1" ]]; then
+  python -u scripts/check_parallel_warmstart_equivalence.py \
+    --baseline_config_name "$BASELINE_CONFIG" \
+    --parallel_config_name "$PARALLEL_CONFIG" \
+    --repo_id "lsnu/twin_dual_push_128_train" \
+    --batch_size 1 \
+    >"$WARMSTART_LOG" 2>&1
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] warm-start equivalence check logged to $WARMSTART_LOG"
+else
+  echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] skipping warm-start equivalence gate (RUN_WARMSTART_CHECK=$RUN_WARMSTART_CHECK)"
+fi
+
+run_train_and_evals "$BASELINE_CONFIG" "$BASELINE_EXP" "$BASELINE_CKPT_ROOT" "$BASELINE_LOG"
+run_train_and_evals "$PARALLEL_CONFIG" "$PARALLEL_EXP" "$PARALLEL_CKPT_ROOT" "$PARALLEL_LOG"
+
+echo "[$(date -u '+%Y-%m-%d %H:%M:%S UTC')] packed dual-push 5k runner finished successfully"