Rewards

See Rewards for the metadata each reward needs and how to write your own.

class ehr2rl.reward.BaseReward

Base class for reward functions.

Subclasses implement compute; shape applies it to every trajectory in a dataset.

abstractmethod compute(trajectory)

Return one reward per timestep for a trajectory.

Parameters:

trajectory (PatientTrajectory) – Trajectory to score. It must not be modified.

Returns:

Shape (T,).

Return type:

ndarray

shape(dataset, policy=None)

Return a copy of a dataset with rewards computed for every trajectory.

Parameters:

  • dataset (EHRDataset) – Input dataset. It is not modified.

  • policy (object | None) – Reserved for policy-dependent rewards. Ignored by all built-in rewards.

Returns:

A new dataset with new trajectories.

Return type:

EHRDataset

class ehr2rl.MortalityReward(survival_reward=1.0, death_penalty=-1.0)

Terminal reward for in-hospital survival.

Reads metadata["died"]; a missing key counts as survival.

Parameters:

  • survival_reward (float) – Reward at the final step for survivors.

  • death_penalty (float) – Reward at the final step for in-hospital deaths.

compute(trajectory)

Place the survival reward or death penalty at the last terminal step.

Parameters:

trajectory (PatientTrajectory) – Trajectory with at least one True terminal.

Returns:

Shape (T,), zero except at the last terminal step.

Return type:

ndarray

class ehr2rl.SofaReward(clip=(-5.0, 5.0))

Dense reward for change in SOFA score; improvement is positive.

Parameters:

clip (tuple[float, float] | None) – (low, high) bounds for each step’s reward, or None for no clipping.

compute(trajectory)

Return sofa[t - 1] - sofa[t] per step, with 0 at the first step.

Parameters:

trajectory (PatientTrajectory) – Trajectory with metadata["sofa_scores"] of shape (T,).

Returns:

Shape (T,), clipped to clip.

Return type:

ndarray

Raises:

ValueError – If sofa_scores is missing or has the wrong shape.

class ehr2rl.ReadmissionReward(window_days=30, readmitted_reward=-1.0, not_readmitted_reward=1.0, censored_value=nan)

Terminal reward for 30-day hospital readmission.

Reads metadata["readmitted_within_30_days"] (True, False, or None when follow-up is censored). Patients who died get 0.

Parameters:

  • window_days (int) – Stored for reference but not used; the window is set by the metadata.

  • readmitted_reward (float) – Reward at the final step when readmitted.

  • not_readmitted_reward (float) – Reward at the final step when not readmitted.

  • censored_value (float) – Reward at the final step when follow-up is censored. The default, NaN, must be replaced or filtered out before training.

compute(trajectory)

Place the readmission reward at the final step.

Parameters:

trajectory (PatientTrajectory) – Trajectory with metadata["readmitted_within_30_days"].

Returns:

Shape (T,), zero except at the final step.

Return type:

ndarray

Raises:

ValueError – If readmitted_within_30_days is missing.

class ehr2rl.CompositeReward(components)

Weighted sum of other reward functions.

Parameters:

components (Sequence[tuple[BaseReward, float]]) – (reward, weight) pairs. At least one is required.

Raises:
compute(trajectory)

Return the weighted sum of every component’s rewards.

Parameters:

trajectory (PatientTrajectory) – Trajectory with whatever metadata the components need.

Returns:

Shape (T,).

Return type:

ndarray

Raises:

ValueError – If a component returns the wrong shape.