run_100k_reference.sh

#!/bin/bash
export CUDA_VISIBLE_DEVICES=7
# VAE Sample Synthesis and Reward Model Training Pipeline
# Author: Leitian Tao
# Date: 2023-11-01
START_STEP=1
cd /u/l/e/leitiantao/public/npos/
# Default parameters
BASE_DIR=""
OUTPUT_DIR="/nobackup2/taoleitian/rm/vae_results/hh_rlhf/llama_instruct_10k"
INPUT_FILE="/nobackup2/taoleitian/neurips/embeddings/hh_rlhf/llama_instruct_10k/train_100k.npy"
MULTI_RESPONSE_FEATURES_PATH="/nobackup2/taoleitian/neurips/embeddings/hh_rlhf/llama_instruct_10k/multi_response/multi_response_embeddings.npy"
MULTI_RESPONSE_REWARDS_PATH="/nobackup2/taoleitian/neurips/embeddings/hh_rlhf/llama_instruct_10k/multi_response/multi_response_rewards.npy"
NUM_SAMPLES=50000
SEEDS=(22)
LATENT_DIM=16
HIDDEN_DIMS=(64)
BATCH_SIZE=128
EPOCHS=20
LR=1e-4
TEMPERATURE=1.0
CONTRASTIVE_WEIGHT=0.01
NOISE_STD=0.01
NUM_VARIANTS=1
N_NOISE=10
TRAIN_SIZE=1000.0
HIDDEN_DIM=512
DROPOUT=0.0
SEED=44
RUN_COMPARISON=true
GPU_ID=0  # Default GPU ID

# 解析命令行参数
while [[ $# -gt 0 ]]; do
  case $1 in
    --gpu)
      GPU_ID="$2"
      shift 2
      ;;
    --base_dir)
      BASE_DIR="$2"
      shift 2
      ;;
    --output_dir)
      OUTPUT_DIR="$2"
      shift 2
      ;;
    --input_file)
      INPUT_FILE="$2"
      shift 2
      ;;
    --multi_response_features_path)
      MULTI_RESPONSE_FEATURES_PATH="$2"
      shift 2
      ;;
    --multi_response_rewards_path)
      MULTI_RESPONSE_REWARDS_PATH="$2"
      shift 2
      ;;
    --num_samples)
      NUM_SAMPLES="$2"
      shift 2
      ;;
    --seeds)
      shift
      SEEDS=()
      while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do
        SEEDS+=("$1")
        shift
      done
      ;;
    --latent_dim)
      LATENT_DIM="$2"
      shift 2
      ;;
    --hidden_dims)
      shift
      HIDDEN_DIMS=()
      while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do
        HIDDEN_DIMS+=("$1")
        shift
      done
      ;;
    --batch_size)
      BATCH_SIZE="$2"
      shift 2
      ;;
    --epochs)
      EPOCHS="$2"
      shift 2
      ;;
    --lr)
      LR="$2"
      shift 2
      ;;
    --temperature)
      TEMPERATURE="$2"
      shift 2
      ;;
    --contrastive_weight)
      CONTRASTIVE_WEIGHT="$2"
      shift 2
      ;;
    --noise_std)
      NOISE_STD="$2"
      shift 2
      ;;
    --num_variants)
      NUM_VARIANTS="$2"
      shift 2
      ;;
    --train_size)
      TRAIN_SIZE="$2"
      shift 2
      ;;
    --hidden_dim)
      HIDDEN_DIM="$2"
      shift 2
      ;;
    --dropout)
      DROPOUT="$2"
      shift 2
      ;;
    --seed)
      SEED="$2"
      shift 2
      ;;
    --run_comparison)
      RUN_COMPARISON="$2"
      shift 2
      ;;
    --start_step)
      START_STEP="$2"
      shift 2
      ;;
    *)
      echo "未知参数: $1"
      exit 1
      ;;
  esac
done

# 设置GPU

# 创建输出目录
OUTPUT_PATH="$BASE_DIR/$OUTPUT_DIR"
mkdir -p "$OUTPUT_PATH"
echo "输出目录: $OUTPUT_PATH"

# 保存配置
echo "保存配置..."
cat > "$OUTPUT_PATH/pipeline_config.txt" << EOF
VAE样本合成和奖励模型训练流程配置
====================================
基础目录: $BASE_DIR
输出目录: $OUTPUT_DIR
输入文件: $INPUT_FILE
多响应特征路径: $MULTI_RESPONSE_FEATURES_PATH
多响应奖励路径: $MULTI_RESPONSE_REWARDS_PATH
样本数量: $NUM_SAMPLES
随机种子: ${SEEDS[@]}
VAE潜在空间维度: $LATENT_DIM
VAE隐藏层维度: ${HIDDEN_DIMS[@]}
批次大小: $BATCH_SIZE
训练轮数: $EPOCHS
学习率: $LR
温度参数: $TEMPERATURE
对比损失权重: $CONTRASTIVE_WEIGHT
噪声标准差: $NOISE_STD
每个样本生成的变体数量: $NUM_VARIANTS
训练样本数量(k): $TRAIN_SIZE
奖励模型隐藏层维度: $HIDDEN_DIM
Dropout比率: $DROPOUT
随机种子: $SEED
运行对比实验: $RUN_COMPARISON
EOF

# 设置随机种子
export PYTHONPATH="$BASE_DIR:$PYTHONPATH"

# Step 1: Sample Selection
# Purpose: Randomly select a subset of samples from the input embeddings
# Input: Input embedding file
# Output: Selected samples saved in seeds_samples directory
SELECTED_SAMPLES_FILE="$OUTPUT_PATH/seeds_samples/${NUM_SAMPLES//1000}k_${SEEDS[0]}.npy"
mkdir -p "$(dirname "$SELECTED_SAMPLES_FILE")"

if [ $START_STEP -le 1 ]; then
  echo "Step 1: Selecting samples..."

  python rm_eval/hh_rlhf/vae/select_samples.py \
    --input_file "$INPUT_FILE" \
    --output_file "$SELECTED_SAMPLES_FILE" \
    --num_samples "$NUM_SAMPLES" \
    --seeds "${SEEDS[@]}" \
    --layer_idx -1

  if [ $? -ne 0 ]; then
    echo "样本选择失败!"
    exit 1
  fi
else
  echo "跳过步骤1: 选择样本"
  SELECTED_SAMPLES_FILE="$OUTPUT_PATH/seeds_samples/${NUM_SAMPLES//1000}k_${SEEDS[0]}.npy"
fi

# Step 2: VAE Training
# Purpose: Train a Variational Autoencoder on the selected samples
# Input: Selected samples from Step 1
# Output: Trained VAE model saved in vae_model directory
HIDDEN_DIMS_PARAM=""
for dim in "${HIDDEN_DIMS[@]}"; do
  HIDDEN_DIMS_PARAM="$HIDDEN_DIMS_PARAM --hidden_dims $dim"
done
if [ $START_STEP -le 2 ]; then
  echo "Step 2: Training VAE model..."
  VAE_OUTPUT_DIR="$OUTPUT_PATH/vae_model"
  mkdir -p "$VAE_OUTPUT_DIR"

  # 构建隐藏层维度参数


  python rm_eval/hh_rlhf/vae/train_vae_dual.py \
    --output_dir "$VAE_OUTPUT_DIR" \
    --train_features_path "$SELECTED_SAMPLES_FILE" \
    --use_shared_vae \
    --latent_dim "$LATENT_DIM" \
    $HIDDEN_DIMS_PARAM \
    --batch_size "$BATCH_SIZE" \
    --epochs 10 \
    --lr "$LR" \
    --temperature "$TEMPERATURE" \
    --contrastive_weight "$CONTRASTIVE_WEIGHT"

  if [ $? -ne 0 ]; then
    echo "VAE训练失败!"
    exit 1
  fi
else
  echo "跳过步骤2: 训练VAE模型"
  VAE_OUTPUT_DIR="$OUTPUT_PATH/vae_model"
fi

# Step 3: Generate Noisy Sample Pairs
# Purpose: Generate noisy sample pairs using the trained VAE
# Input: Trained VAE model and selected samples
# Output: Generated noisy pairs saved in generated_pairs directory
if [ $START_STEP -le 3 ]; then
  echo "Step 3: Generating noisy sample pairs..."
  echo $HIDDEN_DIMS_PARAM
  GENERATED_PAIRS_DIR="$OUTPUT_PATH/generated_pairs"
  mkdir -p "$GENERATED_PAIRS_DIR"

  python rm_eval/hh_rlhf/vae/generate_noisy_pairs.py \
    --model_path "$VAE_OUTPUT_DIR/best_model.pt" \
    --input_features "$SELECTED_SAMPLES_FILE" \
    --output_dir "$GENERATED_PAIRS_DIR" \
    --noise_std "$NOISE_STD" \
    --num_variants "$NUM_VARIANTS" \
    --use_shared_vae \
    --evaluate_vae \
    --n_noise "$N_NOISE" \
    --latent_dim "$LATENT_DIM" \
    $HIDDEN_DIMS_PARAM \
    --include_original_samples true \

  if [ $? -ne 0 ]; then
    echo "样本生成失败!"
    exit 1
  fi
else
  echo "跳过步骤3: 生成带噪声的样本对"
  GENERATED_PAIRS_DIR="$OUTPUT_PATH/generated_pairs"
fi

# Step 4: Train Bradley-Terry Reward Model (with VAE-generated samples)
# Purpose: Train a reward model using the VAE-generated sample pairs
# Input: VAE-generated pairs and selected samples
# Output: Trained reward model saved in reward_model_with_vae directory
if [ $START_STEP -le 4 ]; then
  echo "Step 4: Training Bradley-Terry reward model (with VAE-generated samples)..."
  REWARD_MODEL_DIR="$OUTPUT_PATH/reward_model_with_vae"
  mkdir -p "$REWARD_MODEL_DIR"

  python rm_eval/hh_rlhf/vae/train_bt_vae.py \
    --train_size "$TRAIN_SIZE" \
    --batch_size "$BATCH_SIZE" \
    --epochs "$EPOCHS" \
    --lr "$LR" \
    --hidden_dim "$HIDDEN_DIM" \
    --dropout "$DROPOUT" \
    --seed "$SEED" \
    --train_data_path "$GENERATED_PAIRS_DIR/generated_noisy_pairs.npy" \
    --save_path "$REWARD_MODEL_DIR" \
    --vae_reconstructed_path None \
    --multi_response_features_path "$MULTI_RESPONSE_FEATURES_PATH" \
    --multi_response_rewards_path "$MULTI_RESPONSE_REWARDS_PATH"

  if [ $? -ne 0 ]; then
    echo "使用VAE生成样本的奖励模型训练失败!"
    exit 1
  fi
else
  echo "跳过步骤4: 训练Bradley-Terry奖励模型（使用VAE生成的样本）"
  REWARD_MODEL_DIR="$OUTPUT_PATH/reward_model_with_vae"
fi

# Step 5: Train Bradley-Terry Reward Model (baseline, without VAE-generated samples)
# Purpose: Train a baseline reward model using original samples for comparison
# Input: Original selected samples
# Output: Baseline reward model saved in reward_model_baseline directory
# Note: Only executed if RUN_COMPARISON=true
if [ "$RUN_COMPARISON" = true ] && [ $START_STEP -le 5 ]; then
  echo "Step 5: Training Bradley-Terry reward model (baseline)..."
  REWARD_MODEL_DIR_BASELINE="$OUTPUT_PATH/reward_model_baseline"
  mkdir -p "$REWARD_MODEL_DIR_BASELINE"

  python rm_eval/hh_rlhf/vae/train_bt_vae.py \
    --train_size "$TRAIN_SIZE" \
    --batch_size "$BATCH_SIZE" \
    --epochs "$EPOCHS" \
    --lr "$LR" \
    --hidden_dim "$HIDDEN_DIM" \
    --dropout "$DROPOUT" \
    --seed "$SEED" \
    --train_data_path "$SELECTED_SAMPLES_FILE" \
    --vae_reconstructed_path None \
    --save_path "$REWARD_MODEL_DIR_BASELINE" \
    --multi_response_features_path "$MULTI_RESPONSE_FEATURES_PATH" \
    --multi_response_rewards_path "$MULTI_RESPONSE_REWARDS_PATH"

  if [ $? -ne 0 ]; then
    echo "不使用VAE生成样本的奖励模型训练失败!"
    exit 1
  fi
else
  if [ "$RUN_COMPARISON" = true ]; then
    echo "跳过步骤5: 训练Bradley-Terry奖励模型（不使用VAE生成的样本）"
  fi
fi

# Step 6: Model Performance Comparison
# Purpose: Compare performance between VAE-based and baseline reward models
# Input: Both reward models' evaluation results
# Output: Comparison results saved in comparison directory
# Note: Only executed if RUN_COMPARISON=true
if [ "$RUN_COMPARISON" = true ] && [ $START_STEP -le 6 ]; then
  echo "Step 6: Comparing model performances..."
  COMPARISON_DIR="$OUTPUT_PATH/comparison"
  mkdir -p "$COMPARISON_DIR"

  # 提取两个模型的评估结果
  VAE_MODEL_GOLDEN_REWARD=$(jq '.mean_golden_reward' "$REWARD_MODEL_DIR/gold_reward_results.json")
  VAE_MODEL_GOLDEN_REWARD_STD=$(jq '.std_golden_reward' "$REWARD_MODEL_DIR/gold_reward_results.json")
  
  BASELINE_MODEL_GOLDEN_REWARD=$(jq '.mean_golden_reward' "$REWARD_MODEL_DIR_BASELINE/gold_reward_results.json")
  BASELINE_MODEL_GOLDEN_REWARD_STD=$(jq '.std_golden_reward' "$REWARD_MODEL_DIR_BASELINE/gold_reward_results.json")

  # 计算性能提升
  GOLDEN_REWARD_IMPROVEMENT=$(echo "scale=4; $VAE_MODEL_GOLDEN_REWARD - $BASELINE_MODEL_GOLDEN_REWARD" | bc)

  # 保存比较结果
  cat > "$COMPARISON_DIR/comparison_results.txt" << EOF
Performance Comparison: VAE vs Baseline
====================================
Metric                  Baseline                VAE Model               Improvement
--------------------------------------------------------------------------------
Mean Golden Reward     ${BASELINE_MODEL_GOLDEN_REWARD}±${BASELINE_MODEL_GOLDEN_REWARD_STD}    ${VAE_MODEL_GOLDEN_REWARD}±${VAE_MODEL_GOLDEN_REWARD_STD}    $GOLDEN_REWARD_IMPROVEMENT
EOF

  # 同时保存为JSON格式
  cat > "$COMPARISON_DIR/comparison_results.json" << EOF
{
  "baseline": {
    "mean_golden_reward": $BASELINE_MODEL_GOLDEN_REWARD,
    "std_golden_reward": $BASELINE_MODEL_GOLDEN_REWARD_STD
  },
  "vae_model": {
    "mean_golden_reward": $VAE_MODEL_GOLDEN_REWARD,
    "std_golden_reward": $VAE_MODEL_GOLDEN_REWARD_STD
  },
  "improvement": {
    "mean_golden_reward": $GOLDEN_REWARD_IMPROVEMENT
  }
}
EOF

  echo "Results saved to: $COMPARISON_DIR/comparison_results.txt and comparison_results.json"
  cat "$COMPARISON_DIR/comparison_results.txt"
else
  if [ "$RUN_COMPARISON" = true ]; then
    echo "Skipping Step 6: Model performance comparison"
  fi
fi

# Save experiment completion time
echo "Experiment completed at: $(date)" > "$OUTPUT_PATH/experiment_completed.txt"

echo "================================================"
echo "VAE sample synthesis and reward model training pipeline completed!"
echo "Results saved in: $OUTPUT_PATH"
echo "================================================"