EHR-RobustGym: Benchmarking and Training Agents for Robust Clinical Reasoning

1 INTRODUCTION

• EHR-RobustGym contains 5,486 Clean-Noise pairs across six clinical intents and patient- and population-level queries, with three types of noise.Each pair preserves the clinical goal while changing one necessary condition, without modifying the underlying database.Source: §1 · p. 2
• The environment combines SQL/Python interaction with outcome verification, requiring correct answers on Clean questions and evidence-grounded abstention on Noise questions.A refusal or query failure alone does not count as successful handling of an unsupported request.Source: §1 · p. 2
• Average task success drops from 62.2% on Clean questions to 37.9% on Noise questions across proprietary and large-scale open-weight models.At k = 4, pass^k consistency falls below 50% for most evaluated models; training gains extend to five external EHR benchmarks.Source: §1 · p. 2
• The paper contributes a robustness benchmark, an interactive training environment, and evaluations of 12 LLMs and trajectory-based training.Training gains are reported on EHR-RobustGym and five external EHR benchmarks.Source: §1 · p. 2

2 RELATED WORK

• EHR agent research spans text-to-SQL, iterative code execution, interactive tasks, training, clarification, and longitudinal reasoning over hospital records.Related systems also analyze failed trajectories and use outcome-supervised verification.Source: §2 · p. 3
• Prior robustness work examines schema incompatibility, answerability, ambiguity, value mismatches, confidence-based abstention, and record verification.EHR-RobustGym instead studies schema-valid requests with unavailable records, missing or conflicting values, and unsatisfied relations among records.Source: §2 · p. 3
• EHR-RobustGym pairs unsupported requests with supported Clean requests and requires retrieved evidence to justify abstention.The environment also supports training on interaction trajectories and task outcomes for accurate answering and evidence-grounded abstention.Source: §2 · p. 3

3.1 PROBLEM FORMULATION

• The task formulation requires deriving an answer from database evidence or returning NULL with evidence identifying the unmet requirement.Record-level noise means a required record is absent; value-level noise means a required field is missing or conflicts with the question; query-level noise means available records do not satisfy the requested relations or conditions.Source: §3.1 · pp. 3–4

3.2 DATA CONSTRUCTION

• Clinical requests are grounded in executable queries, and each Noise counterpart changes exactly one necessary condition while preserving the clinical goal and remaining requirements.Pairs are retained only when reference SQL confirms a supported Clean answer and the need for abstention on its Noise counterpart.Source: §3.2 · p. 4
• A language model jointly rewrites each pair while preserving the reference SQL’s meaning and computation, followed by execution and semantic verification.Two reviewers independently audited 344 challenging test pairs, with 92.15% judged valid after adjudication.Source: §3.2 · p. 4
• The benchmark contains 5,486 Clean-Noise pairs across 51 programs and six clinical intents, with 4,478 training and 1,008 test pairs.The test set is balanced across patient- and population-level scopes and three Noise types; its in-domain and structure-disjoint tracks each contain 504 pairs.Source: §3.2 · p. 4

3.3 INTERACTIVE EVALUATION AND TRAINING

• Agents iteratively execute SQL or Python with feedback before submitting an answer and reasoning, while Clean and Noise counterparts appear in separate episodes.Condition labels, pairing information, reference SQL, and reference answers remain hidden during inference.Source: §3.3 · p. 5
• Outcome verification checks Clean answers against reference values, units, and records, and requires Noise abstentions to cite visible tool evidence of an unmet requirement.The scoring setup distinguishes answer matching from abstention validity.Source: §3.3 · p. 5
• Training supports supervised fine-tuning, preference optimization, online PPO and GRPO, and self-improvement from sampled trajectories.Online methods use the same success score, without intermediate or format rewards.Source: §3.3 · p. 5

4.2 EVALUATION RESULTS

• 62.2% average success on Clean questions falls to 37.9% on Noise questions across seven unadapted API-based models.All seven perform worse on Noise requests.Source: §4.2 · p. 7
• Query-level Noise is hardest for all seven Base models; GPT-5.4 reaches 43.2% success, versus 58.3% on Record-level and 65.8% on Value-level Noise.GPT-4.1 returns NULL on 80.6% of Noise requests, but only 35.3% of those responses pass evidence verification, yielding 28.5% success; GPT-5.4 Noise success also falls from 65.5% on patient-level to 46.0% on population-level requests.Source: §4.2 · p. 7

4.3 TRAINING EHR AGENTS WITH EHR-ROBUSTGYM

• SFT improves Clean and Noise success across all five training backbones; subsequent online RL is strongest on Qwen3, with PPO raising Qwen3-14B from 42.4%/34.7% to 70.4%/62.6%.SFT raises Qwen3-32B from 6.0%/0.6% to 48.8%/40.3% and Llama-3.1-8B from near-zero to 52.5%/31.8%; later optimization does not consistently improve both conditions.Source: §4.3 · p. 8
• At k = 4, GPT-5.4 has 82.4% Pass@4 but 44.5% pass^4, while Gemini has 83.9% and 58.3%, respectively.Additional sampling improves task coverage, but the models do not consistently complete the same tasks across repeated rollouts.Source: Figure 4 · p. 7, §4.3 · p. 8
• At 1,000 trajectories, Mixed training reaches 30.1%/29.4% Clean/Noise success, while single-condition training yields near-zero success on the other condition.With both trajectories per pair, increasing Mixed training from 500 to 2,000 trajectories raises Clean/Noise success from 16.9%/21.9% to 39.8%/42.8%; tool-call success alone does not ensure correct answers or evidence-grounded abstention.Source: Figure 4 · p. 7, §4.3 · p. 8
• RS+SFT improves Noise success for all five models, but after the full self-improvement sequence Qwen3-14B rises from 42.4%/34.7% to 57.1%/46.7% Clean/Noise success while Qwen3-8B and Qwen3-32B finish below their SFT Noise scores.Clean success improves after RS+SFT only for Qwen3-4B and Qwen3-32B; later DPO stages yield less consistent gains.Source: Figure 4 · p. 7, §4.3 · p. 9
• After GRPO, Qwen3-14B outperforms Base and SFT on all five external EHR benchmarks, including a rise from 35.4% to 51.2% on TREQS and from 6.8% to 9.0% on EHR-Complex.Across eight general and medical benchmarks, four of five models improve their average over Base after SFT and remain above Base after GRPO; Qwen3-4B is the exception.Source: Figure 6 · p. 8, Figure 5 · p. 8, §4.3 · p. 9
• For GPT-5.4, failures without tool execution errors account for 42.4% of Noise requests versus 20.2% of Clean requests; clinical interpretation and query logic errors make up 25.2% and 27.0% of its annotated error episodes.The analysis covers 12,096 Base trajectories from six models, and trajectories can contain multiple error episodes, including ones later corrected.Source: §4.3 · p. 9, Figure 7 · p. 9

• EHR-RobustGym jointly evaluates accurate answering on supported questions and evidence-grounded abstention on unsupported requests across Record-level, Value-level, and Query-level noise.Source: §5 · p. 9
• Evaluation reveals substantial Clean–Noise robustness gaps and unstable task completion across repeated executions.Source: §5 · p. 9
• Supervised fine-tuning and reinforcement learning improve performance, with gains generalizing to five external EHR benchmarks.Source: §5 · p. 9
• The conclusion positions EHR-RobustGym as a testbed for improving evidence-grounded robustness while preserving accurate answers to supported questions.Source: §5 · pp. 9–10

• AI assistants supported language polishing.Source: AI USE STATEMENT · p. 10
• Language models were used in task construction, trajectory generation, outcome verification, and error analysis.Source: AI USE STATEMENT · p. 10
• The authors state that they reviewed and verified the scientific claims, results, analyses, citations, and final manuscript content, and take responsibility for the work.Source: AI USE STATEMENT · p. 10

• EHR-RobustGym is constructed from de-identified MIMIC-IV v3.1 records that underwent PHI removal, random identifier replacement, and patient-level date shifting.Source: ETHICS STATEMENT · p. 10
• Access to underlying records remains subject to MIMIC-IV credentialing and data-use requirements, and released artifacts must comply with those requirements.Source: ETHICS STATEMENT · p. 10
• The authors describe EHR-RobustGym as a research testbed, not a clinical decision-support system or substitute for professional clinical judgment.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• Section 3 defines the clinical tasks, three noise types, Clean-Noise pairs, and the interactive environment and outcome verification used for evaluation and training.Appendix A covers construction, reference generation, rewriting, and data splits; Appendix B covers challenging-pair selection and human review.Source: REPRODUCIBILITY STATEMENT · p. 10
• Appendix A documents structure-disjoint evaluation and separates patient-level tasks; Appendix B describes human review criteria and adjudication.The statement also points to Appendix E.5 for repeated-run evaluation and Pass@k and pass^k computation.Source: REPRODUCIBILITY STATEMENT · p. 10
• Section 4 describes evaluated models, metrics, and training comparisons, while Appendices C–E provide evaluation protocols, training settings, and additional comparisons.Appendix F gives the error taxonomy and trajectory annotation procedure; Appendix H contains shared agent instructions and response format.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 DATABASE PREPROCESSING

• The conversion renames MIMIC-IV tables and columns while retaining all 342 source columns, their values, data types, and temporal precision; both conditions use the same converted database.Admission age is the only derived field, and fields absent from MIMIC-IV 3.1 are not reconstructed.Source: §A.1 · p. 16, Table 4 · p. 16

A.2 DATASET COMPOSITION

• The benchmark has 4,478 training and 1,008 test Clean-Noise pairs, covering six primary clinical intents across query scopes and noise types.Each pair has one primary intent; 94 laboratory/microbiology–medication pairs have two output categories.Source: §A.2 · pp. 16–17, Table 5 · p. 17, Table 6 · p. 17

A.3 CONSTRUCTION DETAILS

• Each pair changes one necessary condition while preserving its clinical topic, query scope, requested output, and other conditions; construction programs instantiate requests and reference SQL deterministically.References distinguish unsupported conditions from related evidence and distinguish zero eligible population units from one to four; both require NULL, while the latter also withholds exact counts and aggregates.Source: §A.3 · pp. 17–18

A.4 QUESTION VOCABULARY VISUALIZATION

• Figure 8 visualizes question vocabulary by noise type and Clean or Noise condition, with word size representing within-panel frequency.It uses question text only, excludes short and numeric tokens, and applies independent frequency scaling to each panel; the visualization does not establish equal difficulty or absence of leakage.Source: §A.4 · p. 18

A.5 DATA SPLITS, PROVENANCE, AND VALIDATION

• The test set divides evenly between 504 in-domain and 504 structure-disjoint pairs, with 252 patient-level and 252 population-level pairs in each track.The structure-disjoint track reserves six construction programs and separates several program, surface-form, SQL-structure, and semantic-contract identifiers; population-query patients may overlap across partitions.Source: §A.5 · p. 18

B HUMAN VALIDATION OF TASK QUALITY

• The audit selected 344 pairs containing jointly failed questions and reviewed both questions in each pair against evidence, reference outcomes, and pair consistency.Two annotators independently reviewed each pair; a pair was valid only when both questions and their pairing met the stated criteria.Source: B HUMAN VALIDATION OF TASK QUALITY · p. 20
• 92.15% of the 344 audited pairs were valid after adjudication but before correction, and the two annotators had 98.26% raw agreement before adjudication.The 27 pairs with confirmed issues were corrected, replayed, and rechecked; reported benchmark results use the finalized test set.Source: B HUMAN VALIDATION OF TASK QUALITY · pp. 20–21, Table 8 · p. 21

C.1 ENVIRONMENT AND INTERACTION SETTINGS

• Episodes use a task-scoped, read-only DuckDB connection, with persistent Python state within an episode and fresh execution state in each new episode.Model-generated errors and ordinary timeouts remain in trajectories as feedback; infrastructure or judge failures must be resolved before evaluation is reported.Source: §C.1 · p. 21, Table 9 · p. 21

C.2 MODELS AND SAMPLING SETTINGS

• The evaluation includes seven API-based models and five locally deployed open-weight models; the main tables use the first Base run under the stated sampling settings.API-based describes access, not whether model weights are public.Source: §C.2 · p. 21, Table 10 · p. 22

C.3 OUTCOME EVALUATION

• Clean scoring matches the submitted answer to the execution-verified reference, while Noise scoring requires NULL and visible tool evidence supporting abstention.A failed query or wrong-scope empty result alone does not establish abstention validity; population summaries with zero or one to four eligible units require withholding aggregates.Source: §C.3 · pp. 22–23

C.4 REPORTING AND INTERACTION STATISTICS

• Figure 7(a) groups first-run Base trajectories from eight models, using 1,008 questions per condition as the denominator for each outcome or error share.These interaction statistics are descriptive and do not add reward terms to task success.Source: §C.4 · p. 23, Table 11 · p. 23

C.5 CAPABILITY EVALUATION

• The capability evaluation covers eight general and medical benchmarks, with five external EHR benchmarks among the reported subsets.The reported mean weights the eight benchmarks equally rather than by question count.Source: §C.5 · p. 23, Table 12 · p. 24

D.1 TRAINING DATA AND CHECKPOINTS

• The 2,000 selected training demonstrations comprise 1,150 unguided and 850 guided examples, and candidates must pass task evaluation with successful tool-based evidence acquisition.SFT uses verified teacher trajectories; guided candidates receive additional naturalness and self-containment checks.Source: §D.1 · pp. 24–25

D.2 TRAINING CONFIGURATION

• DPO, PPO, and GRPO independently continue from SFT-250 for 250 updates, while online PPO and GRPO collect fresh rollouts through environment interaction.The training stages use separate question IDs, although paired counterparts may appear in different stages.Source: §D.1 · p. 25, §D.2 · p. 25

D.3 SELF-IMPROVEMENT

• Self-improvement samples eight trajectories per question at temperature 0.6, retaining successful trajectories for RS+SFT and successful and unsuccessful trajectories for DPO and iDPO.Each stage trains for one epoch using the shared outcome scorer for selection.Source: §D.3 · p. 25, Table 14 · p. 26

E ADDITIONAL RESULTS

• Tables 15 and 16 use the same evaluations as the main performance tables, with SFT denoting SFT-250 and DPO, PPO, and GRPO denoting independent continuations.Reporting follows the stated evaluation conventions.Source: E ADDITIONAL RESULTS · p. 26

E.1 PERFORMANCE BY NOISE TYPE

• +GRPO Success was 52.7/36.6 for Record-level, 51.5/47.0 for Value-level, and 51.2/34.8 for Query-level conditions, with values reported as Clean / Noise percentages.Each condition and noise type has 336 questions; NR is the NULL rate and AV is abstention validity conditional on NULL.Source: Table 15 · p. 26, §E.1 · p. 27

E.2 PERFORMANCE BY SCOPE, TEST TRACK, AND POPULATION SIZE

• 260 of 504 population Noise questions involve one to four eligible units, requiring evidence-grounded abstention without implying missing records; track score differences do not isolate structural generalization.Table 16 reports success by query scope, test track, and eligible population size.Source: Table 16 · p. 27, §E.2 · pp. 27–28

E.3 TRAINING-DATA SCALING AND SELF-IMPROVEMENT RESULTS

• Scaling separates paired demonstrations from one-per-pair comparisons: equal pair fractions use twice as many Mixed demonstrations, while Figure 4(b) matches 1,000 demonstrations, pair coverage, and 250 updates.Figure 9 reports paired Mixed scaling; its checkpoints train for one epoch, and bars show successful executions divided by tool calls.Source: §E.3 · pp. 28–29, Figure 9 · p. 29
• At an equal budget of 1,000 demonstrations, Mixed reaches 30.1% Clean and 29.4% Noise success, versus Clean-only at 35.4% and 1.1%, and Noise-only at 0.0% and 51.8%.The 500-complete-pair Mixed subset reaches 31.6%/28.3% with the same number of demonstrations but covers fewer pairs.Source: §E.3 · p. 29

E.4 CAPABILITY RETENTION

• Figure 10 compares DPO and PPO after 250 updates from SFT-250 across five models; Table 18 gives eight raw benchmark scores and their equal-weight mean.Radar radii are relative to each model’s Base score, with model-specific radial limits, so polygon areas are not directly comparable.Source: §E.4 · p. 29, Figure 10 · p. 30, Table 18 · p. 30

E.5 REPEATED EVALUATIONS

• Repeated-run scores combine four runs across sampling temperatures, so they describe observed run combinations rather than probabilities for a fixed-temperature distribution.Pass@k measures any-success and pass^k all-success across observed k-run subsets; at k = 1, both are the four-run mean.Source: §E.5 · p. 31
• Across 2,016 test questions, Qwen3-235B-A22B-Instruct-2507 and Qwen3-32B score 20.5% and 8.4% Pass@4, and 5.2% and 0.4% pass^4, respectively.These are observed-run coverage and consistency statistics, not post-training results.Source: §E.5 · p. 31

F TRAJECTORY ERROR ANALYSIS

• Error analysis covers 12,096 Base trajectories from six models; each distinct episode receives one category and its most specific supported subtype, with earlier-error consequences not counted again.Successful trajectories may contain corrected errors, and eventual task success alone does not establish correction.Source: F TRAJECTORY ERROR ANALYSIS · pp. 32–33

G.1 SELECTED MODEL TRAJECTORIES

• Three selected cases show that retrieved values can fail the request’s conditions: Record-, Value-, and Query-level Noise cases all receive the C2.3 temporal and logical constraint error label.The examples illustrate specific mechanisms, not their prevalence; tool calls succeeded without timeouts or truncated feedback.Source: §G.1 · pp. 33, 35

H.1 AGENT SYSTEM PROMPT AND RESPONSE FORMAT

• Agents must query the task-scoped read-only DuckDB connection sequentially, then return one tagged answer with reasoning; population aggregates require at least five qualifying units.The prompt also requires database inspection before answering and a complete result before the final response.Source: §H.1 · pp. 35–36

H.2 EVIDENCE-CHECK INSTRUCTION

• PE is a prompt-only baseline that appends an evidence-check instruction without changing tools or interaction limits.The instruction asks agents to check whether the EHR supports the requested information before answering.Source: §H.2 · p. 37

I LIMITATIONS AND SCOPE

• The benchmark uses de-identified MIMIC-IV v3.1 from one hospital system; broader institutional and clinical-setting generalization needs further evaluation.Free-text notes, imaging, real-time clinician interaction, and prospective deployment remain outside the current scope.Source: I LIMITATIONS AND SCOPE · p. 37