pratimassaravanan/clinical-qwen3-4b-sft-lora

Clinical Trial Recruitment Agent — Qwen3-4B + RL

A clinical trial recruitment agent trained with REINFORCE on a 180-step long-horizon environment. Theme 2 submission for the OpenEnv Hackathon.

Overview

Component	Detail
Base Model	Qwen/Qwen3-4B (4-bit quantized)
Method	Fresh LoRA (r=16, alpha=32) + REINFORCE policy gradient
Environment	Clinical Recruitment Env (8 action types, 37 observation features, 180-step horizon)
Training GPU	NVIDIA L40S (48GB) on Lightning AI
HF Space	pratimassaravanan/clinical-recruitment

Training Pipeline

Phase 1: SFT (Supervised Fine-Tuning)

• Data: 2,000 heuristic-generated traces (observation → JSON action pairs)
• Result: 100% JSON parse rate but complete policy collapse — model only outputs adjust_strategy
• Loss: 3.14 → 0.015 (perfect memorization, zero generalization)
• Finding: SFT teaches format but not observation-conditional behavior

Phase 2: REINFORCE (Policy Gradient RL)

• Key fix: Observation parsing bug (result.get("observation") → result — API returns flat dict)
• Key fix: Heuristic override prevents degenerate actions when candidates exist
• Reward shaping: +0.20 for allocate, +0.15 for recontact, +0.10 for screen, -0.10 for adjust_strategy when productive actions available

Debug Trial Results (3 episodes, 15 steps each):

Episode	Task	Enrolled	Target	Reward	Action Distribution
0	easy_bench	5	80	6.61	screen=6, allocate=9
1	easy_bench	5	80	6.61	screen=6, allocate=9
2	easy_bench	1	80	7.41	screen=14, allocate=1

Improvement over SFT: 0 → 5 enrolled patients, diverse action distribution vs. 100% adjust_strategy collapse.

Phase 3: Full REINFORCE Run (30 episodes)

• 30 episodes across easy/medium/hard benchmarks
• Results uploading upon completion

Key Findings

1. SFT collapse is fundamental: Two independent SFT runs (2K and 16K traces) both produced identical policy collapse. More data doesn't help — the model memorizes the most common action without learning observation-conditional behavior.

2. Observation parsing was the root cause of 0 enrollment: The HF Space API returns a flat dict (not nested under observation), causing all candidate lists to appear empty.

3. Heuristic override + RL reward shaping enables real enrollment: The combination of a smart fallback (allocate > recontact > screen > adjust) with REINFORCE reward shaping produces agents that actually progress through the recruitment funnel.

4. Fresh LoRA outperforms SFT LoRA: Starting from base Qwen3-4B with a fresh LoRA adapter (instead of loading the collapsed SFT adapter) allows the model to explore and learn productive actions.

Repository Structure

clinical-recruitment-env/
├── env.py                  # 180-step clinical trial environment
├── models.py               # Action/Observation Pydantic models  
├── graders.py              # Task-specific scoring (easy/medium/hard)
├── app.py                  # FastAPI server
├── openenv_adapter.py      # OpenEnv protocol adapter
├── train.py                # SFT training script (local, uses Python API)
├── _lightning_reinforce.py  # REINFORCE v3 (Lightning AI, uses HTTP API)
├── _reinforce_v4.py         # REINFORCE v4 (gentler reward shaping)
└── _debug_trial.py          # Debug trial (3 episodes, full logging)

Usage

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B", load_in_4bit=True)
model = PeftModel.from_pretrained(base, "pratimassaravanan/clinical-qwen3-4b-sft-lora/rl_v3_adapter")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

Environment Details

• Actions: screen_patient, recontact, allocate_to_site, adjust_strategy, plan_next_phase, summarize_and_index, retrieve_relevant_history, stop_recruitment
• Observation: 37 features including patient lists, site performance, funnel metrics, world_type
• Reward: Enrollment (+0.50), screening (+0.30), dropout (-0.35), milestone bonuses, hypothesis accuracy (+0.10)
• Tasks: easy_bench (80 target), medium_bench (100 target), hard_bench (150 target)

Compute & Cost

Job	GPU	Duration	Cost
SFT v1 (2K traces)	L40S	~30 min	~$2.50
SFT v2 (16K traces)	L40S	~60 min	~$5.00
RL debug trial	L40S	~5 min	$0.44
RL v3-fixed (30 eps)	L40S	~60 min	~$5.00

License

MIT