Synthetic quickstart

This walkthrough builds a small offline RL dataset from synthetic MIMIC-IV-style data. It needs only the base install: no MIMIC-IV access, no cloud credentials, and no d3rlpy.

The complete script

"""Synthetic quickstart shown in the documentation; needs only the base install."""

from __future__ import annotations

import numpy as np

from ehr2rl import BehaviorPolicy, MortalityReward, make_synthetic_dataset


def main() -> dict[str, int | tuple[int, int]]:
    # 25 deterministic synthetic ICU stays, 24 hourly steps each.
    ds = make_synthetic_dataset(n_patients=25, trajectory_length=24, seed=7)

    trajectory = ds[0]
    print(trajectory.metadata["feature_names"])  # state columns
    print(trajectory.states.shape)  # (T, D) = (24, 4)
    print(trajectory.actions.shape)  # (T, A) = (24, 1)

    # Estimate how often clinicians took each action in each state.
    policy = BehaviorPolicy().fit(ds)
    all_states = np.vstack([t.states for t in ds])
    probabilities = policy.predict_proba(all_states)
    print(probabilities.shape)  # (25 * 24 timesteps, n_actions) = (600, 2)

    # Terminal reward: +1 for survival, -1 for in-hospital death.
    ds = MortalityReward().shape(ds)
    print(ds[0].rewards[-1])

    summary: dict[str, int | tuple[int, int]] = {
        "n_trajectories": len(ds),
        "state_shape": trajectory.states.shape,
        "probability_rows": probabilities.shape[0],
    }
    print(summary)
    return summary


if __name__ == "__main__":
    main()

This is examples/docs_quickstart.py from the repository, shown exactly as the test suite runs it, so the output below is what you should see:

['heart_rate', 'mean_bp', 'lactate', 'creatinine']
(24, 4)
(24, 1)
(600, 2)
1.0
{'n_trajectories': 25, 'state_shape': (24, 4), 'probability_rows': 600}

What each step does

1. Generate trajectories

make_synthetic_dataset returns an EHRDataset: a list of PatientTrajectory objects, one per synthetic ICU stay. Data are deterministic for a given seed.

Each trajectory holds aligned arrays, where T is the number of timesteps:

Field

Shape

Contents

timestamps

(T,)

Unix seconds, one hour apart

states

(T, D)

Heart rate, mean blood pressure, lactate, creatinine

actions

(T, A)

One column: 1 when mean blood pressure is below 65, else 0

rewards

(T,)

All zero until a reward is applied

terminals

(T,)

True only at the last step

metadata carries the feature names, whether the patient died, and a synthetic SOFA score series for SofaReward.

2. Estimate the behavior policy

BehaviorPolicy().fit(ds) learns the probability of each observed action given the state. predict_proba returns one row per state and one column per action class, so stacking all 25 × 24 states gives shape (600, 2).

3. Shape rewards

MortalityReward().shape(ds) returns a new dataset whose trajectories carry +1 at the terminal step for survivors and -1 for in-hospital deaths. The original dataset is not modified.

Next steps