ARCHITECTURE.md

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

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

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