#!/bin/bash
# v10_r2 + asymmetric over-prediction penalty (FORENSICS_OVERPRED_PENALTY).
#
# Motivation: the bw=0.5 sweep (v10_r2_bw05) regressed across the board. The
# compare_eval_dirs.py diff localised the regression to a +1.85pt rise in
# over-prediction rate (+2.12pt on single-seg, 81% of the test set), which tanks
# the precision-aware F1@strict. Multi-seg mIoU actually improved (+0.57), so the
# fix is NOT to touch boundary weight but to discipline segment COUNT at the
# set level — decoupling "boundary precision" (per-token routing via
# hungarian_per_pred_iou) from "don't over-predict" (set-level reward).
#
# This run keeps the v10_r2 reward/routing landscape EXACTLY (boundary_weight
# back to 0.3) and only turns on the existing asymmetric penalty in
# forensics_iou_reward (reward.py:632): fires ONLY when K_pred > K_gt, bounded
# by kappa, so correct multi-segment predictions are never penalised.
#
# SIZING CAVEAT (the v11 lesson, reward.py:609-615): GRPO advantage is z-scored
# WITHIN the group of G=4 rollouts; converged within-group reward_std ~= 0.02.
# v11 died from a ~0.13 additive penalty = ~7x reward_std, which dominated the
# normalised advantage and wrecked the boundary gradient. kappa=0.05 here is
# bounded and chosen to stay a small fraction of that std. Watch reward_std in
# the logs; if over-prediction barely moves, step kappa up to 0.1, not higher.

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

export PYTHONPATH=".:$PYTHONPATH"
export DEBUG_MODE="true"
export LOG_PATH="./qwen2.5_7b_vl_forensics_v10_r2_overpred05.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
# back to v10_r2 value (bw=0.5 regressed):
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 ========================================================
export FORENSICS_STRICT_EDGE_SIGMA=0.5
# ============================================================================

# === only-knob change vs v10_r2: asymmetric over-prediction penalty =========
export FORENSICS_OVERPRED_PENALTY=0.05
# ============================================================================

OUTDIR="${OUTDIR:-outputs_forensics/v10_r2_overpred05}"
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="12370" \
    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