#!/usr/bin/env bash
# Focal POMDP-v3: 3-class generalization of POMDPTemporalHead with focal CE.
# Output: checkpoints/Policy/focal_pomdp_qwen3vl4b_seed{0..2}/best/head.pt
set -euo pipefail
cd "$(dirname "$0")/../.."

CACHE_DIR=data/belief_cache_perframe_qwen3vl4b
LABEL_DIR=data/policy_labels
OUT_BASE=checkpoints/Policy

[[ -f $CACHE_DIR/multisrc_train.pt ]] || { echo "MISSING train cache"; exit 1; }
[[ -f $CACHE_DIR/multisrc_val.pt   ]] || { echo "MISSING val cache"; exit 1; }

for SEED in 0 1 2; do
    EXP=focal_pomdp_qwen3vl4b_seed${SEED}
    echo "=== [focal_pomdp_v3] seed=${SEED} → $OUT_BASE/$EXP ==="
    python -m training.Policy.train_focal_pomdp_v3 \
        --train_cache       $CACHE_DIR/multisrc_train.pt \
        --val_cache         $CACHE_DIR/multisrc_val.pt \
        --label_dir         $LABEL_DIR \
        --output_dir        $OUT_BASE \
        --experiment_name   $EXP \
        --num_epochs 6 \
        --batch_size 128 \
        --lr 2e-4 \
        --focal_gamma 2.0 \
        --label_smoothing 0.05 \
        --use_balanced_sampler \
        --use_ema --ema_decay 0.999 \
        --val_every_n_steps 200 \
        --seed ${SEED}
done

echo
echo "=== focal_pomdp_v3 all seeds done ==="
ls -d $OUT_BASE/focal_pomdp_qwen3vl4b_seed*