README.md

metadata

title: Drug Target Validation Environment
sdk: docker
pinned: false
app_port: 8000
tags:
  - openenv
  - reinforcement-learning
  - drug-discovery
  - pharma

🧬 DrugEnv — Drug Target Validation Environment

DrugEnv — an OpenEnv RL environment that teaches LLMs to do computational drug-target validation.

This repository implements an OpenEnv-compatible reinforcement learning environment in which an agent acts as a computational drug discovery scientist. Given a proposed drug target (gene / protein) and a disease context, the agent must investigate target viability by issuing simulated bioinformatics, clinical, and experimental queries, and finally submit a calibrated go / no-go validation report with a confidence score.

The environment is designed as a partially observable Markov decision process (POMDP) with:

• a hidden ground-truth TargetProfile (expression, druggability, selectivity, toxicity, clinical precedent)
• noisy database / assay outputs governed by DataQualityState
• a single unified experimental credit budget per episode
• visible task metadata, dossier of accumulated findings, and step history
• dense step-wise reward plus terminal reward for decision quality and evidence coverage

Why drug target validation?

Roughly 90% of drug development programs fail in clinical trials, and a large fraction of failures trace back to mistakes during target validation: targets that are not actually disease-driving, are undruggable, lack selectivity, or have hidden toxicity. The cost of progressing a single bad target through Phase III can run into the billions of dollars. Even modest improvements in early-stage decision quality therefore translate into enormous savings and faster cures.

This environment lets you train and benchmark agents on exactly that bottleneck: acquiring the right evidence cheaply and submitting a well-calibrated go / no-go.

How it works

At a high level, each episode looks like this:

1. reset() selects a drug-target-validation scenario and seeds the simulator.
2. The agent receives a ValidationObservation describing the target, indication, remaining credits, accumulated dossier, and step history.
3. The agent submits a DrugTargetAction such as query_expression, druggability_screen, off_target_screen, or submit_validation_report.
4. The rule engine checks credit budget, redundancy, and ordering prerequisites.
5. The transition engine deducts credits and asks the output generator to simulate evidence from the hidden TargetProfile.
6. The reward computer scores the step for novelty, reasoning coherence, credit efficiency, and rule compliance.
7. The environment returns a new observation with an updated EvidenceDossier, latest output, violations, and reward.
8. The episode ends when the agent submits a validation report, exhausts credits, or hits the step limit.

The core mental model

Hidden state

The simulator maintains a FullLatentState that the agent never sees directly:

• TargetProfile — true expression level / tissue specificity / disease over-expression, druggability score, binding-pocket quality, selectivity ratio, off-target genes, toxicity profile, clinical precedent, expected in-vitro and in-vivo behaviour, plus the hidden correct_decision, true_viability_score, key_evidence_dimensions, and any misleading_signals.
• DataQualityState — noise level, false-positive rate, false-negative rate, database coverage.
• CreditState — total / used / remaining experimental credits.
• ValidationProgress — boolean flags for which evidence dimensions have been investigated and whether a report has been submitted.

Visible state

The agent only sees ValidationObservation, which includes:

• target_gene, disease_context, indication
• credits_remaining / credits_total
• dossier — running EvidenceDossier of expression / protein / clinical / safety / literature / experimental findings, plus any flagged_red_flags
• pipeline_history — list of past actions and their summary outputs
• latest_output — typed IntermediateOutput from the most recent step
• rule_violations and step_reward_breakdown for the last step

Action space

Category	Action	Cost (credits)
Expression & omics	query_expression, differential_expression, pathway_enrichment, coexpression_network	2
Protein & structure	protein_structure_lookup, binding_site_analysis, druggability_screen	3
Protein & structure	protein_interaction_network	2
Clinical & safety	clinical_trial_lookup, toxicity_panel, off_target_screen, patient_stratification	3
Literature	literature_search, evidence_synthesis, competitor_landscape	1
Experimental	crispr_knockout, biomarker_correlation	4 / 3
Experimental	in_vitro_assay	5
Experimental	in_vivo_model	8
Meta	flag_red_flag, request_expert_review	0 / 1
Terminal	submit_validation_report	0

submit_validation_report carries two extra fields: final_decision ("go" or "no_go") and confidence in [0, 1]. The episode ends as soon as the report is submitted.

Reward function

Every step receives a decomposed reward:

R_t = evidence_novelty_bonus

• reasoning_coherence_bonus
• credit_efficiency_penalty
• rule_violation_penalty
• [φ(s_{t+1}) − φ(s_t)]

When the episode ends, a terminal reward is added:

R_T = 0.40 * decision_accuracy

• 0.35 * evidence_coverage
• 0.15 * credit_efficiency
• 0.10 * reasoning_coherence

Where:

• decision_accuracy — 1.0 if the final go / no-go matched the hidden correct_decision, scaled by 2 * |confidence - 0.5| so a confidently correct answer is fully rewarded and a confidently wrong answer is fully penalised.
• evidence_coverage — fraction of the scenario's key_evidence_dimensions (e.g. expression, druggability, off_target, clinical, in_vitro) that the agent actually investigated.
• credit_efficiency — 1 − redundant_calls / total_calls.
• reasoning_coherence — fraction of actions whose soft prerequisites (e.g. expression before toxicity, in_vitro before in_vivo) were satisfied.

Hard penalties are applied for: submitting without any evidence, submitting without a decision or confidence, and exhausting credits without ever submitting a report.

Curated scenarios

Name	Difficulty	Correct decision	Why it's interesting
egfr_nsclc_viable	easy	go	Clear viable target — expression + druggability alone are sufficient.
kras_pdac_borderline	medium	go	Historically undruggable; recent inhibitor literature is decisive.
cd33_aml_misleading	hard	no_go	Naive expression query says "go", but off-target + toxicity + clinical reveal the
right answer.
tp53_solid_tumors_clear_fail	easy-medium	no_go	Druggability check alone is sufficient.
ptpn11_juvenile_mml_complex	very hard	go	Requires binding_site_analysis(include_allosteric=True), off-target work,
patient stratification, and an in-vitro assay.

The procedural generator (server/tasks/procedural_generator.py) layers on additional easy / medium / hard scenarios sampled from a pool of 20 real cancer targets and 8 cancer indications.

Setup

# 1. Install dependencies (env runtime only)
pip install -e .

# 2. Or install with training extras (torch + transformers + trl + peft pinned to working set)
pip install -e .[train]

# 3. Run the environment server
PYTHONPATH=. python -m server.app
# server is now available at http://localhost:8000

The legacy uv sync workflow still works if you have uv.lock checked
in locally; the editable pip install path above is the primary
supported route.

## Talking to the environment

from client import DrugTargetEnv
from models import DrugTargetAction

with DrugTargetEnv(base_url="http://localhost:8000") as env:
    result = env.reset()
    print(result.observation.target_gene, "/", result.observation.indication)

    result = env.step(DrugTargetAction(
        action_type="query_expression",
        parameters={"database": "GTEx"},
        reasoning="Establish tissue baseline",
    ))
    print(result.observation.latest_output.summary)

    result = env.step(DrugTargetAction(
        action_type="submit_validation_report",
        reasoning="Sufficient evidence for go",
        final_decision="go",
        confidence=0.85,
    ))
    print("done:", result.done, "reward:", result.reward)

## Running the baseline agent

PYTHONPATH=. python run_agent.py

The script writes a live JSON snapshot to _dashboard_state.json after every step so you can watch the agent's progress.

The agent model is configurable via env vars:

- RUN_AGENT_MODEL_ID (HF model id)
- RUN_AGENT_LORA_ADAPTER (optional LoRA adapter repo/path)

Example (base + LoRA):

RUN_AGENT_MODEL_ID=Qwen/Qwen2.5-0.5B-Instruct \
RUN_AGENT_LORA_ADAPTER=anugrahteesdollar/drugenv-qwen25-05b-lora \
PYTHONPATH=. python run_agent.py

## Quick Training Update (Colab LoRA Adapter)

HF compute credits ran out before completing a full GRPO run, so we ran a short adapter-only LoRA fine-tune on free Colab T4 to
demonstrate measurable improvement.

- Base model: Qwen/Qwen2.5-0.5B-Instruct (4-bit)
- LoRA adapter (HF Hub): anugrahteesdollar/drugenv-qwen25-05b-lora
- Training: ~300 steps, LoRA r=8, alpha=16, dropout=0.05, target modules q_proj, v_proj
- Outcome: improved DrugEnv action JSON parameter schema adherence.

Example (step 0 prompt, EGFR / NSCLC):

Before (base):
{"action_type":"query_expression","parameters":{"term":"EGFR","dossier_id":"NSCLC"},"reasoning":"Query for the presence of EGFR
in the given dataset."}

After (base + LoRA):
{"action_type":"query_expression","parameters":{"database":"GTEx"},"reasoning":"Start by querying the expressible space across
the entire expression space."}

## Reproduce

Three commands cover the env-locally / training-locally / training-on-Space paths:

# 1. Env locally (CPU is fine — the env itself is dependency-light)
pip install -e . && PYTHONPATH=. python -m server.app
# → http://localhost:8000  (also at https://huggingface.co/spaces/anugrahteesdollar/drugenv when deployed)

# 2. Training locally (single GPU, vanilla GRPO)
pip install -e .[train]
PYTHONPATH=. python -m training.training_script \
    --model-id Qwen/Qwen2.5-3B-Instruct \
    --evidence-dir evidence \
    --output-dir runs/grpo-output

# 3. Training on a Hugging Face Space (H200 single-GPU)
#    Push space/training/ to anugrahteesdollar/drugenv-trainer, set PUSH_REPO + HF_TOKEN
#    in the Space variables, then POST /train.
#    → https://huggingface.co/spaces/anugrahteesdollar/drugenv-trainer

The trainer Space's FastAPI control panel (space/training/app.py)
streams a live evidence dashboard while training runs — per-step
training curve, mid-training checkpoint progression, and a before /
after summary card. Default expected hardware: H200 single-GPU
(h200x1); H200 is ≈ 4× A100 throughput, ~$0.05–0.10 per step on
Qwen2.5-3B-class GRPO.

An optional SFT warm-start (training/sft_warmstart.py) is
controlled via the SFT_WARMSTART env var on the Space (default on).
It collects oracle trajectories on the curated scenario library, SFTs
the base model with a small LoRA, and hands the merged checkpoint to
GRPO so the policy starts with a non-zero prior over correct
trajectories.

## Baseline scores

| Difficulty bucket | Random policy | Heuristic policy | Trained Qwen2.5-3B |
|---|---|---|---|
| Easy (egfr_nsclc_viable) | filled in after first training run | filled in after first training run | filled in after first
training run |
| Medium (kras_pdac_borderline) | filled in after first training run | filled in after first training run | filled in after first
training run |
| Hard (cd33_aml_misleading) | filled in after first training run | filled in after first training run | filled in after first
training run |

The trainer Space writes the populated table to
evidence/before_after_metrics.json automatically on every run.

## Evolution note

The deployment scaffolding in this repository — the trainer Space
control panel, the live training-evidence callback, the SFT warm-start
script, and the working dependency pin set — was originally validated
against a particle-physics-themed prototype and then carried forward
when we pivoted to drug discovery. The simulator, scenarios, action
space, reward function, and rules engine are all drug-domain native;
the inheritance is exclusively in the training and evaluation
scaffolding.