#!/bin/bash
# RL GRPO for the gemma-4-26B-A4B NLA (self-contained trainer, bf16, GPU0).
# reward = -log reconstruction MSE; --train-critic REQUIRED (else advantages collapse).
set -u
set -a; source /root/.nla_env; set +a
export PYTHONUNBUFFERED=1
export WANDB_ENTITY=MATS10-CS-JB
export WANDB_RUN_GROUP=g4a
export TOKENIZERS_PARALLELISM=false

DATA=${DATA:-/data/gemma4_nla/or}
CK=/ckpts/gemma4_nla
BASE=google/gemma-4-26B-A4B-it
AV=${AV:-$CK/av/iter_0001000}
AR=${AR:-$CK/ar/iter_0001000}

CUDA_VISIBLE_DEVICES=${GPU:-0} python3 -m nla.train_rl_self_contained \
  --av-ckpt "$AV" --ar-ckpt "$AR" \
  --base-ckpt "$BASE" --quant none \
  --rl-parquet $DATA/rl_shuf.parquet --sidecar $DATA/rl_shuf.parquet \
  --save-dir $CK/rl \
  --num-steps 500 --batch-prompts 16 --group-size 16 \
  --max-new-tokens 150 --temperature 1.0 \
  --lr 1e-5 --kl-beta 0.01 --clip-eps 0.2 \
  --train-critic --critic-lr 5e-5 \
  --logp-micro-batch 2 --max-rows 30000 \
  --save-every 20 --eval-every 20 --eval-n-prompts 20 --eval-skip-rows 35000 \
  --eval-n-hallucination 0 --eval-n-karvonen 0 \
  --max-grad-norm 1.0 \
  --wandb-project cot-oracle --wandb-name g4a_rl_grpo --seed 0