README.md

metadata

title: Narada Env
emoji: 🧬
colorFrom: blue
colorTo: green
sdk: docker
app_port: 7860
pinned: false

Narada — Rare Disease Diagnosis RL Environment

Meta × PyTorch OpenEnv Hackathon × Scaler School of Technology — Grand Finale

Live environment: krishvenky-narada-env.hf.space · Blog: Blog.md · Colab training notebook: Open in Colab

---

Results

GRPO training on Qwen3-1.7B (LoRA rank 16, 200 steps curriculum) vs zero-shot baseline:

Task	Baseline	After GRPO	Gain
monogenic	0.4955	0.572	+15.4%
oligogenic	0.4955	0.561	+13.2%
phenotype_mismatch	0.4955	0.552	+11.4%
Average	0.4955	0.562	+13.3%

Zero-shot baseline vs. trained agent across all three task tiers

Mean reward across 200 training steps (curriculum order). Shaded region = reward_std. Dotted line = zero-shot baseline.

Policy loss across curriculum. Phase boundaries shown as dashed vertical lines.

reward_std > 0 throughout confirms GRPO received real gradient signal — model never collapsed to uniform reward.

---

The Problem

Most rare disease patients wait 4–7 years for a correct diagnosis (Lancet 2024, Nature 2024). The delay is systemic, not a simple workflow gap:

• Rare diseases are clinically heterogeneous — many present with non-specific symptoms common across hundreds of conditions (PMC 2024)
• 60% of patients are initially misdiagnosed with either a different physical illness or a psychological condition; patients consult 8+ specialists (EU) or 17+ specialists (US) before a correct diagnosis (EURORDIS survey)
• "Medical ping-pong" — patients are passed between specialists with poor cross-service communication (NHS Genomics Education)
• 95% of rare diseases have no approved treatment, and poor clinician/patient awareness of symptom profiles compounds the problem (Lancet 2024)

ClinVar contains 2M+ catalogued genetic variants; HPO maps 15,000+ diseases to phenotypic signatures. The bottleneck is reasoning under uncertainty — cross-referencing patient symptoms against thousands of candidate variants while resisting high-salience but causally irrelevant signals.

This is exactly where current LLMs fail: they follow pathogenicity scores, not causal chains.

Narada's scope: The full diagnostic odyssey involves referral pathways, specialist availability, and fragmented care across institutions. Narada targets the variant prioritization substage — given a patient's phenotype profile and a shortlist of candidate genomic variants, which variant is causally responsible? This substage has a ground-truth signal (ClinVar), a tractable action space (graph traversal), and measurable improvement via RL — making it ideal for methodology proof-of-concept work.

---

What Narada Is

A reinforcement learning environment where an LLM agent navigates a 55,000-node gene-disease knowledge graph built from real ClinVar and HPO data. The agent must diagnose a rare disease patient by reasoning through phenotype → disease → gene → variant chains.

Three task tiers, increasing difficulty:

Task	Description	Key Challenge
monogenic	Single causal gene, 3–4 phenotypes	Basic directional reasoning
oligogenic	2 causal genes (one variant each), 5–7 phenotypes	Multi-objective tracking across long trajectory
phenotype_mismatch	Cardiac patient + BRCA1 frameshift decoy	Causal discipline — resist the highest-salience wrong signal

Three-agent system:

• Detective (Qwen3-1.7B, trainable) — navigates the graph, flags the causal variant
• Overseer — local heuristic (no LLM) that scores trajectory quality: penalises hallucinated hops, rewards touching the causal gene, and scales with a concise trail. Added only to successful terminal rewards.
• Adversary (planned) — curriculum case generation targeting Detective failure patterns; reliable adversarial curriculum from agent error logs is an open research problem, deferred to future work

---

Proof-of-Concept Framing

This project uses Qwen3-1.7B as the Detective agent. At this scale, the honest goal is methodology proof-of-concept: can GRPO training on a verifiable graph-navigation task move the needle even on a constrained model? Measurable improvement from a low baseline is a legitimate research contribution.

The environment and training pipeline are designed to generalize to larger models — the same loop applies to Qwen2.5-72B with zero code changes.

---

Graph

Built at runtime from:

• data/hp.obo — 19,389 HPO phenotype terms
• data/clinvar_pathogenic.tsv — 92,000 high-confidence pathogenic variants (GRCh38, criteria provided/expert panel, deduplicated)

Graph stats:

• 55,000+ nodes (phenotype, disease, gene, variant, pathway)
• 70,000+ edge pairs
• 3,268 genes represented

---

Action Space

Action	Effect	Reward
hop(node_id)	Move to connected node	+0.15 relevant / −0.05 irrelevant
flag_causal(variant_id)	Declare diagnosis; oligogenic cases allow multiple flags	+1.0 correct terminal, −0.5 wrong
backtrack()	Return to previous node	+0.05 after wrong direction
request_lab(test)	Get additional phenotype data	−0.10 penalty
summarise_trail()	Compressed visit summary	0.0 (neutral)

Signed raw rewards are mapped into OpenEnv's required score interval (0.01, 0.99). This preserves penalties while keeping the validator-compatible output range.

---

Running Locally

# 1. Clone
git clone https://github.com/KrishVenky/Narada-Env.git
cd Narada-Env

# 2. Create virtualenv (Python 3.10+)
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate

# 3. Install
pip install -r requirements.txt

# 4. Set up environment variables
cp .env.example .env
# Edit .env and fill in your API keys

# 5. Generate filtered ClinVar (one-time, ~3 min)
python scripts/filter_clinvar.py

# 6. Start server
PYTHONPATH=src/envs uvicorn narada.server.app:app --port 7860

# 7. Open browser → http://localhost:7860

---

Contributing

Pull requests welcome. To add a new task tier, reward component, or agent variant:

1. Fork the repo and create a branch off main
2. Environment logic lives in src/envs/narada/ — environment.py for reward design, case_generator.py for case sampling, graph.py for the knowledge graph
3. Add your changes and run the OpenEnv validator before opening a PR:

   openenv validate http://localhost:7860
   

4. Update ARCHITECTURE.md if you change reward values, node types, or the session protocol

Key files:

File	Purpose
src/envs/narada/server/environment.py	Core RL loop, reward computation
src/envs/narada/server/app.py	FastAPI entry point (WebSocket + HTTP debug)
src/envs/narada/graph.py	Knowledge graph build + singleton
src/envs/narada/case_generator.py	Patient case generation per tier
src/envs/narada/models.py	Pydantic schemas (action/observation/state)
src/envs/narada/client.py	Python WebSocket client for agents
training/narada_grpo.ipynb	GRPO training notebook (Colab)
inference.py	Benchmark script (Groq or HF backend)

---

OpenEnv Validation

openenv validate https://krishvenky-narada-env.hf.space

All scores strictly in (0.01, 0.99). [END] line always includes score= field.

---

Baseline Benchmark

Zero-shot evaluation via inference.py (Groq backend, no fine-tuning). Two single-episode samples per model show the core problem: zero-shot LLMs are inconsistent and frequently collapse into summarise_trail loops instead of navigating the graph.

Note: the exact scores below are reference runs from before reward hardening. Re-run inference.py before final submission so the README plots and tables match the current reward mapping.

llama-3.3-70b-versatile (reference runs)

Task	Run 1 Score	Run 2 Score	Notes
monogenic	0.990	0.433	Run 1: solved in 4 steps. Run 2: hit summarise_trail loop after step 3
oligogenic	0.500	0.240	Run 1: WS disconnect mid-episode. Run 2: full summarise_trail timeout
phenotype_mismatch	0.060	0.225	Run 1: looped on wrong gene 9× before giving up. Run 2: pure timeout

Multi-model comparison (zero-shot, all 3 tasks, single run each)

Model	monogenic	oligogenic	phenotype_mismatch	Behavior
llama-3.3-70b-versatile	0.990	0.500	0.060	Hops graph; inconsistent
llama-3.1-8b-instant	0.310	0.425	0.310	Hops but flags wrong variant
mixtral-8x7b-32768	0.233	0.220	0.225	Full summarise_trail timeout
gemma2-9b-it	0.233	0.220	0.225	Full summarise_trail timeout

The pattern is clear: large frontier models (llama-3.3-70b) occasionally navigate the graph correctly but are inconsistent (0.990 vs 0.433 on the same task across runs). Mid-size models (llama-3.1-8b) attempt navigation but misfire on the final flag. Smaller models (mixtral, gemma2) collapse entirely to the summarise_trail loop and never issue a flag_causal. Fine-tuning on the graph-navigation reward signal is intended to make correct phenotype → gene → variant chaining the default, not a lucky outcome that only large models achieve occasionally.

Target post-GRPO (Qwen3-1.7B): consistent flag accuracy > 50% on monogenic, causal path coverage > 60%.

To switch backends: set GROQ_API_KEY for Groq, or HF_TOKEN for HF Inference Router. The script auto-detects which to use.

---

Training

Training notebook: training/narada_grpo.ipynb (Colab, Unsloth + HF TRL GRPO)

Base model: Qwen/Qwen3-1.7B — fits on a free T4 (4-bit quantised, 17.4M trainable LoRA params)

Tasks trained in curriculum order: monogenic → oligogenic → phenotype_mismatch

Architecture

Multi-step outcome GRPO — the model generates a complete 3–5 action diagnostic plan per prompt. The full plan is executed in the environment and the terminal reward (correct flag / wrong flag / timeout) becomes the training signal. This gives 10× more reward variance than single-step GRPO:

Approach	Typical reward range	reward_std	Learning
1-step (old)	0.47–0.56	~0.03	Slow
Multi-step plan	0.28–0.99	~0.20–0.35	Strong

Async-parallel reward — all G=8 completions are evaluated concurrently via asyncio.gather. Each completion's plan runs as an independent WebSocket session. Total overhead ≈ one episode round-trip, not 8×.

Curriculum learning — monogenic cases establish the basic hop→flag behaviour before oligogenic (multi-objective) and phenotype_mismatch (decoy resistance) add harder constraints.

Milestone reward — environment gives a +0.10 bonus the first time the agent visits the actual causal gene node. Creates a two-stage reward landscape: find the gene (+0.10) → flag the variant (+1.0).

Key training config

Param	Value	Why
num_generations	8	More completions per prompt → higher reward_std
temperature	1.1	Forces diverse hop targets within a group
max_completion_length	800	Fits 4–5 JSON action blocks
N_SEEDS_PER_TASK	40	120 total prompts (was 60) — more case diversity
LR	5e-6	Conservative to avoid catastrophic forgetting on 1.7B

---

Multi-Model Benchmark

Run zero-shot comparison across sub-5B models (all runnable on T4):

HF_TOKEN=hf_... python benchmark.py --n_seeds 3

Models compared: Qwen3-0.6B, Qwen3-1.7B, Qwen3-4B, Llama-3.2-3B, Phi-3.5-mini

Results are saved to benchmark_results.json after each run.

---

Graph Export (Neo4j)

Export the 55K-node knowledge graph to Cypher for Neo4j Aura (free tier) or Desktop:

PYTHONPATH=src/envs python scripts/export_neo4j.py
# Generates neo4j_nodes.cypher + neo4j_rels.cypher
# Import into Neo4j Browser, then:
# MATCH (n:NaradaNode) RETURN n LIMIT 100

The live environment also exposes a subgraph JSON endpoint for D3.js visualization:

GET /graph/subgraph?node_id=GENE:BRCA2&depth=2