code/OpenTSLM/grpo/GRPO_INTEGRATION.md

# Planting GRPO on OpenTSLM — runbook

Goal: (A) train OpenTSLM on **our** faithful rationales and compare to their baseline
rationales (downstream SFT experiment), then (B) add a **GRPO** stage with the dual
answer+faithfulness reward on top of the SFT model.

Everything here is GPU-blocked at *run* time (all 8 GPUs busy with ECG agent2 + other
users). The code/setup below is done; launch when a GPU frees.

---

## 0. Environment (no GPU needed to install)

Their stack: Python 3.12, torch ≥2.9.1, transformers ≥4.57, peft, open-flamingo, wfdb.
Keep it separate from our agent pipeline venv (which is 3.13).

```bash
cd ~/Adinath/TimeAgent/OpenTSLM
python3.12 -m venv .venv_otslm
source .venv_otslm/bin/activate
pip install --upgrade pip
pip install -r requirements.txt        # installs torch 2.9.x + open-flamingo + -e .
# HF access for the base LLM:
huggingface-cli login                  # needs a token with Llama-3.2-1B access,
                                        # OR use the open google/gemma-3-270m instead
```

Decision — base LLM: `meta-llama/Llama-3.2-1B` (default, gated) vs `google/gemma-3-270m`
(open, smaller, faster to iterate). Recommend gemma-3-270m for the first end-to-end
smoke run, then Llama-3.2-1B for the paper numbers.

---

## A. Train on our faithful data (SFT: ours vs baseline)

OpenTSLM's HAR loader reads `src/data/har_cot/har_cot_{train,val,test}_cot.csv`
with columns `x_axis,y_axis,z_axis,label,rationale`. `compute_loss` teacher-forces on
`rationale`, so swapping that column is the whole intervention.

1. Let their loader download the originals once (so we have the time series + labels):
   ```bash
   python -c "from opentslm.time_series_datasets.har_cot.har_cot_loader import ensure_har_cot_dataset; ensure_har_cot_dataset()"
   ```
2. Inject our rationales (our faithful HAR CoT lives on anviksha:
   `faithful_har_cot_{train,val,test}.csv`):
   ```bash
   for sp in train val test; do
     python grpo/inject_faithful_data.py \
       --orig src/data/har_cot/har_cot_${sp}_cot.csv \
       --ours /path/to/faithful_har_cot_${sp}.csv \
       --ours-rationale-col our_rationale \
       --out  src/data/har_cot/har_cot_${sp}_cot.csv --backup
   done
   ```
   (`--backup` keeps `*.orig.csv` so you can train the baseline arm too.)
3. Train each arm and compare test macro-F1:
   ```bash
   # ours
   python curriculum_learning.py --model OpenTSLMSP --stages stage3_cot --llm_id google/gemma-3-270m
   # baseline: restore *.orig.csv into place, retrain into a different results dir
   ```
   Results land under `results/<llm_id_safe>/OpenTSLMSP/stage3_cot/`.

The same pattern applies to Sleep (`stage4_sleep_cot`) and ECG (`stage5_ecg_cot`) once
those faithful CSVs are final; ECG uses `--key sample_id` in the injector.

---

## B. Add the GRPO stage (`stage6_grpo`)

The pieces are in `grpo/`:
- `reward.py` — `dual_reward()` = 0.7·answer + 0.3·faithfulness (HAR scorer ready;
  add ECG/Sleep/WESAD scorers from `codebase/agent3_*.py` patterns).
- `grpo_trainer.py` — `GRPOTrainer.grpo_loss(batch)`: samples N rollouts via
  `model.generate`, scores them, group-normalizes advantages, returns a PG loss
  (+ optional KL to a frozen SFT reference).

### Three wiring edits in `curriculum_learning.py`

1. **Carry reward fields through the dataset.** The reward needs the gold label and the
   signal facts per sample. In the GRPO dataset, attach to each item:
   - `item["gold_label"]` = the activity/stage/answer label (already available pre-rationale),
   - `item["facts"]` = the Stage-1 facts dict for that sample (load `signal_facts_*.json`
     keyed by sample_id; add a `sample_id` column to the injected CSVs to join).
   Build a `reward_fn(completion, item)` closure:
   ```python
   from grpo.reward import dual_reward, HAR_SCORER
   reward_fn = lambda c, it: dual_reward(c, it["gold_label"], it["facts"], HAR_SCORER)
   ```

2. **Swap the train step.** In `_train_stage`'s inner loop (currently L1114–1121):
   ```python
   # SFT:
   #   optimizer.zero_grad(); loss = model.compute_loss(batch)
   #   loss.backward(); optimizer.step()
   # GRPO:
   optimizer.zero_grad()
   loss, stats = grpo.grpo_loss(batch)        # grpo = GRPOTrainer(model, reward_fn, cfg, ref_model)
   loss.backward()
   clip_grad_norm_(model.parameters(), 1.0)
   optimizer.step()
   ```
   Gate this on `stage_name == "stage6_grpo"` so SFT stages are untouched. Use a smaller
   LR (e.g. 1e-6–5e-6) and a smaller batch (rollouts multiply cost by N).

3. **Register the stage.** Add `"stage6_grpo"` to `CURRICULUM_STAGES` and a `stage6_grpo`
   method mirroring `stage3_cot` (same HAR dataset, but it must start from the completed
   `stage3_cot` checkpoint — GRPO needs a competent SFT policy first). Initialize
   `ref_model` as a frozen clone of that checkpoint for the KL term (start with
   `kl_coef=0.0`, raise if the policy degenerates).

### GRPO config starting point
```python
from grpo.grpo_trainer import GRPOConfig
GRPOConfig(num_rollouts=8, max_new_tokens=400, temperature=1.0,
           kl_coef=0.0, w_answer=0.7, w_faith=0.3)
```

---

## Order of operations (when a GPU frees)
1. Set up `.venv_otslm`, HF login, smoke-test with `demo/huggingface/03_test_hf_har_cot.py`.
2. SFT arm A (ours) + arm B (baseline) on `stage3_cot`; record test macro-F1 → paper §Downstream.
3. From arm A's checkpoint, run `stage6_grpo`; track `reward_mean / answer_reward /
   faith_reward` from `grpo_loss` stats; evaluate faithfulness + accuracy vs the SFT model.
4. Repeat for Sleep/ECG once their faithful data is final.

## Open risks to watch
- `model.generate` must accept `num_return_sequences`/`do_sample` (they pass `**kwargs`
  to `self.llm.generate`, so it should — verify on the smoke run).
- Rollout cost = N × generate per sample; keep batch small and consider fewer rollouts
  (4–6) for the large splits.
- Reward needs per-sample facts; the cleanest join is a `sample_id` column threaded from
  Stage-1 facts → injected CSV → dataset item.
```