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).

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.

Let their loader download the originals once (so we have the time series + labels):

python -c "from opentslm.time_series_datasets.har_cot.har_cot_loader import ensure_har_cot_dataset; ensure_har_cot_dataset()"

Inject our rationales (our faithful HAR CoT lives on anviksha: faithful_har_cot_{train,val,test}.csv):

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.)

Train each arm and compare test macro-F1:

# 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`

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:
```
from grpo.reward import dual_reward, HAR_SCORER
reward_fn = lambda c, it: dual_reward(c, it["gold_label"], it["facts"], HAR_SCORER)
```

Swap the train step. In _train_stage's inner loop (currently L1114–1121):

# 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).

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

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)

Set up .venv_otslm, HF login, smoke-test with demo/huggingface/03_test_hf_har_cot.py.
SFT arm A (ours) + arm B (baseline) on stage3_cot; record test macro-F1 → paper §Downstream.
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.
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.