{"data":{"kind":"file","path":"README.md","version_id":"x6wpt17ajntpwdhj9mqq32od","entry":{"name":"README.md","path":"README.md","is_directory":false,"size":13063,"modified_at":"2026-10-19T04:23:43.245000","content_hash":"fcb6b969cc814b566427f93b8e06b1b9586698075fc8c4d3511d4caa667319ef"},"entries":[],"content":"# dxenv-diagnostician\n\nCost-aware diagnosis over synthetic patients. The agent works one case at a time: it reads a\ncase sheet (demographics, complaint, vitals, history, allergies), orders tests that cost\nbudget, may prescribe, and ends the case by reporting a **probability distribution** over 149\nconditions. It is scored by a strictly proper rule, so its best strategy is to report\nhonest, calibrated beliefs and to buy a test only when the information is worth the cost.\n\nThe patients are generated from a hand-written generative model, which is what makes a\n**Bayes-optimal ceiling** computable for every patient. This environment measures how well an\nagent learns a synthetic diagnostic mapping under cost. It does **not** measure clinical\nskill. The disease likelihoods are invented, and no real patient data (Synthea or otherwise)\nis used.\n\n## Task and interaction\n\nEach case uses a global menu, identical for every patient, of **73 tests/panels** and\n**72 treatments**, plus the 149 condition labels. A per-patient menu would leak the diagnosis.\nThe agent takes one action per turn as a JSON object:\n\n| action | effect |\n|---|---|\n| `order_test` + `prediction` | reveals the test's analytes and costs budget. The agent must first commit to a coarse prediction of the headline result (`low`/`normal`/`high`, or `normal_categorical`/`abnormal_categorical`) |\n| `prescribe` | records a treatment |\n| `diagnose` | terminal: up to 16 `{condition, probability}` pairs; unassigned mass is spread evenly over the unnamed labels |\n| `abstain` | terminal |\n\nMechanics:\n- Repeat orders are free and return the cached result.\n- An order the remaining budget can't cover is refused, and the turn is still spent.\n- The budget is drawn per patient from {10, 25, 50, 100, 200}, weighted .15/.25/.30/.20/.10.\n- An episode ends on diagnose or abstain, or after 20 turns.\n\n**The policy is stateless.** Each turn is a fresh single-turn model call. Its prompt is\nrebuilt from the typed observation, which has no field that can carry the label. No chat\nhistory is carried over, and each turn is recorded as its own trace.\n\n**Output format (differs from the source).** On the original cluster, a JSON grammar made\ninvalid output impossible. Here the same action grammar is sent as a JSON-schema\n`response_format` on the diagnostician's sampling args, and the wire format is also spelled\nout in the system prompt. Providers that honour structured output (OpenAI-compatible APIs,\nvLLM, Ollama) are then constrained the same way. Ones that ignore it fall back to the prompt.\nA reply wrapped in a markdown fence or a `<think>` block is unwrapped, which changes no\nfield. A reply that still isn't one valid JSON object on the menu ends the case as\n`decode_failure`. The source scores that as a non-decision (−0.05), the same as running out\nof turns. Without the schema, Qwen2.5-7B failed to decode on 5 of 5 smoke episodes. With\nit, 0 of 5 failed.\n\n## Reward\n\nOne reward per episode, `diagnostic`, recorded on **every turn's trace**:\n\n```\nR = diagnosis + test_cost + turn_penalty + treatment + verify\n```\n\n| term | formula | notes |\n|---|---|---|\n| diagnosis | `w(c_true) · [(1 − 1/149) − Σ_i (p_i − 1[i = c_true])²]` | shifted Brier: a uniform report scores 0, a confident correct one about +w, a confident wrong one about −w. `w` ∈ {1.0, 1.8, 3.2, 6.0} by urgency of the **true** condition |\n| abstain / out of turns / decode failure | diagnosis term = −0.05 | strictly worse than reporting the prior |\n| test_cost | `−0.004 · price` for each charged order | repeat and refused orders cost nothing |\n| turn_penalty | `−0.02 · n_turns` | |\n| treatment | `0.5 · coherence · mass_on_truth + 1.0 · correctness − 4.0 · Σ contraindication scales` | see *Coherence fix* below |\n| verify | `0.25 · order cost · (1[prediction right] − P_prior(bucket))` | zero in expectation at chance, always smaller than the order's own cost |\n\nProperties, each checked by `tests/test_exploits.py`:\n- Every test is net negative on its own, even with a correct prediction. Tests only pay\n  through a better diagnosis.\n- Constant predictions earn nothing.\n- Naming severe conditions doesn't pay.\n\nThere is no judge. The reward is computed from hidden ground truth.\n\nMetrics (logged, not rewarded):\n- the separate terms;\n- `source_reward`, the original reward with no coherence fix, for comparison with the\n  original results;\n- `mass_on_truth`, `top1_correct`, `n_tests`;\n- the termination flags `diagnosed`, `abstained`, `out_of_turns`, `decode_failure`;\n- `n_contraindications`, and `final_step` (1.0 on the episode's last trace).\n\n### Coherence fix (this package differs from the original)\n\nThe original treatment term paid `0.5 · coherence`. Coherence is how well the prescriptions\nfit the agent's **own** top diagnosis, whether or not that diagnosis is right. A policy that\norders nothing, reports the population prior, and always prescribes amlodipine (first-line\nfor the most common condition) averaged **+0.46** under the original reward. That beats every\ntrained model in the original work, with no information about the patient. Here the term\nis gated on the probability the report puts on the truth:\n\n```\noriginal:  0.5 · coherence\nthis env:  0.5 · coherence · mass_on_truth\n```\n\n| policy (40 fixture patients, mean R) | original reward | this env |\n|---|---|---|\n| report the prior, order nothing (floor) | −0.019 | −0.019 |\n| prior + one coherent prescription | **+0.464** | −0.028 |\n| same, with 0.3 / 0.6 / 0.9 extra mass moved to the prescribed condition | +0.350 / −0.053 / −0.743 | −0.134 / −0.527 / −1.208 |\n\nThis is the only change to the reward arithmetic. It is applied as a delta on the original\nscore, and the original score is logged as `source_reward`.\n`tests/test_reward_parity.py` replays 62 stored episodes from the source repo and checks\nthree things:\n- the trajectories match;\n- `source_reward` equals the stored reward;\n- the Hub reward equals `source_reward` plus the adjustment, and the adjustment is exactly 0\n  whenever an episode lacks a prescription or a diagnosis.\n\n## Why the reward can be trusted\n\n- **No label in the observation, by type.** The observation model has no field that can\n  carry a condition. The display strings come from the catalog, never from the patient\n  record. The test suite also checks that no case sheet contains its own label.\n- **Fixed menu.** Every patient sees the same tests, treatments and labels.\n- **Proper scoring.** The Brier rule is strictly proper, and the severity weight is keyed on\n  the true condition.\n- **Tests never pay on their own.** This is enforced arithmetically by the source repo's\n  `validate_reward_config`.\n- **A ceiling exists.** The generative model is known, so the source repo computes a\n  per-patient Bayes-optimal ceiling. An episode scoring above it would indicate a leak.\n\n## Data and license\n\nThe data is procedural, and nothing is downloaded or bundled. `load()` regenerates the\nsource repo's frozen 20,000-patient corpus (seed 20260901), which takes about 12 s once per\nprocess. It then re-cuts the split and verifies it against the committed eval hash.\n\n| split | patients | notes |\n|---|---|---|\n| `train` | 14,854 | |\n| `eval` (default) | 3,713 | frozen by hash in the source repo |\n| `holdout` | 1,433 | rheumatologic and obstetric conditions only, which never appear in train or eval |\n\nThe taxonomy, test catalog, costs and treatment tables are small YAML files inside the\n`dxenv` package.\n\n## Configuration\n\n| flag | default | meaning |\n|---|---|---|\n| `--env.taskset.split` | `eval` | `train`, `eval` or `holdout` |\n| `--env.taskset.num-tasks` | whole split | first n patients |\n| `--env.taskset.budget` | sampled | fixed test budget instead of the per-patient draw |\n| `--env.taskset.max-turns` | 20 | turn cap (the source GRPO curriculum started at 12) |\n| `--env.diagnostician.model` / `.sampling.*` | the run's model | the diagnosing model |\n\nReward weights are fixed to the source values. Only the coherence term differs.\n\n## Quickstart\n\n```bash\nprime env install manasganti/dxenv-diagnostician\nuv run vf-eval dxenv-diagnostician -n 5 -r 1 -m <model>\nuv run vf-eval dxenv-diagnostician -n 50 -r 1 -m <model>\npython scripts/episode_summary.py outputs/dxenv-diagnostician--<model>--null--<id>\n```\n\nEach turn's prompt is about 3k tokens: the case sheet, the full menu and the label list.\nAllow about 1k completion tokens per turn, since a diagnosis that names 16 conditions is\nthe longest reply.\n\n**Reading eval output.** vf-eval's headline is a mean over *traces*, and an episode has one\ntrace per turn, so long episodes count more. `scripts/episode_summary.py` gives per-episode\nnumbers. For training, every turn's trace carries its episode's return, so group by episode\nwhen computing advantages.\n\n## Baseline results (vf-eval, this package)\n\nPer-episode means from `scripts/episode_summary.py`. The run: `vf-eval -n 50 -r 1`, T = 1.0,\n1024 max tokens per turn, 20-turn cap, budgets drawn per patient. The model was served\nunquantized (bf16) by vLLM 0.8.5 on one A100 and took 9 minutes.\n\n| model | split | episodes | reward (mean ± SE) | source_reward | top-1 | tests / episode | decode failures |\n|---|---|---|---|---|---|---|---|\n| Qwen2.5-7B-Instruct | eval | 50 | −0.602 ± 0.015 | −0.602 | 0.00 | 4.26 | 0 |\n\nThe untrained 7B model doesn't commit to a diagnosis:\n- **94%** of its episodes ran out of turns, averaging 19.7 of 20 turns;\n- 6% diagnosed;\n- the mean probability it put on the true condition was 0.005.\n\nMost of the score is the turn penalty (−0.40) and test costs (−0.14), which puts it below the\n−0.019 floor of reporting the prior and stopping. It never prescribed, so the coherence fix\ndidn't change its score (`source_reward` = reward). The original work saw the same behaviour\nfrom the same model (see below), and it's the headroom RL is meant to close.\n\n**Model-free reference policies** on the same first 50 patients and budgets a `-n 50` run\nsees. These are the source repo's heuristic policies, scored by this package's reward\n(`python scripts/reference_baselines.py --split <split> -n 50`). The Bayes policies use the\ntrue generative model, so they are references, not competitors.\n\n| policy | eval: reward (± SE) | eval: top-1 | holdout: reward (± SE) | holdout: top-1 | tests / episode |\n|---|---|---|---|---|---|\n| report the prior, order nothing (floor) | −0.019 ± 0.007 | 0.08 | −0.019 ± 0.006 | 0.00 | 0 |\n| vitals-only Bayes, order nothing | +1.221 ± 0.242 | 0.60 | +0.521 ± 0.168 | 0.58 | 0 |\n| greedy Bayes (up to 6 tests) | +1.764 ± 0.272 | 0.78 | +0.783 ± 0.199 | 0.76 | 5.7 |\n| expected Bayes ceiling (diagnosis term, all evidence free) | 2.244 | | 1.397 | | |\n\n## Original results (separate harness)\n\nThese come from the source repo's own harness on VT ARC (vLLM with grammar-constrained\ndecoding, Qwen2.5-7B-Instruct). They were **not** produced by this package or by vf-eval, and\nthey are **scored under the original reward, before the coherence fix**. They are not Hub\nbaselines.\n\nGate B protocol: 200 train-split patients, k = 8 samples, T = 1.0, 20 turns, budget drawn per\npatient.\n\n| policy | mean reward | source |\n|---|---|---|\n| report the prior (blank-record floor) | −0.018 | `runs/phase3/prompted_baseline.json` |\n| vitals-only Bayes (no tests) | +0.675 | same |\n| greedy Bayes (5.76 tests on average) | +1.183 | same |\n| Qwen2.5-7B-Instruct, prompted | −0.689 (best@8 −0.480) | same |\n| + SFT | −0.664 | `docs/design-notes.md` (table only) |\n| + GRPO (99 steps, LoRA) | −0.569 (best@8 +0.034) | same |\n\nThe source repo also reports a mean expected Bayes ceiling of 1.722 on these patients. No\ntrained model cleared the blank-record floor on mean reward.\n\n## Known limitations\n\n- **Synthetic medicine.** The likelihoods are invented, and the Bayes baselines use the same\n  generative model as the data. Strong results mean the agent learned this mapping, not that\n  it can practise medicine.\n- **Treatment term vs proper scoring.** The treatment terms are linear in `mass_on_truth`, so\n  an agent that prescribes has a small incentive to put extra mass on conditions its\n  prescription treats. The diagnosis term dominates: moving mass that way loses more Brier\n  than it gains, as the exploit tests check. Still, a report followed by a prescription is\n  not exactly proper.\n- **Contraindications use ground-truth labs.** A penalty can fire for a lab the agent never\n  ordered. The source's docstring says contraindications are \"checkable by the agent\", which\n  is not quite true.\n- **Structured output depends on the provider.** If a provider ignores `response_format`,\n  format errors end episodes. Report the `decode_failure` rate alongside the reward.\n- vf-eval's trace-mean headline is turn-weighted (see *Reading eval output*).\n\n## Source and citation\n\nEnvironment, generator, observation filter and reward:\n[Manas-Ganti/healthcare-RL](https://github.com/Manas-Ganti/healthcare-RL), pinned at commit\n`5046c92`. Packaged for verifiers 0.4.0.\n","encoding":"utf-8","truncated":false,"total_bytes":13063},"status":null}