drug-tool-rl

A multi-turn, tool-using RL environment for drug-discovery agents. The model answers biomedical questions about drug targets, disease mechanisms, protein function, molecular interactions, and candidate molecules by calling biomedical retrieval tools, then synthesizing a grounded final answer.

Task

Each example is a biomedical prompt paired with a reference tool-use trace from lokahq/drug-tool-sft. Rollouts run an evidence-first loop: call tools, inspect results, reformulate empty searches, and end with one concise synthesis.

Tools

Seven retrieval tools are always available:

search_pubmed - PubMed literature via NCBI eUtils.
search_geo_datasets - NCBI GEO expression datasets.
search_kegg_by_keyword / search_kegg_by_id - KEGG pathways, diseases, drugs, and compounds.
search_uniprot_by_name / search_uniprot_by_id - UniProt protein records and sequences.
search_string_interactions - STRING protein-protein interaction networks.

Three NVIDIA NIM structural-biology tools are optional when include_nvidia_tools=True:

fold_protein - ESMFold structure prediction with rollout-local PDB storage.
dock_ligand - DiffDock protein-ligand docking from raw PDB or a stored fold artifact.
molmim_generate - MolMIM molecule generation or optimization from a seed SMILES.

The environment caches identical tool calls within a rollout and shares rollout-local structural artifacts across stateful tools.

Reward

The only supported reward profile is grounded_tool_synthesis.

signal	weight	description
`key_fact_grounding_reward`	0.33	Fraction of final-answer key facts grounded in tool evidence, down-weighted when the model ignores tool failures.
`tool_result_success_reward`	0.19	Penalizes tool errors, warnings, no-result hits, invalid stateful-tool arguments, duplicate calls, and UniProt entity mismatches.
`tool_selection_reward`	0.14	Tool multiset F1 plus call-order LCS similarity against the reference trace.
`tool_argument_reward`	0.12	Best-match argument-key/value similarity for selected tools against the reference.
`tool_activation_reward`	0.05	Rewards emitting at least one parser-compatible tool call.
`synthesis_completion_reward`	0.05	Rewards a non-truncated final synthesis after tool use.
`efficiency_reward`	0.04	Matches the reference trace's tool-call budget, capped at 8.
`concise_final_answer_reward`	0.04	Gives full credit to concise final answers, decaying by length.
`non_truncated_reward`	0.04	Rewards rollouts that do not truncate.

Additional parser, grounding, hallucination, error, warning, and rate-limit metrics are logged at weight 0 for analysis.

Contract

from drug_tool_rl import load_environment

env = load_environment(
    dataset_name="lokahq/drug-tool-sft",
    dataset_config="openrouter_trinity_v1",
    split="train",
    num_examples=-1,
    max_turns=8,
    include_nvidia_tools=False,
    nvidia_api_key_var="NVIDIA_API_KEY",
    reward_profile="grounded_tool_synthesis",
)

Unsupported reward_profile values raise ValueError.

Environment Variables

NVIDIA_API_KEY or NVCF_RUN_KEY - required only when include_nvidia_tools=True.
NCBI_API_KEY - optional; raises NCBI eUtils rate limits for PubMed/GEO.

Retrieval tools hit public NCBI, KEGG, UniProt, and STRING endpoints with bounded retries. NVIDIA NIM calls are capped by NVIDIA_CALL_DEADLINE so a slow streaming call surfaces as a tool error instead of hanging the eval.

Dependencies

datasets, httpx, rdkit, tenacity, and verifiers.

Published environment handle: lokahq/drug-tool-rl@1.