ehr2rl: A Python Library for Converting MIMIC-IV Electronic Health Records into Offline Reinforcement Learning Datasets

Why Do We Need ehr2rl?

In the field of medical machine learning, there is an awkward gap: supervised learning tools (such as predictive models) are already fairly mature, while offline reinforcement learning (offline RL) libraries assume you already have a clean "state-action-reward" dataset. But in reality, the preprocessing work between electronic health records (EHR) and RL datasets often requires researchers to repeatedly write similar code.

According to a Reddit user, he kept rewriting the same preprocessing pipeline in his research projects: extracting data from EHR tables, aligning vital signs and lab metrics on an hourly basis, defining what counts as an "action," attaching reward signals, and finally converting everything into a format that d3rlpy can accept. Unable to find an existing library to handle this "middle layer" work, he developed ehr2rl.

The goal of this tool is clear: to fill the gap between raw EHR data and offline RL training, so researchers don't have to reinvent the wheel.

Offline reinforcement learning (offline RL) is a branch of reinforcement learning whose core premise is that the agent can only learn a policy from a pre-collected, fixed historical interaction dataset and cannot interact with the environment online to gather new data. This stands in sharp contrast to online RL, which requires real-time trial and error. In medical scenarios, online trial and error is infeasible on ethical and safety grounds, making offline RL a natural choice for learning treatment policies from EHR data. However, offline RL algorithms typically assume the dataset is already organized as "state-action-reward-termination" tuples, whereas raw EHR data is event-stream tabular records. A huge engineering gap exists between the two, and this is exactly the gap ehr2rl attempts to bridge.

Core Features of ehr2rl

ehr2rl builds hourly patient trajectories from MIMIC-IV v3.1 via BigQuery, with each trajectory containing states, actions, rewards, and termination flags. Several design details are worth noting:

• Query safety mechanism: Every query is first executed as a dry run with a hard maximum_bytes_billed cap. This means that even if a cohort query is written incorrectly, it won't silently rack up a huge bill.
• Action construction: Vasopressor and fluid doses are extracted from ICU input events and discretized according to user-defined binning boundaries.
• Replaceable reward functions: Supports mortality, per-step SOFA change, 30-day readmission, or weighted combinations of these rewards.
• Behavior policy estimation: Estimates the observed clinician (behavior) policy from the data, which is needed for off-policy evaluation.
• Direct export: Outputs as a d3rlpy dataset, with an optional JSON sidecar file recording data provenance.

In addition, the project provides a synthetic MIMIC-style data generator so you can try it out even without MIMIC access.

d3rlpy is an open-source offline deep reinforcement learning library that implements a variety of mainstream offline RL algorithms such as CQL, BCQ, and IQL, and provides a unified training and evaluation interface. It requires input data to be an MDPDataset conforming to its format, i.e., a collection of trajectories composed of states, actions, rewards, and termination flags. ehr2rl exports preprocessing results directly as a d3rlpy dataset, meaning users can directly invoke algorithms in d3rlpy for policy training without manually converting formats. This design choice lowers the end-to-end barrier from EHR to policy learning and keeps ehr2rl ecologically compatible with existing offline RL algorithm implementations.

Behavior policy refers to the policy that actually made decisions during the historical data collection process—in medical scenarios, the prescribing and treatment habits of clinicians. The goal of off-policy evaluation (OPE) is to estimate the performance of a new policy (target policy) without conducting online experiments. Since offline data is generated by the behavior policy, there is a distribution shift between the data distribution and the target policy's distribution. OPE methods typically need to explicitly model or estimate the behavior policy in order to compute importance weights or perform distribution correction. The behavior policy estimation feature provided by ehr2rl is precisely intended to supply the necessary inputs for subsequent reliable offline evaluation, rather than merely serving policy training itself.

Key Trade-offs in Design

The author particularly wants feedback on reward design and action discretization. These two issues indeed have no standard answers in clinical offline RL.

Regarding reward design, mortality is a sparse and delayed reward, while SOFA change provides a denser signal but may be inconsistent with long-term outcomes. A weighted mixed reward seems like a compromise, but the choice of weights itself lacks theoretical guidance. Action discretization is equally tricky: the binning boundaries for vasopressor and fluid doses directly affect the results of policy learning, and the current binning approach is more heuristic than principled.

From an engineering perspective, ehr2rl chooses BigQuery as the data access layer, which lowers the barrier to local storage and processing but also binds users to the Google Cloud ecosystem. For researchers with MIMIC-IV access, this is a pragmatic trade-off.

The SOFA (Sequential Organ Failure Assessment) score is a widely used tool in intensive care for assessing organ dysfunction, combining indicators from six systems: respiratory, coagulation, liver, cardiovascular, central nervous system, and renal. Higher scores indicate more severe organ dysfunction. In ehr2rl, the per-step SOFA change is used as a dense reward signal: a decrease in SOFA usually means clinical improvement. Compared with a sparse reward based solely on death as the endpoint, SOFA change provides feedback at every step of the trajectory, which is conducive to learning. However, SOFA is a summary of short-term physiological status and is not fully consistent with long-term functional recovery or quality of survival, so using it as a reward requires careful trade-offs.

Applicable Boundaries and Usage Recommendations

The author explicitly states that ehr2rl is a research tool and is not intended for treatment decisions. This point is crucial for medical AI tools. The application of offline RL in clinical scenarios is still in its early stages, and the reliability of behavior policy estimation and offline evaluation needs careful validation.

For researchers who want to try it, it is recommended to start with synthetic data to understand the workflow of trajectory construction and reward definition, and then move on to real MIMIC-IV data. For action discretization, you can try different binning schemes and observe whether the policy evaluation results are robust. For reward design, it is recommended to start with a single reward (such as mortality) and gradually introduce more complex combinations, avoiding increasing the number of tuning dimensions too early.

The emergence of ehr2rl reflects a shift in medical RL research from "algorithm-centric" to "data-pipeline-centric." A reliable preprocessing layer may do more to advance the field than a new algorithm.

One-line Summary

ehr2rl fills the preprocessing gap between raw EHR data and offline RL datasets, supporting the construction of hourly patient trajectories from MIMIC-IV.

Paragraph Summary

ehr2rl is a research tool for medical offline reinforcement learning, designed to address the pain point of repeatedly hand-writing preprocessing code between electronic health records (EHR) and RL training datasets. It builds hourly patient trajectories from MIMIC-IV v3.1 via BigQuery, containing states, actions, rewards, and termination flags, and includes built-in query safety mechanisms, customizable binning-based action discretization, replaceable reward functions (mortality, SOFA change, readmission, etc.), and behavior policy estimation, ultimately exporting to d3rlpy format. The author openly seeks feedback on reward design and action discretization, and emphasizes that the tool is for research only and not for treatment decisions.