#!/bin/bash
# RL run 2 (2026-10-07): wait for the SFT warm start (runs/sft1) to finish, evaluate it on the fixed
# eval set, then run GRPO r2 under the liveness watchdog. If the warm start cost more than 5 points
# in-distribution versus final.pt (best score 0.894; differences under ~5 points are eval noise),
# r2 starts from final.pt instead. Log: $TA_DATA/logs/rl_pipeline.log
#   setsid nohup bash scripts/rl_pipeline_r2.sh >/dev/null 2>&1 < /dev/null &
set -u
cd "$(dirname "$0")/.."
source env.sh
SFT=$TA_DATA/runs/sft1
OUT=$TA_DATA/rl/r2
L=$TA_DATA/logs/rl_pipeline.log
log() { echo "[$(date '+%F %T')] $*" | tee -a "$L"; }
trainer() { ps -eo pid,args | awk '$3=="scripts/train.py" && $0 ~ /runs\/sft1/ {print $1}' | head -1; }

log "waiting for the SFT warm start ($SFT/final.pt)"
while [ ! -f "$SFT/final.pt" ] || [ -n "$(trainer)" ]; do
  if [ -z "$(trainer)" ] && [ ! -f "$SFT/final.pt" ]; then log "SFT trainer gone without final.pt; see runs/sft1.out"; break; fi
  sleep 30
done
CK=$TA_DATA/runs/main/final.pt
if [ -f "$SFT/final.pt" ]; then
  log "SFT done: $(tail -1 "$SFT/log.jsonl" | cut -c1-160)"
  $TA_PY scripts/eval_agent.py --ckpt "$SFT/final.pt" --tasks 150 --k 8 --verbose 0 > "$TA_DATA/logs/eval_sft1.json" 2>> "$L"
  verdict=$($TA_PY -c "
import json
r = json.load(open('$TA_DATA/logs/eval_sft1.json'))
i, h = r['_in_dist']['pass1'], r['_held_out']['pass1']
print('ok' if i >= 0.894 - 0.05 else 'worse', i, h, {k: v['pass1'] for k, v in r.items() if k in ('multi_hop', 'log_count', 'write_fact')})" 2>&1)
  log "SFT eval (k=8): $verdict"
  if [ "${verdict%% *}" = "ok" ]; then
    CK=$SFT/final.pt
    mkdir -p "$OUT" && cp "$TA_DATA/logs/eval_sft1.json" "$OUT/baseline.json"
  else
    log "SFT cost too much in-distribution; r2 starts from final.pt (new kinds then have no demos)"
    mkdir -p "$OUT" && cp "$TA_DATA/logs/eval_final.json" "$OUT/baseline.json"
  fi
fi
log "starting r2 from $CK under the watchdog"
RL_OUT=$OUT RL_CKPT=$CK bash scripts/rl_watchdog.sh --prefill_rows 8 --vary 0.7 --invented_penalty 0.3
log "r2 pipeline finished: $(tail -2 "$TA_DATA/logs/rl_watchdog.log" | tr '\n' ' ' | cut -c1-300)"