dbs-qwen3b-sft-grpo

DBS Agent: SFT -> GRPO on the Parkinson's Motor Environment

LLM agent (Qwen2.5-3B-Instruct-bnb-4bit + LoRA) trained as an adaptive Deep Brain Stimulation programmer for the Parkinson's Motor OpenEnv environment.

Headline result

Mean grader (averaged over 3 tasks x 1 seeds): 0.664 (base) -> 0.664 (SFT) -> 0.664 (SFT+GRPO) (+0.000, +0.0%)

Success rate: 100% -> 100% (+0 percentage points)

Per-task scores (mean grader, higher is better)

Submission links

• Environment (HF Space): https://huggingface.co/spaces/virustechhacks/parkinsons_Motor
• Trained model (this repo): https://huggingface.co/Dash10107/dbs-qwen3b-sft-grpo
• Blog post: TODO: paste HF blog post URL here
• Video walkthrough: TODO: paste YouTube/Loom video URL here
• Slide deck: TODO: paste Google Slides / pptx URL here

Training pipeline

1. SFT warm-start on rejection-sampled rollouts from the safety_aware heuristic policy ({'beta_suppression': 12, 'tremor_correction': 12, 'full_episode': 6} eps/task across 3 tasks; kept fully when grader >= 0.55, top half kept when grader >= 0.3).
2. GRPO with per-step return-to-go advantages, group normalization, KL anchor to the frozen SFT reference adapter, and a small JSON-format bonus. Curriculum: easy -> medium -> hard, promoted by rolling mean grader.

Plots

Logs (full reproducibility)

• logs/sft/sft_steps.jsonl - one row per SFT optimizer step
• logs/grpo/rollouts.jsonl - one row per GRPO rollout episode
• logs/grpo/updates.jsonl - one row per GRPO update step
• logs/eval/episode_rows.jsonl - one row per eval episode (all policies)
• logs/eval/inference_*.jsonl - full LLM trace (prompt -> response -> action -> obs) for replay
• logs/eval/submission_summary.json - the headline numbers above
• logs/eval/submission_checklist.json - submission-requirement audit

Qualitative replay (base vs trained, same seed)

Qualitative inference replay

Same patient, same seed, two policies. Look at how DBS amplitude and the beta/tremor biomarkers move differently after training.

Task: beta_suppression (seed 0, 30 steps)

• base_qwen: grader = 0.750, success = True
• sft_plus_grpo: grader = 0.750, success = True

  [BASE] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.624  tremor=0.079  side=0.065  reward=+0.646
  [FINAL] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.624  tremor=0.079  side=0.065  reward=+0.646

  [BASE] step= 15  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.189  tremor=0.352  side=0.298  reward=+0.652
  [FINAL] step= 15  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.189  tremor=0.352  side=0.298  reward=+0.652

  [BASE] step= 29  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.210  tremor=0.401  side=0.386  reward=+0.605
  [FINAL] step= 29  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.210  tremor=0.401  side=0.386  reward=+0.605

Task: tremor_correction (seed 0, 48 steps)

• base_qwen: grader = 0.544, success = True
• sft_plus_grpo: grader = 0.544, success = True

  [BASE] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.292  tremor=0.394  side=0.054  reward=+0.592
  [FINAL] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.292  tremor=0.394  side=0.054  reward=+0.592

  [BASE] step= 24  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.079  tremor=0.490  side=0.274  reward=+0.517
  [FINAL] step= 24  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.079  tremor=0.490  side=0.274  reward=+0.517

  [BASE] step= 47  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.008  tremor=0.765  side=0.317  reward=+0.435
  [FINAL] step= 47  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.008  tremor=0.765  side=0.317  reward=+0.435

Task: full_episode (seed 0, 100 steps)

• base_qwen: grader = 0.698, success = True
• sft_plus_grpo: grader = 0.698, success = True

  [BASE] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.679  tremor=0.040  side=0.046  reward=+0.751
  [FINAL] step=  0  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.679  tremor=0.040  side=0.046  reward=+0.751

  [BASE] step= 50  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.078  tremor=0.670  side=0.235  reward=+0.566
  [FINAL] step= 50  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.078  tremor=0.670  side=0.235  reward=+0.566

  [BASE] step= 99  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.002  tremor=0.497  side=0.240  reward=+0.586
  [FINAL] step= 99  amp=1.00mA  pw=130us  freq=130Hz  | beta=0.002  tremor=0.497  side=0.240  reward=+0.586

How to load

from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained('Dash10107/dbs-qwen3b-sft-grpo')
mdl = AutoModelForCausalLM.from_pretrained('Dash10107/dbs-qwen3b-sft-grpo', torch_dtype='auto', device_map='auto')

Reproducing this run

The full Colab notebook is dbs_sft_grpo_colab.ipynb in the env repo. The exact config used is checked into this model repo as cfg.json.