Spaces:
Sleeping
Sleeping
| # Narada — Architecture & Reward Design | |
| ## Overview | |
| Narada is a multi-step reinforcement learning environment where an LLM agent plays a clinical geneticist navigating a gene-disease knowledge graph. The agent observes a patient's HPO phenotype terms and a pool of candidate variants, then navigates the graph to identify the causal variant(s). | |
| --- | |
| ## Knowledge Graph | |
| Built at runtime from two public biomedical datasets: | |
| | Source | File | Content | | |
| |--------|------|---------| | |
| | HPO | `data/hp.obo` | 19,389 phenotype terms + hierarchy | | |
| | ClinVar | `data/clinvar_pathogenic.tsv` | 92,000 high-confidence pathogenic variants (GRCh38, criteria provided / expert panel, deduplicated by AlleleID) | | |
| ### Node types | |
| | Type | ID format | Count | Meaning | | |
| |------|-----------|-------|---------| | |
| | `phenotype` | `HP:XXXXXXX` | ~2,400 | HPO terms (catalog + 5-level ancestors) | | |
| | `disease` | `DIS:<slug>` | ~28,000 | Disease names from ClinVar PhenotypeList | | |
| | `gene` | `GENE:<symbol>` | ~3,268 | Gene symbols | | |
| | `variant` | `VAR:<AlleleID>` | ~43,634 | Individual ClinVar variants | | |
| | `pathway` | `PATH:<name>` | ~14 | Coarse pathway group (cardiac, neurological, …) | | |
| **Total: 55,201 nodes, 70,741 edge-pairs** | |
| ### Edge structure | |
| All edges are undirected (bidirectional). Key connections: | |
| ``` | |
| phenotype ↔ phenotype (HPO parent-child hierarchy) | |
| phenotype ↔ disease (catalog wiring + inverted word-index matching) | |
| disease ↔ gene (via ClinVar PhenotypeList) | |
| gene ↔ variant (one gene → many variants) | |
| gene ↔ pathway (PATHWAY_MAP classification) | |
| ``` | |
| ### Graph build performance | |
| The `_add_hpo_nodes()` method uses an O(N+M) inverted word index instead of O(N×M) brute-force matching. Cold build time: ~2 seconds. Loaded once at server startup, shared across all sessions. | |
| --- | |
| ## Session Lifecycle | |
| ``` | |
| Client Server | |
| │ │ | |
| ├─── WebSocket connect ─────────►│ NaradaEnvironment() created (per session) | |
| │ │ | |
| ├─── {"type":"reset", ...} ─────►│ generate_case() → PatientCase | |
| │◄── {"type":"observation"} ─────┤ _build_observation() → StepResult(reward=0.0) | |
| │ │ | |
| ├─── {"type":"step", action} ───►│ _dispatch_action() → raw reward | |
| │◄── {"type":"observation"} ─────┤ reward mapped to (0.01, 0.99) | |
| │ (repeat) │ | |
| │ │ | |
| ├─── flag_causal(VAR:xxxxx) ────►│ _compute_terminal_reward() + _overseer_score() | |
| │◄── {"type":"observation", │ done=True, final reward returned | |
| │ done:true, reward:X} ──────┤ | |
| │ │ | |
| └─── WebSocket close ───────────►│ session destroyed | |
| ``` | |
| **WORKERS=1** is enforced. All state is in-memory per WebSocket connection; no shared state between sessions. | |
| --- | |
| ## Task Tiers | |
| ### monogenic (easy) | |
| - 1 causal gene, 3–4 HPO phenotypes, 5–8 candidate variants | |
| - Max 15 steps | |
| - Tests: basic phenotype → disease → gene → variant chain reasoning | |
| ### oligogenic (medium) | |
| - 2 causal genes, 5–7 HPO phenotypes, 10–15 candidates | |
| - Max 25 steps | |
| - Tests: multi-objective tracking — the agent must flag both contributing variants; correct intermediate flags are recorded and the episode continues until all are found or a wrong variant is flagged | |
| ### phenotype_mismatch (hard) | |
| - 1 causal gene (cardiac/neurological) | |
| - A high-pathogenicity BRCA1/BRCA2/TP53 frameshift variant in the candidate pool as a decoy | |
| - Max 20 steps | |
| - Tests: causal discipline — pathogenicity score alone is not sufficient; the variant must match the patient's phenotype | |
| --- | |
| ## Reward Design | |
| ### Step-level rewards | |
| | Action | Reward | Condition | | |
| |--------|--------|-----------| | |
| | `hop` | +0.15 | Target node is on the causal path | | |
| | `hop` | −0.05 | Target node is off-path | | |
| | `hop` | −0.10 | Hallucinated hop (node exists but not connected) | | |
| | `backtrack` | +0.05 | Previous node was off-path (recovering) | | |
| | `backtrack` | −0.05 | Previous node was on-path (wrong direction) | | |
| | `request_lab` | −0.10 | Always penalised | | |
| | `summarise_trail` | 0.00 | Neutral | | |
| | Per-step | −0.01 | Applied to every action (efficiency pressure) | | |
| ### Terminal rewards | |
| | Outcome | Reward | | |
| |---------|--------| | |
| | Correct flag (monogenic/mismatch) | +1.0 | | |
| | Correct flag + timing bonus (step < 10) | +1.2 | | |
| | Progress per correct variant (oligogenic) | `0.5 / total` non-terminal | | |
| | All oligogenic variants flagged | `(correct/total) × 0.5` + timing bonus | | |
| | Flagged decoy in mismatch task | −0.5 | | |
| | Wrong flag | −0.5 | | |
| | Timeout (no flag) | `-0.25 + min(0.2, trail_size / max_steps × 0.25)` | | |
| ### Overseer score (additive, 0.0–0.3) | |
| Added only to successful terminal rewards. Computed locally without an LLM call: | |
| | Criterion | Effect | | |
| |-----------|--------| | |
| | Hallucinated hops | −0.05 each | | |
| | Visited < 3 unique nodes | −0.10 | | |
| | Visited causal gene node | +0.05 | | |
| ### Score mapping | |
| Rewards are kept as signed raw values internally, then mapped to the **open interval (0.01, 0.99)** before returning to the client. This preserves the ordering between penalties, neutral moves, and successes while satisfying the OpenEnv validator. `math.isfinite()` guards against NaN/inf. | |
| --- | |
| ## Three-Agent System | |
| ### Detective (trainable) | |
| - Qwen3-1.7B fine-tuned with GRPO | |
| - Navigates the graph, flags the causal variant | |
| - Trained via `training/narada_grpo.ipynb` | |
| ### Overseer (local) | |
| - Heuristic, no LLM call: reads only the trail and hop counters, not the free-form `reasoning` string | |
| - Penalises hallucinated hops and trivial exploration | |
| - Adds 0.0–0.3 only to successful terminal rewards | |
| ### Adversary (exploratory, WIP) | |
| - Intended to generate harder cases targeting Detective failure patterns | |
| - Reliable curriculum generation from agent error logs is an open research problem | |
| - Current implementation falls back to random seed selection | |
| --- | |
| ## Action Space | |
| ```python | |
| class NaradaAction(BaseModel): | |
| action_type: str # hop | flag_causal | backtrack | request_lab | summarise_trail | |
| node_id: Optional[str] # required for hop | |
| variant_id: Optional[str] # required for flag_causal (format: VAR:12345) | |
| test_type: Optional[str] # for request_lab | |
| reasoning: str # agent's stated rationale (logged by Overseer) | |
| ``` | |
| ## Observation Space | |
| ```python | |
| class NaradaObservation(BaseModel): | |
| step: int | |
| max_steps: int | |
| task_type: str | |
| current_node: GraphNode | |
| trail: List[GraphNode] # last 10 visited nodes | |
| patient_phenotypes: List[str] # HPO IDs | |
| phenotype_names: List[str] | |
| candidate_variants: List[Variant] # 5–15 variants | |
| step_reward: float | |
| cumulative_reward: float | |
| done: bool | |
| info: Dict[str, Any] # ground_truth_hint revealed on done=True | |
| ``` | |
| --- | |
| ## OpenEnv Compliance | |
| | Endpoint | Method | Purpose | | |
| |----------|--------|---------| | |
| | `/health` | GET | Liveness check | | |
| | `/metadata` | GET | Environment name + description | | |
| | `/schema` | GET | Pydantic JSON schemas for action/observation/state | | |
| | `/mcp` | POST | JSON-RPC 2.0 tool discovery | | |
| | `/reset` | POST | Start new episode | | |
| | `/step` | POST | Take action | | |
| | `/state` | GET | Current episode metadata | | |
| | `/ws` | WebSocket | Primary transport (persistent session) | | |
| Validation: `python -m openenv.cli validate --url https://krishvenky-narada-env.hf.space` → **6/6 passed** | |