#!/bin/bash
# v11 = v10_r2 + asymmetric over-prediction penalty.
#
# Diagnosis (eval_v10_r2_ckpt956 vs eval_tempsamp_single_span_ckpt1160):
# the +2.62 headline mIoU is diluted because 81% of test videos are
# single-segment, where v10_r2 ties the baseline (+0.05); the real win is on
# the 18.6% multi-segment subset (+13.89). On single-seg, v10_r2 OVER-PREDICTS
# spurious extra segments (single-seg F1@strict 37.73 < baseline 39.31), and
# soft_F1's precision term penalises that too weakly to suppress it.
#
# v11 sets FORENSICS_OVERPRED_PENALTY: forensics_iou_reward subtracts
# kappa*(K_pred-K_gt)/(K_pred+1) ONLY when K_pred > K_gt. Correct multi-segment
# predictions (K_pred==K_gt) are never penalised, so the multi-seg win is
# preserved; the gradient now actively discourages spurious extra segments.
# Everything else is identical to v10_r2 for an apples-to-apples comparison.
#
# --- v10_r2 original header below ---
# v10 + R2 (strict_edge): sub-second boundary precision reward.
#
# Replaces the failed probe-based binary_probing reward from v10_r1 (whose
# reward stayed flat at ~0.08 throughout training — frozen base Qwen had no
# usable forgery signal). Net effect of v10_r1 on ckpt956: mIoU −0.82,
# F1@0.5 −0.53, F1@0.7 +0.08 vs v10_full → noise, not signal.
#
# v10's eval gap: F1@0.7=48 → F1@0.85=26 → F1@0.95=10. strict_boundary @
# τ∈{0.3,0.5,0.7} saturates at IoU>0.7 and gives no gradient into that gap.
# strict_edge adds a per-Hungarian-match boundary-distance term with σ=0.5s
# (product form), which only starts paying out near sub-second edge alignment
# — exactly where strict_boundary is silent.

export WANDB_PROJECT=Forensics-GRPO
export WANDB_NAME="${WANDB_NAME:-AF_v11_overpred}"

export PYTHONPATH=".:$PYTHONPATH"
export DEBUG_MODE="true"
export LOG_PATH="./qwen2.5_7b_vl_forensics_v11_overpred.txt"

export PATH="/mnt/local-fast/zhangt/torch_env/bin:$PATH"
export LD_LIBRARY_PATH="/opt/conda/lib:${LD_LIBRARY_PATH}"

# === v10 baseline (do not edit) ============================================
export FORENSICS_COT=false
export FORENSICS_FORGERY_AWARE=false

export FORENSICS_HUNGARIAN_DECOMP=true
export FORENSICS_HUNGARIAN_MAXK=8
export FORENSICS_BOUNDARY_WEIGHT=0.3
export FORENSICS_BOUNDARY_TAU=2.0

export FORENSICS_SPI_AUG=true
export FORENSICS_SPI_PROB=0.25
export FORENSICS_SPI_CHUNK_S=2.5
export FORENSICS_SPI_SAFETY_S=0.5

export FORENSICS_FBR_AUG=true
export FORENSICS_FBR_PROB=0.15

export FORENSICS_STRICT_TAUS="0.3,0.5,0.7"

unset FORENSICS_COT_CONSIS_WEIGHT
unset FORENSICS_COT_CF_WEIGHT
unset FORENSICS_COT_ALIGN_WEIGHT
unset FORENSICS_K_REQUIRE_TIMESTEP
unset FORENSICS_IOU_REQUIRE_THINK
# ============================================================================

# === R2: strict_edge ========================================================
# σ controls tolerance. 0.5s → bonus decays e-fold per 0.5s of edge error.
# Reward is product over the two boundaries: both must hit to score.
export FORENSICS_STRICT_EDGE_SIGMA=0.5
# ============================================================================

# === v11: over-prediction penalty ==========================================
# kappa for the asymmetric K_pred>K_gt penalty in forensics_iou_reward.
# Sanity-checked: single-seg over-pred reward 0.667->0.533; multi-seg CORRECT
# (K_pred==K_gt) stays 1.000 (win preserved); multi-seg over-pred 0.800->0.700.
export FORENSICS_OVERPRED_PENALTY="${FORENSICS_OVERPRED_PENALTY:-0.4}"
# ============================================================================

OUTDIR="${OUTDIR:-outputs_forensics/v11_overpred}"
MODEL_PATH="${MODEL_PATH:-/mnt/local-fast/zhangt/Qwen2.5-VL-7B-Instruct}"
PREPROCESSED="${PREPROCESSED:-/mnt/local-fast/zhangt/forensics_grpo_cache_uniform3584_fps2.0}"

cd "$(dirname "$0")/.."

torchrun --nproc_per_node="8" \
    --nnodes="1" \
    --node_rank="0" \
    --master_addr="127.0.0.1" \
    --master_port="12371" \
    src/open_r1/grpo_forensics.py \
    --deepspeed scripts/zero3_offload.json \
    --output_dir "$OUTDIR" \
    --model_name_or_path "$MODEL_PATH" \
    --annot_dir /mnt/local-fast/zhangt/annot/annot \
    --video_root /mnt/local-fast/zhangt/video \
    --preprocessed_data_path "$PREPROCESSED" \
    --reward_funcs iou strict_boundary strict_edge \
    --dataset_name activityforensics \
    --max_prompt_length 8192 \
    --max_completion_length 512 \
    --num_generations 4 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --bf16 \
    --torch_dtype bfloat16 \
    --data_seed 42 \
    --gradient_checkpointing true \
    --attn_implementation flash_attention_2 \
    --num_train_epochs 4 \
    --run_name "$WANDB_NAME" \
    --report_to none \
    --save_steps 30 \
    --learning_rate 1e-6 \
    --save_only_model true