training/Policy/run_v4_ablations.sh

#!/usr/bin/env bash
# Run all v4 ablation experiments sequentially.
#
# Ablation matrix:
#   v4_baseline   — v4 architecture, but mono_lambda=0, kl_lambda=0 (sanity)
#   v4_edl        — EDL loss only (no monotonic)
#   v4_edl_mono   — EDL + monotonic constraint (full method)
#   v4_edl_mono_u — EDL + mono + tuned uncertainty threshold
#
# Total time estimate: ~4 runs × 15 min = ~1 hour
# Each run: belief cache mode → ~15 min on GPU
#
# Usage:
#   bash training/Policy/run_v4_ablations.sh
#   bash training/Policy/run_v4_ablations.sh --debug    # smoke test (~5 min total)
set -euo pipefail

ROOT=PROJECT_ROOT
cd "$ROOT"

DEBUG=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG="--debug"
fi

echo "=============================================="
echo "  LKAlert v4 Ablation Suite"
echo "=============================================="
echo ""

# 1. Sanity baseline: EDL architecture but no novel losses
echo "[1/4] v4_baseline (EDL arch, no KL, no mono)"
TAG=baseline KL_LAMBDA=0.0 MONO_LAMBDA=0.0 \
    bash training/Policy/train_policy_v4.sh $DEBUG

echo ""

# 2. EDL only: Evidential loss + KL annealing, no monotonic
echo "[2/4] v4_edl (EDL + KL, no mono)"
TAG=edl KL_LAMBDA=0.1 MONO_LAMBDA=0.0 \
    bash training/Policy/train_policy_v4.sh $DEBUG

echo ""

# 3. Full method: EDL + monotonic constraint
echo "[3/4] v4_edl_mono (EDL + KL + mono_λ=0.1)"
TAG=edl_mono KL_LAMBDA=0.1 MONO_LAMBDA=0.1 \
    bash training/Policy/train_policy_v4.sh $DEBUG

echo ""

# 4. Full method with tighter uncertainty threshold
echo "[4/4] v4_edl_mono_u03 (EDL + mono + u_thr=0.3)"
TAG=edl_mono_u03 KL_LAMBDA=0.1 MONO_LAMBDA=0.1 U_THRESHOLD=0.3 \
    bash training/Policy/train_policy_v4.sh $DEBUG

echo ""
echo "=============================================="
echo "  All ablations complete!"
echo "  Checkpoints in: checkpoints/Policy/policy_warmstart_v4_*/"
echo ""
echo "  Next: run conformal calibration on best variant:"
echo "  python -m training.Policy.conformal_risk \\"
echo "      --sft_checkpoint checkpoints/SFT/sft_v2/best \\"
echo "      --v4_ckpt checkpoints/Policy/policy_warmstart_v4_edl_mono/best \\"
echo "      --label_dir data/policy_labels \\"
echo "      --belief_cache_dir data/belief_cache \\"
echo "      --output_dir eval_results/paper_comparison_v4"
echo "=============================================="