#!/bin/bash
# v10 + Binary Probing (R1) reward.
#
# Adds `binary_probing` to the reward stack of run_grpo_forensics_v10_full.sh,
# keeping every other v10 setting identical so the run is a clean
# "+R1 reward" ablation.
#
# The probing reward uses a FROZEN reference Qwen2.5-VL (loaded once per
# rank). For 8xH100 80GB this typically fits alongside the policy under
# ZeRO-3 offload. If you OOM, switch FORENSICS_PROBE_MODEL to a 3B variant.
#
# Cost note: FORENSICS_PROBE_INTERVAL_STEPS=4 means R1 fires every 4th step,
# averaging ~1.5 extra forward passes per step. Tighten/loosen to taste.

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

export PYTHONPATH=".:$PYTHONPATH"
export DEBUG_MODE="true"
export LOG_PATH="./qwen2.5_7b_vl_forensics_v10_r1.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
# ============================================================================

# === Binary probing (R1) ====================================================
# Path to the FROZEN reference model — usually the same Qwen2.5-VL-7B as
# the policy's init, but a separate weights copy on disk to avoid surprises
# if you ever change the policy ckpt mid-experiment.
export FORENSICS_PROBE_MODEL="${PROBE_MODEL:-/mnt/local-fast/zhangt/Qwen2.5-VL-7B-Instruct}"

# Probe schedule: r1 | r3 | combined
export FORENSICS_PROBE_SCHEDULE="${PROBE_SCHEDULE:-r1}"

# R1 window half-width (sec) and R3 point-clip width (sec).
export FORENSICS_PROBE_DELTA_S=2.0
export FORENSICS_PROBE_POINT_WINDOW_S=1.0

# Cost control. With G=4, M=2, K=3, and 6 probes (R1) or 4 probes (R3) per
# segment: every-4-steps × 2 rollouts × ≤3 segs × 6 probes = ≤36 probes
# averaged to ~9 / step.
export FORENSICS_PROBE_INTERVAL_STEPS=4
export FORENSICS_PROBE_MAX_ROLLOUTS=2
export FORENSICS_PROBE_NUM_GENERATIONS=4
export FORENSICS_PROBE_MAX_SEGS=3
export FORENSICS_PROBE_MAX_PROBES_PER_BATCH=16

# Anti-hacking: drop R1 reward to 0 when K_pred > 2x K_gt (prevents the
# policy from inflating segment count to dilute probe averaging).
export FORENSICS_PROBE_K_HARD_GATE=2.0
# ============================================================================

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